大型語言模型（LLM）往往會(huì)追求更長(zhǎng)的「上下文窗口」，但由于微調(diào)成本高、長(zhǎng)文本稀缺以及新token位置引入的災(zāi)難值（catastrophic values）等問題，目前模型的上下文窗口大多不超過128k個(gè)token。

最近，Microsoft Research的研究人員提出了一個(gè)新模型LongRoPE，首次將預(yù)訓(xùn)練 LLM 的上下文窗口擴(kuò)展到了2048k個(gè)token，在256k的訓(xùn)練長(zhǎng)度下只需要1000個(gè)微調(diào)步驟即可，同時(shí)還能保持原始短上下文窗口的性能。

圖片

論文鏈接：https://arxiv.org/abs/2402.13753

代碼鏈接：https: //github.com/microsoft/LongRoPE

LongRoPE主要包含了三個(gè)關(guān)鍵創(chuàng)新點(diǎn)：

1. 通過高效搜索識(shí)別并利用了位置插值中的兩種非均勻性，為微調(diào)提供了更好的初始化，并在非微調(diào)情況下實(shí)現(xiàn)了 8 倍擴(kuò)展；

2. 引入了漸進(jìn)擴(kuò)展策略，首先微調(diào) 256k 長(zhǎng)度的 LLM，然后在微調(diào)擴(kuò)展的LLM上進(jìn)行第二次位置插值，以實(shí)現(xiàn) 2048k 上下文窗口；

3. 在8k長(zhǎng)度上重新調(diào)整 LongRoPE以恢復(fù)短上下文窗口性能。

在 LLaMA2 和 Mistral 上對(duì)各種任務(wù)進(jìn)行的大量實(shí)驗(yàn)證明了該方法的有效性。

通過 LongRoPE 擴(kuò)展的模型保留了原始架構(gòu)，只對(duì)位置嵌入稍作修改，并且可以重復(fù)使用大部分已有的優(yōu)化。

位置插值的不均勻性

Transformer模型需要明確的位置信息，通常以位置嵌入（position embedding）的形式來表示輸入token的順序。

本文中的位置嵌入表示方法主要來自于RoPE， 對(duì)于位置索引為 n 的標(biāo)記，其相應(yīng)的 RoPE 編碼可簡(jiǎn)化如下：

圖片

其中，d 是嵌入維度，nθi 是標(biāo)記在位置 n 上的旋轉(zhuǎn)角度，θi = θ -2i/d 表示旋轉(zhuǎn)頻率。在 RoPE 中，θ 的默認(rèn)基準(zhǔn)值為 10000。

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

然而，當(dāng)前的非線性在很大程度上依賴于人為設(shè)計(jì)的規(guī)則。

這也自然引出了兩個(gè)問題：

1. 當(dāng)前的位置插值是否是最佳的？

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

圖片

為了回答這些問題，研究人員使用進(jìn)化搜索（evolution search）為L(zhǎng)LaMA2-7B發(fā)現(xiàn)更好的非均勻位置插值。搜索以易錯(cuò)性為指導(dǎo)，使用來自PG19驗(yàn)證集的5個(gè)隨機(jī)樣本。

通過實(shí)證分析，研究人員總結(jié)了幾個(gè)主要發(fā)現(xiàn)。

發(fā)現(xiàn)1：RoPE維度表現(xiàn)出很大的不均勻性，目前的位置插值方法無法有效處理這些不均勻性；

在公式 2 中為每個(gè) RoPE 維度搜索最佳 λ。

圖片

研究人員對(duì)比了PG19和Proof-pile測(cè)試集上使用不同方法的 LLaMA2-7B 在不進(jìn)行微調(diào)的情況下的復(fù)雜度。

圖片

從結(jié)果來看，搜索到的解決方案有明顯改善，表明當(dāng)前的線性（PI，positional interpolation）和非均勻（Dynamic-NTK 和 YaRN）插值方法都不是最佳的。

值得注意的是，YaRN 在 PG19 上的表現(xiàn)不如 PI 和 NTK，因?yàn)槠溥_(dá)不到非微調(diào) LLM 的目標(biāo)上下文窗口長(zhǎng)度。

例如，在 8k 上下文大小的情況下，YaRN 的困惑度在 7k 后達(dá)到峰值。

通過搜索，公式 2 中的重標(biāo)度（rescaled）因子λ變得不均勻，與PI、NTK的公式計(jì)算和YaRN的分組計(jì)算中的固定標(biāo)度s有所不同。

在8k和16k上下文窗口中，這些非均勻因子大大提高了LLaMA2的語言建模性能（即復(fù)雜度），而無需進(jìn)行微調(diào)，主要是因?yàn)橛纱水a(chǎn)生的位置嵌入有效地保留了原始的RoPE，尤其是關(guān)鍵維度，從而降低了LLM區(qū)分近似token位置的難度。

發(fā)現(xiàn)2：輸入序列中初始詞塊的RoPE推斷應(yīng)減少插值；

對(duì)于輸入序列中的初始n個(gè)token，假設(shè)RoPE應(yīng)該做較少的插值，這是因?yàn)闀?huì)獲得較大的注意力分?jǐn)?shù)，從而對(duì)注意力層至關(guān)重要，正如在Streaming LLM和 LM-Infinite 中觀察到的那樣。

