KD에 살짝의 변화를 줘보자!! 😜 - Knowledge Distillation of Large Language Models 논문 리뷰

The overview of this paper

 이전의 KD는 주로 black-box model API를 모방하기 위해 white-box 분류 모델 또는 small model을 학습시키는데 적용된다. white-box 생성 LLM으로부터 어떻게 효과적으로 distill 하는지는 아직 under-explore 되어 있다.

 

 이 논문에서는 forward KLD를 reverse KLD로 대체함으로써 생성적 larger LM으로부터 smaller LM을 distill 하는 MiniLLM을 소개하였다. 이것은 student model이 teacher 분포의 low-probability 영역을 과도하게 평가하는 것으로부터 모델을 보호하기 때문에 생성적 LM에 더욱 적합한 LM이다. MiniLLM은 전반적으로 높은 퀄리티, 낮은 bias 노출, 더 낮은 calibration, 더 좋은 long-text 생성 성능을 가지는 더욱 정확한 응답을 생성해 낸다.

 

 

Table of Contents

1. Introduction

2. Methods

3. Experiments

 

 

1. Introduction

 KD는 많은 양의 computing 자원에 대한 필요를 감소시킨다. 이 KD에는 다음의 2가지 종류가 있다.

 

1. black-box KD: teacher 예측만 사용 가능
2. white-box KD: teacher 파라미터만 사용 가능

 

 student model은 white-box KD로부터 더 많은 것을 배울 수 있지만, 지금까지는 white-box KD w/ 생성 LLM이 별로 탐구되지 않았다.

 

 이 논문에서는 LLM의 white-box KD를 조사하였다. 논문에서는 기존의 KLD objective는 forward KLD가 사용되었지만, 이는 분류 task에서는 $p(y|x)$와 $q_{\theta}(y|x)$가 적은 수의 모드를 가지기 때문에 문제가 없었다. 하지만 open text generation task는 output space가 더욱 복잡하고 $p(y|x)$는 $q_{\theta}(y|x)$가 표현할 수 있는 모드보다 더 많은 모드를 가진다. 따라서 forward KLD를 최소화하는 것은 합당하지 않다.

 

 이러한 문제를 완화하기 위해 reverse KLD를 최소화하는 것을 제안하였다. $KL[q_{\theta}||p]$를 최소화하는 것은 $q_{\theta}$가 $p$의 주요 모드를 검색하게 해주고, $p$의 비어 있는 공간에는 낮은 확률값을 지정해준다. 또한 $min_{\theta} KL[q_{\theta}||p]$를 최적화하기 위해 Policy Gradient를 사용해서 objective의 기울기를 구하였다. 논문에서는 모델을 학습시키는데 높은 분산, reward 해킹, 생성 길이 bias 등의 문제를 겪는 것을 발견하였고, 이를 다음과 같이 해결하였다.

 

1. 분산을 줄이기 위해 sigle-step 정규화 진행
2. reward 해킹을 완화하기 위해 teacher-mixed 샘플링 진행
3. 길이 bias를 제거하기 위해 length 정규화 진행

 

그림 1. 평가 세트에서 평균 GPT-4 피드백 스코어 측면에서 seqKD와 MiniLLM의 비교

 

2. Methods

 논문에서는 KD를 teacher 모델 분포와 student model 분포 간의 차이를 최소화하는 최적화 문제로 고안하였다. 기존에 사용되던 forward KLD $KL[p||q_{\theta}]$는 $q_{\theta}$가 $p$의 모든 모드를 커버할 정도로 충분히 표현력이 있을 때 language generation task에서 $p$의 비어 있는 영역을 과대 평가하는 모습을 보였다.

 

2-1. MiniLLM: Knowledge Distillation with Reverse KLD

 

 논문에서는 student & teacher 분포 간의 reverse KLD를 최소화하는 것을 MiniLLM의 learning objective로 삼았다.

 

 

 $q_{\theta}$는 $p$의 거대한 모드에 대해 높은 확률을 지정하고, 작은 것들에 대해서는 무시하는 모습을 보여준다. 그림 2처럼 forward KLD 최소화는 $p$의 영확률 위치에 큰 확률을 부여하지만, reverse KLD는 $p$의 주요 모드에 집중하는 것을 알 수 있다.

 

