[바미] 지식 증류(Knowledge Distillation)로 성능은 챙기고, 모델은 줄이기

들어가기전에

실서비스에 모델을 붙이려면, 성능만큼 중요한 게 무게(추론 비용)입니다. 모델이 무거우면 바로 병목이 되버리죠.

 

큰 모델(Teacher)은 정확하지만 느리고 비싸고, 작은 모델(Student)은 빠르고 싸지만 성능이 아쉬운 경우가 많은데 이 둘 사이를 이어주는 대표적인 방법이 Knowledge Distillation(지식 증류)입니다.

이를 한 문장으로 정리하면 다음과 같이 표현해 볼 수 있어요.

큰 모델(Teacher)의 예측 지식을 작은 모델(Student)이 따라 배우도록 학습시키는 방법.

 

그럼 왜 굳이 증류가 필요할까요?

서비스 관점에서 모델이 좋다는 것은 단순히 정확도만 의미하지 않습니다.

• 모바일/엣지 환경에서 돌아가야 함
• 응답 속도가 중요함(지연 시간, latency)
• 비용이 중요함(GPU/서버 비용)
• 트래픽이 올라가도 안정적으로 버텨야 함

그런데 큰 모델은 보통 이 조건들과 충돌하기 때문에 큰 모델의 성능을 최대한 유지하면서, 더 작은 모델로 바꾸고 싶을 때 증류가 등장합니다.

 

그래서 정답(label)만 보지 말고 선생님의 답안지를 함께 봅니다.

일반적인 학습은 아래처럼 진행됩니다.

• 입력 x → 모델이 예측 p(y∣x) → 정답 레이블 y와 비교 → 손실(loss) 최소화

지식 증류에서는 여기에 하나가 더 들어옵니다.

• 입력 x → Teacher가 예측한 분포 pT​(y∣x)
• 입력 x → Student가 예측한 분포 pS​(y∣x)
• 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를 따라갈 수 있는 한계가 있습니다.

그럼 간단한 예시를 살펴봅시다.

큰 모델로 라벨을 만들어 작은 모델을 학습을 예로 들어 증류를 개념적으로 설명하면 이런 식입니다.

 

1. Teacher 모델을 충분히 좋은 성능으로 학습하고, 확보합니다.
2. 학습 데이터(또는 추가 데이터)에 대해 Teacher의 확률 분포(logits)를 뽑아 저장합니다.
3. Student 모델은 정답(label) + Teacher 분포(soft label)를 함께 보고 학습합니다.

 

실제로는 모델, 프레임워크에 따라 구현이 다양하지만 전체 흐름은 이 구조를 크게 벗어나지 않습니다.

마치며

지식 증류는 큰 모델이 낸 결과를 작은 모델이 따라 배우게 한다는 단순한 아이디어지만 서비스 관점에서 굉장히 실용적인 기술입니다.

성능만 보고 Teacher를 선택하는 게 아니라 배포 환경(속도, 비용, 디바이스)을 고려해서 Student를 설계하고, Teacher의 예측 분포를 학습 신호로 활용하는 방식이라고 보시면 됩니다.