摘要

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

背景知識(shí)

旋轉(zhuǎn)位置編碼

位置插值

附加符號(hào)

方法

然而，盡管PI在所有RoPE維度上實(shí)現(xiàn)了均等的拉伸，但我們發(fā)現(xiàn)PI所描述的理論插值界限在預(yù)測(cè)RoPE與LLM內(nèi)部嵌入之間的復(fù)雜動(dòng)態(tài)時(shí)不足夠準(zhǔn)確。在接下來(lái)的小節(jié)中，我們分別描述了我們?cè)赑I中找到的主要問(wèn)題，并解決了這些問(wèn)題，以便向讀者提供每種方法的背景、起源和理論依據(jù)，以此來(lái)獲得完整的YaRN方法。

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

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

根據(jù)符號(hào)表示法，"NTK-aware"插值方案只需簡(jiǎn)單地應(yīng)用基變換公式，即：

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

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

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

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

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

在斜坡函數(shù)的幫助下，我們將新的波長(zhǎng)定義為：

α和β的值應(yīng)該根據(jù)具體情況進(jìn)行調(diào)整。例如，我們經(jīng)實(shí)驗(yàn)證明，對(duì)于Llama模型家族而言，α和β的合適取值為α = 1和β = 32。將λd轉(zhuǎn)換為θd后，該方法可以描述為：

使用本節(jié)中描述的技術(shù)，我們發(fā)布了一種名為“NTK-by-parts”插值的改進(jìn)方法。這個(gè)改進(jìn)的方法比之前的PI和“NTK-aware”插值方法在未微調(diào)模型和微調(diào)模型上的表現(xiàn)都更好。由于該方法避免了在旋轉(zhuǎn)領(lǐng)域中具有不均勻分布的維度進(jìn)行外推，它避免了之前方法中的所有微調(diào)問(wèn)題。

動(dòng)態(tài)縮放 - “動(dòng)態(tài) NTK”插值

當(dāng)使用RoPE插值方法來(lái)擴(kuò)展上下文大小而無(wú)需進(jìn)行微調(diào)時(shí)，我們希望模型在更長(zhǎng)的上下文大小下能夠逐漸降級(jí)，而不是在設(shè)置的比所需值更高的比例s時(shí)在整個(gè)上下文大小上完全降級(jí)。回想一下，s = L′/L表示PI中的比例，其中L是訓(xùn)練的上下文長(zhǎng)度，L′是新擴(kuò)展的上下文長(zhǎng)度在“動(dòng)態(tài)NTK”方法中，我們動(dòng)態(tài)計(jì)算比例s如下：

在推理過(guò)程中，當(dāng)上下文大小超過(guò)訓(xùn)練的上下文限制L時(shí)，動(dòng)態(tài)改變尺度允許所有模型以平滑方式退化，而不是立即失敗。

在使用動(dòng)態(tài)尺度調(diào)整和kv緩存時(shí)需要注意，因?yàn)樵谀承?shí)現(xiàn)中，RoPE嵌入是被緩存的。正確的實(shí)現(xiàn)應(yīng)該在應(yīng)用RoPE之前緩存kv嵌入，因?yàn)槊總€(gè)標(biāo)記的RoPE嵌入在s改變時(shí)也會(huì)改變。

增加長(zhǎng)距離平均最小余弦相似度-YaRN

即使我們解決了在前面描述的局部距離的問(wèn)題，更大的距離也必須在閾值α處進(jìn)行插值以避免外推。直觀來(lái)說(shuō)，這似乎不是一個(gè)問(wèn)題，因?yàn)槿志嚯x不需要高精度來(lái)區(qū)分token的位置。然而，我們發(fā)現(xiàn)，由于平均最小距離隨著令牌數(shù)的增加而變得更加接近4，它使得注意力softmax分布變得“更尖銳”（即降低了注意力softmax的平均熵）。換句話說(shuō)，由于插值導(dǎo)致遠(yuǎn)距離衰減的效果減弱，網(wǎng)絡(luò)“更關(guān)注”較多的token。這種分布偏移導(dǎo)致LLM的輸出出現(xiàn)退化，這與之前的問(wèn)題無(wú)關(guān)。

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

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

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

外推和遷移學(xué)習(xí)

在Code Llama中，使用了一個(gè)包含16k個(gè)上下文的數(shù)據(jù)集，并將比例因子設(shè)置為s≈88.6，這對(duì)應(yīng)于一個(gè)上下文大小為355k。他們表明，網(wǎng)絡(luò)在訓(xùn)練過(guò)程中從未見(jiàn)過(guò)這些上下文尺寸的情況下，可以推斷出多達(dá)100k個(gè)上下文。YaRN也支持使用比數(shù)據(jù)集長(zhǎng)度更高的比例因子s進(jìn)行訓(xùn)練。由于計(jì)算約束，我們僅測(cè)試s=32，通過(guò)對(duì)具有64k上下文的相同數(shù)據(jù)集進(jìn)行200個(gè)步驟的微調(diào)來(lái)進(jìn)一步調(diào)優(yōu)s=16模型。

