지식 증류(Knowledge Distillation)의 개념과 원리 완벽 정리

Bean·2025년 4월 30일

DeepLearning KnowledgeDistillation ModelCompression NeuralNetworks SoftOutput crossentropy kldivergence

인공지능

목록 보기

20/123

1. 지식 증류(knowledge distillation) 설명

지식 증류(knowledge distillation)는 큰 모델(teacher)의 지식을 작은 모델(student)로 전달하여, 작은 모델도 높은 성능을 내도록 만드는 모델 압축 기법입니다. 성능과 효율성 사이의 균형을 맞출 때 자주 사용됩니다.

1.1. 기본 개념

Teacher Model (교사 모델)
- 크고 정확한 모델 (예: ResNet-152, HRNet 등)
- 연산량이 많고 느리지만 학습된 정보가 풍부함
Student Model (학생 모델)
- 작고 빠른 모델 (예: MobileNet, ShuffleNet 등)
- 일반적으로 teacher보다 성능이 낮지만, 경량화되어 실시간 처리에 적합
지식 전이 방식
- student는 단순히 ground truth(정답 라벨)만 배우는 게 아니라, teacher의 예측 결과도 함께 학습함
- 이때 teacher의 soft output(확률 분포, 중간 feature 등)을 참고

1.2. # Soft Output이란?

Soft output은 신경망이 마지막에 출력하는 확률 분포를 말합니다.
쉽게 말해, 모델이 어떤 클래스일 가능성이 얼마나 되는지를 나타내는 "부드러운" 예측 결과입니다.

예시로 설명해볼게요:

사람의 이미지를 보고 모델이 "고양이냐 강아지냐"를 예측한다고 해봅시다.

모델의 마지막 출력 (softmax 전):

[2.5, 0.3]  ← raw logits (고양이, 강아지)

softmax를 적용하면 확률 분포가 나옵니다:
```
[0.92, 0.08] ← soft output
```
모델의 hard output은 가장 높은 확률을 가진 클래스 (여기선 "고양이")입니다.

Soft Output vs Hard Output

구분	설명	예시
Soft Output	각 클래스에 대한 확률	`[0.92, 0.08]`
Hard Output	가장 높은 확률을 갖는 클래스만 선택	`"고양이"`

지식 증류에서 Soft Output이 중요한 이유

Teacher 모델의 soft output은 다음과 같은 추가 정보를 담고 있습니다:

단순한 정답 외의 정보도 전달 (ex. 강아지도 조금 비슷하게 보였다는 것)
학생 모델이 더 일반화 능력이 뛰어난 예측을 하도록 도움

이렇게 soft output을 참고해서 학습하면, student 모델도 teacher처럼 뉘앙스 있는 예측을 할 수 있게 됩니다.

2. 지식 증류 학습 구조

학습 시 두 종류의 loss를 동시에 사용합니다:

정답 라벨에 대한 loss (cross-entropy):
- student가 실제 정답을 잘 예측하는지 평가
teacher output에 대한 loss (distillation loss):
- student가 teacher와 비슷한 예측 분포를 따르도록 유도
- 보통 KL-divergence나 MSE 등을 사용

이때, softmax 출력에 "temperature T"를 적용하여 분포를 더 부드럽게 만듭니다. 이로 인해 student는 단순히 정답만 맞추는 게 아니라, 각 클래스 간의 미묘한 차이도 학습할 수 있습니다.

KL-divergence란?

KL-divergence(Kullback–Leibler divergence, 쿨백-라이블러 발산)는 두 확률 분포 간의 차이를 측정하는 지표입니다.
지식 증류(knowledge distillation)에서도 teacher와 student 모델의 출력 확률 분포를 맞추는 데 사용됩니다.

"두 확률 분포가 얼마나 다른지를 측정하는 수치"

즉, teacher가 예측한 분포 ( P )와 student가 예측한 분포 ( Q )가 있을 때,
student의 분포 ( Q )가 teacher의 분포 ( P )를 얼마나 잘 흉내내고 있는지**를 KL-divergence로 계산합니다.

수식

D_{KL}(P \parallel Q) = \sum_{i} P(i) \log \frac{P(i)}{Q(i)}

$P(i)$ : teacher가 예측한 확률 분포 (정답처럼 취급)
$Q(i)$ : student가 예측한 확률 분포
$D_{KL}$ : 이 값이 작을수록 Q가 P와 유사하다는 뜻

특징

항상 0 이상 (0이면 두 분포가 완전히 같음)
비대칭: $D_{KL}(P \parallel Q) \neq D_{KL}(Q \parallel P)$
teacher가 예측한 정보를 기준으로 얼마나 차이 나는지 측정함

지식 증류에서의 예시

Teacher의 soft output (P):

[0.7, 0.2, 0.1]

Student의 soft output (Q):

[0.6, 0.3, 0.1]

이 두 분포 간의 차이를 KL-divergence로 계산해서 loss로 사용하면, student가 teacher처럼 "부드러운 확률 분포"를 예측하도록 유도할 수 있습니다.

3. 지식의 종류

Soft label distillation:
teacher의 출력 확률 분포를 학습 (가장 일반적인 방식)
Feature-based distillation:
중간 layer의 feature map을 맞춤
Attention-based distillation:
teacher의 attention map(어디를 집중해서 보는지)을 모방
Relational knowledge:
여러 샘플 간의 상호 관계(거리 등)를 student가 따라 하도록 학습

4. 장점

작은 모델의 성능 향상
추론 속도 개선
모바일·임베디드 환경에 적합
기존 데이터셋만으로도 teacher의 일반화 능력을 student에게 전달 가능

5. 예시

예를 들어 Human Pose Estimation Task에서 대형 HRNet이 사람의 자세를 매우 정확히 예측할 수 있다면, 그 예측을 참고해 작은 SimpleBaseline 모델도 비슷한 결과를 내도록 학습시키는 방식입니다. 이렇게 하면 작은 모델도 정확도는 최대한 유지하면서도 연산량은 훨씬 줄어듭니다.