<Distillation, DPO> Zephyr: Direct Distillation of LM Alignment (2023.10)

관심있는 NLP 논문을 읽어보고 간단히 정리했습니다. 

혹시 부족하거나 잘못된 내용이 있다면 댓글 부탁드립니다 🙇‍♂️

---

usechatgpt init success

[H4 Team]
- distilled Direct Preference Optimization (dDPO)를 적용하여 sLLM이 user intent에 잘 align 될 수 있도록 함
- fine-tuning 동안 추가적인 sampling 없이 Mistral-7B 모델을 학습한 Zephyr-7B 모델을 공개

1. Introduction

• capable teacher model의 output을 사용하는 distilled supervised fine-tuning (dSFT) 등의 방법론이 연구
• 그럼에도 불구하고 여전히 smaller LM들은 not "intent aligned"
• MT-Bench, AlpacaEval과 같이 user intent를 제대로 반영하는지 확인하는 벤치마크가 등장

 

---

2. Realted Works

LLaMA 모델의 등장 이후..

• Efficient Fine-Tuning, Longer Prompt Context, Retrieval Augmented Generation (RAG), Quantization
• MPT, RedPajama-INCITE, Falcon, LLaMA-2 등 다양한 sLLM 등장

 

---

3. Contributions

1. AI Feedback 데이터셋으로부터의 Direct Preference Optimization (DPO)을 이용하여 Zephyr-7B 모델 개발
2. MT-Bench에서 LLaMA-2-Chat-70B 모델을 outperform 할 정도의 성능

 

---

4. Method

larger teacher model: π_{T} / student model: π_{θ}

학습은 세 단계로 구성됨

4.1. Distilled Supervised Fine-Tuning (dSFT)

• teacher 모델이 instruction과 repsonse를 생성하고, sutdent 모델은 이에 대해 directly 학습하는 방식
• J개의 seed prompt (x_{1}, x_{2}, ... x_{J})가 주어지면, 이를 바탕으로 iterative self-prompting을 수행
  • teacher 모델로부터 response를 생성: y_{1} ~ π_{T}(· | x_{1}) 이런식으로..
    원래는 위첨자 표기를 사용하여 생성 종료시점까지 response를 반복적으로 생성함
  • 최종적으로 C = {(x_{1}, y_{1}), ... , (x_{J}, y_{J})} 데이터셋이 구축
• objective function

4.2. AI Feedback through Perferences (AIF)

• LLM을 사람의 선호에 align하기 위해 제공하는 Human feedback (HF) 대신 다른 모델로부터 생성된 결과에 대한 teacher model의 AI preference를 사용
• UltraFeedback 방식을 사용
  • 네 개의 other model로부터 생성된 결과를 준비
  • GPT-4 (teacher model)이 네 개의 결과물에 대한 score를 측정
  • 그중에서 score가 가장 높은 것을 y_{w}, 나머지 세 개 중에 무작위로 선택한 것을 y_{l}로 정함
  • 결과적으로 feedback dataset D는 (x, y_{w}, y_{l})의 triplet 형식을 갖추게 됨

4.3. Distilled Direct Preference Optimization (dDPO)

• π_{dSFT}가 y_{l}보다 y_{w}에 대해 더 높은 선호 랭킹을 가질 수 있도록 likelihood를 maximize
• 기존 reward function을 preference model에 직접 plug. 각 AIF triple (x, y_{x}, y_{l})에 대해 반복
  
  1. dSFT 모델로부터 (x, y_{x}) & (x, y_{l})에 대한 확률을 계산
  2. dDPO 모델로부터 (x, y_{x}) & (x, y_{l})에 대한 확률을 계산
  3. 아래 objective function으로 loss를 계산하여 backpropagate. 반복 

 

---

5. Experiments

• Model & Baseline
  • Mistral 7B에 대해 fine-tuning
  • Transformer Reinforcement Learning (TRL) library, DeepSpeed ZeRO-3, FlashAttention-2
• Datsets
  • UltraChat, UltraFeedback
• Benchmarks
  • MT-Bench (multi-turn), AlpacaEval (single-turn), Open LLM Leaderboard
• Results
  • dDPO improves Chat Capabilities
  • dDPO improves Academic Task Performance: ARC, Hella Swag, MMLU, Truthful QA 등에 대해서도 준수한 성능을 보이는 것을 확인

 

---

7. Insights

scalability만 확인할 수 있다면 굉장히 효과적인 방식이 아닐까 싶은 생각이 든다.

기존에는 항상 reward 모델을 따로 학습시키고 이에 대해 model이 출력하는 결과를 바탕으로 optimize하는 방식이었는데, 이 과정을 간소화하면서도 preference를 잘 주입시킨 것 같다는 생각이 든다.

 

다만, 이 방식이 왜 reward 모델을 학습시키고 이를 이용하는 것보다 유효한 것인지에 대한 설명이 부족한 것 같다.

또한 학습 데이터에 대한 과적합이 있음에도 불구하고 벤치마크 성능 테스트에서 가장 좋은 스코어를 달성했다는 것이 잘 이해되지 않는다.

그렇다면 학습 데이터의 분포가 테스트 데이터의 분포와 유사하다는 것 이외에 다른 설명이 가능할까 싶기도 하고..

 

개인적으로는 teacher 모델을 여러 개 쓰는 방법은 어떨까 궁금증이 생긴다.

 

 

 

출처 : https://arxiv.org/abs/2310.16944

Zephyr: Direct Distillation of LM Alignment