모델 압축·배포 최적화 시리즈 1편: Knowledge Distillation

"큰 모델이 배운 지식을 작은 모델에 전수할 수 없을까?"
Knowledge Distillation(지식 증류)은 Teacher-Student 방식으로 모델을 압축하는 대표적인 기법입니다.

---

Knowledge Distillation

 

Knowledge Distillation이란?

Knowledge Distillation(KD)은 큰 모델(Teacher) 이 예측한 soft output(로짓, 확률분포)을
작은 모델(Student) 이 모방하도록 학습시키는 방법입니다.

이 방식으로 Student 모델은:

• 라벨 정보(정답)뿐 아니라,
• Teacher가 가진 클래스 간 유사성, 결정 경계 정보까지 학습

결과적으로 작고 가벼운 모델이 Teacher 모델 수준의 성능을 재현할 수 있습니다.

---

왜 필요한가?

• 대형 모델은 학습/추론에 비용이 많이 듦 → 모바일/엣지 디바이스 배포 어렵
• 단순한 weight pruning, quantization은 성능이 급격히 떨어질 수 있음
• KD는 성능 유지 + 모델 크기 감소 두 마리 토끼를 잡는 방법

---

동작 원리

Distillation Loss:

 

L = (1 - α) * CE(y_pred_student, y_true) + α * T² * KLDiv(softmax(z_student/T), softmax(z_teacher/T))

• CE: Cross-Entropy Loss (정답 라벨)
• KLDiv: KL Divergence (teacher vs student soft label)
• T: Temperature (softening factor, T>1이면 더 부드러운 확률분포)
• α: 가중치 (soft loss vs hard loss 비율)

즉, Student는 정답(label) + Teacher 예측 분포를 동시에 맞추도록 학습.

---

장점

1. 작은 모델도 큰 모델의 일반화 능력 습득
2. 추론 속도·메모리 사용량 감소
3. 라벨 데이터 적은 경우에도 효과적 (teacher의 soft label이 데이터 증강 역할)

---

단점 & 주의점

Teacher 모델 학습/추론 비용 발생 (사전 준비 필요)
Teacher 품질이 낮으면 Student도 그 수준에 맞춰 학습 → 성능 한계
Temperature, α 등 하이퍼파라미터 튜닝 필요

---

PyTorch 예시 코드 (간단)

 

import torch.nn.functional as F

def distillation_loss(student_logits, teacher_logits, labels, T=2.0, alpha=0.5):
    # Hard loss
    hard_loss = F.cross_entropy(student_logits, labels)

    # Soft loss (KL Divergence between softened distributions)
    p_s = F.log_softmax(student_logits / T, dim=1)
    p_t = F.softmax(teacher_logits / T, dim=1)
    soft_loss = F.kl_div(p_s, p_t, reduction="batchmean") * (T * T)

    return (1 - alpha) * hard_loss + alpha * soft_loss

---

실제 활용 예시

• 모바일 배포: ResNet-50 teacher → MobileNet student
• 대형 LLM 압축: GPT-3 → Distilled GPT-3 (6B → 1.3B)
• 분류 태스크: 데이터 부족 시 teacher 예측값으로 pseudo-labeling

---

Distillation + 다른 기법

KD는 다른 경량화 기법과 함께 쓰면 더 강력합니다.

조합	효과
KD + Quantization	모델 작아지면서 성능 유지
KD + Pruning	필요 없는 weight 제거 후 재학습
KD + LoRA/Adapter	student 모델에도 PEFT 적용 가능

---

실무 적용 팁

• Temperature: 2~4 권장 (soft label smooth하게)
• Teacher Freezing: teacher는 inference mode로 유지 (gradient 불필요)
• Student capacity가 너무 작으면 KD 효과 한계 → 작은 모델이라도 capacity 적당히 확보

---

결론

Knowledge Distillation은 "모델의 지식을 다른 모델로 전수"하는 기법으로,
작은 모델로도 큰 모델 수준의 성능을 내게 해줍니다.
모바일 배포나 실시간 추론 환경에서 필수적으로 고려해야 할 방법입니다.

 

 

 

다음 편은 아래에 있습니다!

https://machineindeep.tistory.com/68

모델 압축·배포 최적화 시리즈 2편: Quantization (양자화)