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