그림 2. forward KLD와 reverse KLD의 가우시안 분포

 

 논문에서는 LLM을 reverse KLD를 최소화시키는 KD method를 그림 3처럼 MiniLLM으로 이름 지었다. seqKD와 달리 MiniLLM은 teacher 분포 $p$로부터 샘플링된 모든 $y$에 대해 $q_{\theta}$를 맞추도록 강요하지 않는다. 그 대신에, student model이 자신의 능력 내에서 teacher가 선호하는 샘플을 생성하도록 권장하며, 이는 성취 가능성이 더 높다.

 

그림 3. seqKD(왼쪽)와 MiniLLM(오른쪽)의 비교. seqKD는 student가 모든 teahcer-generated 샘플을 기억하도록 강요하는 반면, MiniLLM은 자기 자신의 생성과 teacher의 피드백을 사용해서 student가 개선될 수 있도록 허락해 줌

 

2-2. Optimization with Policy Gradient

 

Gradient Derivation.  논문에서는 objective function의 기울기가 Policy Gradient 정리를 사용해서 다음과 같이 얻어질 수 있다고 말한다.

 

 

 여기서 $T = |\mathbf{y}|$이고 $R_{t} = \sum_{t'=t}^{T} log \frac{p(y_{t'}|\mathbf{y}_{<t'}, x)}{q_{\theta}(y_{t'}|\mathbf{y}_{<t'}, x)}$는 각 스텝 생성의 퀄리티를 측정하는 $r_{t'} = log \frac{p(y_{t'}|\mathbf{y}_{<t'}, x)}{q_{\theta}(y_{t'}|\mathbf{y}_{<t'}, x)}$의 축적이다. 직관적으로 teacher 분포 하에서 high probability를 가지기 위해 $p(y_{t'}|\mathbf{y}_{<t'}, x)$가 증가되길 원하지만, 동시에 $q_{\theta}(y_{t'}|\mathbf{y}_{<t'}, x)$를 낮춤으로써 다양성을 유지하기를 원한다. 하지만, 아직 몇 가지 문제를 가지고 있는데, 논문에서는 이 문제를 완화하기 위한 3가지 전략을 제시하였다.

 

Single-Step Regularization.  single-step 생성 퀄리티 $r_{t}$는 분산을 학습하는데 중요하다. $r_{t}$에 더 잘 집중하기 위해 $R_{t}$로부터 $r_{t}$를 분할하기 위해 $\bigtriangledown \mathbf{J}(\theta)$를 재작성하고 $\mathbb{E}_{y_{t} ~ q_{\theta}(t)}[r_{t}]$의 기울기를 정규화로 직접적으로 계산하였다.

 

 

 이러한 정규화는 학습 중에 분산을 줄여주고 수렴을 가속화시켜 주는 single-step 생성 퀄리티의 더욱 정확하고 효율적인 측정을 준다.

 

Teacher-Mixed Sampling.  논문에서는 reward 해킹을 줄이기 위해 teacher & syudent 분포를 각 time step에서 섞었다.

 

 

 여기서 $\alpha$는 하이퍼파라미터로 teacher mix-in의 강도를 조절한다. $\tilde{p}$로부터의 샘플링이 teacher의 도움으로 low-quality 생성을 억제하고, reward 해킹을 완화시켜 준다. 논문에서는 기울기의 unbiased 측정자를 얻기 위해 $(\bigtriangledown \mathbf{J})_{Main}$과 $(\bigtriangledown \mathbf{J})_{Reg}$를 importance smapling과 함께 재작성하였다. 

 

 

 여기서 $w_{t}$는 importance weight인데, $w_{t}$를 사용하는 것은 하이퍼 파라미터에 대해 매우 민감하고, 느리게 수렴하게 된다. 그래서 측정자의 분산을 낮추기 위해 $w_{t} \approx \frac{q_{\theta}(y_{t'}|\mathbf{y}_{<t'}, x)}{\tilde_{p}(y_{t'}|\mathbf{y}_{<t'}, x)}$로 근사하였다.

 

