안녕하세요
최근 캡스톤 디자인을 하면서 정리한 내용들을 티스토리에 정리해보려고 합니다..
정리용이라 약간 미흡할 수 있으니 양해 바랍니다.
저희 캡스톤 디자인 주제에서 사용하는 모델은 on-device 형태로 진행해야하기 때문에 경량화가 필수불가결적인데요.
그래서 저흰 지식 증류라는 기법을 사용하게되었습니다.
그리하여, 오늘은 딥러닝 학습 방법 중 하나인 Knowledge Distillation, 지식 증류 (또는 모델 증류)에 대해 알아보겠습니다.
1. 정의
Knowledge Distillation, 지식 증류란 복잡하고 큰 모델 (Teacher Model)이 학습한 지식을 상대적으로 간단하고 작은 모델 (Student Model)에 전달하는 기법입니다.
즉, “성능은 비슷하게 유지하면서, 더 가볍고 빠른 모델을 만들기 위한 학습 방법”입니다.
이 방법은 작은 모델이 큰 모델의 성능을 최대한 따라가도록 훈련할 때 사용됩니다.
다만 Transfer learning (전이 학습)과는 다른 개념입니다.
Transfer learning은 parameter를 일부 복사하여 하위 task에 대해 fine-tuning을 하는 방식이지만,
knowledge distillation은 parameter 복사 없이 teacher model의 예측값 자체를 label로 사용하여 학습합니다.
Teacher model의 좋은 성능과 지식을 copy 하는 것으로, 모델 경량화 기법이라 볼 수 있습니다.
전이학습과 지식증류를 구체적으로 비교하면 다음과 같습니다.
| Knowledge Distillation | Transfer Learning | |
| 목적 | Teacher 모델의 지식을 Student 모델로 전달 | 기존 모델의 지식을 새로운 작업에 활용 |
| 효과 | 모델 압축 및 경량화 | 새로운 작업에서 빠른 학습 (데이터 부족 시 유용) |
| 모델 | Teacher와 Student 모델 구조가 다를 수 있음 | 기존 모델 구조를 유지하거나 약간 수정 |
| 데이터 | Teacher 모델과 동일한 데이터로 학습 | 새로운 작업에 맞는 데이터 필요 |
2. 방법
이제 구체적으로 knowledge distillation을 어떻게 하는지 알아보도록 하겠습니다.
Teacher model은 학습 데이터의 복잡한 패턴과 높은 수준의 지식을 학습합니다.
이를 직접 Student model에 적용하기 어렵기 때문에, knowledge distillation에서는 다음과 같은 과정이 진행됩니다:
1) Teacher Model 학습
크고 복잡한 모델(예: ResNet-152, GPT-3 등)을 먼저 학습합니다.
이 모델은 높은 정확도와 강력한 추론 능력을 가지고 있습니다.
2) Soft Target 생성
Teacher Model은 입력 데이터를 예측할 때, 단순히 정답 레이블만 예측하지 않고, 클래스 간 확률 분포 (Soft Target)도 생성합니다.
T가 클수록 label이 더 soft 해진다는 것을 알 수 있습니다.
- T=1: 평소 softmax
- T>1: 분포가 평탄해짐 → teacher의 “헷갈림”이 더 잘 드러남
예를 들어, 동물 분류 문제에서 Teacher Model의 출력이
다음과 같은 확률 분포를 보인다고 가정합시다.
여기서 soft target은 고양이가 맞긴 한데 개/말이랑도 좀 비슷함 이라는 유사도 정보를 담고 있습니다.
이 분포는 단순히 Hard Target(데이터셋의 true label)보다 깊은 의미를 지닙니다.
예를 들어, 어떤 고양이 이미지는 개와 유사한 모습일 수도 있습니다.
Hard Target은 이런 정보를 전혀 반영하지 않고 “그냥 고양이(1) vs 나머지(0)”로만 보지만,
Soft Target은 “고양이 0.7, 개 0.1.5, 말 0.05…”처럼 클래스 간 유사도와 모델의 불확실성까지 함께 표현합니다.
마치 학생이 교과서만 보고 공부하는 것이 아니라 선생님의 심도있는 지식과 인사이트를 습득하는 것과 비슷합니다.
3) Student Model 훈련
Student Model은 Teacher Model이 제공한 Soft Target과 원래의 정답 레이블(Hard Target)을 조합하여 학습합니다.
따라서 distillation loss는 두 개의 loss가 결합된 형태인데, 1) soft loss 2) hard loss 가 일정한 가중치로 더해집니다.
먼저 soft loss는 일반적으로 Kullback-Leibler Divergence 함수를 통해 Soft Target과 Student 모델의 예측값 간 차이를 최소화합니다. 이때 Soft target은 모델의 출력값을 temperature으로 나눈 뒤 softmax 하여 만들어집니다.
Hard loss는 일반적으로 볼 수 있는 Cross Entropy Loss를 활용한 함수입니다.
Supervised learning 할 때처럼 Student 모델 예측값과 데이터셋 true label (hard target) 간 차이를 최소화합니다.
실무에서는 soft loss와 hard loss의 비율(α)은 하이퍼파라미터라서,
실험으로 적절한 값을 찾는 것이 일반적입니다.
Hinton 논문이나 후속 연구들에서 soft loss 비중을 상대적으로 크게 주는 경우가 많지만,
데이터/모델에 따라 최적 비율은 달라집니다.
4) 다양한 방식
앞서 설명한 방식은 Teacher 모델의 output에 기반하여 학습하는 'response-based KD'라 불립니다.
특히 Hinton 등이 제안한 conventional(=vanilla) KD의 전형적인 형태입니다. 다음은 vanilla KD에 대한 설명입니다.
- Teacher와 Student의 아키텍처는 다를 수 있지만, 출력 차원(클래스 수)은 같아야 함
Temperature-softmax로 soft target을 만들고, KL-divergence 기반 KD loss와 CE loss(정답 라벨)를 가중합하는 구조 - 보통 “정답 라벨 CE loss + 교사 분포 KD loss”를 섞어서 사용
- 정답(one-hot)만 쓰는 것보다, 교사가 다른 클래스에 주는 “애매한 확률” 정보까지 학생에 전달해 일반화 성능을 높이기 위해서 사용
이외에도 Teacher 모델의 지식을 전달하는 확장된 기법이 존재합니다.
TA-KD (Teacher Assistant KD)
- 중간 크기 보조 교사를 두는 KD
- 거대한 Teacher → 중간 크기 Assistant → 작은 Student 순서로 단계적 증류
- 각 단계에서 Vanilla KD와 비슷한 방식으로 distill
- 아주 큰 교사와 아주 작은 학생을 바로 연결하면 capacity gap 때문에 학습이 불안정해지는데, 중간 모델을 두어 난이도를 낮추려는 전략
feature-based KD
- Teacher 모델의 중간 레이어에서 추출한 정보를 Student 모델이 학습하도록 하는 방법
- Teacher 모델의 특정 레이어의 출력 (Feature Map)과 Student 모델의 해당 레이어 출력을 비교하는데, MSE 또는 L2 Loss로 그 차이를 줄임
hidden states-based distillation
- 출력층만이 아니라, 중간 레이어의 hidden states 을 맞추는 KD
- Transformer 각 층(or 선택된 몇 층)의 hidden states $h^l_t$, $h^l_s$를 L2, cosine, MSE 등으로 정렬
- 내부 표현 자체를 가이드하면, 아웃풋만 맞출 때보다 더 풍부한 구조 정보가 전달되어 수렴 안정성과 성능이 좋아지는 경우가 많음
feature-based / hidden-states-based KD
- 사실상 같은 축
- CNN에서는 feature map, Transformer에서는 hidden states라고 부르는 정도라 “Representation-based KD”이라 통칭해도 됨
RKD (Relational KD)
- Teacher 모델이 학습한 데이터들 간 관계 정보를 전달하는 것
- 즉 입력 데이터들 간 관계를 Teacher 모델이 학습했을 때, 이를 Student 모델도 학습하도록 유도Distance-wise, angle-wise loss를 사용하여 학습
logit-based KD
- Teacher의 Logit 값을 활용하여 Student를 학습시키는 방법 (그림 참고)
- Logit
- Softmax 함수를 적용하기 전의 예측값
- 클래스 간 유사도와 상대적 가능성 정보를 포함
블랙박스 KD
- 교사 모델 내부(파라미터, 히든스테이트)에 접근할 수 없고, 입·출력만 사용할 수 있는 KD
- teacher API (예: GPT-4, 상용 LLM)로 입력을 던지고, 응답 텍스트/점수/라벨만 받음
- 그 응답을 pseudo-label/soft label로 써서 학생을 지도학습
- 상용 LLM, 폐쇄형 모델을 teacher로 쓸 때 내부 로짓·히든스테이트를 볼 수 없기 때문
화이트박스 KD
- 교사 내부에 완전 접근 가능한 KD
- teacher 로짓, 히든스테이트, 어텐션 등 내부 정보를 모두 꺼내와 loss로 사용
- 더 많은 신호(중간 표현, 어텐션 패턴 등)를 학생에게 줄 수 있어서, 같은 데이터로 더 좋은 학생 성능을 얻기 쉬움
feature-based KD나 hidden-states-based KD, attention distillation 등은
Teacher의 내부 표현에 접근해야 하므로 white-box KD에 해당합니다.
반대로, 상용 LLM API처럼 내부를 볼 수 없는 경우에는
입·출력만 이용하는 black-box KD만 사용할 수 있습니다.
사용해볼만한 코드
https://github.com/SforAiDl/KD_Lib
A Pytorch Knowledge Distillation library for benchmarking and extending works in the domains of Knowledge Distillation, Pruning, and Quantization
https://github.com/arcee-ai/DistillKit
An Open Source Toolkit For LLM Distillation
https://github.com/modelscope/easydistill
a toolkit on knowledge distillation for large language models
https://github.com/jongwooko/distillm-2
Official PyTorch implementation of DistiLLM-2: A Contrastive Approach Boosts the Distillation of LLMs (ICML 2025 Oral)
참고자료: https://velog.io/@tony3ynot/Knowledge-Distillation-%EC%A7%80%EC%8B%9D-%EC%A6%9D%EB%A5%98
'AI > 캡스톤디자인' 카테고리의 다른 글
| 「2025년도 확장형 캡스톤디자인 성과보고회:STaRISE」 AIShield 장려상 수상 후기 + 프로젝트 회고 (텍스트 기반 보이스피싱 탐지) (0) | 2026.03.02 |
|---|