為了驗(yàn)證這一點(diǎn)，研究人員使用PI和NTK將上下文窗口擴(kuò)展到 8k 和 16k，保留前 n（0,2, ..., 256）個(gè)token，不進(jìn)行插值。當(dāng)n=0 時(shí)，則恢復(fù)到原來的 PI 和 NTK。

圖片

上表中可以觀察到兩個(gè)結(jié)果：

1. 保留起始token而不進(jìn)行位置插值確實(shí)能提高性能。

2. 最佳起始token數(shù)n取決于目標(biāo)擴(kuò)展長(zhǎng)度。

發(fā)現(xiàn)3：在微調(diào)和非微調(diào)設(shè)置中，非均勻位置插值都能有效擴(kuò)展 LLM 上下文窗口。

雖然已經(jīng)證明，在不進(jìn)行微調(diào)的情況下，搜索到的非均勻位置插值能顯著提高8k和16k擴(kuò)展性能，但更長(zhǎng)的擴(kuò)展需要微調(diào)。

因此使用搜索到的RoPE對(duì)LLaMA2-7B的64k上下文窗口大小進(jìn)行了微調(diào)。

圖片

從結(jié)果中可以看到，在微調(diào)LLaMA2-7B之前和之后，該方法都明顯優(yōu)于PI和YaRN，主要原因是有效地使用了非均勻位置插值、最小化信息損失，以及為微調(diào)提供了更好的初始化。

受上述發(fā)現(xiàn)的啟發(fā)，研究人員提出了LongRoPE，首先引入了一種高效的搜索算法，以充分利用這兩種不均勻性，并將LLM上下文窗口擴(kuò)展到 200 萬個(gè)token。

圖片

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

實(shí)驗(yàn)結(jié)果

研究人員將LongRoPE應(yīng)用于LLaMA2-7B和Mistral-7B模型上，并從三個(gè)方面對(duì)其性能進(jìn)行了評(píng)估：

1. 長(zhǎng)文檔中擴(kuò)展上下文 LLM 的困惑度；

2. 密鑰（passkey）檢索任務(wù)，該任務(wù)衡量模型從大量無關(guān)文本中檢索簡(jiǎn)單密鑰的能力；

3. 4096上下文窗口的標(biāo)準(zhǔn)LLM基準(zhǔn)；

在256k范圍內(nèi)進(jìn)行長(zhǎng)序列語言建模。

在Proof-pile和PG19上通過不同插值方法擴(kuò)展的 LLaMA2 和 Mistral 的困惑度。

從實(shí)驗(yàn)結(jié)果中可以得出兩個(gè)關(guān)鍵的結(jié)論：

1. 從 4k 到 256k 的評(píng)估長(zhǎng)度來看，擴(kuò)展模型展現(xiàn)出整體困惑度下降的趨勢(shì)，表明模型有能力利用更長(zhǎng)的上下文；

2. 即使在上下文窗口長(zhǎng)度為 16 倍的情況下（這通常是在較短上下文長(zhǎng)度下保持性能所面臨的挑戰(zhàn)），我們的 LongRoPE-2048k 模型在 256k 上下文長(zhǎng)度內(nèi)的性能仍優(yōu)于最先進(jìn)的基線模型。

圖片

圖片

超過2000k的長(zhǎng)序列語言建模

為了評(píng)估超長(zhǎng)文檔的有效性，研究人員使用了Books3數(shù)據(jù)集。

為了評(píng)估效率，隨機(jī)選擇20本書，每本長(zhǎng)度超過2048k個(gè)token，并使用256k的滑動(dòng)窗口。

圖片

從結(jié)果中可以看出，LongRoPE成功地將LLaMA2-7B和Mistral-7B的上下文窗口擴(kuò)展到2048k，同時(shí)還在8k-128k的較短長(zhǎng)度內(nèi)實(shí)現(xiàn)了與基線相當(dāng)或更好的困惑度。

還可以觀察到2048k LLaMA2和Mistral之間的顯著性能差異：Mistral在較短的長(zhǎng)度上優(yōu)于基線，但困惑度在超過256k長(zhǎng)度時(shí)達(dá)到7。

LLaMA2的性能與預(yù)期一致：隨著時(shí)間的延長(zhǎng)，困惑感會(huì)有所減少，在1024k和2048k時(shí)略有增加。

此外，在LLaMA2上，LongRoPE-2048k在256k比128k的微調(diào)長(zhǎng)度下表現(xiàn)更好，主要是由于次級(jí)延伸比（secondary extension ratio）更?。?倍對(duì)16倍）。

相比之下，Mistral在微調(diào)128k的窗口大小方面表現(xiàn)更好，主要原因是對(duì)于Mistral的128k和256k微調(diào)，研究人員遵循YaRN的設(shè)置使用16k訓(xùn)練長(zhǎng)度，影響了Mistral在微調(diào)后進(jìn)一步擴(kuò)展上下文窗口的能力。

參考資料：

https://arxiv.org/abs/2402.13753