我們知道在使用 GPT 和 LLaMA 等大型語言模型時，輸入的 prompt 存在字符數(shù)限制，比如 ChatGPT 目前的輸入字符限制是 4096 個字符。這會限制上下文學(xué)習(xí)和思維鏈等技術(shù)的發(fā)揮空間，畢竟用戶只能提供有限的示例數(shù)量。近日，Nous Research、EleutherAI 和日內(nèi)瓦大學(xué)的一個研究團隊提出了一種擴展上下文窗口的方案 YaRN ，并在實驗中取得了優(yōu)于其它所有方法的效果，而且他們還發(fā)布了使用 YaRN 微調(diào)過的 LLaMA 2 7B/13B 模型，其上下文窗口為 64k 和 128k。   

基于 Transformer 的大型語言模型（LLM）已經(jīng)展現(xiàn)出執(zhí)行上下文學(xué)習(xí)（ICL）的強大能力，并且?guī)缀跻呀?jīng)成為許多自然語言處理（NLP）任務(wù)的不二選擇。Transformer 的自注意力機制可讓訓(xùn)練高度并行化，從而能以分布式的方式處理長序列。LLM 訓(xùn)練所用的序列的長度被稱為其上下文窗口。

Transformer 的上下文窗口直接決定了可以提供示例的空間量，從而限制了其 ICL 能力。

如果模型的上下文窗口有限，那么為模型提供穩(wěn)健示例的空間就更少，而這些穩(wěn)健示例正是執(zhí)行 ICL 所用的。此外，當(dāng)模型的上下文窗口特別短時，摘要等其它任務(wù)也會受到嚴(yán)重妨礙。

就語言本身的性質(zhì)來說，token 的位置對有效建模來說至關(guān)重要，而自注意力由于其并行性，并不會直接編碼位置信息。Transformer 架構(gòu)引入了位置編碼來解決這個問題。

原始的 Transformer 架構(gòu)使用了一種絕對正弦位置編碼，之后其被改進成了一種可學(xué)習(xí)的絕對位置編碼。自那以后，相對位置編碼方案又進一步提升了 Transformer 的性能。目前，最流行的相對位置編碼是 T5 Relative Bias、RoPE、XPos 和 ALiBi。

位置編碼有一個反復(fù)出現(xiàn)的限制：無法泛化在訓(xùn)練期間看到的上下文窗口。盡管 ALiBi 等一些方法有能力做一些有限的泛化，但還沒有方法能泛化用于顯著長于其預(yù)訓(xùn)練長度的序列。

已經(jīng)出現(xiàn)了一些試圖克服這些限制的研究成果。比如，有研究提出通過位置插值（PI）來稍微修改 RoPE 并在少量數(shù)據(jù)上微調(diào)來擴展上下文長度。

兩個月前，Nous Research 的 Bowen Peng 在 Reddit 分享了一種解決思路，即通過納入高頻損失來實現(xiàn)「NTK 感知型插值」。這里的 NTK 是指 Neural Tangent Kernel（神經(jīng)正切核）。

其聲稱經(jīng)過 NTK 感知型擴展的 RoPE 能讓 LLaMA 模型的上下文窗口大幅擴展（超過 8k），同時還無需任何微調(diào)，對困惑度造成的影響也極小。

近日，他與另外三位合作者的相關(guān)論文出爐了！

• 論文：https://arxiv.org/abs/2309.00071
• 模型：https://github.com/jquesnelle/yarn

在這篇論文中，他們對 NTK 感知型插值做出了兩點改進，它們分別側(cè)重于不同的方面：

• 動態(tài) NTK 插值法，無需微調(diào)就能用于預(yù)訓(xùn)練模型。
• 部分 NTK 插值法，當(dāng)使用少量更長上下文的數(shù)據(jù)微調(diào)后，模型能取得最佳表現(xiàn)。

研究者表示，在這篇論文誕生前，就已經(jīng)有研究者將 NTK 感知型插值和動態(tài) NTK 插值用于一些開源模型。比如 Code Llama（使用 NTK 感知型插值）和 Qwen 7B（使用動態(tài) NTK 插值）。

在這篇論文中，基于之前有關(guān) NTK 感知型插值、動態(tài) NTK 插值和部分 NTK 插值的研究成果，研究者提出了 YaRN（Yet another RoPE extensioN method），一種可以高效擴展使用旋轉(zhuǎn)位置嵌入（Rotary Position Embeddings / RoPE）的模型的上下文窗口的方法，可用于 LLaMA、GPT-NeoX 和 PaLM 系列模型。研究發(fā)現(xiàn)，只需使用原模型預(yù)訓(xùn)練數(shù)據(jù)規(guī)模大約 0.1% 量的代表性樣本進行微調(diào)，YaRN 就能實現(xiàn)當(dāng)前最佳的上下文窗口擴展性能。

