💻 🧐

배경

머신러닝 디자인패턴을 읽던 중 4장 distilling 에 관한 기법이 나와 찾아보게 되었습니다.

요약

지금까지으 대규모 머신러닝 시스템은 학습과 배포단계에서 같으느 모델을 사용했는데 이때문에 추론 레벨에서 리소스가 커진다는 단점이 있었습니다. 그래서 거대한 모델로부터 지식을 하나의 작은 모델로 전이하는 방법을 통해 이 제약을 극복하려고 합니다. 이것이 “증류” 의 표현입니다.

기존의 hard label 은 [1,0,0] 처럼 정확한 확률을 알려주었는데, 실제로는 고양이를 닮은 개가 있을수도 있으니 [0.6,0.4,0] 같은 label 도 의미있는 지식일 수 있다는것이 아이디어 입니다.

그래서 정리하자면 거대한 모델에서 산출된 마지막 레이어의 값들을 학습에 활용하는 soft target 으로 활용하자는 것입니다.

Distillation

기존의 신경망은 클래스 분류 태스크를 수행할 때 output layer 에 softmax 를 취해 logit 을 변환합니다.

일반적으로 T 는 1로 세팅되어 있습니다. T를 높게 할수록 클래스 확률값이 soft 하게 출력됩니다.

T 가 커질수록 클래스 값들이 soft 해지는 것을 알 수 있습니다.

논문에서 제안한 “증류” 의 방식은 거대한 모델을 teacher model 로 특수한 목적으로 만들 작은 모델을 Student(distilled model) 이라 합니다.

t를 최대로 하고, 기존의 transfer dataset 을 teacher 모델에 넣고 soft label (1) 을 얻습니다. 그리고 student 모델에 inference 해서 soft prediction(2) 을 얻습니다. 그리고 student model 에서 hard prediction(3)결과를 얻습니다. 그리고 (1) 와 (2) 의 cross entropy , (2) 와 (3) 의 cross entropy 를 가중합 하는 방식이 성능이 좋다고 합니다.

cross entropy gradient 의 는 아래와 같이 나타냅니다. 여기서 Vi 는 큰 모델의 결과값을 나타내고 pi 는 soft label 의 확률값을 의미합니다. 그러니까 qi 와 pi 의 cross entropy 를 계산하는 문제가 logit 간의 차이를 근사하는 문제로 변환시키고 있음을 의미합니다.

증류에서 Teacher Model 과 Student 의 모델 output 차이를 활용해 gradient 계산을 하려는 움직임입니다.

여기서 T가 충분히 크다면 테일러 근사를 활용해 아래와 같이 나타낼 수 있습니다

그리고 logit 의 평균이 0 이라고 가정한다면 0 으로 변환되니

student model output 에 대한 Cross entropy 변화율 즉 gradient 는 nt^1 에 반비례하게 됩니다. 그래서 결론은 “T가 충분히 큰 상황에서 logit 들의 평균이 0 으로 주어졌다면, Distillation 은 1/nt^2 을 최소화 하는 문제가 됩니다.

반대로 T 가 낮다면 gradient 를 최대화 시키니 모델이 너무 작다면 중간정도의 temperature 를 사용하는것이 좋다고 주장합니다

기여도

본 논문의 가장 큰 기여는 다음 세 가지로 정리할 수 있습니다.

1. Soft Target을 활용한 지식 전달 방식의 정식화

기존의 분류 모델은 정답 레이블(one-hot 벡터)만을 학습 신호로 사용하였습니다. 그러나 본 논문은 모델의 출력 확률 분포 전체가 지식이라는 관점을 제시하였습니다.

특히 Teacher 모델이 출력한 softmax 확률 분포에는 다음과 같은 정보가 포함되어 있습니다.

• 클래스 간 유사도
• 모델이 헷갈리는 정도
• 데이터 분포에 대한 암묵적 구조

이를 Student 모델이 학습하도록 만들면, 단순히 정답을 맞추는 것이 아니라 Teacher의 판단 구조 자체를 모방하게 됩니다.

이는 이후 모든 Knowledge Distillation 연구의 출발점이 되었습니다.

2. Temperature 기반 확률 분포 제어 메커니즘 제안

논문은 softmax 함수에 temperature TTT를 도입하여 확률 분포를 조절하는 방식을 제안하였습니다.

• T=1일 경우 일반적인 softmax
• T>1일 경우 분포가 평탄해짐

Temperature를 높이면 모델의 확신(confidence)이 낮아지고, 클래스 간 상대적 관계 정보가 더 잘 드러납니다. 이로 인해 Student 모델은 단순 정답이 아니라 확률 구조 자체를 학습하게 됩니다.

이러한 수학적 근사 분석을 통해, KD가 결국 logit 간 차이를 줄이는 문제로 수렴한다는 것을 보인 점 역시 중요한 기여입니다.

3. 앙상블 모델의 압축 방법 제시

논문은 여러 개의 모델을 앙상블하여 얻은 고성능 Teacher 모델을 단일 Student 모델로 압축할 수 있음을 보였습니다.

이는 다음과 같은 실용적 의미를 가집니다.

• 학습 단계에서는 대규모 모델 사용
• 배포 단계에서는 경량 모델 사용
• 추론 속도 향상 및 리소스 절감

즉, 학습과 배포의 구조적 분리를 가능하게 한 기법이라는 점에서 산업적 의미가 매우 큽니다.

실험결과

MNIST

MNIST 실험에서는 다음과 같은 설정을 사용하였습니다.

• Teacher 모델: 여러 개의 모델을 앙상블한 고성능 네트워크
• Student 모델: 상대적으로 작은 네트워크

결과적으로 Student 모델은 다음과 같은 특징을 보였습니다.

1. 일반적인 hard label 학습 대비 더 낮은 error rate
2. Teacher의 성능에 근접한 정확도 달성
3. 과적합 감소 효과

특히 데이터가 충분하지 않은 상황에서도 soft target을 사용하면 일반화 성능이 개선되는 경향을 보였습니다.

이는 soft target이 일종의 regularizer 역할을 수행한다는 것을 의미합니다.

Speech Recognition

음성 인식 실험에서는 대규모 acoustic model을 Teacher로 사용하였습니다.

Teacher는 매우 복잡하고 큰 네트워크였으며, 직접 배포하기에는 비효율적이었습니다.

Distillation을 적용한 Student 모델은:

• 파라미터 수 감소
• 추론 속도 개선
• 정확도는 Teacher에 근접

특히 soft target 기반 학습이 hard label 기반 학습보다 안정적인 수렴 특성을 보였습니다.

이는 KD가 classification뿐 아니라 sequence 기반 문제에도 적용 가능함을 보여주는 사례입니다.

참교자료 :

[논문 리뷰] Distilling the Knowledge in a Neural Network​