LLM上下文視窗突破200萬!無架構變化+複雜微調,輕鬆擴充8倍

2024.04.29

大型語言模型(LLM)往往會追求更長的「上下文視窗」,但由於微調成本高、長文本稀缺以及新token位置引入的災難值(catastrophic values)等問題,目前模型的上下文視窗大多不超過128k個token

最近,Microsoft Research的研究人員提出了一個新模型LongRoPE,首次將預訓練LLM 的上下文窗口擴展到了2048k個token,在256k的訓練長度下只需要1000個微調步驟即可,同時還能保持原始短上下文視窗的效能。

圖片圖片

論文連結:https://arxiv.org/abs/2402.13753

程式碼連結:https: //github.com/microsoft/LongRoPE

LongRoPE主要包含了三個關鍵創新點:

1. 透過高效搜尋識別並利用了位置插值中的兩種非均勻性,為微調提供了更好的初始化,並在非微調情況下實現了8 倍擴展;

2. 引入了漸進擴展策略,首先微調256k 長度的LLM,然後在微調擴展的LLM上進行第二次位置插值,以實現2048k 上下文視窗;

3. 在8k長度上重新調整LongRoPE以恢復短上下文視窗效能。

在LLaMA2 和Mistral 上對各種任務進行的大量實驗證明了該方法的有效性。

透過LongRoPE 擴充的模型保留了原始架構,只對位置嵌入稍作修改,並且可以重複使用大部分現有的最佳化。

位置內插的不均勻性

Transformer模型需要明確的位置訊息,通常以位置嵌入(position embedding)的形式來表示輸入token的順序。

本文中的位置嵌入表示方法主要來自於RoPE,而對於位置索引為n 的標記,其對應的RoPE 編碼可簡化如下:

圖片圖片

其中,d 是嵌入維度,nθi 是標記在位置n 上的旋轉角度,θi = θ -2i/d 表示旋轉頻率。在RoPE 中,θ 的預設基準值為10000。

受NTK 和 YaRN 的啟發,研究人員注意到這兩個模型可以從非線性嵌入中獲得性能提升,特別是在考慮RoPE 各維度的不同頻率以進行專門的內插法和外推法時。

然而,當前的非線性在很大程度上依賴人為設計的規則。

這也自然引出了兩個問題:

1. 目前的位置內插是否是最佳的?

2. 是否存在尚未探索的非線性?

圖片圖片

為了回答這些問題,研究人員使用進化搜尋(evolution search)為LLaMA2-7B發現更好的非均勻位置插值。搜尋以易錯性為指導,使用PG19驗證集的5個隨機樣本。

透過實證分析,研究人員總結了幾個主要發現。

發現1:RoPE維度表現出很大的不均勻性,目前的位置內插法無法有效處理這些不均勻性;

在公式2 中為每個RoPE 維度搜尋最佳λ。

圖片圖片

研究人員比較了PG19和Proof-pile測試集上使用不同方法的LLaMA2-7B 在不進行微調的情況下的複雜度。

圖片圖片

從結果來看,搜尋到的解決方案有明顯改善,顯示目前的線性(PI,positional interpolation)和非均勻(Dynamic-NTK 和YaRN)插值方法都不是最佳的。

值得注意的是,YaRN 在PG19 上的表現不如PI 和NTK,因為其無法達到非微調LLM 的目標上下文視窗長度。

例如,在8k 上下文大小的情況下,YaRN 的困惑度在7k 後達到峰值。

透過搜索,公式2 中的重標度(rescaled)因子λ變得不均勻,與PI、NTK的公式計算和YaRN的分組計算中的固定標度s有所不同。

在8k和16k上下文視窗中,這些非均勻因子大大提高了LLaMA2的語言建模性能(即複雜度),而無需進行微調,主要是因為由此產生的位置嵌入有效地保留了原始的RoPE,尤其是關鍵維度,從而降低了LLM區分近似token位置的難度。

發現2:輸入序列中初始詞塊的RoPE推論應減少內插;

對於輸入序列中的初始n個token,假設RoPE應該做較少的插值,這是因為會獲得較大的注意力分數,從而對注意力層至關重要,正如在Streaming LLM和LM-Infinite 中觀察到的那樣。

為了驗證這一點,研究人員使用PI和NTK將上下文視窗擴展到8k 和16k,保留前n(0,2, ..., 256)個token,不進行插值。當n=0 時,則恢復到原來的PI 和NTK

圖片圖片

上表中可以觀察到兩個結果:

1. 保留起始token而不進行位置內插確實能提升效能。

2. 最佳起始token數n取決於目標擴展長度。

發現3:在微調和非微調設定中,非均勻位置插值都能有效擴展LLM 上下文視窗。

虽然已经证明,在不进行微调的情况下,搜索到的非均匀位置插值能显著提高8k和16k扩展性能,但更长的扩展需要微调。

因此使用搜尋到的RoPE對LLaMA2-7B的64k上下文視窗大小進行了微調。

圖片圖片

從結果中可以看到,在微調LLaMA2-7B之前和之後,該方法都明顯優於PI和YaRN,主要原因是有效地使用了非均勻位置插值、最小化資訊損失,以及為微調提供了更好的初始化。

受上述發現的啟發,研究人員提出了LongRoPE,首先引入了一種高效的搜尋演算法,以充分利用這兩種不均勻性,並將LLM上下文視窗擴展到200 萬個token

圖片圖片

具體形式化演算法參見原文。

實驗結果

研究人員將LongRoPE應用於LLaMA2-7B和Mistral-7B模型上,並從三個方面對其性能進行了評估:

1. 長文檔中擴展上下文LLM 的困惑度;

2. 密鑰(passkey)檢索任務,該任務衡量模型從大量無關文字中檢索簡單密鑰的能力;

3. 4096上下文視窗的標準LLM基準;

在256k範圍內進行長序列語言建模

在Proof-pile和PG19上透過不同內插法擴展的LLaMA2 和Mistral 的困惑度。

從實驗結果可以得到兩個關鍵的結論:

1. 從4k 到256k 的評估長度來看,擴展模型展現出整體困惑度下降的趨勢,顯示模型有能力利用更長的脈絡;

2. 即使在上下文視窗長度為16 倍的情況下(這通常是在較短上下文長度下保持性能所面臨的挑戰),我們的LongRoPE-2048k 模型在256k 上下文長度內的性能仍優於最先進的基線模型。

圖片圖片

圖片圖片

超過2000k的長序列語言建模

為了評估超長文件的有效性,研究人員使用了Books3資料集。

為了評估效率,隨機選擇20本書,每本長度超過2048k個token,並使用256k的滑動視窗。

圖片圖片

從結果中可以看出,LongRoPE成功地將LLaMA2-7B和Mistral-7B的上下文視窗擴展到2048k,同時也在8k-128k的較短長度內實現了與基線相當或更好的困惑度。

还可以观察到2048k LLaMA2和Mistral之间的显著性能差异:Mistral在较短的长度上优于基线,但困惑度在超过256k长度时达到7

LLaMA2的表現與預期一致:隨著時間的延長,困惑感會減少,在1024k和2048k時略有增加。

此外,在LLaMA2上,LongRoPE-2048k在256k比128k的微調長度下表現更好,主要是由於次級延伸比(secondary extension ratio)更小(即8倍對16倍)。

相較之下,Mistral在微調128k的視窗大小方面表現較好,主要原因是對於Mistral的128k和256k微調,研究人員遵循YaRN的設定使用16k訓練長度,影響了Mistral在微調後進一步擴展上下文視窗的能力。