ALiBi Positional Bias/zh

Article
Topic area	transformers
Prerequisites	Transformer, Self-Attention, Positional Encoding

概述

線性偏置注意力（ALiBi）是一種向 Transformer 注意力中注入位置信息的方法，它通過在注意力分數上添加一個固定的、與距離相關的偏置，而不是將位置嵌入到 token 表示中。該方法由 Press、Smith 和 Lewis 於 2021 年提出。ALiBi 用一個靜態的、按頭獨立的偏置矩陣替代了正弦位置編碼和可學習位置嵌入，其元素隨查詢位置與鍵位置之間的距離線性增長。每個注意力頭對遠距離 token 的關注會受到不同程度的懲罰，懲罰速率由每個頭固定的斜率控制。

ALiBi 的核心特性是能夠外推到遠超訓練時所見的序列長度。標準的位置編碼在評估長度超過訓練上下文的輸入時性能會急劇下降，而採用 ALiBi 在長度 1024 上訓練的模型，可以在長度 2048、4096 甚至更長的序列上進行評估，僅出現輕微的困惑度下降，且無需任何重新訓練或插值。這一特性使 ALiBi 成為早期長上下文語言建模的主力，並成為 BLOOM、MPT 和 Replit Code 中首選的位置方法。

動機

原始的 Transformer 使用正弦位置編碼，將其與輸入 token 的嵌入相加。可學習的絕對位置嵌入（如 BERT 和 GPT-2 中所用）將固定的正弦曲線替換為可訓練的向量。這兩種方法都存在一個根本性的局限：超出訓練範圍的位置要麼不在嵌入表中（可學習方法），要麼落在模型從未需要解釋的正弦曲線區域（正弦方法）。經驗上，超出範圍位置的困惑度會急劇上升。

Press、Smith 和 Lewis 觀察到，Transformer 的位置信息最終是通過注意力分數進入網絡的。與其讓位置信息通過輸入嵌入傳遞，為什麼不直接在它被使用的地方進行編碼呢？ALiBi 完全移除了位置嵌入，轉而通過查詢與鍵之間相對距離的函數來對注意力 logits 施加偏置。由於偏置是關於距離的簡單線性函數，沒有任何可學習參數，因此在前所未見的距離上對其求值是良定義的，並且表現平滑。

公式

標準的縮放點積注意力對於查詢 $ Q $、鍵 $ K $ 和值 $ V $ 計算如下：

$ {\displaystyle \operatorname{Attention}(Q, K, V) = \operatorname{softmax}\!\left(\frac{Q K^\top}{\sqrt{d_k}}\right) V.} $

ALiBi 在softmax內部、按行歸一化之前為每個頭添加一個偏置矩陣 $ B^{(h)} $，從而修改上述公式為：

$ {\displaystyle \operatorname{Attention}^{(h)}(Q, K, V) = \operatorname{softmax}\!\left(\frac{Q K^\top}{\sqrt{d_k}} + B^{(h)}\right) V.} $

在因果語言建模設置中，對於頭 $ h $，查詢位置 $ i $ 與鍵位置 $ j $（其中 $ j \le i $）之間的偏置為：

$ {\displaystyle B^{(h)}_{ij} = -m_h \cdot (i - j),} $

其中 $ m_h > 0 $ 是每個頭固定的（非可學習的）斜率，$ (i - j) \ge 0 $ 是因果距離。位置 $ j > i $ 按慣例被屏蔽為 $ -\infty $。對於雙向模型，偏置是對稱的：$ B^{(h)}_{ij} = -m_h \cdot |i - j| $。

其關鍵特性是：不向輸入添加任何位置嵌入；偏置僅依賴於距離 $ i - j $，而非絕對位置；且斜率 $ m_h $ 不可學習。因此，每個頭都實現了對鄰近 token 呈指數衰減的偏好，衰減速率由其斜率控制。

斜率方案

ALiBi 按幾何方式在各注意力頭之間分配斜率。對於一個具有 $ n $ 個頭的注意力層，斜率為：

$ {\displaystyle m_h = 2^{-8 h / n}, \quad h = 1, 2, \ldots, n.} $

對於 $ n = 8 $ 個頭，該方案給出的斜率為 $ 1/2, 1/4, 1/8, \ldots, 1/256 $。對於 $ n = 16 $，斜率以幾何方式在相同範圍內插值，每個斜率都是其在 8 頭方案中相鄰斜率的幾何平均值。斜率較大的頭幾乎只關注近鄰 token；斜率較小的頭則偏置近乎平坦，可以關注遙遠的上下文。該方案由作者經驗性地選定，並在不同模型規模上保持不變；儘管簡單，至今仍未被顯著改進。

關鍵的是，這些斜率是常數，而不是參數。它們不通過 Backpropagation 更新，不出現在優化器狀態中，也不為模型增加任何參數。因此，就參數量而言，ALiBi 比可學習位置嵌入或 Rotary Position Embedding 都更省。

長度外推

原始 ALiBi 論文最顯著的成果是長度外推。一個具有 13 億參數的 Transformer，在 1024 token 上下文上訓練，在 3072 token 上下文上評估時，保持的困惑度與直接在 3072 token 上訓練時所能達到的相比僅有微小差距。而正弦編碼和可學習位置嵌入在相同的分布外長度下評估時，困惑度會上升一個數量級或更多。

