[KD] Distilling the Knowledge in a Neural Network 리뷰

es.Seong·2025년 2월 12일
Knowledge distillationreview
0

Model Compression

목록 보기
1/2

논문 제목
Distilling the Knowledge in a Neural Network (NIPS 2014)
URL: https://arxiv.org/abs/1503.02531
인용수 : 23,389회 (25.2.11 기준)

요새 LLM의 주요 경량화 방법이 양자화 (Quantization)와 지식 증류 (Knowledge Distillation, KD)이다. 모델의 성능을 향상시키기 위해서는 파라미터나 연산량이 커지는게 거의 대다수기 때문에 모바일이나 각종 장비에 넣기엔 제약사항이 많다. 제일 중요한건 보통 이런 서비스를 사용하는 기업은 시간=돈 이라 경제적 이익도 가져다 줄 수 있는게 큰 것 같다. 그래서 지식 증류 공부도 하고 내 연구에 접목이 가능한지 알아보기 위해 서베이랑 제안 논문부터 차례로 읽어보고 정리해보았다,,

요약

  • ML 알고리즘에서 성능을 향상하는 방법 중 가장 간단한 방법은 동일한 데이터로 여러 개의 모델을 학습한 후 이들의 예측을 평균내는 것이다. (모델 앙상블)
  • 하지만 개별 모델이 큰 신경망이라면 앙상블을 활용하여 예측을 수행하는 것은 computationally expensive이며, 특히 배포가 어렵다는 문제가 있다.
  • 본 논문에서는 acoustic model 성능 개선을 위해 앙상블 모델의 지식을 단일 모델로 distilling할 수 있음을 보였다. 또한, 하나 이상의 full model과 여러 개의 specialist models로 구성된 새로운 유형의 앙상블을 도입했습니다.
  • specialist models은 전체 모델이 헷갈려하는 fine-grained classes를 구별하도록 학습된다.
  • Mixture of Experts(MoE)와 다르게 specialist models은 빠르게 학습할 수 있고, 병렬로 훈련할 수 있다고 함.

Introduction

  1. 대형 모델은 훈련 시에 복잡한 구조를 학습할 수 있도록 크고 정교할 필요가 있다. 하지만 배포 시에는 latency와 computational resources가 작은 모델이 필요하다.
  2. "Distillation(증류)"는 대형 모델의 학습된 지식을 작은 모델로 전이하여, 배포에 적합한 모델을 반드는 방법.
  3. 대형 모델이 정답 외의 오답 클래스에도 확률을 할당하는데, 이러한 상대적 확률은 모델의 일반화 능력을 나타내는 중요한 정보다.

    보통 모델이 예측할 때 단순하게 Target Label을 출력하는게 아니라, 확률기반함수(Softmax 등)를 사용하여 각 Label에 대한 확률을 할당.
    ex)
    고양이 : 90%, 강아지 5%, 호랑이 1%, 햄스터 1%, 비행기 0.3%, 택시 0.2% 등등
    이렇게 된다고 가정해보자.
    이 결과를 출력한 학습된 모델은 동물의 feature를 잘 학습했다고할 수 있다.
    만약 비행기나 택시의 확률이 더 높았다면 모델의 구분 능력이 많이 떨어진다고 할 수 있다.
    따라서 오답 클래스들 간의 확률 차이(상대적 확률)를 보면 모델이 어떤 특징을 학습했는지를 알 수 있고, 이것이 모델의 일반화 능력과 직결된다.

  1. soft targets을 활용하면 작은 모델이 대형 모델의 일반화 방식을 모방하여 학습할 수 있으며, 더 적은 데이터로도 효과적인 학습이 가능하다.

  2. 높은 temperature(TT)값을 사용한 소프트맥스 출력을 사용하여 확률분포를 soft하게 만들면, 작은 모델의 distillation이 효과적으로 이루어 진다고함. (2.Distillation에서 설명)

Distillation

신경망은 일반적으로 sotfmax 출력층을 사용하여 각 클래스의 logit ziz_i를 확률 qiq_i로 변환한다.

qi=exp⁡(zi/T)∑jexp⁡(zj/T)q_i = \frac{\exp(z_i / T)}{\sum_j \exp(z_j / T)}

여기서 TT는 temperature라고 하며, 일반적으로 1로 설정된다. TT값을 높이면 클래스 간 확률 분포가 더욱 soft해진다.
distillation의 가장 심플한 형태에서는 지식을 전이하기 위해 transfer set에서 생성된 각 샘플에 대해 높은 T값으로 설정한 softmax를 적용한 soft targets을 사용하여 작은 모델을 학습시킨다. 훈련 과정에서도 동일한 high temperature를 사용하지만 학습된 후에는 일반적인 T=1을 사용한다.

