LLM學習筆記-長度外推技術

老张哈哈哈發表於2024-09-19

長度外推為在不需要對模型進行額外訓練的情況下,模型可以處理更長的序列。本篇文章主要介紹目前大模型用到的一些長度外推技術,包括以RoPE為基礎進行位置插值、NTK-aware、動態NTK、NTK-by-parts 和YaRN。關於RoPE,可參見我的上一篇部落格LLM學習筆記-位置編碼篇

位置插值

回想一下Transformer結構中使用RoPe進行query, key, value(下文用q、k、v表示)的計算

xm表示m位置處的詞嵌入, θd是d維的旋轉角度。為了在預測階段擴充套件視窗長度,一個非常樸素的方法就是在預測階段給長文字的位置乘上一個縮放因子1/s,將長文字的位置縮放到訓練階段視窗長度之內。縮放因子1/s一般為訓練階段文字長度L/預測階段文字長度L', 即 1/s=L/L'

優點

  • 實現簡單
  • 相容性強,與當前的位置編碼技術,比如正餘弦編碼、RoPE高度相容

缺點

  • 仍需要進行微調。根據Meta的論文,需要大約1000步長文字資料微調才能得到一個能處理長文字的模型
  • 效果有限。位置插值均等壓縮長文字中的每個位置,擾亂了模型的區域性解析度,鄰近token的位置關係被影響,導致模型無法捕捉長文字中相鄰token的關係,進而影響最終效果

NTK-aware RoPE

根據NTK理論,如果輸入維度較低且相應的嵌入缺乏高頻分量,深度神經網路將難以學習高頻資訊。位置插值利用縮放因子s對RoPE embedding的所有維度進行均等縮放會損害embedding中的高頻資訊。NTK-aware RoPE的提出者認為位置插值效果有限的原因就在於此。NTK-aware RoPE對所有維度採用非均勻縮放,對於高頻(低維)部分縮放程度小一點,對於低頻(高維)部分縮放程度大一點。總結來說就是‘高頻外推,低頻內插’,具體透過對θ中的base,10000進行縮放來完成。

原來Description, 現在改為Description。為實現低頻內插,則需當i = d/2-1時,有Description,解得Description。 將該結果帶入到表示式中可得Description。 當i為0時,最右邊式子的第二項為1,不會進行縮放;當i為d/2-1時,最右邊式子的第二項為1/s,從而實現了高頻外推,低頻內插。

優點

  • 實現免訓練外推
  • 實現簡單
  • 相容性強,與當前的位置編碼技術,比如正餘弦編碼、RoPE高度相容

缺點

  • 外推長度有限

NTK-by-parts

NTK-by-parts在NTK-aware基礎上對embedding的不同維度進一步的區別對待, 不對高頻做任何改動,只內插低頻部分。 NTK-by-parts引入了波長的概念Description,[關於波長,可參見我的上一篇部落格LLM學習筆記-位置編碼篇]。對於RoPE embedding而言,給定一個視窗長度L,embedding中某些維度的波長要大於L,這意味著該維度可能訓練的不夠充分,關於該部分蘇劍林部落格裡解釋的很好,我參考他的部落格內容簡單解釋一下。

加入RoPE之後的qk內積可以用複數表示為: Description
其中Description表示為單位圓上的一個點,當m-n增大時,該點開始旋轉,轉速與θ有關。θ越大轉速越快,因此m-n在增大到L-1的過程中該點會轉很多圈,即該單位圓上的所有點都被訓練到。θ越小轉速越慢,m-n從0增大到L-1時,可能一個單位圓還沒有轉完只轉了部分弧度,該單位圓上的點訓練不充分,因此當m-n超出視窗大小時,該點可能會落到沒有經過訓練的點上,從而有無法預估的表現,這時就需要用內插法把該點壓縮回訓練的弧裡。而該點在單位圓上轉的圈數與維度對應的波長有關 。 Ti為週期,計算方法與波長一樣。

NTK-by-parts的做法就是將視窗長度L與維度對應的波長進行比較,如果波長遠小於視窗長度,說明該維度對應的點轉了很多圈,不用進行改動;如果波長大於等於視窗長度,則需要進行內插;如果波長與視窗長度的關係在前兩者之間,則兼而有之。具體操作上,NTK-by-parts引入了兩個超參α和β,對於某個維度而言,如果該維度對應的轉數r(d) > β,則無需對改動,如果該維度對應的轉數r(d) < α,則需要進行內插,實驗中α=1, β=32。定義斜坡函式γ如下
Description
對頻率進行如下縮放
Description

Dynamic NTK

Dynamic NTK主要提出了動態的縮放因子。對於自迴歸模型而言,其所處理的序列是不斷變化的,每一步比前一步增加一個token。 在這種情況下有兩種插值的方法

  1. 在整個推理階段,使用固定的縮放因子s=L'/L, 其中L'是預測階段文字長度
  2. 在每輪前向傳播中,位置編碼更新縮放因子s = max(1, l'/L), 其中l'是當前序列的長度

方案1的問題在於,當序列長度小於L'時,模型效能可能會有一點小折扣;而當序列長度大於L'時,模型效能則會突然下降,而使用方案2可以使模型效能下降的更少。論文中把方案二這種在推理階段不斷調整縮放因子的方式稱為Dynamic 縮放,再結合NTK感知方法,論文稱之為‘動態NTK’插值

YaRN

Yarn實際上是將attention scale 和 NTK-by-parts結合到一起。具體來說,是在attention 計算的時候引入了一個溫度因子t,


t的值是根據實驗得出, 實驗結果顯示引入這個縮放因子後,ppl有所下降。

優點

  • YaRN可以直接和修改注意力機制的庫相容,比如Flash attention
  • 效果拔群,在不到原始訓練資料0.1%進行微調,YaRN在上下文視窗的擴充套件上優於其他測試的方法

總結

本篇部落格主要總結了基於RoPE的長度外推技術,主要採取縮放的方法,1. 按token位置進行縮放,例如位置插值,2. 按embedding維度進行縮放,例如NTK-aware, NTK-by-parts, Dynamic-NTK等,採用不同的方法對embeeding的維度進行非均勻縮放,保證高頻不縮放或者縮放程度小一點,低頻進行縮放

Ref

YaRN: Efficient Context Window Extension of Large Language Models
Transformer升級之路:16、“覆盤”長度外推技術
大模型長度擴充套件綜述:從直接外推ALiBi、插值PI、NTK-aware插值(對此介紹最詳)、YaRN到S2-Attention

相關文章