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

大型語言模型（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在微調後進一步擴展上下文視窗的能力。