transfer set의 정답 라벨을 알고 있는 경우 이 방법은 더욱 개선될 수 있다고 한다. 정답 라벨을 활용하여 soft target을 수정하는 것이지만, 본 논문에서는 두 개의 다른 objective functions를 단순히 weighted average해서 사용하는 것이 더 효과적이라고 한다.

첫 번째 objective function은 soft target과의 cross entropy이며, CE는 대형 모델에서 soft target을 생성할 때 사용한 것과 동일한 T값을 사용하여 계산한다.

두 번째 objective function은 hard targets과의 cross entropy이며, 이 CE는 작은 모델의 T=1 softmax일 때 계산된다.

본 논문에서는 두 번째 objective function에서 상대적으로 낮은 가중치를 부여할 때 가장 좋은 결과를 얻었다고 한다. soft target에서 생성된 그래디언트의 크기가 1/T21/T^2 비율로 축소되기 때문에, soft targe과 hard target을 함께 사용하는 경우 그래디언트를 T2T^2배 스케일링 해야함.

이를 통해 distillation 과정에서 사용되는 T가 변경되더라도 hard target과 soft target의 상대적 기여도가 거의 변하지 않는다고 함.

Matching logits is a special case of distillation

transfer set의 각 케이스는 각 distilled model의 logit ziz_i에 대한 cross-entropy gradient에 기여함.
만약 복잡한 모델이 logit viv_i를 생성하고, 이로부터 soft target 확률 pip_i를 만든다고 가정한다면, T에서 훈련될 때 기울기는 아래 수식으로 정의된다.

∂C∂zi=1T(qi−pi)=1T(ezi/T∑jezj/T−evi/T∑jevj/T)\frac{\partial C}{\partial z_i} = \frac{1}{T} \left( q_i - p_i \right) = \frac{1}{T} \left( \frac{e^{z_i/T}}{\sum_j e^{z_j/T}} - \frac{e^{v_i/T}}{\sum_j e^{v_j/T}} \right)

만약 T가 logit 값들에 비해 매우 크다면, 다음과 같은 근사식을 사용할 수 있다.

∂C∂zi≈1T(1+zi/TN+∑jzj/T−1+vi/TN+∑jvj/T)\frac{\partial C}{\partial z_i} \approx \frac{1}{T} \left( \frac{1 + z_i/T}{N + \sum_j z_j/T} - \frac{1 + v_i/T}{N + \sum_j v_j/T} \right)

각 transfer case에 대해 logit의 평균이 0이 되도록 조정한다고 가정한다면 ∑jzj=∑jvj=0\sum_j z_j = \sum_j v_j = 0
다음과 같이 단순화 된다.

∂C∂zi≈1NT2(zi−vi)\frac{\partial C}{\partial z_i} \approx \frac{1}{N T^2} (z_i - v_i)

정리하면 T값이 매우 높다면 distillation이 다음과 같은 최소화 문제와 동일하다고 함.

12(zi−vi)2\frac{1}{2} (z_i - v_i)^2

각각의 transfer case에 대해 logits가 개별적으로 zero-mean으로 된 경우에 해당.
T가 낮아질수록, 증류 과정은 평균보다 훨씬 작은 egative logits에 대한 매칭을 덜 집중하게 된다.
이러한 logits들은 복잡한 모델을 사용할 떄 사용된 cost function에 의해 거의 제약을 받지 않아서 매우 noisy할 가능성이 있다고 한다. 따라서 이를 무시하는 것이 잠재적으로는 유리할 수 있다.
하지만 negative logits는 복잡한 모델이 학습한 지식에 대한 유용한 정보를 포함하고 있을 수 있다. 이를 판단하기 위해서는 실험적으로 검증해야하는 문제라고 저자는 언급했다.

Distilled model이 너무 작아서 복잡한 모델이 학습한 모든 지식을 캡쳐할 수 없는 경우, , intermediate temperatures가 가장 좋은 성능을 보인다는 것을 보여주고, 이는 large negative logits을 무시하는 것이 실제로 도움이 될 수 있다고 함.

이 이후의 섹션들은 실험에 대한 디테일 및 결과들이기 때문에 간단히 요약 정리만 하였습니다.

Preliminary experiments on MNIST

distillation work가 잘되는지 보기 위해서 두 개의 은닉층을 가진 하나의 대형 신경망을 1200개의 ReLU 활성화 유닛을 사용하여 훈련했으며, 60,000개의 훈련 샘플을 모두 사용했다고 한다.이 신경망은 드롭아웃과 가중치 제약을 사용하여 강하게 정규화되었다고 한다.
이 신경망은 67개의 테스트 오류를 기록했고, 두 개의 은닉층과 800개의 ReLU 유닛을 가지며 정규화를 적용하지 않은 작은 신경망은 146개의 오류를 기록했으나, distilled model이 큰 신경망에서 생성된 soft target을 T=20에서 매칭하는 추가 작업을 통해 정규화되었을 때, 74개의 테스트 오류를 기록했다고 한다.
이는 soft target이 distilled model로 많은 지식을 전달할 수 있으며, transfer set에 어떠한 변환된 데이터가 포함되지 않더라도 일반화 지식이 잘 전달될 수 있음을 보였다.

