Survey: Positional Encodings in Transformers: ALiBi, RoPE, and Sinusoidal

2025-01-13
| ICLR 2022

Ofir Press, Noah A. Smith, Mike Lewis

當前挑戰

在語言與視覺 Transformer 的長序列推論與擴展過程中,常見的位置編碼通常會遇到幾項瓶頸:

傳統絕對位置編碼缺乏長度外推能力

在原始 Transformer 使用的 Sinusoidal PE 中,位置是透過預先計算的固定 \sin/\cos 函數直接加到輸入 embedding 上。這種作法雖然完全不帶額外參數量、計算極快,但只要推論時序列長度超過訓練時設定的最大長度 L,超出範圍的位置編碼分佈模型從未見過,會導致 Attention 分數計算徹底失準,Perplexity 呈現爆炸性惡化。

四大位置編碼在推論長度延伸時的 Perplexity 表現
不同位置編碼在長度外推(Extrapolation)下的 Perplexity 變化。Sinusoidal PE 一旦超過訓練長度 L,困惑度直接飆升;而 ALiBi 則展現了極佳的長度外推能力,在訓練長度 1024 下可直接平滑外推至 2048 甚至更長的輸入。

相對位置計算與矩陣旋轉帶來運算及記憶體開銷

為了改善絕對位置無法外推的缺點,旋轉位置編碼(RoPE)將位置資訊以旋轉矩陣的形式注入 Query 與 Key 的內積中,保證了點積結果嚴格依賴相對距離。然而,RoPE 需要在每一層 Attention 計算時頻繁對高維向量進行旋轉變換,這在大規模批次推論時會引入可觀的計算量與記憶體頻寬開銷。

可學習相對偏置矩陣大幅增加訓練參數與記憶體負擔

T5 採用的 Relative Position Bias 透過在 Attention Logit 上加上一個可學習的偏置矩陣(Learnable Bias Matrix)來捕捉相對距離。這種設計雖然靈活,但偏置矩陣需要佔用額外的模型參數,且無法預先離線計算,在長序列訓練時對記憶體與計算資源造成不小負擔。

小筆記:Perplexity(困惑度)與推論長度
在自迴歸 Decoder 評測中,Perplexity 定義為預測 Ground Truth token 機率對數平均的指數:

    \[l = \frac{1}{L} \sum_{i=1}^L \log_2 p(y_i \mid y_{<i}), \quad \text{Perplexity} = 2^{-l}\]


Perplexity 越小代表模型對預測結果越有把握(理論最佳值為 1)。當推論序列長度大於訓練長度時,評估 Perplexity 是否能維持穩定,就是檢驗位置編碼外推能力的最直接標準。


主要貢獻

針對長度外推與計算效率的矛盾,ALiBi(Attention with Linear Biases)提出了一種捨棄位置 Embedding、直接在 Attention Matrix 上施加線性距離懲罰的思路:

解決: 傳統絕對位置編碼缺乏長度外推能力

ALiBi 完全不往 Token Embedding 裡加任何位置向量,而是在計算 Attention Softmax 之前,根據 Query 與 Key 之間的相對距離 |i - j| 直接在點積結果上扣除一個靜態的線性懲罰項:

    \[\text{Attention}(Q, K, V) = \text{Softmax}\left( \frac{Q K^T}{\sqrt{d_k}} - m \cdot |i - j| \right) V\]


其中 m 是預先設定好、固定不需訓練的斜率(Slope)。因為懲罰項只跟相對距離 |i - j| 成線性正比,即便推論時序列拉長到未見過的長度,模型依然能按照既有的距離衰減規律分配注意力,完美解決外推崩潰的問題。

解決: 相對位置計算與矩陣旋轉帶來運算及記憶體開銷

相較於 RoPE 每一層都要做旋轉矩陣運算,ALiBi 的距離矩陣形式極其簡單,只需要在計算 Attention Mask 時直接加上一個靜態的偏置矩陣即可。這讓 ALiBi 在維持長度外推能力的同時,推論速度(Words Per Second, WPS)與記憶體佔用幾乎與最快的 Sinusoidal PE 持平。

推論速度與記憶體開銷對比
各類位置編碼的吞吐量(WPS)與記憶體佔用比較。ALiBi 同時兼顧了高速推論、低記憶體消耗與出色的長度外推表現。

解決: 可學習相對偏置矩陣大幅增加訓練參數與記憶體負擔

ALiBi 中的懲罰斜率 m 是幾何級數設定的非學習常數(例如在 8 個 Head 下分別設為 2^{-8/8}, 2^{-7/8}, \dots, 2^{-1/8}),完全零額外學習參數。模型不需要在訓練過程中耗費容量去學習位置權重,既降低了過擬合風險,又節省了優化器狀態所佔用的顯存。


延伸探討

四大位置編碼機制綜合對比

機制 核心作法 參數量 長度外推能力 推論速度 / 顯存開銷 代表模型
Sinusoidal PE 絕對位置 \sin/\cos 加進 Embedding 0 無外推能力(超出即崩潰) 最快 / 顯存極低 原始 Transformer
RoPE 旋轉矩陣作用於 Query / Key 0 具備外推能力(支援 2D 擴展) 偏慢 / 顯存中等 LLaMA, GPT-3, VARSR
T5-Bias Attention 上加可學習偏置矩陣 需額外參數 具備外推能力 較慢 / 顯存最高 T5
ALiBi Attention 上直接扣靜態線性懲罰 -m \cdot |i-j| 0 極佳外推能力(Train 1024, Test 2048+) 極快 / 顯存極低 BLOOM, MPT

心得

ALiBi 最令人驚艷的地方在於它的「大膽做減法」——研究人員花了多年時間試圖設計更複雜的幾何旋轉或連續函數來表示位置,但 ALiBi 證明了在很多場景下,Attention 真正需要的只是一個隨著相對距離單調遞減的先驗懲罰

不過 ALiBi 也並非萬靈丹。它底層依賴「距離越遠關聯度越低」的單調衰減假設,這在長篇文本理解中非常合理,但在影像或 2D 結構處理上(例如影像超解析度或物體偵測),空間關係是非單調且多維度的,這時能自然擴展至 2D 坐標軸的 RoPE(如 VARSR 中使用的 SaRoPE)往往展現出更強的空間幾何適配性。