方法

旋轉(zhuǎn)位置嵌入（Rotary Position Embeddings / RoPE）最早由論文《RoFormer: Enhanced transformer with rotary position embedding》引入，也是 YaRN 的基礎(chǔ)。

簡單來說，RoPE 可以寫成如下形式：

對于使用固定上下文長度預(yù)訓(xùn)練的 LLM，如果使用位置插值（PI）來擴展上下文長度，則可以表示為：

可以看出 PI 對所有 RoPE 維度都會做同等延展。研究者發(fā)現(xiàn) PI 論文中描述的理論插值界限不足以預(yù)測 RoPE 和 LLM 內(nèi)部嵌入之間的復(fù)雜動態(tài)。下面將描述研究者發(fā)現(xiàn)并解決的 PI 的主要問題，以便讀者了解 YaRN 中各種新方法的背景、起因和解決理由。

高頻信息丟失 ——NTK 感知型插值

如果只從信息編碼的角度看 RoPE，根據(jù)神經(jīng)正切核（NTK）理論，如果輸入維度較低且對應(yīng)的嵌入缺乏高頻分量，那么深度神經(jīng)網(wǎng)絡(luò)難以學(xué)習(xí)高頻信息。

為了解決在對 RoPE 嵌入插值時丟失高頻信息的問題，Bowen Peng 在上述 Reddit 帖子中提出了 NTK 感知型插值。這種方法不會對 RoPE 的每個維度進行同等擴展，而是通過更少地擴展高頻和更多地擴展低頻來將插值壓力分散到多個維度。

在測試中，研究者發(fā)現(xiàn)在擴展未經(jīng)微調(diào)的模型的上下文大小方面，這種方法優(yōu)于 PI。但是，這種方法有一個重大缺點：由于它不只是一種插值方案，某些維度會被外推入一些「界外」值，因此使用 NTK 感知型插值進行微調(diào)的效果不及 PI。

更進一步說，由于存在「界外」值，理論上的擴展因子就無法準(zhǔn)確地描述真實的上下文擴展程度。在實踐中，對于給定的上下文長度擴展，必須將擴展值 s 設(shè)置得比期望的擴展值高一點。

相對局部距離的損失 —— 部分 NTK 插值

對于 RoPE 嵌入，有一個有趣的觀察：給定一個上下文大小 L，存在某些維度 d，其中的波長 λ 長于預(yù)訓(xùn)練階段見過的最大上下文長度（λ > L），這說明某些維度的嵌入可能在旋轉(zhuǎn)域中的分布不均勻。

PI 和 NTK 感知型插值會平等地對待所有 RoPE 隱藏維度（就好像它們對網(wǎng)絡(luò)有同樣的效果）。但研究者通過實驗發(fā)現(xiàn)，網(wǎng)絡(luò)會給某些維度不同于其它維度的待遇。如前所述，給定上下文長度 L，某些維度的波長 λ 大于或等于 L。由于當(dāng)一個隱藏維度的波長大于或等于 L 時，所有的位置配對會編碼一個特定的距離，因此研究者猜想其中的絕對位置信息得到了保留；而當(dāng)波長較短時，網(wǎng)絡(luò)僅可獲得相對位置信息。

當(dāng)使用擴展比例 s 或基礎(chǔ)變化值 b' 來拉伸所有 RoPE 維度時，所有 token 都會變得與彼此更接近，因為被一個較小量旋轉(zhuǎn)過的兩個向量的點積會更大。這種擴展會嚴(yán)重?fù)p害 LLM 理解其內(nèi)部嵌入之間小的局部關(guān)系的能力。研究者猜測這種壓縮會導(dǎo)致模型對附近 token 的位置順序感到困惑，從而損害模型的能力。

為了解決這個問題，基于研究者觀察到的現(xiàn)象，他們選擇完全不對更高頻率的維度進行插值。

他們還提出，對于所有維度 d，r < α 的維度按擴展度 s 線性插值（與 PI 一樣，避免出現(xiàn)外推）；r > β 的維度就完全不插值（總是外推）。

使用這一小節(jié)描述的技術(shù)，一種名為部分 NTK 插值的方法誕生了。這種改進版方法優(yōu)于之前的 PI 和 NTK 感知型插值方法，其適用于無微調(diào)和已微調(diào)模型。因為該方法避免了對旋轉(zhuǎn)域分布不均勻的維度進行外推，因此就避免了之前方法的所有微調(diào)問題。  