Distilled net의 두 은닉층 각각에 300개 이상의 유닛이 포함된 경우, T>=8이면 결과가 거의 유사했지만, 각 층의 유닛수를 30개로 극단적으로 줄인다면 2.5<=T<=4 사이인 경우가 그 외의 T 값에서 보다 훨씬 좋은 성능을 보였다.

이 후 내용은 transfer set에 특정 클래스(논문에서는 3)을 포함하지 않았을 때 distilled model이 어떻게 일반화 하는지에 대해서 설명했다.
핵심 요약하면 다음과 같다.

숫자 3을 전이 데이터셋에서 완전히 제거했음에도 불구하고, 증류된 모델은 테스트에서 3을 일부 인식할 수 있었다.
-> distillation는 훈련 중 특정 클래스를 본 적이 없어도 해당 클래스를 일부 인식할 수 있도록 학습할 수 있다.

Experiments on speech recognition

본 논문에서는 음성 인식에서 DNN 기반 음향 모델의 앙상블 효과를 실험했으며, 지식 증류(distillation)를 통해 단일 모델이 앙상블 모델의 성능을 따라잡을 수 있는지 실험했다고 한다.

8개 은닉층(각 2560개 유닛), 최종 14,000개 HMM 타겟을 갖는 소프트맥스 층을 사용한 모델을 학습했으며, 2000시간의 영어 음성 데이터(7억 개 훈련 샘플)로 학습한 결과, 58.9%의 프레임 정확도 및 10.9%의 WER을 달성했다고 한다.
증류된 모델(distilled model)이 단일 모델보다 더 나은 성능을 보였으며 앙상블 모델과 유사한 지표를 보였다.

Results

10개의 개별 모델을 동일한 아키텍처와 훈련 절차로 학습했으며, 서로 다른 초기화 값을 사용하여 충분한 다양성을 확보한 결과, 개별 모델보다 앙상블 모델의 성능이 월등히 뛰어남.
Distillation 과정에서 T={1,2,5,10}을 실험해서 최적의 결과를 나타냈다고 한다.
앙상블 모델은 WER(Word Error Rate)에서 추가적인 개선을 보였지만, 이 WER 향상 역시 증류된 모델로 효과적으로 전이된 것을 알 수 있다.

Training ensembles of specialists on very big datasets

모델의 앙상블을 학습하는 것은 parallel computation의 이점을 활용하는 매우 간단한 방법이며, 일반적으로 앙상블 모델이 테스트 시 과도한 연산량을 요구한다는 반론은 distillation을 사용 선에서 처리가 가능하다고 한다.
그러나 앙상블에 대한 또 다른 중요한 반론이 존재하는데 각 개별 모델이 large neural networks이고, 데이터셋이 매우 크다면, 훈련 시 요구되는 computation이 지나치게 많아질 수 있으며, 이는 병렬화가 용이하더라도 여전히 문제로 남을 수 있다.

그래서 논문에서는 specialist model이 헷갈리기 쉬운 부분을 집중적으로 학습하면 앙상블하는데 필요한 전체 연산량을 줄일 수 있다고 한다. 하지만 specialist model의 문제는 fine-grained distinctions에 집중하면 과적합이 너무 쉽게 발생한다는 것이었는데, 이건 soft target으로 방지했다고 한다.

The JFT dataset

JFT는 100 million 개의 라벨링된 이미지와 15,000개의 라벨을 포함하는 Google 내부 데이터셋이라고 한다.

찾아보니 데이터 크기가 현재는 더 증가된 것 같다. 300M, 3B 두 버전이 대표적이다.

https://paperswithcode.com/dataset/jft-300m

Specialist Models

클래스의 개수가 매우 많을 경우, cumbersome model이 모든 데이터를 학습한 generalist model 하나와 여러 개의 specialist models로 구성된 앙상블 형태로 설계하는 것이 합리적이다.
각 전문가 모델은 서로 혼동되기 쉬운 특정 클래스 그룹에 대한 데이터로 집중적으로 학습됨.
이렇게하면 전문가 모델의 softmax 출력 크기를 크게 줄일 수 있으며, 중요치 않은 모든 클래스를 dustbin class로 통합해서 단순화할 수 있다고 한다.

Assigning classes to specialists

