딥러닝 모델의 경량화 기술 중 가장 각광받는 지식 증류(Knowledge Distillation)는 거대한 Teacher 모델의 지식을 가볍고 빠른 Student 모델로 전이하는 핵심 프로세스입니다. 이 과정에서 가장 중요한 물리량이 바로 온도(Temperature, $T$)입니다. 단순히 Softmax를 적용하는 것을 넘어, 왜 특정 온도에서 지식 전이가 더 잘 일어나는지, 그리고 실무 데이터셋에 따라 어떻게 $T$를 설정해야 하는지에 대한 정답을 찾기란 쉽지 않습니다. 본 포스팅에서는 지식 증류의 수학적 배경부터 시작하여, Soft Label의 엔트로피를 조절하는 온도 $T$의 최적 설정 기준, 그리고 실무 개발자가 바로 적용할 수 있는 파이토치(PyTorch) 기반의 실전 예제 7가지를 상세히 공유합니다.
1. 지식 증류에서 온도(Temperature)의 정의와 역할
일반적인 분류 모델의 Softmax 함수는 클래스 간의 점수 차이를 극대화하여 가장 높은 점수를 가진 클래스에 확률을 몰아주는 경향이 있습니다. 하지만 지식 증류에서는 Teacher 모델이 가진 "정답이 아닌 클래스들 사이의 관계(Dark Knowledge)"를 학생에게 전달하는 것이 목표입니다. 온도 $T$가 포함된 Softmax 수식은 다음과 같습니다: $$q_i = \frac{\exp(z_i / T)}{\sum_j \exp(z_j / T)}$$ 여기서 $T$가 커질수록 확률 분포는 부드러워지며(Soft), $T$가 1에 가까워질수록 일반적인 Hard Max에 가까워집니다. 즉, 온도는 지식의 '부드러움'을 결정하는 조절 나사와 같습니다.
2. 온도 설정에 따른 모델 성능 차이 및 요약 비교
온도를 어떻게 설정하느과에 따라 Student 모델이 학습하는 정보의 양과 질이 달라집니다. 다음은 온도 설정 기준에 따른 특징을 비교한 표입니다.
| 구분 | 저온 (Low T, $T \approx 1$) | 최적 온도 (Mid T, $T=2 \sim 5$) | 고온 (High T, $T > 10$) |
|---|---|---|---|
| 확률 분포 특징 | 가장 높은 클래스에 집중 (Peak) | 클래스 간 상관관계 보존 (Soft) | 모든 클래스가 균일해짐 (Uniform) |
| 전이되는 지식 | 정답 레이블 위주의 정보 | 오답 클래스 간의 유사성 정보 | 배경 노이즈 및 무의미한 정보 |
| Student 성능 | Teacher와 유사한 예측 편향 | 일반화 성능 최대화 | 수렴 속도 저하 및 성능 하락 |
| 주요 용도 | 간단한 데이터셋 (MNIST 등) | 복잡한 시각 데이터 (ImageNet 등) | 특수한 보안/프라이버시 증류 |
3. 실무 적용을 위한 Python & PyTorch 실전 예제 7가지
개발자가 현업에서 지식 증류를 구현할 때 직면하는 다양한 상황별 해결 코드를 제안합니다. 모든 코드는 복사하여 즉시 사용 가능합니다.
Example 1: 기본 Distillation Loss 함수 구현
Kullback-Leibler Divergence(KLD)를 이용한 기본적인 Soft Label 손실 함수입니다.
import torch
import torch.nn as nn
import torch.nn.functional as F
def distillation_loss(student_logits, teacher_logits, T=3.0):
"""
T: Temperature 하이퍼파라미터
"""
soft_targets = F.softmax(teacher_logits / T, dim=1)
soft_prob = F.log_softmax(student_logits / T, dim=1)
# T^2을 곱해주는 이유는 그래디언트 스케일을 맞추기 위함입니다.
return F.kl_div(soft_prob, soft_targets, reduction='batchmean') * (T ** 2)
Example 2: 데이터셋 복잡도에 따른 동적 온도 조절 방법
학습 초기에는 높은 온도로 전체적인 구조를 익히고, 후반부에는 온도를 낮추는 스케줄링 기법입니다.
def get_temperature(epoch, total_epochs):
# Annealing Temperature: 5.0에서 1.0으로 서서히 감소
start_t = 5.0
end_t = 1.0
return start_t - (start_t - end_t) * (epoch / total_epochs)
# 학습 루프 내 적용 예시
# current_T = get_temperature(epoch, 100)
# loss = distillation_loss(s_out, t_out, T=current_T)
Example 3: 다중 클래스 불균형 해결을 위한 Weighted Temperature
특정 클래스의 샘플 수가 적을 때 해당 클래스에 더 민감하게 반응하도록 온도를 개별 적용합니다.
def weighted_distillation_loss(s_logits, t_logits, class_weights, T_base=3.0):
# class_weights: [batch_size] 크기의 텐서
# 각 샘플의 클래스 중요도에 따라 온도를 미세 조정
per_sample_T = T_base * class_weights
# 개별 연산을 위해 반복문 또는 브로드캐스팅 활용
loss = 0
for i in range(s_logits.size(0)):
s_p = F.log_softmax(s_logits[i:i+1] / per_sample_T[i], dim=1)
t_p = F.softmax(t_logits[i:i+1] / per_sample_T[i], dim=1)
loss += F.kl_div(s_p, t_p, reduction='batchmean') * (per_sample_T[i] ** 2)
return loss / s_logits.size(0)
Example 4: 정답(Ground Truth)과 Soft Label의 결합 비율 최적화
단순 증류보다 실제 라벨을 함께 학습하는 것이 훨씬 효율적입니다.
def total_criterion(s_out, t_out, labels, alpha=0.5, T=4.0):
# alpha: 정답 라벨과 Soft Label의 반영 비중
distill_loss = distillation_loss(s_out, t_out, T)
student_loss = F.cross_entropy(s_out, labels)
return alpha * distill_loss + (1 - alpha) * student_loss
Example 5: 검증 데이터 성능 기반의 자동 온도 검색 (Grid Search)
최적의 $T$를 찾기 위한 간단한 실험 스크립트 구조입니다.
temp_candidates = [1.5, 2.0, 3.0, 5.0, 10.0]
best_acc = 0
best_T = 1.0
for T in temp_candidates:
# 모델 초기화 및 학습 루틴 실행
# acc = train_student(T)
print(f"Testing Temperature: {T}, Accuracy: {acc}")
if acc > best_acc:
best_acc = acc
best_T = T
print(f"Optimal Temperature found: {best_T}")
Example 6: Feature-based Distillation에서의 온도 활용
Logit이 아닌 중간 Feature Map을 증류할 때 채널 간의 어텐션을 부드럽게 만드는 기법입니다.
def feature_distillation_loss(s_feature, t_feature, T=2.0):
# Feature Map을 평탄화한 후 Softmax 적용
s_map = F.softmax(s_feature.view(s_feature.size(0), -1) / T, dim=1)
t_map = F.softmax(t_feature.view(t_feature.size(0), -1) / T, dim=1)
return F.mse_loss(s_map, t_map)
Example 7: Large-Scale 모델(BERT/ViT)을 위한 고온 증류 전략
모델이 너무 크면 Logit 값의 분산이 매우 큽니다. 이때는 $T$를 높게 잡아야 정보가 소실되지 않습니다.
# Large Model의 경우 T=10~20 사이에서 안정적인 수렴을 보이는 경우가 많음
# 이를 위해 안정적인 Log-Softmax를 직접 구현하여 수치적 불안정성 제거
def stable_large_model_distillation(s_logits, t_logits, T=15.0):
s_logits = s_logits - torch.max(s_logits, dim=1, keepdim=True)[0]
t_logits = t_logits - torch.max(t_logits, dim=1, keepdim=True)[0]
return distillation_loss(s_logits, t_logits, T=T)
4. 결론: 최적의 온도를 찾는 3단계 가이드라인
지식 증류 프로젝트를 성공시키기 위한 온도 설정 프로세스는 다음과 같이 요약할 수 있습니다.
- 초기 설정: 먼저 $T=3.0$ 또는 $T=4.0$으로 시작하십시오. 이는 대부분의 컴퓨터 비전과 NLP 과업에서 표준적으로 사용되는 수치입니다.
- Teacher 분석: 만약 Teacher 모델의 신뢰도가 너무 높아 특정 클래스 확률이 거의 1에 육박한다면, 온도를 5.0 이상으로 높여 분포를 강제로 펴주어야 합니다.
- 검증(Validation): Student 모델이 학습 데이터에 과적합된다면 온도를 낮추고, 반대로 학습이 너무 느리거나 Teacher를 따라가지 못한다면 온도를 높여 '부드러운 가이드'를 제공하십시오.
본 가이드를 통해 여러분의 모델 경량화 프로젝트가 한 단계 더 높은 성능에 도달하기를 바랍니다.
내용 출처 및 참고 자료
- Hinton, G., Vinyals, O., & Dean, J. (2015). "Distilling the Knowledge in a Neural Network." arXiv:1503.02531.
- PyTorch Official Documentation: "Knowledge Distillation Tutorial"
- Gou, J., Yu, B., Maybank, S. J., & Tao, D. (2021). "Knowledge Distillation: A Survey." International Journal of Computer Vision.
