摘要

旋轉位置編碼（RoPE）已被證明可以有效地在基于Transformer的語言模型中編碼位置信息。然而，這些模型在超過它們訓練的序列長度后無法推廣。我們提出了YaRN（另一種RoPE擴展方法），這是一種計算高效的方法，可以擴展此類模型的上下文窗口，所需token減少10倍，訓練步驟減少2.5倍。使用YaRN，我們展示了LLaMA模型可以有效地利用和推斷出比其原始預訓練允許的上下文長度長得多的上下文長度，并且在上下文窗口擴展中達到了SOTA。此外，我們證明YaRN表現出了超越微調數據集有限上下文的能力。

背景知識

旋轉位置編碼

位置插值

附加符號

方法

然而，盡管PI在所有RoPE維度上實現了均等的拉伸，但我們發現PI所描述的理論插值界限在預測RoPE與LLM內部嵌入之間的復雜動態時不足夠準確。在接下來的小節中，我們分別描述了我們在PI中找到的主要問題，并解決了這些問題，以便向讀者提供每種方法的背景、起源和理論依據，以此來獲得完整的YaRN方法。

高頻信息丟失 - “NTK 感知”插值

為了解決在插值RoPE嵌入時丟失高頻信息的問題，[4]中開發了"NTK-aware"插值。與同樣乘以因子s的方式相比，我們通過在多個維度上縮放高頻率較小并且低頻率較大的方式來分散插值壓力。可以通過多種方式獲得這樣的轉換，但最簡單的方式是對θ的值進行基變換。我們希望最低頻率的縮放與線性位置縮放盡可能一致，并且最高頻率保持不變，因此我們需要找到一個新的基b′，使得最后一個維度與具有縮放因子s的線性插值的波長匹配。由于原始的RoPE方法跳過奇數維度，以便將cos(2πx/λ)和sin(2πx/λ)分量連接為單個嵌入，所以最后一個維度d ∈ D為|D|?2，解b′的值如下：

根據符號表示法，"NTK-aware"插值方案只需簡單地應用基變換公式，即：

在測試中，與PI方法相比，這種方法在擴展非精調模型的上下文大小方面表現更好。然而，這種方法的一個主要缺點是，由于它不僅僅是一個插值方案，某些維度會略微外推到“超出界限”的值，因此使用“NTK-aware”插值進行精細調節得到的結果比PI方法更差。此外，由于“超出界限”的值，理論縮放因子s并不能準確描述實際上的上下文擴展比例。在實踐中，為了給定上下文長度的擴展，必須將縮放值s設得比預期的縮放值更高。

相對局部距離的損失 - “NTK-by-parts”插值

RoPE嵌入的一個有趣觀察是，給定上下文大小L，在一些維度d中，波長比預訓練期間觀察到的最大上下文長度要長（λ > L），這表明某些維度的嵌入可能在旋轉域中分布不均勻。在PI和 “NTK-aware” 插值的情況下，我們將所有的RoPE隱藏維度視為相等。然而，實驗中發現網絡對一些維度與其他維度有所不同。如前所述，給定上下文長度L，某些維度的波長λ大于或等于L?？紤]到當隱藏維度的波長大于或等于L時，所有位置對編碼唯一的距離，我們假設絕對位置信息被保留下來，而當波長較短時，網絡只能獲得相對位置信息。當我們通過縮放因子s或使用基礎變換b′拉伸所有RoPE維度時，所有的標記都會彼此更加接近，因為較小旋轉角度下兩個向量的點積較大。這種縮放嚴重影響了LLM理解其內部嵌入之間的小而局部關系的能力。假設這種壓縮導致模型對接近的標記的位置順序感到困惑，從而損害了模型的能力。為了解決這個問題，根據我們的觀察，選擇根本不插值較高頻率的維度。

為了確定我們所需要的維度d，在給定某個上下文長度L下，我們可以按照如下的方式擴展方程（14）：

我們還提出對于所有滿足r < α的維度d，我們會線性插值一個比例s（就像à一樣，避免任何外推），而對于r>β的維度d，則不進行插值（總是進行外推）。定義斜坡函數γd為：

在斜坡函數的幫助下，我們將新的波長定義為：

α和β的值應該根據具體情況進行調整。例如，我們經實驗證明，對于Llama模型家族而言，α和β的合適取值為α = 1和β = 32。將λd轉換為θd后，該方法可以描述為：

使用本節中描述的技術，我們發布了一種名為“NTK-by-parts”插值的改進方法。這個改進的方法比之前的PI和“NTK-aware”插值方法在未微調模型和微調模型上的表現都更好。由于該方法避免了在旋轉領域中具有不均勻分布的維度進行外推，它避免了之前方法中的所有微調問題。

動態縮放 - “動態 NTK”插值

當使用RoPE插值方法來擴展上下文大小而無需進行微調時，我們希望模型在更長的上下文大小下能夠逐漸降級，而不是在設置的比所需值更高的比例s時在整個上下文大小上完全降級?；叵胍幌?，s = L′/L表示PI中的比例，其中L是訓練的上下文長度，L′是新擴展的上下文長度在“動態NTK”方法中，我們動態計算比例s如下：

在推理過程中，當上下文大小超過訓練的上下文限制L時，動態改變尺度允許所有模型以平滑方式退化，而不是立即失敗。

在使用動態尺度調整和kv緩存時需要注意，因為在某些實現中，RoPE嵌入是被緩存的。正確的實現應該在應用RoPE之前緩存kv嵌入，因為每個標記的RoPE嵌入在s改變時也會改變。