Length Normalization.  long sequence는 작은 $R_{t+1}$을 가지는 경향이 있는데, 이는 모델이 짧은 응답을 생성하도록 북돋아준다. 그래서 논문에서는 length normalization을 $R_{t+1}$에 추가하였다.

 

 

In Summary.  앞서 나열된 전략을 묶어서, 다음과 같이 최종 최적화 기울기를 얻을 수 있었다.

 

 

2-3. Training Algorithm

 

 MiniLLM의 training algorithm은 알고리즘 1에 나타나 있다. 논문에서는 가장 낮은 validation loss를 가지는 training data에서 fine-tune 된 checkpoint로부터 student model을 초기화하였다. 그리고 $(\bigtriangledown \mathbf{J})_{Main}$에 PPO clipping 전략을 추가해서 학습 안전성을 개선시켰다. 또한 학습 효율성을 개선시키기 위해 PPO에서 value network와 KL 정규화를 사용하지는 않았다. 그리고 pre-training corpus에 language modeling loss $L_{PT}$를 추가하였다. 

 

알고리즘 1. MiniLLM: Knowledge Distillation of LLMs

 

3. Experiments

3-1. Experimental Setup

 

 거대 모델을 instruction-response 데이터셋 $D$에서 fine-tune 함으로써 teacher $p$를 만들었다. 그다음에 teacher의 지도와 함께 $D$에서 smaller student model을 distill 하기 위해 서로 다른 KD method를 비교하였다.

 

Base Models.  3가지 종류의 모델을 다양한 사이즈를 사용해서 distill 하였다.

 

• OPT(1.3B, 2.7B, 6.7B) - teacher model: OPT-13B
• GPT-2(120M, 340M, 760M) - teacher model: GPT-2-1.5B
• LLaMA(7B) - teacher model: LLaMA-13B

 

Training.  training data로 databricks-dolly-15k를 사용하였다. 그리고 $D_{PT}$에 대해 GPT-2는 OpenWebText를 사용하고 다른 모델에는 RoBERTa training corpus를 사용하였다.

 

Evaluation.  논문에서는 5개의 instruction-following 데이터셋에서 학습된 모델을 평가하였다.

 

• DollyEval
• SelfInst
• VicunaEval
• S-NI
• UnNI

 

 또한 model-generated 응답을 평가하기 위해 2개의 metric을 사용하였다.

 

• ROUGE-L
• GPT-4

 

Baselines.  논문에서는 메인 실험에서 총 3개의 baseline을 고려하였다.

 

• SFT w/o KD: student model을 golden response로 supervise 된 $D$에서 바로 fine-tune
• KD: student model을 teacher-generated data에서 fine-tune
• SeqKD: student model을 teacher-generated data에서 fine-tune

 

3-2. Results

 

 평가 결과는 표 1에 나타나 있다.

 

표 1. 평가 결과 betst score는 볼드체, teacher model을 능가하는 student model은 * 표시함.

 

 이 결과를 통해 논문에서 관찰한 점은 다음과 같다.

 

1. SFT를 forward KLD를 최소화하는 KD & seqKD와 비교함으로써 기존의 KD method가 teacher model로부터 대부분의 경우에 성공적으로 지식을 distill 하는 것을 확인할 수 있었다.
2. MiniLLM과 baseline의 GPT-4 피드백 score를 비교함으로써 논문의 method로 distill 된 모델이 대부분의 경우에 baseline을 능가한다. 이것은 MiniLLM이 전반적으로 높은 성능을 보여준다는 것을 알 수 있었다. 또한 다른 데이터셋보다 DollyEval에서 더 잘 작동하는 모습을 보여줬는데, 이는 좋은 OOD 일반화를 가진다는 것을 의미한다.
3. ROUGE-L score는 MiniLLM이 ground-truth 응답과 가장 많은 오버랩을 가지는 가장 정확한 응답을 생성한다는 것을 보여준다.
4. MiniLLM의 개선은 base model의 사이즈가 120M으로부터 13B로 다양해짐에도 일관되었다. 이러한 경향은 그림 1에 나타나 있고, 이것은 훌륭한 scalability와 generalization을 보여준다.

 

