References
Introduction
일반적으로 지식 증류는 소형 모델에서도 대형 모델만큼의 성능을 내고자 하는 의도로 수행하게 된다.
그러나 막상 증류를 해보고 나면 생각했던 것만큼 성능이 잘 나오지 않았던 경험이 종종 있었을 것이다.
왜일까?
많은 연구에서 다양한 이유를 제시하고 있다. Student - Teacher 간 갭이 너무 크거나, 둘 간의 모델 구조 차이가 너무 많이 나거나,
여러 이유가 있지만, 이 논문에서는 작은 모델로의 증류 시 사용하는 projection 과정의 비효율을 지적한다.
증류를 해야 하는데 깔때기가 허름하면 무슨 소용이랴.
VkD는 개선된 projection 방법론을 제안하여 지식 증류의 효율을 더 끌어올리고자 하였다.
왜 지식 증류에 projection이라는 것을 사용하며,
기존의 projection 방법론이 무엇이었고,
그 점을 어떻게 개선하였는지
아래에서 간단명료하게 설명하도록 하겠다.
Research Gap: feature Knowledge Distillation과 projection의 허점
일반적인 지식 증류는 모델의 끝단, 즉 두 모델 간 logit의 값을 비교하여, student의 답이 teacher가 내놓은 답과 최대한 같아지도록 페널티를 부여하여 학습한다(이때 페널티에 주로 KL-divergence를 사용한다).
이게 제프리 힌튼이 2015년 처음 제안한 지식 증류 방법론이었으나,
당시 방법론은 image classification을 기준으로 두었기에 태스크나 modality가 달라지면 그 효과가 제한적이었다.
그래서 이후 등장한 방법론이 logit 단위의 비교 대신 모델의 중간 feature 출력을 비교하는 것이었다.
일반적으로 feature knowledge distillation은 student / teacher 각 모델의 k번째 레이어 출력을 비교하여, MSE loss등을 기반으로 student가 teacher와 유사해지도록 유도한다.
여기서 드는 의문,
'작은 모델을 증류해야 한다면 항상 teacher와 student의 dimension이 맞지 않을 텐데, 둘을 어떻게 비교하지?'
Teacher의 hidden dimension이 576인데, student는 384라면?
당연히 MSE 계산이 불가능해진다. 이 문제를 해결하기 위해 projection을 사용한다.
일반적으로 student의 차원이 teacher의 차원과 같아지도록, 레이어 중간에 projection을 붙여서 student도 차원이 576이 되도록 유도한다.
Projection layer는 학습 중에만 사용하고, 지식 증류 파이프라인이 종료되고 난 후(추론 시점)에는 떼어내 원래 student 구조로 복원된다.
이때 student의 차원을 키워주는 linear projection layer는 원래 모델에 loss를 전달해야 하므로, learnable하다.
이게 핵심 문제다.
Projector는 단순히 원래 student의 feature geometry를 잘 보존하면서 크기만 키우는 역할을 해야 하는데,
loss를 받아서 학습하는 일종의 '모델'이다 보니, 이 레이어 자체가 teacher 출력과 잘 닮도록 유도되어 버린다는 거다.
원래는 student가 배워야 할 정보를 linear layer가 배워 버리니, projection layer가 없는 실 사용 시점에는 증류가 잘 안 되어있는 것이다.
특히나 증류 효율을 높이려고 linear를 MLP 구조로 무겁게 만드는 경우도 존재하는데, 이 경우 projector의 task 학습은 더욱 심해질 수밖에 없다.
그럼 어떻게 해야 할까?
Method: Orthogonal Projection
답은 'projection layer가 최대한 정보를 새로 배우지 않도록 한다'이다.
즉 projector - model 간 정보 차가 적어야 한다.
이를 실현하기 위해 orthogonal projection 이라는 방법론을 제안한다.
선형대수에서 말하는 선형 종속, 선형 독립 개념에 대해 알고 있다면 이해가 편하다.
일반적으로 선형 종속인 원소들은 동일한 성질을 공유하고, 독립인 원소들은 서로 다른 원소라고 본다.
한 행렬 내 선형 독립인 원소가 많다는 것은 선형대수에서의 rank 수가 많다는 것으로 해석 가능하고,
이는 행렬, 측 feature 내부에 유의미한 정보들이 많다는 것을 의미한다.
이러한 점에서 Orthogonal, 즉 직교 형태의 projector를 만드는 이유는,
Projection 자체를 'matrix의 회전 연산'으로 만들어 원래 student의 정보를 왜곡시키지 않게 하고, 이로 인해 어떤 방향이 압축되고 겹쳐서 내부 feature이 선형 종속으로 변하는 현상을 막기 위함이다.
이해를 위한 그림을 아래에 첨부한다.
Orthogonal projection은 다음과 같은 절차를 통해 만들어진다.
최대한 쉬운 설명으로 작성하고자 노력했는데, 한번에 이해가 잘 될지 모르겠다.
즉 요약하자면
1. 144*288 사이즈의 projector가 필요하더라도 의도적으로 288*288 사이즈의 반대칭 행렬 W를 만들고,
2. Exponential로 orthogonal 성질을 만든 후,
3. 이 행렬을 필요한 만큼(144*288) 위에서부터 잘라 쓰고,
4. 페널티는 잘린 만큼의 행렬(projector, 144*288)에 주어진다 - loss가 일부에만 들어가더라도 영향은 행렬 전체에 미치므로.
5. 그리고 loss로 학습한 행렬은 새로운 연산에선 또 다시 288*288 사이즈를 만들고, 위 과정을 반복한다.
이라고 이야기할 수 있겠다.
아래는 논문에서 제시한 Figure인데,
이와 같이 일반 linear projector는 기존 feature를 찌그러뜨리고 일부 변형하는 반면,
VkD는 기존 feature의 geometry를 훼손하지 않고 회전만 유도할 수 있었다고 한다.
여기에 별도로, teacher의 feature에 task-specific한 regularization을 걸어 주어 증류의 효율을 향상시키기도 하였다.
간단히 정리하면 classification task에는 feature에 standardization을,
generation에는 feature에 whitening을 걸어 주었을 때 성능이 추가로 향상되었다는 이야기다.
Result & Conclusion
우선 Vision task(Image Classification)에서 baseline 모델과, 경량화된 모델들을 나열하며 파라미터 대비 성능을 보여주었다.
VkD를 적용한 student 모델 백본은 DeiT였으며, teacher는 84M 파라미터에 82.6%의 성능을 보여준 RegNetY-160이었다.
VkD 뒤에 붙은 Ti, S는 각각 Tiny(6M), Small(22M)을 의미한다.
기존의 방법론들을 모두 뛰어넘는 성능을 보여주었다. 특히 VkD-S에서는 teacher 1/4 파라미터 수준으로 teacher만큼의 성능을 보여주었다.
이번에는 Object Detection이다. (* 표에 나오는 Swin은 Swin-Transformer를 의미)
Distillation을 사용하지 않았을 때에 비해, 본 방법론 적용 시 성능이 기존대비 향상하였음을 확인할 수 있다.
결과적으로,
간단한 구조의 projector를 통해 feature KD의 효율 개선이 가능함을 보여주었다.
더 나은 지식 증류를 위해 남은 문제는 무엇일까?
개인적으로는 'student-teacher 간의 비교에 MSE를 사용하는 것이 가장 효율적인 방법일까' 에 대한 의문이 든다.
task/modality에 따라 단순 MSE보단 더 나은 방법들이 생겨날 수 있지 않을까?
향후 또 흥미로운 논문 리뷰로 찾아오도록 하겠다.
'[Efficient AI] > [Efficient AI] Knowledge Distillation' 카테고리의 다른 글
| Distill-Whisper: Robust Knowledge Distillation via Large-Scale Pseudo Labelling (0) | 2026.07.17 |
|---|---|
| Knowledge Distillation(지식 증류) - Overview (0) | 2026.02.25 |