論文題目：PoSE: Efficient Context Window Extension of LLMs via Positional Skip-wise Training

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

代碼鏈接：https://github.com/dwzhu-pku/PoSE

一、研究簡介

大型語言模型（LLMs）通常有一個預(yù)定義的上下文窗口大小，這限制了它們在長輸入的場景中的使用。為了使 LLMs 適應(yīng)更長的輸入，通常需要用目標(biāo)長度的樣本對其進(jìn)行微調(diào)（全長微調(diào)），由此導(dǎo)致訓(xùn)練成本十分昂貴。

舉例來說，在 Positional Interpolation [1] 這份工作中，將 LLaMA 的上下文窗口從 2048 拓展到 8192 使用了 32 張 A100，對于更大的上下文窗口則使用了 128 張 A100。

為了將訓(xùn)練長度與目標(biāo)長度解耦合，以實現(xiàn)高效的上下文窗口擴(kuò)展，我們提出了一種稱為位置跳躍式訓(xùn)練（Positional Skip-wisE training, PoSE）的方法，在原始的上下文窗口中模擬更長的訓(xùn)練樣本。

如下圖所示，我們將原始的上下文窗口分成幾塊，然后引入不同的 bias 項來調(diào)整每個塊的位置編碼。對于每一條訓(xùn)練樣本，這些 bias 項和塊的長度都會發(fā)生變化，因此通過大量的訓(xùn)練，模型能適應(yīng)目標(biāo)長度內(nèi)的所有位置。

實驗結(jié)果表明，PoSE 有以下三方面的優(yōu)勢：

• 訓(xùn)練的時空效率：由于只需要按照原始的上下文長度進(jìn)行訓(xùn)練，PoSE避免了由目標(biāo)上下文長度增加帶來的平方級別的計算復(fù)雜度，使得訓(xùn)練對于內(nèi)存和時間的開銷都大大減小。
• 能支持極長的上下文：通過解耦合訓(xùn)練長度和目標(biāo)長度，我們僅使用2k的訓(xùn)練窗口就成功將 LLaMA 拓展到 128k。
• 兼容所有基于 RoPE 的模型和位置插值策略：PoSE 的有效性在 LLaMA、GPT-J、Baichuan 等多種基礎(chǔ)模型，和 Linear、NTK、YaRN 等多種插值策略上得到了驗證。

二、技術(shù)背景

旋轉(zhuǎn)位置編碼 RoPE：RoPE 是當(dāng)下主流的位置編碼方式，被 LLaMA、GPT-J 等大語言模型所采用。給定一個 維的隱向量和位置 ，RoPE 通過如下方式編碼位置信息：

其中 。此前的絕對位置編碼多是直接作用在輸入向量上，與之不同的是，RoPE 是在作用在每一層的 query 和 key 向量上。RoPE 可以看作是一種相對位置編碼，給定位置處的 query 向量和位置處的 key 向量，注意力分?jǐn)?shù)可以寫成如下函數(shù)：

上下文窗口擴(kuò)展：給定一個以為原始上下文窗口長度的大語言模型，我們的目標(biāo)是其支持的上下文長度拓展到 ，使得在個輸入內(nèi)能較好地保持原有的性能。

位置插值（PI）：為了將 LLM 的上下文窗口從? 拓展到 ，一種直接的做法是使用長的輸入文本 ，設(shè)定其位置編碼為 ， 對 LLM 進(jìn)行微調(diào)。

然而，實踐表明 [1] [2]，這部分的位置在前向傳播時會產(chǎn)生災(zāi)難性的離群值，從而導(dǎo)致訓(xùn)練無法達(dá)到預(yù)期的效果。這主要是因為模型在預(yù)訓(xùn)練時只見過 

這些位置，無法很好的泛化到外推出去的這部分位置。

為了解決這個問題，Position Interpolation [1] 這份工作首先提出用“內(nèi)插”代替“外推”，設(shè)定縮放因子，并將上述注意力公式修改為  （也就是將位置編碼線性修改為  ）。

這種方式可以減少離群值的出現(xiàn)，將上下文窗口拓展到了 32k。在此基礎(chǔ)上，NTK 提出通過修改 來進(jìn)行位置插值，取得了更好的效果。YaRN 則根據(jù)不同的維度，對上述線性插值和 NTK 進(jìn)行了整合。