增加長距離平均最小余弦相似度-YaRN

即使我們解決了在前面描述的局部距離的問題，更大的距離也必須在閾值α處進行插值以避免外推。直觀來說，這似乎不是一個問題，因為全局距離不需要高精度來區分token的位置。然而，我們發現，由于平均最小距離隨著令牌數的增加而變得更加接近4，它使得注意力softmax分布變得“更尖銳”（即降低了注意力softmax的平均熵）。換句話說，由于插值導致遠距離衰減的效果減弱，網絡“更關注”較多的token。這種分布偏移導致LLM的輸出出現退化，這與之前的問題無關。

由于在將RoPE嵌入插值到更長的上下文大小時，注意力softmax分布的熵減小，我們的目標是反向減小熵（即增加注意力logits的“溫度”）。這可以通過將中間注意力矩陣乘以一個大于1的溫度t，然后應用softmax來實現，但由于RoPE嵌入被編碼為旋轉矩陣，我們可以簡單地通過一個常量因子√t來調整RoPE嵌入的長度。這種“長度縮放”技巧使我們避免了對注意力代碼的任何修改，從而顯著簡化了與現有的訓練和推斷流程集成的過程，并且具有O(1)的時間復雜度。

由于我們的RoPE插值方案不均勻地插值RoPE維度，因此很難計算所需的溫度比例t相對于比例s的解析解。幸運的是，我們通過最小化困惑度實驗發現，所有的Llama模型都大致遵循同一擬合曲線：

上面的方程是通過將 perplexity 最低的√t 與比例擴展因子 s 對擬合得到的，在 LLaMA 7b、13b、33b 和 65b 模型上，使用在 3.2 中描述的插值方法。我們還發現，這個方程也適用于Llama 2 模型 (7b、13b 和 70b)，只有輕微的差異。這表明增加熵的這個屬性在不同的模型和訓練數據中是普遍且可推廣的。

外推和遷移學習

在Code Llama中，使用了一個包含16k個上下文的數據集，并將比例因子設置為s≈88.6，這對應于一個上下文大小為355k。他們表明，網絡在訓練過程中從未見過這些上下文尺寸的情況下，可以推斷出多達100k個上下文。YaRN也支持使用比數據集長度更高的比例因子s進行訓練。由于計算約束，我們僅測試s=32，通過對具有64k上下文的相同數據集進行200個步驟的微調來進一步調優s=16模型。

我們展示了s=32模型在訓練過程中使用僅有64k上下文成功推斷出128k上下文的情況。與以前的“盲目”插值方法不同，YaRN在增加比例因子s時的遷移學習效率更高。這證明了從s=16到s=32的成功遷移學習，而網絡無需重新學習插入的嵌入，因為s = 32模型在整個上下文尺寸上等效于s=16模型，盡管僅在s=32上進行了200個步驟的訓練。

評估

長序列語言建模

為了評估長序列語言建模的性能，我們使用了GovReport和Proof-pile數據集，這兩個數據集都包含許多長序列樣本。對于所有評估，我們只使用了這兩個數據集的測試集進行計算。所有困惑度評估都是使用來自的滑動窗口方法進行計算的，其中S=256。

首先，我們評估了模型在上下文窗口增大時的性能。我們從Proof-pile中選擇了10個至少為128k標記長度的隨機樣本，并在從2k標記的序列長度到128k標記的情況下截斷了每個樣本的前2k個步驟，并評估了每個樣本的困惑度。圖1顯示了與其他上下文窗口擴展方法相比的結果。特別是，我們與Together.ai的公開可用的32k PI模型 和"NTK-aware" Code Llama進行了比較。這些結果總結在表1中。

表2顯示了在32k上下文窗口下，對至少含有16k個標記的50個未截斷的GovReport文檔進行評估的最終困惑度。

我們觀察到該模型在整個上下文窗口中表現出色，優于所有其他上下文窗口擴展方法。特別值得注意的是YaRN（s=32）模型，在128k中繼續呈現出下降的困惑度，盡管微調數據僅限于長度為64k的標記。這表明該模型能夠推廣到未見過的上下文長度。

標準化基準

Hugging Face Open LLM排行榜對一系列LLM模型在四個公共基準測試上進行了比較。具體來說，有25-shot ARC-Challenge [9]，10-shot HellaSwag，5-shot MMLU和0-shotTruthfulQA。

為了測試模型性能在上下文擴展下的退化情況，我們使用這套測試套件評估了我們的模型，并將其與Llama 2基準模型的已建立分數以及公開可用的PI和“NTK-aware”模型進行了比較，結果總結在表3中。

我們觀察到YaRN模型和它們相應的Llama 2基線之間存在最小的性能下降。我們還觀察到在YaRN s=16和s=32模型之間的平均分數下降了0.49%。因此我們得出結論，從64k到128k的迭代擴展導致的性能損失可以忽略不計。

總結

綜上所述，我們展示了YaRN改進了所有現有的RoPE插值方法，并可以作為PI的替代方案，其沒有缺點且實施起來所需的工作量很小。經過微調的模型在多個基準測試中保留其原有的能力，同時能夠關注非常大的上下文范圍。此外，在較短的數據集上進行微調的YaRN能夠高效地進行外推，并且能夠利用遷移學習實現更快的收斂，這對于計算限制的情況至關重要。最后，我們展示了YaRN進行外推的有效性，其中它能夠"訓練短，測試長"。