我們展示了s=32模型在訓(xùn)練過(guò)程中使用僅有64k上下文成功推斷出128k上下文的情況。與以前的“盲目”插值方法不同，YaRN在增加比例因子s時(shí)的遷移學(xué)習(xí)效率更高。這證明了從s=16到s=32的成功遷移學(xué)習(xí)，而網(wǎng)絡(luò)無(wú)需重新學(xué)習(xí)插入的嵌入，因?yàn)閟 = 32模型在整個(gè)上下文尺寸上等效于s=16模型，盡管僅在s=32上進(jìn)行了200個(gè)步驟的訓(xùn)練。

評(píng)估

長(zhǎng)序列語(yǔ)言建模

為了評(píng)估長(zhǎng)序列語(yǔ)言建模的性能，我們使用了GovReport和Proof-pile數(shù)據(jù)集，這兩個(gè)數(shù)據(jù)集都包含許多長(zhǎng)序列樣本。對(duì)于所有評(píng)估，我們只使用了這兩個(gè)數(shù)據(jù)集的測(cè)試集進(jìn)行計(jì)算。所有困惑度評(píng)估都是使用來(lái)自的滑動(dòng)窗口方法進(jìn)行計(jì)算的，其中S=256。

首先，我們?cè)u(píng)估了模型在上下文窗口增大時(shí)的性能。我們從Proof-pile中選擇了10個(gè)至少為128k標(biāo)記長(zhǎng)度的隨機(jī)樣本，并在從2k標(biāo)記的序列長(zhǎng)度到128k標(biāo)記的情況下截?cái)嗔嗣總€(gè)樣本的前2k個(gè)步驟，并評(píng)估了每個(gè)樣本的困惑度。圖1顯示了與其他上下文窗口擴(kuò)展方法相比的結(jié)果。特別是，我們與Together.ai的公開可用的32k PI模型 和"NTK-aware" Code Llama進(jìn)行了比較。這些結(jié)果總結(jié)在表1中。

表2顯示了在32k上下文窗口下，對(duì)至少含有16k個(gè)標(biāo)記的50個(gè)未截?cái)嗟腉ovReport文檔進(jìn)行評(píng)估的最終困惑度。

我們觀察到該模型在整個(gè)上下文窗口中表現(xiàn)出色，優(yōu)于所有其他上下文窗口擴(kuò)展方法。特別值得注意的是YaRN（s=32）模型，在128k中繼續(xù)呈現(xiàn)出下降的困惑度，盡管微調(diào)數(shù)據(jù)僅限于長(zhǎng)度為64k的標(biāo)記。這表明該模型能夠推廣到未見(jiàn)過(guò)的上下文長(zhǎng)度。

標(biāo)準(zhǔn)化基準(zhǔn)

Hugging Face Open LLM排行榜對(duì)一系列LLM模型在四個(gè)公共基準(zhǔn)測(cè)試上進(jìn)行了比較。具體來(lái)說(shuō)，有25-shot ARC-Challenge [9]，10-shot HellaSwag，5-shot MMLU和0-shotTruthfulQA。

為了測(cè)試模型性能在上下文擴(kuò)展下的退化情況，我們使用這套測(cè)試套件評(píng)估了我們的模型，并將其與Llama 2基準(zhǔn)模型的已建立分?jǐn)?shù)以及公開可用的PI和“NTK-aware”模型進(jìn)行了比較，結(jié)果總結(jié)在表3中。

我們觀察到Y(jié)aRN模型和它們相應(yīng)的Llama 2基線之間存在最小的性能下降。我們還觀察到在YaRN s=16和s=32模型之間的平均分?jǐn)?shù)下降了0.49%。因此我們得出結(jié)論，從64k到128k的迭代擴(kuò)展導(dǎo)致的性能損失可以忽略不計(jì)。

總結(jié)

綜上所述，我們展示了YaRN改進(jìn)了所有現(xiàn)有的RoPE插值方法，并可以作為PI的替代方案，其沒(méi)有缺點(diǎn)且實(shí)施起來(lái)所需的工作量很小。經(jīng)過(guò)微調(diào)的模型在多個(gè)基準(zhǔn)測(cè)試中保留其原有的能力，同時(shí)能夠關(guān)注非常大的上下文范圍。此外，在較短的數(shù)據(jù)集上進(jìn)行微調(diào)的YaRN能夠高效地進(jìn)行外推，并且能夠利用遷移學(xué)習(xí)實(shí)現(xiàn)更快的收斂，這對(duì)于計(jì)算限制的情況至關(guān)重要。最后，我們展示了YaRN進(jìn)行外推的有效性，其中它能夠"訓(xùn)練短，測(cè)試長(zhǎng)"。