三、方法描述

盡管上述 Linear / NTK / YaRN 等插值方式能一定程度上解決位置外推的問題，他們?nèi)匀恍枰媚繕?biāo)長度的訓(xùn)練樣本來訓(xùn)練模型（即全長微調(diào)）。

隨著目標(biāo)長度的增加，平方級別的計算復(fù)雜度帶來的開銷依舊是難以承受的。因此，在插值技術(shù)的基礎(chǔ)上，我們提出調(diào)整原始的上下文窗口中的位置編碼，來模擬更長的訓(xùn)練樣本，從而實現(xiàn)高效的上下文擴(kuò)展。

位置編碼的調(diào)整主要有兩個考量：

• 為了避免推理時遇到 out-of-distribution 的相對位置，調(diào)整后的位置編碼應(yīng)覆蓋 所有這些相對位置；
• 用調(diào)整后的位置編碼來微調(diào) LLM 不應(yīng)該損害其原有性能，因此調(diào)整后的位置編碼的結(jié)構(gòu)應(yīng)該和預(yù)訓(xùn)練時盡可能接近。

第一步：我們將原上下文窗口  分成 N 個塊 ，每個塊的長度為，滿足 。記 的起始位置編碼為 ，則這個塊的位置編碼如下：

第二步：我們從離散均勻分布 中采樣出跳躍偏置項 ，并施加到 的位置編碼上：

為了避免塊之間位置編碼的重合，我們施加了 這一限制。值得注意的是，對于每條數(shù)據(jù)，我們會重新采樣每個塊的大小和跳躍偏置項。直觀上來說，通過這種方式，我們擴(kuò)大了原上下文窗口能覆蓋的相對位置范圍，并且位置編碼的不連續(xù)只發(fā)生在塊之間，因此盡可能地保持了預(yù)訓(xùn)練階段的位置編碼結(jié)構(gòu)。

第三步：選定每個塊內(nèi)的內(nèi)容。給定輸入文本，我們用類似的方法來抽取每個塊內(nèi)的填充的內(nèi)容：

我們也嘗試了其它  的賦值方式，如，此時塊間的內(nèi)容也是連續(xù)的；或如，此時調(diào)整后的位置編碼恰好對應(yīng)訓(xùn)練數(shù)據(jù)在原始文本中的位置。實驗結(jié)果表明，這幾種賦值方式并沒有明顯的差別。

第四步：位置插值及超參初始化。我們使用位置插值來使訓(xùn)練更穩(wěn)定。 和 設(shè)置成 0，N 設(shè)置為 2 。

四、實驗分析

1. 實驗設(shè)置

訓(xùn)練過程：我們主要使用 LLaMA-7B 作為基模型，對于所有設(shè)定都只訓(xùn)練 1000 步，訓(xùn)練時長度為 2k，batch size 為 64。我們使用 8 張 V100 進(jìn)行訓(xùn)練，1 張 A100 進(jìn)行推理。對于我們的方法和各個 baseline，我們都默認(rèn)采用線性插值來使訓(xùn)練更穩(wěn)定。

Baseline：

• 全長微調(diào)（Full-length fine-tuning）
• 隨機(jī)位置（RandPos）：給定目標(biāo)長度  和原始長度 ，從  中隨機(jī)采樣  個位置，按升序排列，作為位置編碼。

2. 主要結(jié)果

語言模型：

我們使用滑動窗口的方式來計算困惑度 PPL。在 GovReport 和 Proof-Pile 兩個數(shù)據(jù)集上，PoSE 的性能和 Full-length 十分接近，遠(yuǎn)超未做窗口擴(kuò)展的版本（Original）和隨機(jī)位置的版本（RandPos）。且隨著窗口長度從 2k 增加到 32k，PPL 呈下降趨勢，說明拓展后的模型能充分利用更長的上下文信息。

密碼檢索：

在密碼檢索任務(wù)上，利用 PoSE 拓展到 16k 和 32k 的模型能分別在 16k 和 32k 的上下文內(nèi)取得接近 100% 的密碼檢索準(zhǔn)確率，說明模型能關(guān)注到目標(biāo)長度內(nèi)的每個位置。

