[바미] 지식 증류(Knowledge Distillation)로 성능은 챙기고, 모델은 줄이기
들어가기전에
실서비스에 모델을 붙이려면, 성능만큼 중요한 게 무게(추론 비용)입니다. 모델이 무거우면 바로 병목이 되버리죠.
큰 모델(Teacher)은 정확하지만 느리고 비싸고, 작은 모델(Student)은 빠르고 싸지만 성능이 아쉬운 경우가 많은데 이 둘 사이를 이어주는 대표적인 방법이 Knowledge Distillation(지식 증류)입니다.
이를 한 문장으로 정리하면 다음과 같이 표현해 볼 수 있어요.
큰 모델(Teacher)의 예측 지식을 작은 모델(Student)이 따라 배우도록 학습시키는 방법.
그럼 왜 굳이 증류가 필요할까요?
서비스 관점에서 모델이 좋다는 것은 단순히 정확도만 의미하지 않습니다.
- 모바일/엣지 환경에서 돌아가야 함
- 응답 속도가 중요함(지연 시간, latency)
- 비용이 중요함(GPU/서버 비용)
- 트래픽이 올라가도 안정적으로 버텨야 함
그런데 큰 모델은 보통 이 조건들과 충돌하기 때문에 큰 모델의 성능을 최대한 유지하면서, 더 작은 모델로 바꾸고 싶을 때 증류가 등장합니다.
그래서 정답(label)만 보지 말고 선생님의 답안지를 함께 봅니다.
일반적인 학습은 아래처럼 진행됩니다.
- 입력 → 모델이 예측 → 정답 레이블 와 비교 → 손실(loss) 최소화
지식 증류에서는 여기에 하나가 더 들어옵니다.
- 입력 → Teacher가 예측한 분포
- 입력 → Student가 예측한 분포
- Student가 Teacher의 예측 분포를 따라가도록 학습
여기서 포인트는 정답만 맞히는 것이 아니라, Teacher가 어떤 선택지를 얼마나 유력하게 봤는지(확률 분포)까지 배우는 데 있습니다.
예를 들어 분류 문제에서 정답이 고양이라고 해도 Teacher는 아래와 같이 예측할 수 있습니다.
- 고양이 0.70
- 강아지 0.20
- 여우 0.08
- 기타 0.02
Student는 이 분포를 보면 '강아지/여우도 헷갈릴 수 있는 케이스구나!' 같은 미묘한 정보(일종의 힌트)를 얻게 되는데 이게 증류의 핵심 가치입니다.
여기서 Soft Label이랑 Temperature가 나옵니다
증류 설명에서 거의 항상 같이 나오는 용어가 있습니다.
Soft Label(소프트 라벨)
정답이 one-hot(딱 하나만 1)인 하드 라벨이 아니라, Teacher가 내놓은 확률 분포 전체를 말합니다.
Temperature(온도, T)
Teacher의 확률 분포를 더 부드럽게 만들기 위해 사용하는 값입니다. 소프트맥스(softmax)에 Temperature를 적용하면 분포가 완만해져서, Student가 선택지들 간의 관계를 더 잘 배우는 경우가 많습니다.
- 𝑇가 커질수록 분포가 더 평평해집니다(덜 확신하는 형태)
- 𝑇가 작아질수록 분포가 더 뾰족해집니다(더 확신하는 형태)
실무에서는 보통 𝑇 = 2 ∼ 5정도를 후보로 두고 실험하는 경우가 많습니다.(물론 정답은 없고 데이터/모델에 따라 다릅니다).
그럼 실제로는 어떻게 학습할까요?
지식 증류는 보통 아래 두 학습 목표를 같이 씁니다.
- 정답(label)을 맞추는 손실 (일반적인 CE loss)
- Teacher 분포를 따라가는 손실 (KL divergence 등)
즉 Student는 정답도 맞추고, Teacher처럼 예측하도록 학습됩니다.
이를 간단히 애기하면 아래와 같이 표현할 수 있겠네요!
정답을 맞추되, 선생님의 풀이 방식도 따라 해라.
그럼 어디에 많이 쓰일까요?
대규모 모델 → 경량 모델로 배포
- 서버 비용이나 응답 속도가 중요한 서비스
- 모바일/온디바이스 모델 배포
NLP/추천/검색에서 빠른 랭킹 모델 만들기
- Teacher로 좋은 성능을 뽑아놓고
- Student는 빠르게 서빙 가능한 구조로 가져가는 패턴이 자주 나옵니다.
이 뿐만 아니라 LLM 시대에도 여전히 많이 씁니다.
LLM 자체를 작은 LLM으로 증류하기도 하고, LLM의 결과를 이용해 작은 분류기/리랭커를 훈련하는 형태로도 많이 사용되죠.
장점과 주의할 점
장점
- 작은 모델 성능을 꽤 끌어올릴 수 있습니다.
- 추론 비용/지연 시간을 줄이면서 정확도 손실을 완화할 수 있습니다.
- 레이블이 부족한 상황에서 Teacher가 추가 신호를 제공하는 효과도 있습니다.
주의점
- Teacher가 틀린 경우, Student도 그 방향을 배울 수 있습니다.
(Teacher 품질이 중요합니다.) - 데이터 분포가 바뀌면 Teacher/Student 모두 재학습이 필요할 수 있습니다.
- Student 구조가 너무 단순하면 Teacher를 따라갈 수 있는 한계가 있습니다.
그럼 간단한 예시를 살펴봅시다.
큰 모델로 라벨을 만들어 작은 모델을 학습을 예로 들어 증류를 개념적으로 설명하면 이런 식입니다.
- Teacher 모델을 충분히 좋은 성능으로 학습하고, 확보합니다.
- 학습 데이터(또는 추가 데이터)에 대해 Teacher의 확률 분포(logits)를 뽑아 저장합니다.
- Student 모델은 정답(label) + Teacher 분포(soft label)를 함께 보고 학습합니다.
실제로는 모델, 프레임워크에 따라 구현이 다양하지만 전체 흐름은 이 구조를 크게 벗어나지 않습니다.
마치며
지식 증류는 큰 모델이 낸 결과를 작은 모델이 따라 배우게 한다는 단순한 아이디어지만 서비스 관점에서 굉장히 실용적인 기술입니다.
성능만 보고 Teacher를 선택하는 게 아니라 배포 환경(속도, 비용, 디바이스)을 고려해서 Student를 설계하고, Teacher의 예측 분포를 학습 신호로 활용하는 방식이라고 보시면 됩니다.