Knowledge Distillation (지식 증류): Teacher-Student 모델 기반 On-Device AI 경량화 기법

1. On-Device AI 모델 경량화를 위한 Knowledge Distillation 도입 배경

NPU(Neural Processing Unit) 및 DSP와 같은 엣지 하드웨어 환경에서는 메모리 용량과 연산 자원이 극도로 제한됩니다. 대규모 파라미터를 가진 거대 AI 모델(Teacher Model)은 추론 정밀도가 높지만, Latency(지연 시간)가 길고 메모리 점유율이 높아 On-Device AI 실시간 추론에 적합하지 않습니다.

이러한 문제를 해결하기 위해 거대 모델의 지식을 경량화된 소형 모델(Student Model)로 전수하는 Knowledge Distillation(지식 증류) 기술이 활용됩니다. Knowledge Distillation은 단순한 Hard Label 학습을 넘어, Teacher Model이 출력하는 확률 분포인 Soft Label을 Student Model에 학습시킵니다. 본 포스팅에서는 Teacher-Student 모델 아키텍처 구조, Temperature Hyperparameter 조절 기법, 그리고 Distillation Loss 함수 정의 및 구현 방식을 심층 분석합니다.

2. Knowledge Distillation (지식 증류) 핵심 요약

• Teacher-Student 모델 아키텍처 (Teacher-Student Architecture): 대용량 Teacher Model의 숨겨진 지식(Dark Knowledge)을 Soft Target 형태의 확률 분포로 변환하여 소형 Student Model에 전수합니다.
• Softmax 온도 스케일링 (Softmax Temperature Scaling): 온도 파라미터 $T$를 활용하여 출력 확률 분포의 엔트로피를 조절하고, 클래스 간의 숨겨진 연관 관계 정보를 추출합니다.
• 복합 손실 함수 (Combined Loss Function): Kullback-Leibler (KL) Divergence Loss와 Cross-Entropy Loss의 가중합(Weighted Sum)을 계산하여 Student Model의 가중치를 최적화합니다.

3. Knowledge Distillation 아키텍처 및 Softmax Temperature 동작 원리 분석

3.1 Teacher-Student 모델 아키텍처 구조 비교

Knowledge Distillation 메커니즘은 두 개의 독립된 신경망 구조로 구성됩니다. Teacher Model은 고성능을 보장하지만 많은 파라미터를 가지며, Student Model은 수용량이 적고 소형화된 신경망입니다.

항목 (Metric)	Teacher Model	Student Model
모델 크기 및 파라미터 (Model Size)	대용량 (High-Capacity, Deep NN)	경량화 (Lightweight, Shallow NN)
타겟 하드웨어 (Target Hardware)	Server / Cloud GPU Cluster	On-Device NPU / Edge HW
출력 데이터 형태 (Output Label)	Soft Probabilities (Soft Targets)	Logits & Predicted Class
손실 함수 연산 (Loss Function)	학습 완료 상태 (Frozen Weights)	KL Divergence + Cross-Entropy

3.2 Temperature Parameter (T)를 활용한 Softmax 확률 분포 스케일링

표준 Softmax 함수는 가장 높은 값을 가지는 클래스의 확률을 1에 가깝게 만들고, 나머지 클래스의 확률을 0으로 수축시킵니다. 이는 클래스 간의 숨겨진 상관관계(Dark Knowledge) 정보를 손실시킵니다. Softmax 함수 내에 Temperature 변수 $T$를 적용하여 확률 분포를 평탄하게 스케일링합니다.

$$\mathbf{p}_i = \frac{\exp(z_i / T)}{\sum_{j} \exp(z_j / T)}$$

• $z_i$: 신경망의 마지막 레이어에서 출력되는 Logit 값
• $T$: Temperature 스케일링 파라미터 ($T > 1$ 일 때 확률 분포가 부드러워짐)

$T \to \infty$ 일수록 모든 클래스의 확률 분포가 균일해지며, $T = 1$ 일 경우 표준 Softmax 함수와 동일하게 동작합니다.