時空效率：

在時空效率方面，全長微調(diào)的訓(xùn)練時長和內(nèi)存消耗隨目標(biāo)長度的增加而迅速增長，相比之下，PoSE 需要的訓(xùn)練時間和內(nèi)存較為穩(wěn)定。并且在每個時間步上，性能和全長微調(diào)都很接近。

兼容性：

兼容性方面，PoSE 可以適配 LLaMA、LLaMA2、GPT-J、Baichuan2 等各種基于 RoPE 的基礎(chǔ)模型，以及 Linear、NTK、YaRN 等各種插值策略，展現(xiàn)出較好的普適性。其中 NTK 在最后階段會有一個 PPL 的突增，這主要是因為給定縮放因子 ，NTK 實際實現(xiàn)的縮放倍數(shù)會略小于  [3]。YaRN 解決了這個缺陷，取得了三者中最好的效果。

超長上下文拓展的潛力：

只使用 2k 的訓(xùn)練長度和 1000 步的訓(xùn)練步數(shù)，我們嘗試了將 LLaMA 模型拓展到 128k。實驗表明，在使用 YaRN 的情況下，模型在 128k 的窗口下仍然能保持較低的 PPL。

原窗口內(nèi)的語言能力：

最后，我們分析了經(jīng)由 PoSE 訓(xùn)練過后的模型在原窗口內(nèi)的語言能力。可以看出，和全長微調(diào)以及原始模型相比，PoSE 模型能力的損失非常微小，這說明 PoSE 在拓展上下文窗口的同時較好地保持了模型的基礎(chǔ)能力。

五、總結(jié)與討論

本文提出了一種位置跳躍式訓(xùn)練（PoSE）來高效的拓展大語言模型的上下文窗口。通過調(diào)整位置編碼，PoSE 在原始的上下文窗口中模擬更長的訓(xùn)練樣本，以達(dá)到解耦合訓(xùn)練長度和目標(biāo)長度的目的。

實驗結(jié)果表明 PoSE 在和全長微調(diào)保持同等性能的情況下，大大縮小了訓(xùn)練所需的時空開銷，并表現(xiàn)出良好的普適性和超長上下文擴(kuò)展的潛力。我們相信 PoSE 將大大降低上下文窗口拓展的成本，使更多人可以參與到相關(guān)的研究中來，從而推動長上下文建模領(lǐng)域的快速發(fā)展。

PoSE 完成于 2023 年 9 月，我們相信這種位置跳躍的思路是 Long Context 的有效解決方案。結(jié)合近幾個月來 Long Context 相關(guān)研究的進(jìn)展，我們認(rèn)為 PoSE 可能有以下一些方面值得進(jìn)一步探究：

• 應(yīng)用范圍的拓展：融合到預(yù)訓(xùn)練或者 SFT 階段中。PoSE 的實驗主要是對基礎(chǔ)模型進(jìn)行輕量級的 post-pretrain，如果能直接融入預(yù)訓(xùn)練中，可能可以更好的解決推理時長文本位置 out-of-distribution 的問題；如果能用于 SFT 中，則可以將模型更好地適配到具體的下游任務(wù)或者 alignment 要求上。
• 效果的提升：探索更優(yōu)的 skip 策略和合適的訓(xùn)練數(shù)據(jù)配比。一方面，實驗中只將窗口分成了兩塊，且跳步由隨機(jī)采樣決定，如何在更多場景設(shè)計更科學(xué)合理的 skip 結(jié)構(gòu)指的關(guān)注。

另一方面，PoSE 簡單從 The Pile 數(shù)據(jù)集采樣了部分超過 2k 長的樣本作為訓(xùn)練數(shù)據(jù)，并沒有特別關(guān)注數(shù)據(jù)的來源、長度等配比。根據(jù) Fu et al. (2024) [4] 的結(jié)論，優(yōu)化訓(xùn)練數(shù)據(jù)的分布對于長上下文建模能力的獲取以及模型原始能力的保持都能有較大的幫助。

本文轉(zhuǎn)載自PaperWeekly

原文鏈接：??https://mp.weixin.qq.com/s/H4X15Rcv77slWBK3Nm7axg??