動態(tài)縮放 —— 動態(tài) NTK 插值

當(dāng)使用 RoPE 插值方法無微調(diào)地擴展上下文大小時，我們希望模型在更長的上下文大小上慢慢地劣化，而不是在擴展度 s 超過所需值時在整個上下文大小上完全劣化。

在動態(tài) NTK 方法中，擴展度 s 是動態(tài)計算的。

在推理過程中，當(dāng)上下文大小被超過時，就動態(tài)地更改擴展度 s，這樣可讓所有模型在達到訓(xùn)練的上下文限制 L 時緩慢地劣化而不是突然崩潰式劣化。   

增加用于長距離的平均最小余弦相似度 ——YaRN

即便解決了前面描述的局部距離問題，為了避免外推，也必須在閾值 α 處插值更大的距離。直覺來看，這似乎不應(yīng)該是個問題，因為全局距離無需高精度也能區(qū)分 token 位置（即網(wǎng)絡(luò)只需大概知道 token 是在序列的開頭、中間或末尾即可）。

但是，研究者發(fā)現(xiàn)：由于平均最小距離隨著 token 數(shù)量的增加而變得更近，因此它會使注意力 softmax 分布變得更尖（即減少了注意力 softmax 的平均熵）。換句話說，隨著長距離衰減的影響因插值而減弱，網(wǎng)絡(luò)會「更加關(guān)注」更多 token。這種分布的轉(zhuǎn)變會導(dǎo)致 LLM 輸出質(zhì)量下降，這是與之前問題無關(guān)的另一個問題。

由于當(dāng)將 RoPE 嵌入插值到更長的上下文大小時，注意力 Softmax 分布中的熵會減少，因此研究者的目標(biāo)是逆轉(zhuǎn)這種熵減（即增加注意力 logit 的「溫度」）。這可以通過在應(yīng)用 softmax 之前將中間注意力矩陣乘以溫度 t > 1 來完成，但由于 RoPE 嵌入被編碼為一個旋轉(zhuǎn)矩陣，就可以簡單地按常數(shù)因子 √t 來擴展 RoPE 嵌入的長度。這種「長度擴展」技巧讓研究可以不必修改注意力代碼，這能大幅簡化與現(xiàn)有訓(xùn)練和推理流程的集成，并且時間復(fù)雜度僅有 O (1)。

由于這種 RoPE 插值方案對 RoPE 維度的插值不均勻，因此很難計算相對于擴展度 s 所需的溫度比例 t 的解析解。幸運的是，研究者通過實驗發(fā)現(xiàn)：通過最小化困惑度，所有 LLaMA 模型都遵循大致相同的擬合曲線：

研究者是在 LLaMA 7B、13B、33B 和 65B 上發(fā)現(xiàn)這個公式的。他們發(fā)現(xiàn)這個公式也能很好地適用于 LLaMA 2 模型（7B、13B 和 70B），差別很細(xì)微。這表明這種熵增特性很常見，可以泛化到不同的模型和訓(xùn)練數(shù)據(jù)。

這種最終修改方案產(chǎn)出了 YaRN 方法。新方法在已微調(diào)和未微調(diào)場景中都勝過之前所有方法，而且完全不需要修改推理代碼。只需要修改一開始用于生成 RoPE 嵌入的算法。YaRN 如此簡單，使其可以在所有推理和訓(xùn)練庫中輕松實現(xiàn)，包括與 Flash Attention 2 的兼容性。

實驗

實驗表明 YaRN 能成功擴展 LLM 的上下文窗口。此外，他們僅訓(xùn)練了 400 步就得到了這一結(jié)果，這差不多只有模型的原始預(yù)訓(xùn)練語料庫的 0.1%，與之前的研究成果相比有大幅下降。這說明新方法具有很高的計算效率，沒有額外的推理成本。

為了評估所得到的模型，研究者計算了長文檔的困惑度，并在已有基準(zhǔn)上進行了評分，結(jié)果發(fā)現(xiàn)新方法勝過所有其它上下文窗口擴展方法。

首先，研究者評估了上下文窗口增大時模型的性能表現(xiàn)。表 1 總結(jié)了實驗結(jié)果。

表 2 展示了在 50 個未截斷的 GovReport 文檔（長度至少為 16k token）上的最終困惑度。

為了測試使用上下文擴展時模型性能的劣化情況，研究者使用 Hugging Face Open LLM Leaderboard 套件評估了模型，并將其與 LLaMA 2 基準(zhǔn)模型以及公開可用的 PI 和 NTK 感知型模型的已有分?jǐn)?shù)進行了比較。表 3 總結(jié)了實驗結(jié)果。