Knowledge Distillation

지식 증류는 작은 "학생" 신경망을 학습시켜 더 크고 정확한 "교사" 모델의 동작을 재현하도록 하는 모델 압축 기법입니다. 학생 모델은 정답 레이블 대신 또는 정답 레이블에 더해 교사 모델의 출력 확률이나 내부 특성으로부터 학습하며, 이러한 풍부한 타깃을 통해 레이블만으로 학습한 동일 아키텍처보다 높은 정확도를 얻을 수 있습니다. 그 결과 교사 모델의 성능을 상당 부분 유지하면서 휴대전화, 임베디드 보드 및 기타 엣지 AI 하드웨어에서 사용할 수 있을 만큼 빠르고 작은 모델을 얻습니다. Buciluǎ, Caruana, Niculescu-Mizil은 2006년에 앙상블의 지식을 하나의 모델에 압축할 수 있음을 보였고, Hinton, Vinyals, Dean은 2015년에 이 접근법을 지식 증류라고 명명하고 대부분의 방법이 여전히 사용하는 온도 조정 소프트 타깃을 소개했습니다.

지식 증류의 작동 방식#

레이블은 이미지에 늑대가 있다는 사실만 알려줍니다. 학습된 교사의 소프트맥스 출력은 더 많은 정보를 담고 있습니다. 이미지가 허스키와 어느 정도 비슷하고, 고양이와는 조금 비슷하지만, 자동차와는 전혀 비슷하지 않다는 식입니다. 흔히 다크 지식이라고 부르는 잘못된 클래스 간 상대 확률이 소프트 타깃으로 학생에게 전달됩니다. 소프트맥스 전에 로짓을 1보다 큰 온도 T로 나누면 분포가 평탄해져 작은 확률도 학습에 활용할 만큼 충분한 가중치를 갖게 됩니다. 증류 손실은 교사와 학생의 온도 조정 분포 간 쿨백-라이블러 발산에 T²을 곱한 값입니다. 이렇게 하면 그래디언트 크기가 일반 레이블 손실의 그래디언트 크기와 비슷하게 유지됩니다.

import torch
import torch.nn.functional as F

# Logits for one wolf photo over the classes [wolf, husky, cat, car]
teacher_logits = torch.tensor([[6.0, 4.0, 1.0, -2.0]])
student_logits = torch.tensor([[4.0, 1.0, 1.5, 0.0]])

# Dividing logits by a temperature T > 1 softens the teacher's distribution
for T in (1.0, 4.0):
    probs = F.softmax(teacher_logits / T, dim=1)
    print(f"T={T}: teacher probabilities {[round(p, 3) for p in probs[0].tolist()]}")

# Distillation loss: KL divergence between the softened distributions, scaled by T^2
T = 4.0
soft_targets = F.softmax(teacher_logits / T, dim=1)
student_log_probs = F.log_softmax(student_logits / T, dim=1)
loss = F.kl_div(student_log_probs, soft_targets, reduction="batchmean") * T**2
print(f"Distillation loss at T={T}: {loss.item():.3f}")

T=1일 때 교사는 늑대 클래스에 0.875를 부여하고 고양이에는 0.006만, 자동차에는 0.0을 부여합니다. T=4에서는 같은 로짓이 0.493, 0.299, 0.141, 0.067로 바뀌어 학생에게 늑대가 고양이보다 허스키와 더 비슷하고, 고양이가 자동차보다 더 비슷하다는 점을 알려줍니다. 여기서는 1.244인 이 손실을 학생이 최소화하며, 일반적으로 정답 레이블을 사용한 정규 태스크 손실과 가중합을 계산합니다. 교사는 학습 내내 고정된 상태를 유지하며 순전파만 수행합니다.

증류 유형#

• 응답 기반: 학생은 위의 소프트 타깃 방식처럼 교사의 최종 출력을 따라 하며, 교사에게서 로짓 외의 정보를 받을 필요가 없습니다. DistilBERT는 사전 학습 중 언어 모델링 및 은닉 상태 손실과 이 방법을 결합하여 BERT의 언어 이해 능력을 97% 유지하면서 모델 크기를 40% 줄이고 속도를 60% 높였습니다.
• 특성 기반: 학생은 교사의 중간 특성 맵을 따라 하며, 학습 가능한 작은 프로젝터를 사용해 채널 수의 차이를 연결합니다. FitNets는 이러한 "힌트"를 도입했으며, 교사보다 파라미터 수가 약 10.4배 적으면서 CIFAR-10에서 교사보다 뛰어난 성능을 보인 학생 모델을 학습했습니다.
• 시퀀스 수준: 언어 모델 교사가 텍스트를 생성하고, 학생은 해당 출력으로 파인튜닝됩니다. DeepSeek은 Qwen2.5 및 Llama 3를 기반으로 한 R1 증류 모델을 DeepSeek-R1을 사용해 선별한 800,000개 샘플로 파인튜닝했습니다.

특성 기반 증류는 단일 클래스 분포가 아니라 수천 개의 박스와 점수를 출력하는 검출기에 적합합니다. Ultralytics YOLO는 학습에 distill_model 교사 모델이 제공되면 이 방법을 사용합니다. 고정된 대형 YOLO26 모델이 검출 헤드에 입력되는 3개의 넥 특성 맵을 제공하고, 학생은 교사 신뢰도로 가중치를 부여한 L2 손실을 적용한 1×1 컨볼루션 프로젝터를 통해 해당 특성 맵을 따라 합니다. 손실은 기본값이 6.0인 dis로 조정됩니다. COCO val2017에서 증류된 YOLO26 nano 모델은 교사 없이 학습한 동일 모델의 40.9보다 높은 41.5 mAP_50-95를 기록하며, 추론 비용은 동일합니다.

from ultralytics import YOLO

# 학생: 배포할 소형 YOLO26 nano 모델
model = YOLO("yolo26n.pt")

# 교사: 학습을 안내하고 이후에는 폐기하는 대형 YOLO26 체크포인트
results = model.train(data="coco8.yaml", epochs=5, distill_model="yolo26s.pt")

증류와 관련 용어 비교#

• 전이 학습: 새 태스크의 시작점으로 모델 자체의 가중치를 재사용합니다. 증류는 같은 태스크에서 다른, 대개 더 작은 아키텍처가 모델을 모방하도록 학습합니다.
• 모델 프루닝: 학습된 네트워크에서 가중치, 채널 또는 레이어를 제거합니다. 정확도를 회복하기 위해 프루닝한 모델을 원본 모델에서 증류해 파인튜닝하는 경우가 많습니다.
• 모델 양자화: INT8처럼 낮은 정밀도로 가중치와 활성화 값을 저장합니다. 양자화는 동일한 아키텍처의 크기를 줄이고, 증류는 더 작은 아키텍처를 선택하므로 두 방법을 함께 사용할 수 있습니다.
• 모델 앙상블: 추론 시 여러 모델의 출력을 평균합니다. 앙상블을 하나의 학생 모델로 증류하는 것이 이 기법의 원래 목적이었습니다.
• 데이터셋 증류: 이름은 비슷하지만 모델이 아니라 데이터셋을 압축합니다. Wang 외 연구진은 MNIST 학습 이미지 6만 개를 클래스당 하나씩, 합성 이미지 10개로 압축했습니다. 이 이미지를 사용하면 초기화가 고정된 네트워크가 몇 번의 경사 하강 단계만으로 원래 정확도에 근접합니다.