Knowledge Distillation(지식 증류) - Overview

 

개요

지식 증류는 일반적으로 이미 학습된 큰 모델(Teacher라 부름)에서, 동일 task를 수행하기 위한 더 작은 모델(Student라 부름)을 효율적으로 훈련하기 위해 사용하는 기법이다.

 

2015년에 제프리 힌튼이 처음으로 제안한 방법이다.(https://arxiv.org/abs/1503.02531)

 

이 방법론은 말 그대로 Teacher의 지식을 증류하여 Student에게 전달하는 것인데, 큰 메커니즘은 아래와 같다.

• Teacher와 Student에게 동일한 샘플을 보게 하고,
• Student는 task로부터의 loss(예를 들어, 분류 문제였으면 통상적으로 Cross-entropy, 회귀면 MSE 등이 사용될 수 있을 것이다.)를 학습함과 동시에,
• Teacher가 예측한 답안과 오차로부터의 loss(= teacher의 representation과 얼마나 다른가에 대한 penalty인데, MSE도 사용될 수 있고 방법은 다양하다.)를 추가로 학습한다.

 

이때, 당연히 Teacher는 loss를 업데이트하지 않고 freeze 해둔다. 이미 학습이 완료되었던 시점의 teacher 지식을 증류하는 게 목적이니까!

 

일반적으로 동일 모델을 그냥 학습하는 것보다, 지식 증류를 사용하면 Teacher의 표현 방식 & 결정 경계를 모방하게 되어 더 우수한 성능으로 학습할 수 있다. 특히 적은 파라미터로 큰 파라미터를 가진 모델의 성능을 유도할 수 있어 사용한다.

 

방법론

(Knowledge Distillation을 편의상 KD라고 기재하겠다. 중간중간 예시는 음성 도메인으로 설명한다.)

KD 방법은 크게 응답 기반(response-based), 특징 기반(feature-based), 관계 기반(relation-based) 지식 전달로 구분할 수 있다.

• Response-based KD
  • Teacher 모델의 최종 출력(logit)을 소프트 타겟으로 사용하여 학생 모델을 훈련하는 전통적인 접근
  • Student는 Teacher의 예측 분포를 모방하도록 학습하기 위해, teacher-student 간 logit 차이를 최소화하는 distillation loss를 사용
  • X+AI Example
    • ASR에서는 이런 방식으로 음소/문자 레벨의 확률 분포 (soft label)를 전달하여 학생 모델이 교사의 출력 패턴 을 학습하도록 함
• Feature-based KD
  • Teacher의 중간 표현(feature activations)을 학생이 모방하도록 하는 접근
  • Teacher와 student 모델의 특정 은닉층 출력을 맞추는 레이어 정렬(loss)을 적용하거나, MSE 손실 등으로 feature vector 간 차이를 최소화한다.
  • ASR 모델의 hidden layers는 음성 신호의 음향 특성과 언어 패턴을 추상화한 표현을 학습하므로, 이러한 중간 특징 지식을 전달하면 학생이 단순 출력만 모방하는 것보다 더 깊은 수준의 지식을 얻을 수 있다.
• Relation-based KD
  • 데이터 간 또는 특징 간 관계를 지식으로 활용하는 방법
  • Teacher 모델이 학습한 내부 구조나 표현 공간에서 샘플들 간의 유사도/상관관계 등을 학생에게 전달한다.
    ex. Teacher 출력의 pairwise 유사도 행렬이나 특징 맵들 간의 상호상관을 학생이 모방하도록 학습시킬 수 있음
  • FYI. 포스텍 연구실에서 제안한 방법이다. 관심 있다면 깊게 읽어 보시길! (https://arxiv.org/abs/1904.05068)

 

또한 지금까지는 Single teacher distillation을 기준으로 설명했는데,

뿐만 아니라 self-distillation, multi-teacher distillation 또한 가능하다.

• Self-Distillation
  • 동일 모델의 깊은 층(teacher 역할)과 얕은 층(student 역할) 간에 지식을 전달하거나
  • 이전 epoch의 모델을 교사로 삼아 현재 epoch의 모델을 지도하는 형태 등 방법론은 다양하다.
  • FYI. ASR 분야에서 자기 증류의 대표적인 예는 Moriya et al.(2020)의 연구 (https://www.isca-archive.org/interspeech_2020/moriya20_interspeech.html)
    

기타 지식증류 관련 주의사항/인사이트 등

• Capacity gap
  • Teacher에 비해 student size(width가 작거나 layer 개수가 적은 상황 등)가 작으면 teacher의 decision boundary를 정확히 근사하지 못하는 문제가 발생한다.
  • 이를 부분 해결하기 위해 (teacher-student 사이)중간 사이즈의 assistent teacher를 하나 더 두는 Teacher-Assistant KD라는 방법론도 존재한다.
  • 논문 참고 (https://arxiv.org/abs/1902.03393)

 

• Depth와 Width는 물론 student가 더 작아질 수밖에 없지만, 양상이 유사해야 함 - 레이어 수가 적고 width가 큰 teacher를 레이어 많고 width 적은 student에 distill하면 KD 반응 양상이 달라질 수 있음

 

• Acrhitectural mismatch & learning pattern
  • Transformer teacher에서 CNN/RNN student에 distill을 고려할 경우, feature-based KD는 구조적으로 불안정할 수밖에 없음(중간 feature를 처리하는 layer 구조가 달라졌으니까 데이터에 다르게 반응하는 것)
    • 그럼 어떻게 해야 할까? feature 자체가 아니라 sample 간 relation(distance, angle)을 distill하자. (Relation distillation)
  • 동일한 맥락으로, 설령 logits KD를 수행하더라도 구조적 차이로 인해 teacher의 inductive bias를 온전히 받아들이지 못할 수 있음(예를 들어, CNN teacher는 locality에 강하겠지만, Transformer student에게 globality를 잘 전달하기 어려움) ⇒ KD는 student 구조가 허용하는 정보만 전달된다

 

• NIPS 논문 "What Knowledge Gets Distilled in Knowledge Distillation"(2023)에서 정리한 인사이트
  1. 지역화 지식 (Localization Knowledge)
     • 결과: 증류를 거친 학생은 이미지 분류 시 교사와 유사한 영역에 집중하는 경향을 보인다.
     • 증거: 클래스 활성화 맵(Class Activation Map, CAM)을 사용하여 각 네트워크가 이미지의 어느 부분에 집중하는지 시각화하고, 교사와 학생의 CAM 간 코사인 유사도를 측정했다. 그 결과, 독립적으로 학습된 학생과 비교했을 때, KL과 CRD 방법으로 증류된 학생의 CAM은 교사의 CAM과 훨씬 더 높은 유사도를 보였다. 이는 교사의 클래스 확률을 모방하는 것만으로도 '어디를 봐야 하는지'에 대한 정보가 간접적으로 전달됨을 시사한다.
  2. 적대적 취약성 (Adversarial Vulnerability)
     • 결과: 교사 네트워크를 속이기 위해 특별히 제작된 적대적 이미지는, 독립적으로 학습된 학생보다 지식 증류를 받은 학생을 속일 확률이 더 높다.
     • 증거: 반복적 FGSM(Iterative FGSM) 공격 기법을 사용해 교사 모델을 속이는 적대적 이미지를 생성한 후, 이 이미지들을 학생 모델에 적용하여 '속임률(fooling rate)'을 측정했다. ResNet50→ResNet18, VGG19→VGG11과 같은 CNN→CNN 증류 환경에서, 모든 증류 방법(KL, Hint, CRD)을 거친 학생들은 독립 학생보다 높은 속임률을 기록했다. 이는 학생이 교사의 가중치나 그래디언트에 직접 접근하지 않음에도 불구하고, 교사가 속는 특정한 방식을 상속받는다는 놀라운 결과를 보여준다. 단, 이 효과는 ViT→ResNet18과 같이 아키텍처가 상이한 경우(트랜스포머→CNN)에는 뚜렷하게 나타나지 않았다.
  3. 데이터 변환에 대한 불변성 (Invariance to Data Transformations)
     • 결과: 교사가 특정 데이터 증강(augmentation)을 통해 명시적으로 학습한 불변성(invariance)은, 해당 증강 기법 없이 학습된 학생에게도 이전될 수 있다.
     • 증거: 색상 지터(color jitter) 증강을 통해 '색상 불변성'을 학습한 교사를 준비했다. 학생은 색상 지터 없이 표준적인 증강만으로 학습시켰다. 이후 원본 이미지와 색상이 변경된 이미지에 대한 예측 '일치율(agreement score)'을 평가한 결과, KL 및 CRD로 증류된 학생들은 독립 학생에 비해 월등히 높은 일치율을 보이며 교사의 색상 불변성을 상속받았음을 확인했다. Hint 기법의 경우, 모방하는 특징 맵이 네트워크의 더 깊은 층에 위치할수록 불변성 이전 효과가 강해지는 경향을 보였다.
  4. 보지 못한 도메인에 대한 지식 (Knowledge about Unseen Domains)
     • 결과: 특정 도메인(예: ImageNet)에서만 증류를 수행하더라도, 학생은 이전에 본 적 없는 새로운 도메인(예: 스케치, 실루엣 이미지)에 대해 교사와의 예측 일치도가 높아진다.
     • 증거: ImageNet 데이터로 교사와 학생을 학습 및 증류시킨 후, 스케치, 스타일화 이미지 등 5개의 새로운 도메인 데이터셋에서 교사와 학생의 예측이 얼마나 일치하는지 '합의 점수(consensus score)'를 측정했다. 실험 결과, 대부분의 경우 지식 증류를 통해 독립 학생 대비 합의 점수가 일관되게 상승했다. 이는 증류된 지식이 학습 데이터에만 국한되지 않고 더 넓은 범위로 일반화될 수 있음을 시사한다.