Llama 2 - Open Foundation and Fine-Tuned Chat Models

本文為 “Llama 2: Open Foundation and Fine-Tuned Chat Models” (2023.07) 的論文摘要

論文全文參考

Open Source

Description

Goal & Contributions

• NLP Foundation model (commercial use)
  • 與 LLaMA 1 的差異為
    1. 訓練資料的增加
    2. context window 擴增 (2048 → 4096)
    3. 依照 ChatGPT 的完整訓練 (pre-train → SFT → RLHF)
  • 開源 pre-train 版本為 7B, 13B 70B, RLHF版本會多個 Chat 的 tag

LLaMA 2-Chat training pipeline

Methodology

Pre-train

• 使用了與 LLaMA 1 類似的架構, 包含 RMSNorm, SwiGLU, RoPE
  • RMSNorm (root mean square): 改進 LayerNorm 公式, 刪除對結果不顯著的中心化過程 (- mean), 提升 40% 計算速度
  • SwiGLU: GLU 的變形, GLU 與其他 activation function 差異為有一個類 LSTM 的 gate control, 可以更好的控制做資訊篩選
  • RoPE (rotary positional embeddings): 蘇神提出的, 在形式上是絕對位置編碼, 但在計算 attention 時可以變為相對位置編碼, 細節可以參考這裡
  

  LLaMA 2 architecture
• 使用了與 LLaMA 1 不同的架構, 包含 擴展 context length, GQA (grouped-query attention)

  • 擴展 context length: 在 pre-train 階段將文本訓練長度從 2k 擴增為 4k, 與上一篇的 FOT 架構不同, 屬於暴力 train 的方法

  • GQA: 在近期很多大模型均有用到這種共享類的 attention layer (e.g, PaLM, Falcon), 這種方式好處是可以加速推論速度, 藉由共享不同 attention heads 間的 key 和 value, 達到參數的精簡

    • Falcon 使用的是 multi-query attention, 所有的 query 共享 key 會有最快的推論速度
    

    attention layer comparison

  • Note: LLaMA 2 僅在 34B (尚未 release), 70B 有使用到

• 訓練細節
  • 訓練時使用了 AdamW 的 optimizer
  • 初始學習率設定在 $3*10^{-4}$ 並使用 cosine schedule
  • 權重衰減率 0.1 且有用上 gradient clipping = 1
• Meta 花了很大的功夫在處理資料歧視問題, 並聲稱沒用到 Facebook 中用戶的資訊

Fine-tune (SFT)

• 與 LLaMA 1 最大的差異為多了 SFT + RLHF, LLaMA 1 是一個 pre-train model 僅有文字接龍能力, 無法讀懂人類指令, 因此有許多延伸模型 (e.g., Alpaca, Vicuna)
• 在 Meta 之前的 LIMA 有提到可以用少量的高質量和多元性資料, 就可以使得 SFT model 有很好的表現, 在 LLaMA 2 使用了 27,540 個標註資料
• 如何決定要停止標註?
  • 透過手動地抽樣檢查來確保標註資料的正確性 (抽查 180 筆)
  • SFT 的預測結果可以與人類標註資料相匹配時, 就可以 move on 到 RLHF 的標註上
• 訓練細節
  • 初始學習率設定在 $2*10^{-5}$ 並使用 cosine schedule
  • 訓練文本長度控制在 4,096, 為了可以盡可能滿足這項設定, 在訓練集會將 prompt 和 answer 綁在一起, 中間會用特殊 token 區隔
  • batch size = 64, 訓練了 2 個 epochs

Fine-tune (RLHF)

• Reward Model

  • 做 RLHF 訓練前需要先做一個 reward model, 用來判斷 LLM 輸出的好壞
  • LLaMA 2 在 reward model 採取和 chatGPT 類似的方式, 標註 1 對 1 模型輸出結果的優劣來做為 reward model 的訓練資料
  • 為了有更多樣化的訓練資料會使用不同的 LLM 以及 temperature 參數來輸出結果
  • 標註的 label 會分為 significantly better, better, slightly better, or negligibly better unsure.
  • 評估標準為 helpfulness 和 safety
    • safety 會在分為以下三個 label
      1. the preferred response is safe and the other response is not
      2. both responses are safe, and 3) both responses are unsafe
      3. both responses are unsafe
    • 有趣的是 helpfulness 和 safety 有時候會是 trade-off 關係, 這時候會訓練兩個 reward model
  • Reward Model 的初始模型為 pre-train chat, 這個作法是為了讓 Reward Model 知道 SFT 模型是不是在 pre-train chat 時有訓練過這個知識
  • 訓練細節
    • 經實驗發現, 訓練太久會造成 over-fitting, 因此只訓練 1 個 epoch
    • 最大學習率設定為 $1*10^{-8}$ (7B, 13B, 34B), $5*10^{-6}$ (70B) 並使用 cosine schedule
    • batch size = 1024

• Iterative Fine-Tuning

  • 有了 Reward model 後可以進行大量的 RLHF 訓練 (RLHF v1, RLHF v2, …, RLHF v5)
  • RLHF finee-tune 使用了兩種演算法 PPO (Proximal Policy Optimization), Rejection Sampling fine-tuning
  • 在初期 (v1~v4) 均使用 Rejection Sampling, 後期則是先找出 top Rejection Sampling 在使用 PPO

• Multi-Turn Consistency

  • 使用了 GAtt (ghost attention) 解決多輪對話會遺忘 instruction prompt 的問題, 這個 idea 參考了 Context Distillation 的想法
  • 作法為擷取部分的多輪對話數據, 來產生新的 fine-tune 資料

Conclusion

• 此篇論文主要介紹 LLaMA 2 - chat 的訓練過程, 雖然預測表現尚且不如 ChatGPT, 但已有相同甚至更好的訓練方法和架構, 相信隨著訓練資料的加入會越來越接近
• LLaMA 2 在中文使用上還是非常不又善 (pre-train 資料: 英文占 89.7 %, 中文 0.13%), 論文中提到的原因是因為非英語公開資料取得不易, 想要在中文領域應用需要一定量級的中文 pre-train, 可關注 Chinese LLaMA 2
• 隨著開源 LLM 的普及 training 相關技術沒有太多的未知變數, 可以開始研究如何加速 inference 速度以及與 retrieval 方法作結合