3.3 PyTorch 기반 Knowledge Distillation 손실 함수 구현 및 분석

Knowledge Distillation의 전체 Loss($L_{total}$)는 Student Model의 자체 Ground Truth 손실($L_{student}$)과 Teacher Model의 Soft Label을 복제하는 Distillation Loss($L_{distill}$)의 가중합으로 구성됩니다.

$$L_{total} = \alpha \cdot T^2 \cdot D_{KL}\left(\text{Softmax}\left(\frac{z_s}{T}\right) \parallel \text{Softmax}\left(\frac{z_t}{T}\right)\right) + (1 - \alpha) \cdot H(y, \text{Softmax}(z_s))$$

import torch
import torch.nn as nn
import torch.nn.functional as F

class KnowledgeDistillationLoss(nn.Module):
    def __init__(self, temperature=4.0, alpha=0.7):
        super(KnowledgeDistillationLoss, self).__init__()
        self.temperature = temperature
        self.alpha = alpha
        self.kl_div = nn.KLDivLoss(reduction='batchmean')
        self.cross_entropy = nn.CrossEntropyLoss()

    def forward(self, student_logits, teacher_logits, labels):
        # Compute Soft Targets with Temperature scaling
        soft_student = F.log_softmax(student_logits / self.temperature, dim=1)
        soft_teacher = F.softmax(teacher_logits / self.temperature, dim=1)
        
        # Calculate Distillation Loss using KL Divergence
        distillation_loss = self.kl_div(soft_student, soft_teacher) * (self.temperature ** 2)
        
        # Calculate Standard Student Cross-Entropy Loss
        student_loss = self.cross_entropy(student_logits, labels)
        
        # Total Weighted Loss Computation
        total_loss = (self.alpha * distillation_loss) + ((1.0 - self.alpha) * student_loss)
        return total_loss

4. Knowledge Distillation 학습 및 On-Device 배포 팁

1. Temperature ($T$) 및 Alpha ($\alpha$) Hyperparameter 튜닝: 일반적으로 Temperature 값은 $T \in [2.0, 6.0]$ 범위에서 최적의 performance를 보입니다. Alpha 값은 Soft Label의 비중을 높이기 위해 $0.7 \sim 0.9$ 사이로 설정하는 것이 권장됩니다.
2. Gradient Scaling 수용: KL Divergence 역전파 계산 시 $1/T^2$의 미분 항이 추가되므로, Distillation Loss 항목에 $T^2$을 곱해주어야 Student Gradient의 크기가 적절한 스케일을 유지합니다.

5. Knowledge Distillation 구현 시 흔히 하는 실수 및 디버깅 기법

1. Log Softmax 대신 표준 Softmax 입력 전달 오류
   • 발생 원인: PyTorch의 nn.KLDivLoss 함수는 첫 번째 입력인 Target Student Prediction으로 Log-Probability 형태를 요합니다.
   • 해결 방법: Student Logit에 F.softmax()가 아닌 F.log_softmax()를 적용하여 입력해야 정확한 Divergence Gradient가 계산됩니다.
2. Teacher Model의 Weight가 학습 중에 업데이트되는 현상
   • 발생 원인: Teacher Model을 Instantiation할 때 eval() 모드를 설정하지 않거나 torch.no_grad() 블록을 누락하는 경우 발생합니다.
   • 해결 방법: Teacher Model의 텐서 연산 구간 전체를 with torch.no_grad():로 감싸서 메모리 낭비와 불필요한 Gradient 업데이트를 방지해야 합니다.

6. On-Device AI 모델을 위한 Knowledge Distillation 결론

Knowledge Distillation은 거대 모델의 추론 성능을 유지하면서도 엣지 디바이스에 적합한 소형 신경망을 구축할 수 있는 효율적인 기술입니다. Temperature Scaling 기법을 통해 클래스 간 Soft Label 지식을 성공적으로 전달하고, $T^2$ 가중치 조절을 준수하여 손실 함수를 구현하면 On-Device NPU 환경에 최적화된 모델 경량화를 달성할 수 있습니다.