其機制非常直觀。偏置 $ -m_h (i - j) $ 是關於距離的線性函數，對於任意非負的 $ i - j $ 都有良定義，包括大於訓練中任何所見值的距離。模型從不需要解釋它從未遇到過的位置輸入；它只需對一個其形狀與始終見過的完全相同的線性偏置應用softmax。在遠距離上，注意力分布會自然地集中到鄰近的 token 上，比例由每個頭的斜率決定，而模型在訓練過程中學會利用這一點。

在實際中，外推並非免費。質量會隨長度逐漸下降，而非常長的評估（超過訓練長度數倍時）會丟失遠距離 token 的信息，因為所有頭的偏置都會變得很大且為負。儘管如此，這種平緩的退化使 ALiBi 成為訓練時上下文長度受限但推理時上下文可能增長的場景下的一個可靠選擇。

與其他位置方法的比較

正弦位置編碼 將位置的固定函數加到輸入嵌入上。它們是絕對的（每個位置都有一個獨立的編碼），無參數，原則上對任何位置都有定義；但實際上無法外推，因為訓練好的網絡沒有學會解釋超出訓練範圍的高頻分量。

可學習的絕對位置嵌入為每個位置分配一個可訓練的向量。它們完全無法外推；超出表的位置根本不存在。

相對位置編碼（T5 偏置，Shaw 等人）根據查詢與鍵之間分桶的相對距離，向注意力分數添加一個可學習的標量偏置。ALiBi 與之密切相關，但使用固定的線性偏置而非可學習的分桶偏置，這既消除了參數，又使外推到未見過的距離成為可能。

旋轉位置嵌入（RoPE）按與位置成正比的角度對 $ Q $ 和 $ K $ 中的維度對進行旋轉。在現代大型語言模型（LLaMA、Mistral、Qwen、Gemma）中，RoPE 現在比 ALiBi 更流行，因為它在分布內通常能給出略好的質量，並通過 NTK 感知或 YaRN 縮放支持基於插值的上下文擴展。開箱即用時 RoPE 的外推不如 ALiBi 乾淨，但結合位置插值技巧後，它在長上下文微調方面已超過 ALiBi。

變體與擴展

已經提出了若干改進。

對稱 ALiBi 對雙向編碼器（如 BERT 類模型）使用 $ B_{ij} = -m_h |i - j| $。原始論文聚焦於因果模型，但對稱變體已被用於僅編碼器和編碼器-解碼器場景。

可學習斜率讓 $ m_h $ 可訓練，而非固定不變。這通常不會改善質量，並部分破壞了外推保證，因為學到的斜率可能朝著依賴訓練分布的方向漂移。

Sandwich 和 KERPLE 用其他關於距離的單調函數（對數函數、核化函數）替代線性偏置，在保留外推性質的同時，有時還能給出分布內更好的困惑度。這些方法在概念上是 ALiBi 的衍生物，並占據相同的架構位置。

動態 NTK ALiBi 在推理時調整斜率方案，以進一步擴展有效上下文，其思路類似於針對 RoPE 的 NTK 感知縮放。

實現細節

ALiBi 不需要修改輸入流水線，不需要位置嵌入層，也不需要修改 $ Q $ 或 $ K $。實現通常會按每個序列長度預先計算一次偏置矩陣 $ B^{(h)} $，並將其加到注意力 logits 上。在 FlashAttention 及類似的融合內核中，偏置是根據行列索引在線計算的，只需要一個大小為 $ n $ 的小型按頭斜率緩衝區。

對於使用鍵值緩存的推理，偏置僅依賴於已緩存鍵相對於當前查詢的絕對位置，而這些位置是已知的。因此，上下文長度不斷增長的流式推理無需任何特殊處理，只需將偏置擴展到新的位置即可。

ALiBi 與 Layer Normalization、解碼器與編碼器架構以及標準的注意力掩碼都兼容。由於沒有位置信息混入殘差流，因此輸入嵌入表也比使用可學習位置嵌入的模型略小。

局限性

ALiBi 存在若干已知的弱點。首先，儘管外推過程平緩，但在位置精度至關重要的基準上（例如精確匹配的片段檢索、字符級任務、某些代碼補全場景），其分布內質量通常略低於經過良好調優的 RoPE。其次，斜率方案是固定的，不針對任務或層級進行自適應調整；部分注意力頭實際上幾乎未被使用，因為其斜率過於陡峭，無法關注緊鄰環境之外的內容。第三，線性距離偏置是一種強烈的偏向局部性的歸納偏置；需要關注單個遠距離 token 的任務（例如從長上下文開頭檢索某個鍵）會隨著上下文增長而變得更困難，因為偏置會主動抑制遠處的鍵。

這些權衡解釋了為何在近期的大型語言模型中，ALiBi 在一定程度上被 RoPE 加位置插值方法取代，儘管它仍是長度外推中一種乾淨、無參數的基線方法，並且仍在 2023 年之前訓練的生產系統中被廣泛使用。