3-3. Analysis

 

Exposure Bias.  forward KLD는 exposure bias를 겪는다. MiniLLM에서는 학습 중에 student model로부터 샘플을 수집하는데 이것이 training과 evaluation 간의 미스매치를 완화시켜 준다.

 

Calibration.  RL-trained model은 안 좋은 calibration을 보여준다. 그래서 MiniLLM과 KD baseline들의 calibration을 테스트하였다. 그 결과, KD와 seqKD로 학습된 모델은 teacher model보다 안 좋은 calibration을 보여준다. 이는 forward KLD를 최소화하기 때문에 student와 teacher 모델 간의 상당한 분포 차이를 이끈다. 반면에 MiniLLM은 타깃 분포의 주된 부분을 정확하게 학습하는데 주목해서, 이것이 student & teacher 간의 ECE score 갭을 좁힌다.

 

Performance on DIfferent Response Length.  서로 다른 범위의 golden response 응답 길이가 주어졌을 때 모델의 성능에 대해 연구하였다. 그림 4에서는 ground-truth 응답의 길이로 분할된 3개의 S-NI 하위 집합에서 SFT 모델에 대한 다양한 KD 모델의 ROUGE-L socre를 보여준다. 논문에서는 짧은 응답($\leq 5$)을 예측하는 prompt에서는 낮은 score를 달성하는 모습을 발견하였다. 이는 training set의 대부분은 긴 문장으로 이루어져 있기 때문이다. 그리고 output space가 비교적 작아서 reverse & forward KLD는 비슷한 성능을 가진다. longer response($\geq 6$)의 prompt에 대해 MiniLLM은 기존 KD approach보다 장점을 가졌다.

 

그림 4.&nbsp;distilled model의 ROUGE-L score

 

3-4. Ablations

 

Effect of Optimization Strategies.  최적화를 안정화시키고 가속화하기 위해 제안된 3개의 전략에 대해 ablation을 진행하였다. Teacher-Mixed 샘플링과 Length Normalization은 training을 안정화시키기 위해 필수적이라는 것을 발견하였다. 비록 이러한 전략 없이도 reverse KD는 감소하지만, 그럴 경우에 모델은 reward 해킹을 겪을 뿐만 아니라 표 2처럼 낮은 생성 성능을 이끈다. 그림 5로부터 Single-Step 정규화는 training 프로세스의 분산을 효과적으로 줄여준다는 것을 발견하였다.

 

표 2.&nbsp;서로 다른 MiniLLM 전략이 적용될 때 validation set와 test set에서의 성능

 

그림 5.&nbsp;서로 다른 최적화 전략이 적용될 때 MiniLLM training 중에 student와 teacher 간의 forward KLD

 

Effect of Teacher-Mix-in Strength $\alpha$.  그림 6에서 서로 다른 teacher-mix-in strength $\alpha$에서 MiniLLM의 성능을 비교하였다. 그 결과 $\alpha = 0.2$가 가장 적절하였다.

 

그림 6.&nbsp;알파 값에 따른 ROUGE-L score 변화

 

Effect of Adding Pre-training Loss.  표 3에서 pre-training loss를 추가하는 것의 효과를 연구하였다. PT loss를 추가하는 것이 instruction-following task에서의 성능을 거의 변화 없이 유지하면서 NLP task에서의 능력을 보존하는데 도움을 줬다.

 

표 3. pre-training loss를 추가하는 것의 효과