Knowledge Distillation(지식 증류)
1. 왜 Knowledge Distillation(지식 증류)이 필요한가요?
딥러닝 모델은 일반적으로 크기가 커질수록 성능이 좋아지는 경향이 있습니다.
하지만 큰 모델은 다음과 같은 문제가 있습니다.
- 파라미터(Parameter)가 많아 메모리를 많이 사용합니다.
- 연산량(FLOPs)이 커서 추론(Inference)이 느립니다.
- 모바일(Mobile), 드론(Drone), 엣지 디바이스(Edge Device)처럼 자원이 제한된 환경에서 사용하기 어렵습니다.
- 서버 비용과 전력 소모가 커집니다.
이 문제를 해결하기 위해 Efficient AI(효율적인 AI)에서는 대표적으로 다음 방법들을 사용합니다.
| 방법 | 핵심 아이디어 | 예시 |
| Pruning(가지치기) | 중요하지 않은 Weight(가중치)나 Channel(채널)을 제거 | 0에 가까운 Weight 제거 |
| Quantization(양자화) | FP32 값을 INT8 등 낮은 정밀도로 변환 | 32bit → 8bit |
| Knowledge Distillation(지식 증류) | 큰 모델의 지식을 작은 모델에 전달 | Teacher → Student |
여기서 Knowledge Distillation(지식 증류)은 큰 Teacher Model(교사 모델)이 학습한 판단 방식과 표현을 작은 Student Model(학생 모델)이 따라 배우도록 하는 방법입니다.
loss = alpha * student_loss + beta * distillation_loss 와 같은 형식
2. Knowledge Distillation(지식 증류)의 기본 구조
Knowledge Distillation(지식 증류)의 기본 구조는 다음과 같습니다.
- Teacher Model(교사 모델): 크고 성능이 좋은 모델입니다.
- Student Model(학생 모델): 작고 빠르지만 성능이 낮을 수 있는 모델입니다.
- Distillation Loss(증류 손실): Student가 Teacher의 출력을 따라가도록 만드는 손실입니다.
- Classification Loss(분류 손실): 실제 정답 Label(레이블)을 맞추기 위한 일반적인 손실입니다.
전체 학습 목표는 보통 다음과 같습니다.
Total Loss = α × Classification Loss + β × Distillation Loss즉, Student Model(학생 모델)은 다음 두 가지를 동시에 학습합니다.
- 실제 정답 Label(레이블)을 맞추는 방법
- Teacher Model(교사 모델)이 각 Class(클래스)를 어떻게 판단했는지
3. Hard Label(하드 레이블)과 Soft Label(소프트 레이블)
일반적인 분류 학습에서는 정답이 One-hot(원-핫) 형태입니다.
고양이 이미지 정답 = [1, 0, 0]이런 정답을 Hard Label(하드 레이블)이라고 부를 수 있습니다.
하지만 Teacher Model(교사 모델)은 단순히 정답 하나만 알려주는 것이 아니라, 각 Class(클래스)에 대한 확률분포(Probability Distribution)를 출력합니다.
Teacher 출력 예시 : [0.70, 0.20, 0.10]
고양이: 0.70, 강아지: 0.20, 토끼: 0.10이런 출력을 Soft Label(소프트 레이블) 또는 Soft Target(소프트 타깃)이라고 합니다.
Soft Label(소프트 레이블)에는 다음 정보가 들어 있습니다.
- 정답 Class(클래스)에 대한 확신도
- 오답 Class(클래스) 사이의 상대적 유사도
- Teacher Model(교사 모델)이 학습한 데이터 구조
예를 들어 고양이 이미지를 보고 Teacher가 다음처럼 판단했다고 하겠습니다.
고양이: 0.70
강아지: 0.25
자동차: 0.05이 경우 Student는 단순히 “고양이가 정답이다”만 배우는 것이 아니라, 고양이는 강아지와 더 비슷하고 자동차와는 덜 비슷하다는 관계도 함께 배울 수 있습니다.
4. Temperature(온도)란 무엇인가요?
Knowledge Distillation(지식 증류)에서 중요한 개념이 Temperature(온도)입니다.
Softmax(소프트맥스)는 Logits(로짓)을 확률분포로 바꾸는 함수입니다.
softmax(z_i) = exp(z_i) / sum(exp(z_j))Temperature(온도)를 적용하면 다음처럼 계산합니다.
softmax(z_i / T)여기서 T가 Temperature(온도)입니다.
| Temperature(온도) | 결과 |
| T = 1 | 일반 Softmax(소프트맥스) |
| T > 1 | 확률분포가 부드러워짐 |
| T가 매우 큼 | Class 간 차이가 완만해짐 |
예를 들어 Logits(로짓)이 다음과 같다고 하겠습니다.
Cat: 5
Dog: 1T=1이면 Cat 확률이 매우 높게 나옵니다.T=10이면 Cat과 Dog의 차이가 완화되어 더 부드러운 분포가 됩니다.
이렇게 부드러운 분포를 사용하면 Student Model(학생 모델)이 Teacher Model(교사 모델)의 세부 판단 정보를 더 잘 배울 수 있습니다.
5. PyTorch Tutorial(파이토치 튜토리얼)의 실험 구성
PyTorch Tutorial(파이토치 튜토리얼)에서는 CIFAR-10 Dataset(데이터셋)을 사용합니다.
코드 : 원본 코드를 한글로 번역 후 수행한 결과입니다.
- CIFAR-10은 10개 Class(클래스)를 가진 이미지 분류 데이터셋입니다.
- 입력 이미지는 RGB 3채널, 32×32 크기입니다.
- Teacher Model(교사 모델)은 더 깊은 CNN(합성곱 신경망)을 사용합니다.
- Student Model(학생 모델)은 더 작은 CNN(합성곱 신경망)을 사용합니다.
5.1 기본 Import(임포트)
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import torchvision.transforms as transforms
import torchvision.datasets as datasets5.2 Device(디바이스) 설정
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using {device} device")GPU가 있으면 CUDA(쿠다)를 사용하고, 없으면 CPU를 사용합니다.
5.3 CIFAR-10 Dataset(데이터셋) 로딩
transforms_cifar = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]
),
])
train_dataset = datasets.CIFAR10(
root="./data",
train=True,
download=True,
transform=transforms_cifar
)
test_dataset = datasets.CIFAR10(
root="./data",
train=False,
download=True,
transform=transforms_cifar
)
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=128,
shuffle=True,
num_workers=2
)
test_loader = torch.utils.data.DataLoader(
test_dataset,
batch_size=128,
shuffle=False,
num_workers=2
)6. Teacher Model(교사 모델)과 Student Model(학생 모델)
6.1 DeepNN(깊은 신경망): Teacher Model(교사 모델)
Teacher Model(교사 모델)은 더 깊고 넓은 CNN(합성곱 신경망)입니다.
class DeepNN(nn.Module):
def __init__(self, num_classes=10):
super().__init__()
self.features = nn.Sequential(
nn.Conv2d(3, 128, kernel_size=3, padding=1),
nn.ReLU(),
nn.Conv2d(128, 64, kernel_size=3, padding=1),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2),
nn.Conv2d(64, 64, kernel_size=3, padding=1),
nn.ReLU(),
nn.Conv2d(64, 32, kernel_size=3, padding=1),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2),
)
self.classifier = nn.Sequential(
nn.Linear(2048, 512),
nn.ReLU(),
nn.Dropout(0.1),
nn.Linear(512, num_classes)
)
def forward(self, x):
x = self.features(x)
x = torch.flatten(x, 1)
x = self.classifier(x)
return x6.2 LightNN(가벼운 신경망): Student Model(학생 모델)
Student Model(학생 모델)은 더 작고 빠른 CNN(합성곱 신경망)입니다.
class LightNN(nn.Module):
def __init__(self, num_classes=10):
super().__init__()
self.features = nn.Sequential(
nn.Conv2d(3, 16, kernel_size=3, padding=1),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2),
nn.Conv2d(16, 16, kernel_size=3, padding=1),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2),
)
self.classifier = nn.Sequential(
nn.Linear(1024, 256),
nn.ReLU(),
nn.Dropout(0.1),
nn.Linear(256, num_classes)
)
def forward(self, x):
x = self.features(x)
x = torch.flatten(x, 1)
x = self.classifier(x)
return x6.3 파라미터 수 비교
PyTorch Tutorial(파이토치 튜토리얼)의 예시에서는 Teacher Model(교사 모델)이 Student Model(학생 모델)보다 훨씬 많은 파라미터를 가집니다.
| 모델 | 역할 | 특징 |
| DeepNN | Teacher Model(교사 모델) | 크고 성능이 좋음 |
| LightNN | Student Model(학생 모델) | 작고 빠름 |
total_params_deep = sum(p.numel() for p in nn_deep.parameters())
total_params_light = sum(p.numel() for p in nn_light.parameters())
print(f"DeepNN parameters: {total_params_deep:,}")
print(f"LightNN parameters: {total_params_light:,}")DeepNN parameters: 1,186,986
LightNN parameters: 267,738
7. Baseline: CrossEntropyLoss(교차 엔트로피 손실)만 사용
먼저 Teacher Model(교사 모델)과 Student Model(학생 모델)을 각각 일반적인 CrossEntropyLoss(교차 엔트로피 손실)로 학습합니다.
def train(model, train_loader, epochs, learning_rate, device):
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)
model.train()
for epoch in range(epochs):
running_loss = 0.0
for inputs, labels in train_loader:
inputs = inputs.to(device)
labels = labels.to(device)
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
print(f"Epoch {epoch+1}/{epochs}, Loss: {running_loss / len(train_loader):.4f}")
평가 함수는 다음과 같습니다.
def test(model, test_loader, device):
model.to(device)
model.eval()
correct = 0
total = 0
with torch.no_grad():
for inputs, labels in test_loader:
inputs = inputs.to(device)
labels = labels.to(device)
outputs = model(inputs)
_, predicted = torch.max(outputs, dim=1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
accuracy = 100 * correct / total
print(f"Test Accuracy: {accuracy:.2f}%")
return accuracy기준 실험의 의미는 다음과 같습니다.
- Teacher Model(교사 모델)의 성능을 먼저 확인합니다.
- Student Model(학생 모델)이 혼자 학습했을 때 성능을 확인합니다.
- 이후 Knowledge Distillation(지식 증류)을 적용했을 때 Student 성능이 좋아지는지 비교합니다.
8. Output Distillation(출력 증류): Teacher의 Logits(로짓)을 따라가기
이제 본격적으로 Knowledge Distillation(지식 증류)을 적용합니다.
8.1 핵심 아이디어
Teacher Model(교사 모델)의 Logits(로짓)을 Temperature(온도)로 부드럽게 만든 뒤 Soft Target(소프트 타깃)을 만듭니다.
Student Model(학생 모델)도 같은 Temperature(온도)를 적용하여 Log Probability(로그 확률)를 만듭니다.
soft_targets = F.softmax(teacher_logits / T, dim=-1)
soft_prob = F.log_softmax(student_logits / T, dim=-1)
이후 다음 KL Divergence(KL 발산) 형태의 손실을 계산합니다.
soft_targets_loss = torch.sum(
soft_targets * (soft_targets.log() - soft_prob)
) / soft_prob.size(0) * (T ** 2)
이 코드는 수식으로 보면 다음과 같습니다.
KL(Teacher || Student)
= sum Teacher_prob × (log(Teacher_prob) - log(Student_prob))8.2 왜 Student는 log_softmax(로그 소프트맥스)를 사용하나요?
PyTorch에서 KLDivLoss(KL 발산 손실)를 사용할 때는 보통 다음 형태를 사용합니다.
student_log_probs = F.log_softmax(student_logits / T, dim=-1)
teacher_probs = F.softmax(teacher_logits / T, dim=-1)
loss_kd = F.kl_div(
student_log_probs,
teacher_probs,
reduction="batchmean"
) * (T ** 2)| 대상 | 처리 |
| Teacher Logits(교사 로짓) | softmax()로 Probability(확률) 변환 |
| Student Logits(학생 로짓) | log_softmax()로 Log Probability(로그 확률) 변환 |
| Loss(손실) | KLDivLoss 또는 직접 구현한 KL Divergence |
8.3 Distillation 학습 코드
def train_knowledge_distillation(
teacher,
student,
train_loader,
epochs,
learning_rate,
T,
soft_target_loss_weight,
ce_loss_weight,
device
):
ce_loss = nn.CrossEntropyLoss()
optimizer = optim.Adam(student.parameters(), lr=learning_rate)
teacher.eval()
student.train()
for epoch in range(epochs):
running_loss = 0.0
for inputs, labels in train_loader:
inputs = inputs.to(device)
labels = labels.to(device)
optimizer.zero_grad()
# Teacher Model(교사 모델)은 학습하지 않습니다.
with torch.no_grad():
teacher_logits = teacher(inputs)
# Student Model(학생 모델)은 학습 대상입니다.
student_logits = student(inputs)
# Teacher Soft Target(교사 소프트 타깃)
soft_targets = F.softmax(teacher_logits / T, dim=-1)
# Student Log Probability(학생 로그 확률)
soft_prob = F.log_softmax(student_logits / T, dim=-1)
# Distillation Loss(증류 손실)
soft_targets_loss = torch.sum(
soft_targets * (soft_targets.log() - soft_prob)
) / soft_prob.size(0) * (T ** 2)
# Classification Loss(분류 손실)
label_loss = ce_loss(student_logits, labels)
# Total Loss(전체 손실)
loss = (
soft_target_loss_weight * soft_targets_loss
+ ce_loss_weight * label_loss
)
loss.backward()
optimizer.step()
running_loss += loss.item()
print(f"Epoch {epoch+1}/{epochs}, Loss: {running_loss / len(train_loader):.4f}")8.4 사용 예시
torch.manual_seed(42)
student_kd = LightNN(num_classes=10).to(device)
train_knowledge_distillation(
teacher=nn_deep,
student=student_kd,
train_loader=train_loader,
epochs=10,
learning_rate=0.001,
T=2,
soft_target_loss_weight=0.25,
ce_loss_weight=0.75,
device=device
)
test_accuracy_light_kd = test(student_kd, test_loader, device)9. KLDivLoss(KL 발산 손실)를 사용할 때 주의할 점
Output Distillation(출력 증류)에서는 Teacher와 Student의 최종 출력 분포를 맞추기 때문에 KLDivLoss(KL 발산 손실)가 자연스럽습니다.
9.1 KLDivLoss의 입력 형태
loss = F.kl_div(input, target, reduction="batchmean")| 인자 | 기대 형태 | KD에서의 의미 |
input |
Log Probability(로그 확률) | Student 예측 |
target |
Probability(확률) | Teacher 기준 |
reduction="batchmean" |
전체 KL 합 / Batch Size(배치 크기) | 샘플별 KL 평균 |
예시는 다음과 같습니다.
teacher_probs = F.softmax(teacher_logits / T, dim=-1)
student_log_probs = F.log_softmax(student_logits / T, dim=-1)
loss_kd = F.kl_div(
student_log_probs,
teacher_probs,
reduction="batchmean"
) * (T ** 2)9.2 Histogram(히스토그램) 출력일 때
최종 출력이 Histogram(히스토그램) 형태라면, 먼저 Row-wise L1 Normalization(행 단위 L1 정규화)을 통해 확률분포로 바꾸는 것이 좋습니다.
eps = 1e-8
teacher_probs = F.normalize(teacher_hist + eps, p=1, dim=-1)
student_probs = F.normalize(student_hist + eps, p=1, dim=-1)
student_log_probs = torch.log(student_probs)
loss_hist = F.kl_div(
student_log_probs,
teacher_probs,
reduction="batchmean"
)주의할 점은 LayerNorm(레이어 정규화)처럼 평균을 빼고 표준편차로 나누는 방식은 KLDivLoss(KL 발산 손실)에 적합하지 않다는 점입니다. LayerNorm(레이어 정규화)은 음수 값을 만들 수 있고, 합이 1인 확률분포를 보장하지 않습니다.
10. Feature Distillation(특징 증류): 중간 표현을 따라가기
Knowledge Distillation(지식 증류)은 최종 출력만 맞추는 것이 아닙니다. MIT Efficient AI Lecture 10에서는 “무엇을 Match(매칭)할 것인가?”를 중요한 주제로 다룹니다.
대표적인 대상은 다음과 같습니다.
| Match 대상 | 설명 | 사용 가능한 Loss |
| Logits(로짓) | 최종 출력 전 값 | KLDivLoss |
| Probability(확률분포) | Softmax 출력 | KLDivLoss |
| Feature Map(특징 맵) | 중간 Layer(레이어)의 출력 | MSELoss, CosineEmbeddingLoss |
| Attention Map(어텐션 맵) | 모델이 주목한 영역 | MSELoss, L2 Loss |
| Sparsity Pattern(희소 패턴) | ReLU 활성화 여부 | L1 Loss |
| Relational Information(관계 정보) | 샘플 간 거리/각도 관계 | Distance Loss |
11. CosineEmbeddingLoss(코사인 임베딩 손실): Hidden Representation(은닉 표현)
PyTorch Tutorial(파이토치 튜토리얼)에서는 Teacher와 Student의 Hidden Representation(은닉 표현)을 맞추기 위해 CosineEmbeddingLoss(코사인 임베딩 손실)를 사용합니다.
11.1 왜 CosineEmbeddingLoss(코사인 임베딩 손실)를 사용하나요?
중간 Feature(특징)는 확률분포가 아닙니다. 따라서 KLDivLoss(KL 발산 손실)를 사용하는 것보다 벡터 방향을 비교하는 CosineEmbeddingLoss(코사인 임베딩 손실)가 더 자연스럽습니다.
| 비교 대상 | 추천 Loss |
| 최종 출력 Probability Distribution(확률분포) | KLDivLoss |
| Histogram(히스토그램) 분포 | KLDivLoss |
| Hidden State(은닉 상태) | CosineEmbeddingLoss |
| Context Vector(문맥 벡터) | CosineEmbeddingLoss |
| Feature Map(특징 맵) 값 자체 | MSELoss |
11.2 Teacher와 Student의 차원 맞추기
Teacher Model(교사 모델)의 Hidden Representation(은닉 표현)과 Student Model(학생 모델)의 Hidden Representation(은닉 표현)은 차원이 다를 수 있습니다.
이 경우 다음 방법을 사용할 수 있습니다.
- Pooling(풀링)으로 Teacher 차원을 줄입니다.
- Projection Layer(투영 레이어)를 사용합니다.
- Regressor(회귀기)를 Student 쪽에 붙여 Teacher 차원으로 변환합니다.
PyTorch Tutorial(파이토치 튜토리얼)에서는 Teacher의 Flatten Vector(평탄화 벡터)를 Average Pooling(평균 풀링)하여 Student와 차원을 맞춥니다.
11.3 Teacher Model(선생님 모델)
class ModifiedDeepNNCosine(nn.Module):
def __init__(self, num_classes=10):
super().__init__()
self.features = nn.Sequential(
nn.Conv2d(3, 128, kernel_size=3, padding=1),
nn.ReLU(),
nn.Conv2d(128, 64, kernel_size=3, padding=1),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2),
nn.Conv2d(64, 64, kernel_size=3, padding=1),
nn.ReLU(),
nn.Conv2d(64, 32, kernel_size=3, padding=1),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2),
)
self.classifier = nn.Sequential(
nn.Linear(2048, 512),
nn.ReLU(),
nn.Dropout(0.1),
nn.Linear(512, num_classes)
)
def forward(self, x):
x = self.features(x)
flattened_conv_output = torch.flatten(x, 1)
logits = self.classifier(flattened_conv_output)
# Teacher의 2048차원을 Student의 1024차원에 맞추기 위해 평균 풀링을 사용합니다.
hidden = F.avg_pool1d(flattened_conv_output, kernel_size=2)
return logits, hidden11.4 Student Model(학생 모델)
class ModifiedLightNNCosine(nn.Module):
def __init__(self, num_classes=10):
super().__init__()
self.features = nn.Sequential(
nn.Conv2d(3, 16, kernel_size=3, padding=1),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2),
nn.Conv2d(16, 16, kernel_size=3, padding=1),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2),
)
self.classifier = nn.Sequential(
nn.Linear(1024, 256),
nn.ReLU(),
nn.Dropout(0.1),
nn.Linear(256, num_classes)
)
def forward(self, x):
x = self.features(x)
hidden = torch.flatten(x, 1)
logits = self.classifier(hidden)
return logits, hidden11.5 Cosine Loss 학습 코드
def train_cosine_loss(
teacher,
student,
train_loader,
epochs,
learning_rate,
hidden_rep_loss_weight,
ce_loss_weight,
device
):
ce_loss = nn.CrossEntropyLoss()
cosine_loss = nn.CosineEmbeddingLoss()
optimizer = optim.Adam(student.parameters(), lr=learning_rate)
teacher.eval()
student.train()
for epoch in range(epochs):
running_loss = 0.0
for inputs, labels in train_loader:
inputs = inputs.to(device)
labels = labels.to(device)
optimizer.zero_grad()
with torch.no_grad():
_, teacher_hidden = teacher(inputs)
student_logits, student_hidden = student(inputs)
# target=1은 두 벡터가 비슷해지도록 학습하라는 의미입니다.
target = torch.ones(inputs.size(0)).to(device)
hidden_loss = cosine_loss(
student_hidden,
teacher_hidden,
target
)
label_loss = ce_loss(student_logits, labels)
loss = (
hidden_rep_loss_weight * hidden_loss
+ ce_loss_weight * label_loss
)
loss.backward()
optimizer.step()
running_loss += loss.item()
print(f"Epoch {epoch+1}/{epochs}, Loss: {running_loss / len(train_loader):.4f}")11.6 사용 예시
modified_teacher = ModifiedDeepNNCosine(num_classes=10).to(device)
modified_teacher.load_state_dict(nn_deep.state_dict())
modified_student = ModifiedLightNNCosine(num_classes=10).to(device)
train_cosine_loss(
teacher=modified_teacher,
student=modified_student,
train_loader=train_loader,
epochs=10,
learning_rate=0.001,
hidden_rep_loss_weight=0.25,
ce_loss_weight=0.75,
device=device
)
test_accuracy_cosine = test(modified_student, test_loader, device)12. Intermediate Regressor(중간 회귀기): FitNets 방식
CosineEmbeddingLoss(코사인 임베딩 손실)로 단순히 Hidden Representation(은닉 표현)을 맞추는 것은 항상 좋은 결과를 보장하지 않습니다.
이유는 다음과 같습니다.
- Teacher와 Student는 구조가 다릅니다.
- 같은 의미를 표현하더라도 Feature 순서나 채널 구성이 다를 수 있습니다.
- Student가 Teacher의 표현을 1:1로 따라가는 것이 오히려 부담이 될 수 있습니다.
- 차원이 다르면 Loss를 직접 계산하기 어렵습니다.
그래서 FitNets 계열에서는 Regressor(회귀기) 또는 Projection Layer(투영 레이어)를 사용하여 Student의 Feature Map(특징 맵)을 Teacher의 Feature Map(특징 맵) 크기에 맞춥니다.
이 그림은 Teacher와 Student의 중간 Feature Map(특징 맵)을 맞추는 구조입니다.
Teacher Feature Map
↑
MSE Loss
↓
Student Feature Map 또는 Regressor를 통과한 Student Feature MapFitNets에서는 이 부분을 Hint-based Training(힌트 기반 학습) 이라고 볼 수 있습니다.
수식으로는 대략 다음과 같습니다.
$$L_{hint} = || F_T - r(F_S) ||^2$$
FitNets는 단순히 Teacher(교사 모델)의 최종 출력만 따라 하는 방식이 아닙니다.
- Teacher의 중간 Layer 출력 = Hint
- Student의 중간 Layer 출력 = Guided Layer
보통 학습을 두 단계로 설명합니다.
| 단계 | 학습내용 | Loss |
| 1단계 | Student의 중간 표현이 Teacher의 Hint를 따라가도록 학습 | MSE / L2 Hint Loss |
| 2단계 | Student가 최종 분류도 잘하도록 학습 | CE Loss + KD Loss |
그런데 Teacher와 Student는 구조가 다르기 때문에 중간 Feature의 shape이 다를 수 있습니다.
예를 들면 Teacher [batch, 256, H, W], Student [batch, 64, H, W] 인 경우 이 상태에서는 바로 비교할 수 없습니다.
그래서 Student 중간 Feature 뒤에 작은 변환층을 붙입니다.
이 Intermediate Regressor(중간 회귀기)는 보통 1x1 Convolution(1x1 합성곱) 또는 Linear Layer(선형층)로 구현됩니다.
FitNets의 핵심은 1단계의 중간 Hint 학습이지만, 최종적으로는 일반 Knowledge Distillation(지식 증류)처럼 Teacher의 출력 분포도 활용할 수 있습니다.
12.1 Feature Map Shape(특징 맵 형태) 비교
sample_input, _ = next(iter(train_loader))
sample_input = sample_input.to(device)
student_feature = nn_light.features(sample_input)
teacher_feature = nn_deep.features(sample_input)
print("Student feature shape:", student_feature.shape)
print("Teacher feature shape:", teacher_feature.shape)
예시 형태는 다음과 같습니다.
Student feature shape: [batch, 16, 8, 8]
Teacher feature shape: [batch, 32, 8, 8]Student는 Channel(채널)이 16이고 Teacher는 Channel(채널)이 32입니다. 따라서 Student Feature를 Teacher Feature와 비교하려면 채널 수를 맞춰야 합니다.
12.2 Regressor(회귀기)를 추가한 Student
class ModifiedLightNNRegressor(nn.Module):
def __init__(self, num_classes=10):
super().__init__()
self.features = nn.Sequential(
nn.Conv2d(3, 16, kernel_size=3, padding=1),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2),
nn.Conv2d(16, 16, kernel_size=3, padding=1),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2),
)
# Student의 16채널 Feature를 Teacher의 32채널 Feature에 맞춥니다.
self.regressor = nn.Conv2d(16, 32, kernel_size=1)
self.classifier = nn.Sequential(
nn.Linear(1024, 256),
nn.ReLU(),
nn.Dropout(0.1),
nn.Linear(256, num_classes)
)
def forward(self, x):
feature = self.features(x)
projected_feature = self.regressor(feature)
flattened = torch.flatten(feature, 1)
logits = self.classifier(flattened)
return logits, projected_feature12.3 Teacher도 Feature를 반환하도록 수정
class ModifiedDeepNNRegressor(nn.Module):
def __init__(self, num_classes=10):
super().__init__()
self.features = nn.Sequential(
nn.Conv2d(3, 128, kernel_size=3, padding=1),
nn.ReLU(),
nn.Conv2d(128, 64, kernel_size=3, padding=1),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2),
nn.Conv2d(64, 64, kernel_size=3, padding=1),
nn.ReLU(),
nn.Conv2d(64, 32, kernel_size=3, padding=1),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2),
)
self.classifier = nn.Sequential(
nn.Linear(2048, 512),
nn.ReLU(),
nn.Dropout(0.1),
nn.Linear(512, num_classes)
)
def forward(self, x):
feature = self.features(x)
logits = self.classifier(torch.flatten(feature, 1))
return logits, feature12.4 Feature Map Matching(특징 맵 매칭) 학습 코드
def train_mse_regressor(
teacher,
student,
train_loader,
epochs,
learning_rate,
feature_loss_weight,
ce_loss_weight,
device
):
ce_loss = nn.CrossEntropyLoss()
mse_loss = nn.MSELoss()
optimizer = optim.Adam(student.parameters(), lr=learning_rate)
teacher.eval()
student.train()
for epoch in range(epochs):
running_loss = 0.0
for inputs, labels in train_loader:
inputs = inputs.to(device)
labels = labels.to(device)
optimizer.zero_grad()
with torch.no_grad():
_, teacher_feature = teacher(inputs)
student_logits, student_feature = student(inputs)
feature_loss = mse_loss(student_feature, teacher_feature)
label_loss = ce_loss(student_logits, labels)
loss = (
feature_loss_weight * feature_loss
+ ce_loss_weight * label_loss
)
loss.backward()
optimizer.step()
running_loss += loss.item()
print(f"Epoch {epoch+1}/{epochs}, Loss: {running_loss / len(train_loader):.4f}")13. MIT Efficient AI lec10의 “What to Match?” 정리
MIT Efficient AI Lecture 10에서는 Knowledge Distillation(지식 증류)을 단순히 “Teacher의 최종 출력만 따라가는 방법”으로 보지 않습니다. 핵심 질문은 다음입니다.
Teacher의 어떤 지식을 Student에게 전달할 것인가?13.1 Matching Intermediate Weights(중간 가중치 매칭)
Teacher와 Student의 중간 Layer(레이어) Weight(가중치)를 직접 비교합니다.
- L2 Distance(L2 거리)를 사용할 수 있습니다.
- Student의 Channel(채널)이 부족하면 1×1 Convolution(합성곱)으로 Projection(투영)합니다.
- 단, 구조가 다른 두 모델의 Weight를 직접 맞추는 것은 항상 자연스럽지는 않습니다.
13.2 Matching Intermediate Features(중간 특징 매칭)
Teacher와 Student의 Feature Map(특징 맵)을 비교합니다.
- MSELoss(평균제곱오차 손실)
- CosineEmbeddingLoss(코사인 임베딩 손실)
- MMD(Maximum Mean Discrepancy, 최대 평균 불일치)
등을 사용할 수 있습니다.
13.3 Matching Attention Maps(어텐션 맵 매칭)
Teacher와 Student가 입력의 어느 영역에 주목하는지를 비교합니다.
- CNN(합성곱 신경망)에서는 Feature Map(특징 맵)을 기반으로 Attention Map(어텐션 맵)을 만들 수 있습니다.
- 좋은 Teacher는 중요한 위치에 더 강한 Attention(주의)을 가질 수 있습니다.
- Student가 비슷한 Attention Map(어텐션 맵)을 갖도록 학습할 수 있습니다.
13.4 Matching Sparsity Pattern(희소 패턴 매칭)
ReLU(렐루)를 통과한 Activation(활성값)은 양수 영역과 0 영역으로 나뉩니다.
ρ(x) = 1[x > 0]이 패턴은 모델이 입력 공간을 어떻게 나누는지와 관련이 있습니다. 따라서 Teacher와 Student의 Activation Boundary(활성 경계)를 맞추는 것도 Knowledge Distillation(지식 증류)의 한 방법입니다.
13.5 Matching Relational Information(관계 정보 매칭)
Teacher와 Student의 개별 출력만 비교하는 것이 아니라, 데이터 샘플 사이의 관계를 맞춥니다.
예를 들어 다음을 맞출 수 있습니다.
- 샘플 간 Distance(거리)
- 샘플 간 Angle(각도)
- 클래스 간 구조
- Batch 내부 관계
일반 KD는 보통 이렇게 학습합니다.
Teacher(x1) ≈ Student(x1)
Teacher(x2) ≈ Student(x2)
Teacher(x3) ≈ Student(x3)반면 Relational KD는 이렇게 봅니다.
Teacher에서 x1과 x2가 가까우면, Student에서도 x1과 x2가 가까워야 한다.
Teacher에서 x1과 x3가 멀면, Student에서도 x1과 x3가 멀어야 한다.
14. Task별 Knowledge Distillation(지식 증류) 응용
MIT Efficient AI Lecture 10의 summary02에서는 다양한 Task(작업)에서 Knowledge Distillation(지식 증류)을 어떻게 적용하는지 다룹니다.
14.1 Object Detection(객체 탐지)
Object Detection(객체 탐지)은 Classification(분류)보다 더 복잡합니다.
- Class Prediction(클래스 예측)
- Bounding Box Regression(바운딩 박스 회귀)
- Localization(위치 추정)
을 동시에 수행하기 때문입니다.
따라서 Detection Distillation(탐지 증류)에서는 다음 정보를 전달할 수 있습니다.
- Class Probability(클래스 확률)
- Bounding Box(바운딩 박스)
- Feature Map(특징 맵)
- Region Proposal(영역 제안)
- Localization Information(위치 정보)
14.2 Semantic Segmentation(의미론적 분할)
Semantic Segmentation(의미론적 분할)은 픽셀 단위로 Class(클래스)를 예측합니다. 따라서 단일 이미지 전체에 대한 Class(클래스)만 맞추는 것이 아니라, 공간적 구조(Spatial Structure)를 함께 맞춰야 합니다.
- Pixel-wise Prediction(픽셀 단위 예측)
- Structured KD(구조적 지식 증류)
- Feature Map Matching(특징 맵 매칭)
14.3 GAN(생성적 적대 신경망)에서의 KD
GAN(Generative Adversarial Network, 생성적 적대 신경망)에서도 Knowledge Distillation(지식 증류)을 사용할 수 있습니다.
- 큰 Generator(생성자)를 작은 Generator(생성자)로 압축합니다.
- Discriminator(판별자)의 피드백을 활용할 수 있습니다.
- 생성 품질을 유지하면서 추론 비용을 줄이는 것이 목표입니다.
14.4 NLP(자연어 처리)에서의 KD
NLP(Natural Language Processing, 자연어 처리)에서는 Transformer(트랜스포머) 기반 대형 모델을 작은 모델로 압축할 때 Knowledge Distillation(지식 증류)이 자주 사용됩니다.
전달할 수 있는 정보는 다음과 같습니다.
- Logits(로짓)
- Hidden State(은닉 상태)
- Attention Map(어텐션 맵)
- Token-level Probability(토큰 단위 확률)
- Sequence-level Probability(문장 단위 확률)
15. Distillation Scheme(증류 방식)의 종류
MIT Efficient AI Lecture 10에서는 다양한 Distillation Scheme(증류 방식)도 소개합니다.
| 방식 | 설명 |
| Offline Distillation(오프라인 증류) | 미리 학습된 Teacher를 고정하고 Student를 학습 |
| Online Distillation(온라인 증류) | 여러 모델이 동시에 학습하며 서로 지식을 전달 |
| Self-Distillation(자가 증류) | 같은 모델 내부 또는 이전 세대 모델을 Teacher로 사용 |
| BYOT(Bring Your Own Teacher) | 하나의 네트워크 내부에서 보조 Head를 Teacher처럼 활용 |
16. Knowledge Distillation과 Quantization(양자화)의 관계
Knowledge Distillation(지식 증류)과 Quantization(양자화)은 서로 다른 기술이지만, 함께 사용할 수 있습니다.
16.1 Quantization(양자화)이란?
Quantization(양자화)은 모델의 Weight(가중치)나 Activation(활성값)을 낮은 정밀도로 표현하는 기술입니다.
예를 들어:
FP32 → INT8처럼 32비트 부동소수점 값을 8비트 정수로 바꿉니다.
장점은 다음과 같습니다.
- 모델 크기 감소
- 메모리 사용량 감소
- 추론 속도 향상 가능
- 엣지 디바이스 배포에 유리
하지만 단점도 있습니다.
- 정밀도가 낮아져 성능이 떨어질 수 있습니다.
- 작은 모델일수록 손실이 더 크게 나타날 수 있습니다.
- 민감한 Layer(레이어)는 양자화 후 성능 하락이 발생할 수 있습니다.
16.2 KD와 Quantization을 함께 쓰는 이유
Student Model(학생 모델)을 작게 만들고 Quantization(양자화)까지 적용하면 성능이 더 떨어질 수 있습니다. 이때 Teacher Model(교사 모델)의 Soft Label(소프트 레이블)이나 Feature(특징)를 사용하면 성능 하락을 줄이는 데 도움이 될 수 있습니다.
Teacher Model
↓ Knowledge Distillation
Small Student Model
↓ Quantization
Efficient On-device Model즉, Knowledge Distillation(지식 증류)은 Quantization(양자화)이나 Pruning(가지치기) 이후의 성능 회복에도 활용될 수 있습니다.
17. Loss 선택 가이드: KLDivLoss vs CosineEmbeddingLoss vs MSELoss
Teacher-Student 구조에서 어떤 Loss(손실)를 사용할지는 “무엇을 맞출 것인가?”에 따라 달라집니다.
| 비교 대상 | 데이터 성격 | 추천 Loss |
| 최종 Class Probability(클래스 확률) | 확률분포 | KLDivLoss |
| 최종 Histogram(히스토그램) | 분포 | KLDivLoss |
| Logits(로짓) | 점수 벡터 | KLDivLoss 또는 MSELoss |
| Hidden State(은닉 상태) | 표현 벡터 | CosineEmbeddingLoss |
| Context Vector(문맥 벡터) | 표현 벡터 | CosineEmbeddingLoss |
| Feature Map(특징 맵) | 텐서 값 | MSELoss |
| Attention Map(어텐션 맵) | 공간 중요도 | MSELoss 또는 L2 Loss |
| Sparsity Pattern(희소 패턴) | 0/1 패턴 | L1 Loss |
정리하면 다음과 같습니다.
출력 분포를 맞춘다 → KLDivLoss
중간 표현 방향을 맞춘다 → CosineEmbeddingLoss
값 자체를 맞춘다 → MSELoss18. Transformer Encoder(트랜스포머 인코더)에 적용한다면?
Transformer Encoder(트랜스포머 인코더) 기반 모델에 Knowledge Distillation(지식 증류)을 적용한다면 다음처럼 생각할 수 있습니다.
| 위치 | 출력 형태 | 추천 Loss |
| Encoder Hidden State(인코더 은닉 상태) | [batch, seq_len, hidden_dim] |
CosineEmbeddingLoss, MSELoss |
| CLS Token(분류 토큰) | [batch, hidden_dim] |
CosineEmbeddingLoss |
| Attention Map(어텐션 맵) | [batch, head, seq, seq] |
MSELoss, KLDivLoss |
| Final Logits(최종 로짓) | [batch, num_classes] |
KLDivLoss |
| Histogram Output(히스토그램 출력) | [batch, bins] |
KLDivLoss |
예를 들어 Encoder Context Vector(인코더 문맥 벡터)는 확률분포가 아니라 표현 벡터입니다. 따라서 CosineEmbeddingLoss(코사인 임베딩 손실)가 자연스럽습니다.
반대로 최종 출력이 Histogram(히스토그램) 형태라면 Row-wise Normalization(행 단위 정규화)을 통해 확률분포로 만든 뒤 KLDivLoss(KL 발산 손실)를 사용할 수 있습니다.
19. 전체 학습 코드 예시: CE + KD + Cosine
최종 출력은 KLDivLoss(KL 발산 손실)로 맞추고, 중간 표현은 CosineEmbeddingLoss(코사인 임베딩 손실)로 맞추는 통합 예시는 다음과 같습니다.
def train_student_with_kd_and_feature_loss(
teacher,
student,
train_loader,
epochs,
learning_rate,
temperature,
ce_weight,
kd_weight,
feature_weight,
device
):
ce_loss = nn.CrossEntropyLoss()
cosine_loss = nn.CosineEmbeddingLoss()
optimizer = optim.Adam(student.parameters(), lr=learning_rate)
teacher.eval()
student.train()
for epoch in range(epochs):
running_loss = 0.0
for inputs, labels in train_loader:
inputs = inputs.to(device)
labels = labels.to(device)
optimizer.zero_grad()
with torch.no_grad():
teacher_logits, teacher_feature = teacher(inputs)
student_logits, student_feature = student(inputs)
# 1. 실제 정답 Label(레이블)에 대한 손실
loss_ce = ce_loss(student_logits, labels)
# 2. Teacher 출력 분포를 따라가는 KL Loss
teacher_probs = F.softmax(teacher_logits / temperature, dim=-1)
student_log_probs = F.log_softmax(student_logits / temperature, dim=-1)
loss_kd = F.kl_div(
student_log_probs,
teacher_probs,
reduction="batchmean"
) * (temperature ** 2)
# 3. 중간 Feature 방향을 맞추는 Cosine Loss
target = torch.ones(inputs.size(0)).to(device)
loss_feature = cosine_loss(
student_feature,
teacher_feature,
target
)
# 4. 전체 손실
loss = (
ce_weight * loss_ce
+ kd_weight * loss_kd
+ feature_weight * loss_feature
)
loss.backward()
optimizer.step()
running_loss += loss.item()
print(f"Epoch {epoch+1}/{epochs}, Loss: {running_loss / len(train_loader):.4f}")이 구조는 다음 의미를 가집니다.
loss_ce: 정답 Label(레이블)을 맞춥니다.loss_kd: Teacher의 최종 판단 분포를 따라갑니다.loss_feature: Teacher의 중간 표현 방향을 따라갑니다.
결과
Teacher accuracy: 75.88%
Student accuracy without teacher: 70.37%
Student accuracy with CE + KD: 70.67%
Student accuracy with CE + CosineLoss: 70.57%
Student accuracy with CE + RegressorMSE: 70.88%
20. 참고 자료
- PyTorch Tutorial: Knowledge Distillation Tutorial
- MIT Efficient AI Lecture 10
'03.딥러닝 > OnDevice' 카테고리의 다른 글
| 05. 양자화(Quantization) 심화 (0) | 2026.06.15 |
|---|---|
| 04 양자화(Quantization) (0) | 2026.06.15 |
| 03 네트워크 프루닝 마인드맵 (Network Pruning Mindmap) (0) | 2026.06.13 |
| 02 신경망 프루닝(Network Pruning) 심화편 (0) | 2026.06.12 |
| 01 신경망 프루닝(Network Pruning) (0) | 2026.06.11 |