ICML 2025 | 清華、上海AI Lab等提出傅里葉位置編碼,多項任務遠超RoPE
長文本能力對語言模型(LM,Language Model)尤為重要,試想,如果 LM 可以處理無限長度的輸入文本,我們可以預先把所有參考資料都喂給 LM,或許 LM 在應對人類的提問時就會變得無所不能。
但是,LM 通常只在較短窗長下進行訓練,可能產生過擬合,只學習到指定范圍內的位置關系,但是無法理解沒學習過的位置關系。為了緩解這個問題,當下最流行的便是引入具有周期性的旋轉位置編碼(Rotary Position Embedding,RoPE)。由于周期性編碼每間隔一定距離就會出現數值重復,所以 LM 可以使用在少數幾個周期內學習到的經驗泛化到更多的周期當中。
但奇怪的是,使用 RoPE 的 LM 依然難以直接在訓練長度之外起效,必須依靠其他算法(如 YARN)來輔助其進行外推。 那么,到底是什么限制了 RoPE 的周期延拓,進而限制了 LM 的長度外推呢?
于是,清華大學講席教授、上海AI Lab主任/首席科學家周伯文教授的團隊對這一問題進行了深入探索,使用傅里葉分析工具解讀了使用 RoPE 的 Transformer 模型長文本泛化能力不足的原因之一是 RoPE 帶來的周期性延拓受到了頻譜破壞的影響。進一步地,該文章提出的傅里葉位置編碼(Fourier Position Embedding,FoPE)大幅提升了Transformer的長文本泛化能力。
- 論文標題:Fourier Position Embedding: Enhancing Attention’s Periodic Extension for Length Generalization
- arXiv 鏈接:https://arxiv.org/pdf/2412.17739
- 代碼鏈接:https://github.com/TsinghuaC3I/Fourier-Position-Embedding
研究亮點
發現 —— 頻譜損壞限制周期延拓
作者們通過觀察 RoPE 的公式可以發現,它為 Hidden States 的每一維都指定了單一的頻率,并假設這一維度的語義信息按照這個波長影響其他位置的語義。所以,RoPE 周期延拓性的起效前提是 “Hidden States 的每一維只存在單一頻率的語義”。如果每一維明明存在不同頻率的語義,卻仍然按照單一頻率的波長來估計這部分語義的傳遞規律,RoPE 所帶來的周期延拓將產生混亂,進而無法實現長文本泛化。
遺憾的是,在使用 RoPE 的 LM 中,這個假設只在 LM 的第一層中成立,但在后面的所有層中都不成立。因為后面的所有層中,每一維中都摻雜了除主頻之外的其他頻率分量,這個現象可以被稱作頻譜損壞(Spectrum Damage)。頻譜損壞主要有三個來源:① 線性函數;②激活函數;③時域截斷。
線性函數
如果我們假設線性層的權重為,其對輸入
進行處理得到輸出
,那么每一維輸出將由每一維輸入的線性組合構成:
考慮到每一維輸入所對應的頻率是不一樣的,那么將摻雜輸入中的所有頻率分量。
激活函數
根據高等數學的知識,給定一個含有兩個頻率的函數,以及任意一個與時間無關的非線性函數
,
被函數
作用后得到的輸出將存在多個頻率分量,這些頻率將是輸入頻率的線性組合:
這個結論可以通過泰勒展開進行簡單證明,也可以任意地推廣到存在更多頻率的情況。可以看到,經過線性層之后,每一維本就摻雜了多種頻率。在經過激活函數之后,這種摻雜會變得更加嚴重。
時域截斷
給定一個被截斷為長度N的單頻率函數
通過傅里葉變換可以得到這個函數的頻譜是:
其中,是截斷長度與周期長度相除后向下取整,
是在
處無限大但積分為 1、在其他位置取值均為零的奇異函數。可以看到,時域截斷會讓主要頻率
的強度向周圍的頻率擴散,降低信噪比。如果該函數的周期
大于截斷長度N,這個信噪比將變得很低,會極大地抑制對只要頻率的學習。巧合的是,RoPE 中其實使用了大量頻率低、周期長的分量,所以這個問題尤為嚴重。
算法 —— 頻域魯棒性是長文本泛化關鍵
在以往的研究中,大家普遍認為只有 Attention 才會影響長度外推。但從上面的分析可以看出,整個模型中的線性層、激活函數和時域截斷也都會對長度外推產生影響,并且是不利影響(也就是上文提到的頻譜損壞)。為了改善頻譜損壞對長文本泛化的不利影響,這篇論文提出了傅里葉位置編碼(FoPE,Fourier Position Embedding)來提升模型的頻域魯棒性和周期延拓性,進而提升長文本泛化。
FoPE 的核心思想是 “打不過就加入”。考慮到線性層和激活函數可以帶來更強的表征能力,時域截斷又是受到硬件限制無法改變,FoPE 索性就仍然保留了各層中的頻譜損壞,轉而提出了對于頻譜損壞更加魯棒的位置編碼。魯棒性的提升主要源于兩方面:① 既然每一維中不可避免的混雜其他頻率的分量,那就干脆在一開始就把每一維都建模成一個傅里葉級數(Fourier Series)。即使這樣的建模不會避免頻譜破壞,FoPE 卻可以在每一維中解碼出更多頻率的信息(利用三角函數的正交性);② 既然極低頻的分量周期過長,會導致這些頻率分量的周期特性無法被學習到,那就將他們裁剪成頻率為 0 的直流分量。考慮到直流分量的良好性質(既可以看作周期無限短,又可以看作周期無限長),這個新加入的頻率既保證了周期性,又可以讓每個詞匯的信息向無限遠的詞匯傳遞;
綜上,FoPE 的公式可以寫作:
實驗
進一步地,文章在困惑度、大海撈針準確率以及很多下游任務 Benchmark 對不同方法進行了對比,實驗發現 FoPE 在這些任務上都有穩定的表現,在絕大多數遠超過使用 RoPE 的模型。
潛在影響
論文中使用傅里葉工具得到的分析結論和算法可能存在更廣泛的潛在價值,有潛力應用在更多的領域和任務:① AI 領域內:長視頻生成、kv-cache 壓縮、多模型協同等;② AI 領域外:語義通信、光計算和腦機接口。
作者簡介:華爾默,清華大學博士生,研究方向是基礎模型的架構設計與訓練算法設計,在 ICML、ICLR、NeurIPS、ACL、EMNLP、COLM、AAAI 等頂級會議上發表過論文。