Knowledge Distillation에 대해 알아보자

반갑습니다.
이번 글에서는 Knowledge Distillation에 대해서 소개해보겠습니다.

Knowledge Distillation은 최근에도 많이 사용되고 있는 기술입니다.
API가 아주 저렴한 GPT 4o-mini도 GPT 4o 모델을 KD하여 만든 모델이죠.
그럼 자세히 알아가봅시다 !!

Knowledge Distillation?

고성능의 Teacher 모델로부터 지식을 전달 받아 Student 모델을 학습 시키는 기법을 의미합니다.
성능 저하를 최소화하면서 모델을 압축하는 방법 중 하나입니다.

목적이 뭐냐면
Teacher 모델보다 훨씬 작은 파라미터를 갖는 Student 모델의 성능을 Teacher 모델만큼 올려보겠다는 겁니다.
Teaching을 하는 것이니 KD를 하려면 고성능의 Teacher 모델이 일단 준비가 돼 있어야 합니다.

Knowledge Distillation의 분류

Response-based

Teacher 모델의 Logit을 통해 학습하는 Logit-based 방법이 있고, Output을 통해 학습하는 Output-based 방법이 있습니다.

1. Logit Based

Teacher 분류 모델의 Output 확률분포가 [0.8, 0.1, 0.1] 라면
Student 모델이 자신의 Output과 비교하여 Loss를 계산한 뒤 Backward를 거칩니다.

즉 Teacher의 예측과 Student의 예측이 유사해지도록 학습을 진행한다고 생각하시면 되겠습니다.

일반적으로 Student 모델은 Teacher 모델의 Output과 비교한 KL Divergence Loss 한개, 정답인지 아닌지에 대한 Cross-entropy Loss 한개 총 두개의 로스가 더해져서 학습이 진행됩니다.

이 때, Teacher 모델의 Logit을 Softmax에 통과시킬 때, Temperature로 나눈 뒤 통과시킵니다.

저렇게 한 쪽에 치우친 분포를 학습한다는 것은..
결국 Cross-entropy Loss를 두 번 곱해서
Backward 하는 것과 다를 게 없으니 Teacher 모델의 확률 분포를 좀 더 완만하게 학습하는 것이지요.

2. Feature-based

Teacher 모델의 중간 레이어 Feature/Representation을 사용합니다.

Logit-based KD에서는 마지막 Output을 사용했다면 Feature-based KD에서는 Student 모델이 Teacher 모델의 가운데 Feature map을 학습합니다.

$f_T$ 와 $f_S$의 차원은 맞지 않을 수 있기에 Regressor Layer R을 사용하여 크기를 맞춘 뒤, MSE Loss를 적용합니다.

일반적으로 초기 레이어는 너무 저수준의 특징을, 후반 레이어는 너무 고수준의 특징을 표현하기에 중간 레이어를 활용합니다.

Transparency(투명도)

위 방법은 White Box 모델에서 사용 가능한 KD 방법입니다.
AI모델은 크게 세가지로 분류되는데

1. Black Box
   Teacher 모델의 Output만 알 수 있습니다. (GPT)

2. White Box
   Teacher 모델의 내부 구조, 파라미터를 알 수 있습니다. (llama)

3. Gray Box
   Teacher 모델의 Output 및 최종 Logit값 등 제한된 정보를 알 수 있습니다.

이 중, 당연히 White Box Model이 좋겠지만 Black Box Model들의 성능이 너무 좋다보니 Black Box Model들의 지식을 가져오고자 하는 시도가 생겼습니다.

Imitation Learning

다른 Agent의 행동을 관찰하고 모방하여 자신의 정책을 학습하는 기계 학습 방법입니다.
NLP 분야에서 많이 활용되는 방법이며, Black Box 모델로부터 지식을 배울 수 있습니다.

방법은 간단합니다.
다양한 질문 셋을 Black Box LLM에게 던져 답을 얻은 뒤, 그 QA 데이터셋을 SLM(Small Language Model)이 학습합니다.

이러한 Imitation Learning은 사실상 유일한 Black Box 모델의 KD 접근법입니다.
장점으로는 지식을 인간이 해석할 수 있다는 점이 있습니다.
하지만 QA 데이터셋을 통해 학습하는 것이다 보니 QA 데이터셋의 품질에 따라 성능이 많이 좌우됩니다.

'
'
'
이로서 Knowledge Distillation에 대한 설명을 마칩니다.
도움이 되었길 바랍니다.
감사합니다 !!