NeurIPS 2024 | Transformer長度外推，全新位置編碼DAPE大幅提升模型性能
AIxiv專欄是機器之心發佈學術、技術內容的欄目。過去數年，機器之心AIxiv專欄接收報導了2000多篇內容，覆蓋全球各大高校與企業的頂級實驗室，有效促進了學術交流與傳播。如果您有優秀的工作想要分享，歡迎投稿或者聯繫報導。投稿郵箱：liyazhou@jiqizhixin.com；zhaoyunfeng@jiqizhixin.com
本篇論文已被NeurIPS 2024接收，論文第一作者鄭傳陽來自香港中文大學，共同作者包括新加波國立大學高伊杭，諾亞實驗室石涵、任曉哲、蔣欣、李震國，香港中文大學 黃敏斌、 李靖瑤，香港大學熊璟，香港浸會大學吳國寶，香港中文大學李煜
在當今的人工智能領域，Transformer 模型已成為解決諸多自然語言處理任務的核心。然而，Transformer 模型在處理長文本時常常遇到性能瓶頸。傳統的位置編碼方法，如絕對位置編碼（APE）和相對位置編碼（RPE），雖然在許多任務中表現良好，但其固定性限制了其在處理超長文本時的適應性和靈活性。
為了應對這一挑戰，提出了一種全新的位置編碼方法：Data-Adaptive Positional Encoding（DAPE）。DAPE 通過動態調整位置編碼，使其能夠根據輸入上下文和學習到的固定先驗進行自適應調整。這種創新方法不僅保留了局部和反局部信息，還在模型訓練長度和長度泛化方面顯著提升了模型性能。相關研究成果已被 NeurIPS 2024 收錄。
-
論文地址: https://arxiv.org/abs/2405.14722
-
代碼地址: https://github.com/chuanyang-Zheng/DAPE
背景與挑戰
Transformer 模型的成功離不開其強大的序列處理能力，但在超出其訓練長度時，其性能往往會顯著下降。這主要是由於傳統的位置編碼方法（如 APE 和 RPE）在處理長文本時的固定性和缺乏適應性，導致模型難以有效捕捉長距離的依賴關係。最近的一些工作（e.g. Kerple, FIRE, BiPE）指出 transformer 通過合適的位置編碼可以提升模型長度外推的能力，但是在外推長度達到訓練長度 (512) 16 倍 (8192) 的時候，依然出現了 perplexity 的上升。相反的，DAPE 做到了在 128 長度上訓練，在 8192 乃至 16384 上拿到了更低的困惑度（perplexity）。
方法
Additive Relative Position Encoding
對於大多數這些加性相對位置編碼（RPE）方法，(softmax 之前的) 注意力 logits 的計算可以通過以下公式統一表示：
其中，偏置矩陣
項定義為
的
生成，
由位置編碼函數
不同的b的公式和參數化方法導致了各種 RPE 的變體。一些支持任意序列長度的方法包括 T5 的 RPE，ALiBi，Kerple，Sandwich，以及 FIRE。加性 RPE 的示例包括：
(1) ALiBi:
為超參數；
，其中標量
(2) Kerple:
是兩個可學習的參數；
和
，其中
(3) FIRE:
是一個轉換函數，旨在為局部位置分配更多的模型能力。
從數據中學習，
由參數
，其中位置編碼函數
之前方法的局限
這些位置編碼的共同特徵是它們是預定義且靜態的。具體來說，它們在各種任務和模型中都是固定的，這可能導致它們無法有效適應不同的輸入長度和上下文。為瞭解決這個問題，近期的研究提出了相對位置編碼的函數插值方法（FIRE），它利用神經網絡從輸入位置到位置偏置的隱式映射進行學習。儘管 FIRE 使用多層感知機（MLP）來學習位置嵌入，但這些嵌入在訓練完成後在不同任務中仍然是固定的。從直觀上看，所學習的靜態位置編碼（如 Kerple 和 FIRE）是所有訓練樣本的平均最優解。因此，儘管它們通常是有效的，但對於任何特定實例來說，它們本質上是次優的。這種靜態特性限制了它們在訓練上下文以外的各種實際場景中的靈活性和適用性。
Data-Adaptive Positional Encoding
本文受靜態位置編碼局限性的啟發，提出了一種數據自適應位置編碼（DAPE）方法。DAPE 根據語義信息（如當前的注意力值）和位置信息動態調整位置編碼。由於 MLP 具有普適逼近能力，本文採用
來根據注意力動態調整位置編碼。我們注意到 DAPE 與所有加性相對位置編碼兼容，並在可解釋性和易於實現方面具有優勢。所提出的 DAPE 結合了語義信息和位置信息，使得位置編碼能夠根據輸入數據進行自適應調整。這種適應性使 DAPE 能夠克服靜態編碼的局限性，通過對每個具體輸入數據的動態調整，實現相對最優的性能。據我們所知，這是在 Transformer 架構中首次引入的基於數據語義依賴的自適應位置編碼方法。
在這裏，我們使用注意力
是一個隱式函數，它將語義和位置信息整合為所需的位置編碼。因此，結合 DAPE 的 softmax 前注意力 logit 公式如下：
，其中
來表示注意力語義信息，使用位置偏置矩陣B（例如 ALiBi, Kerple 和 FIRE）來捕捉位置信息。然後，數據自適應 PE 可表示為
這裏
根據輸入上下文動態調整位置嵌入，確保編碼方法既具有自適應性又依賴於輸入數據。
，因為它具有普適逼近性。所有參數在訓練過程中直接從數據中學習。這種架構允許
神經網絡來參數化
是逐元素函數。實際上，我們採用一個兩層
在自然語言任務中，DAPE 的設計旨在捕捉 token 之間複雜的關係。Arora et al. 指出 aassociate recall 佔據了 Transformer 模型、基於 RNN 的模型和卷積模型之間困惑度（perplexity）差異的大部分。比如，我們考慮一個在長段落中 「Hakuna」 總是緊跟 「Matata」 的一致配對。這種模式表明模型對位置信息的依賴減少，而更注重增強詞嵌入的相似性，從而使得 「Hakuna」 可以有效地與前面的 「Matata」 聯繫起來。同樣，在涉及長上下文理解和搜索的任務中，注意力機制應該優先考慮語義相似性，而不是被與位置編碼相關的信息所掩蓋，因為在較長距離上位置編碼的相關性可能較低。因此，Transformer 應能夠保存信息而不受位置距離的過度影響。相反，一個滿意的 PE 應該結合語義和位置信息。因此，基於語義依賴的位置編碼方法是更優的，預計能夠提升模型性能。
實驗結果
相比於之前的方法
在訓練長度內更好的效果。DAPE 增強了在序列長度內部的表現，表明其較低的困惑度可能來自對整個句子的充分利用，而不是忽視長距離信息。與 ALiBi、Kerple 和 FIRE 相比，改進後的版本 DAPE-ALiBi、DAPE-Kerple 和 DAPE-FIRE 在序列長度內部的表現始終顯著更好。隨著序列長度的增加，ALiBi 往往從全局注意力過渡到幾乎局部的注意力，這就是為什麼 ALiBi 在訓練長度內的表現比大多數基線差，但在超出訓練長度後表現更好的原因。結果表明 DAPE 在序列長度內部的優越表現具有統計顯著性，p 值小於 0.05。因此，在不同訓練長度 (長度 128，512 以及 1024) 中的表現表明，DAPE 較低的困惑度是由於它有效利用了整個序列，而不是僅關注局部部分並忽視長距離信息。
在長度外推上更好的效果。與 ALiBi、Kerple 和 FIRE 相比，DAPE 顯著提升了長度外推（length extrapolation）性能。在不同長度的訓練和評估中，DAPE-Kerple 明顯超越 Kerple 等競爭對手。在 Arxiv 數據集上，訓練長度為 128 時，DAPE-Kerple 在評估長度為 8192 時達到了驚人的低困惑度 5.00，而 Kerple 的困惑度為 31.93。同樣，在 Books3 數據集上，訓練長度為 512 時，DAPE-Kerple 在相同的擴展評估長度下的困惑度為 17.88，遠遠優於 Kerple 的 39.31。這些結果證明，DAPE 通過其語義適應性和靈活性，持續提升了超出訓練長度的性能，超越了靜態位置編碼方法。
在更大模型上上保持更好的結果
DAPE 在更大模型上有效。隨著模型規模的增長（如圖 4 所示），DAPE 在性能指標上持續展現出提升。當模型規模從 125M 增加到 350M 時，DAPE-ALiBi 在評估序列長度為 8192（訓練長度為 512）時的困惑度顯著下降，從 3.82 降至 3.57。這些數值明顯小於原始 ALiBi 的困惑度，ALiBi 從 4.54 降至 4.21，表明了 DAPE 的強勁性能提升。此外，DAPE-Kerple 大幅減少了 Kerple 的困惑度，從最初的 22.76 降至令人印象深刻的 3.43。在 2.7B 和 6.7B 的模型上，DAPE-Kerple 依然取得了最低的 perplexity。這些結果證實了 DAPE 即使在模型規模增大的情況下仍能保持其有效性，並繼續表現出色，主要得益於其採用了語義自適應的位置編碼方法。
不同 hidden dimension 情況下的表現
即使是較小的 hidden dimension
，DAPE 仍然具有有效性。
設置為 4 時，DAPE-ALiBi 在評估長度為 128 時的困惑度為 8.25，在評估長度為 8192 時為 5.67，均優於 ALiBi 的 8.31 和 5.85。不論 hiddien dimension 設置為 4、16、32 或 64，DAPE 的性能在所有評估長度上都優於原始 ALiBi。這表明即使使用較小的
也能提升性能。實驗在 ALiBi 和 DAPE-ALiBi 上進行。如附錄圖 6 所示，當訓練長度為 128，且
關於偏置矩陣 Bias Matrix 的消融實驗
我們進一步對函數 f 進行了消融研究，證明 f 有助於增強偏置矩陣。DAPE（動態位置編碼）改進了偏置矩陣，使得最終的注意力矩陣
，這表明上下文自適應是重要的。
優於簡單的
，這表明通過提高偏置矩陣的表達能力可以獲得更好的效果；2). DAPE 的
優於
。對於未見過的位置，偏置矩陣 B 部分可以一定程度上處理（FIRE 將問題轉化為插值），但不夠準確，因此 DAPE 通過注意力得分幫助增強偏置矩陣 B。實驗結果表明兩點：1).
用於計算
得到了提升，而
在 CHE 基準上的表現
DAPE 在需要位置信息的任務中表現更好。DAPE（與 Kerple 和 FIRE 結合）在 11 項需要位置信息的任務中有 10 項表現最佳，並在 Solve Equation 任務中取得了第二好的表現。這凸顯了 DAPE 通過語義適應性處理需要位置信息的任務的有效性。
可視化結果
DAPE 展現 local pattern 和 anti-local pattern. 我們在圖 1 中繪製了第 8192 個位置的查詢 token 的學習位置編碼偏置，涵蓋了所選層中的所有注意力頭。我們想強調 DAPE 的兩個特點。首先，與固定的局部歸納偏置（如 Kerple 和 ALiBi）相比，DAPE 的偏置矩陣在不同的注意力頭中，能夠學習到既包含局部注意力模式，又包含 「反局部」 注意力模式 (DAPE Bias Head-8)，強調更遠的 key（類似於 FIRE）。其次，與為所有注意力固定的靜態偏置相比，DAPE 的偏置矩陣可以根據不同的注意力值動態調整。
代碼實現
未來展望
通過引入語義和位置信息的結合，DAPE 極大地提升了 Transformer 模型在長文本處理上的表現。同時，應將繼續優化 DAPE 的方法，提高其計算效率和適應性，探索其在更多實際應用中的潛力。