전문가 모델이 담당할 클래스를 정하기 위해, 네트워크의 예측 패턴을 기반으로 클러스터링을 수행했다.
정확한 정답 라벨을 필요로 하지 않으며, on-line K-means를 활용해 효과적으로 클래스 그룹을 만들었다고 한다.
다른 클러스터링 방법을 사용해도 유사한 결과를 얻었다함.

Performing inference with ensembles of specialists

specialist models이 증류될 때 어떤 일이 발생하는지 조사하기 전에, 우리는 전문가 모델을 포함하는 앙상블이 어떻게 성능을 발휘하는지 확인하고자 했다.

플로우는
1. 일반화 모델이 먼저 확률이 높은 클래스를 선택하고, 해당 클래스를 담당하는 전문가 모델들이 추가로 활용된다.
2. 최종적으로 전문가 모델과 일반화 모델의 확률 분포가 KL 발산을 최소화하도록 결합된다.

3. 이 접근법을 통해 전문가 모델의 장점을 활용하면서도, 일반화 모델이 커버할 수 없는 부분을 보완한다함.

Results

specialist models은 baseline full network로부터 빠르게 학습된다함. JFT의 경우 몇 주가 걸렸지만 specialist models은 며칠만에 끝났다고 한다.
전문가 모델을 추가하면 전체 성능이 향상되며, 특정 클래스를 많이 담당할수록 정확도가 더 높아진다. 그리고 병렬 학습이 가능하며, 빠르게 학습할 수 있어 실용적이다.

Soft Targets as Regularizers

소프트 타겟은 하드 타겟보다 훨씬 더 많은 정보(패턴, 확률 분포)를 포함하며, 모델이 더 일반화(generalization)되도록 돕는다고 한다.
적은 양의 데이터로도 소프트 타겟을 사용하면 효과적인 학습이 가능하며, 과적합을 줄일 수 있으며, 소프트 타겟을 활용하면 데이터의 규칙성을 유지하면서도 새로운 모델로 쉽게 지식을 전이할 수 있다.

Using soft targets to prevent specialists from overfitting

  • 전문가 모델이 특정 클래스에 과적합되지 않도록, 비전문가 클래스에 대한 소프트 타겟을 활용하면 효과적이다.
  • 일반화 모델의 가중치로 초기화한 후, hard targets과 함께 소프트 타겟을 사용하면 과적합을 방지하면서 지식을 유지할 수 있다.
  • 위 표처럼 데이터 3%만 사용한 실험에서, soft targets을 활용하면 적은 데이터를 사용했음에도 과적합을 줄일 수 있음을 확인함.

Relationship to Mixtures of Experts

데이터의 subsets에서 훈련된 specialists의 방식은 Mixture of Experts(MoE)와 유사한 점이 있다. MoE에서는 gating network가 각 샘플을 어느 전문가 모델에 할당할지를 결정하는 확률을 계산.
전문가 모델들이 자신에게 할당된 샘플을 학습하는 동시에, 게이팅 네트워크는 각 샘플을 어떤 전문가에게 할당할지를 학습하는데, 이 과정에서 각 샘플에 대해 전문가들의 discriminative performance을 비교하여 가장 적합한 전문가를 선택하는 방식을 사용함.

전문가 모델들의 판별 성능을 활용하여 학습된 할당 방식을 결정하는 것은 단순히 입력 벡터를 클러스터링하여 전문가를 지정하는 것보다 훨씬 우수한 접근 방식이지만, 이 방법은 훈련을 병렬화하기 어렵게 만드는 문제가 있다.

  1. 각 전문가 모델의 weighted training set은 다른 모든 전문가 모델들과의 관계에 따라 지속적으로 변경됨.
  2. 게이팅 네트워크는 동일한 샘플에 대해 여러 전문가들의 성능을 비교하여, 전문가 할당 확률을 조정하는 방법을 학습해야 함.

두 가지 이유 때문에 MoE는 가장 유용할 가능성이 높은 대규모 데이터셋에서 잘 사용되지 않고, Multiple specialists 학습은 MoE보다 훨씬 쉽게 병렬화할 수 있음.

Discussion

  1. distilling works는 앙상블 모델이나 대형 모델에서 작은 모델로 지식을 효과적으로 전이할 수 있다.
  2. MNIST에서는 일부 클래스가 전이 데이터셋에 없어도 증류가 효과적으로 작동했다.
  3. 대형 음향 모델에서도 앙상블의 성능을 단일 모델로 증류하여 유지할 수 있었다.
  4. 전문가 모델을 활용하면 대형 신경망의 성능을 크게 향상시킬 수 있다.
profile
es.Seong
Graduate student at Pusan National University, majoring in Artificial Intelligence
다음 포스트

[KD] Knowledge Distillation: A Survey 정리 & 번역

0개의 댓글

관련 채용 정보