大模型的位置編碼和外推性問題一樣嗎

作者?|?王嘉寧?

整理?|?NewBeeNLP

現(xiàn)如今很多大模型都開始支持超過4096長度的推理，例如GPT-4支持超過30k，ChatGLM2-6B也支持最長為32K的文本。但是由于顯存資源的限制，這些大模型在真正在訓(xùn)練過程中不一定要訓(xùn)練這么長的文本，通常在預(yù)訓(xùn)練時(shí)只會(huì)設(shè)計(jì)到4k左右。

因此 如何確保在模型推理階段可以支持遠(yuǎn)遠(yuǎn)超過預(yù)訓(xùn)練時(shí)的長度 ，是目前大模型的核心問題之一，我們將這一問題歸為 大模型的外推性 。

大模型的外推性目前主要在這兩個(gè)方面考慮，也是提升最有效的兩個(gè)角度：

尋找或設(shè)計(jì)合適的位置編碼；

設(shè)計(jì)局部注意力機(jī)制。

本文從這兩方面進(jìn)行深度探討大模型的位置編碼和外推性問題。

一、位置編碼基礎(chǔ)介紹

對(duì)于一個(gè)token ，其表征向量記作 ，對(duì)于一個(gè)句子 則表示為 。那么可以通過一個(gè)映射函數(shù)將這個(gè)句子中的token表征為 ：

其中 和 表示第 和第 個(gè)token。

1.1 絕對(duì)位置編碼

在Transformer中，采用正余弦函數(shù)來表示絕對(duì)位置，公式如下：

這種編碼方式又叫做 Sinusoidal編碼 。直覺上看，第個(gè)位置的表征向量維度是 ，這個(gè)向量的奇數(shù)位置元素使用余弦值，偶數(shù)位置元素使用正弦值。
可視化圖如下所示：

相鄰的位置編碼向量很相似，較遠(yuǎn)的位置編碼向量差異很大，說明基于正余弦函數(shù)的絕對(duì)位置可以表征位置的相關(guān)性；

不需要顯式地學(xué)習(xí)位置，提高效率。

最后 映射函數(shù)可以定義如下所示。即輸入表征為token的表征和其對(duì)應(yīng)的絕對(duì)位置表征。

該表征通常是直接將位置表征與Word表征直接相加。

1.2 相對(duì)位置編碼

（1）顯式的相對(duì)位置

對(duì)于第 和第 個(gè)位置的token，其相對(duì)位置可以表示為，即兩個(gè)token之間的相對(duì)距離，且由最大最小值做約束（相對(duì)位置不能超過最大值或小于最小值）。

因此，相比于絕對(duì)位置，相對(duì)位置只需要有表征向量即可，即在計(jì)算兩個(gè)token之間的attention值時(shí)，只需要在attention計(jì)算過程中注入這兩個(gè)位置對(duì)應(yīng)的相對(duì)位置 對(duì)應(yīng)的相對(duì)位置表征向量即可：

這樣一來，只需要有限個(gè)位置編碼，就可以表達(dá)出任意長度的相對(duì)位置（因?yàn)檫M(jìn)行了截?cái)啵?，不管是選擇可訓(xùn)練式的還是三角函數(shù)式的，都可以達(dá)到處理任意長度文本的需求。

該表征通常是直接將位置表征與Word表征直接相加。
參考論文：《Self-Attention with Relative Position Representations》

（2）Transformer-XL（XLNet）

將第 和第 個(gè)位置的QK計(jì)算進(jìn)行了分解。引入了一些可學(xué)習(xí)的參數(shù)：

其中 表示待學(xué)習(xí)的第 個(gè)token的相對(duì)位置向量， 表示待學(xué)習(xí)的第 個(gè)token的相對(duì)位置向量。通過分解讓相對(duì)位置注入在attention的計(jì)算過程中。

該表征通常是在Attention計(jì)算過程中融入絕對(duì)位置。

（3）Transformer-XL的改進(jìn)

第二和第四項(xiàng)使用相對(duì)位置表征 來替換絕對(duì)位置表征。同時(shí)加入新的可訓(xùn)練參數(shù) 和 。

該表征方法在T5模型中被首次使用，參考論文：《Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer》

該表征通常是在Attention計(jì)算過程中融入相對(duì)位置。

（4）可訓(xùn)練偏置項(xiàng)

Transformer-XL中的位置表征是將QK分解為4項(xiàng)，而后面3個(gè)夠項(xiàng)都有跟位置有關(guān)的參數(shù)，可以直接將后面的3項(xiàng)之和抽象為一個(gè)偏置：

進(jìn)一步改造還可以得到：

以及：

這三種方法均是對(duì)Transformer-XL的表征形式進(jìn)行改進(jìn)。

二、RoPE旋轉(zhuǎn)位置編碼

RoPE（Rotary Position Embedding）的出發(fā)點(diǎn)就是“ 通過絕對(duì)位置編碼的方式實(shí)現(xiàn)相對(duì)位置編碼 ”，或者可以說是實(shí)現(xiàn) 相對(duì)位置編碼和絕對(duì)位置編碼的結(jié)合 。

這樣做既有理論上的優(yōu)雅之處，也有實(shí)踐上的實(shí)用之處，比如它可以拓展到線性Attention中就是主要因?yàn)檫@一點(diǎn)。

2.1 復(fù)數(shù)的引入

假設(shè) 和 是對(duì)應(yīng)位置的二維行向量（即每個(gè)位置都有兩個(gè)向量來表征位置），因此這個(gè)二維向量可以用復(fù)數(shù)來代替（包括實(shí)部和虛部），因此他們的內(nèi)積可以作為其對(duì)應(yīng)的Attention值。

內(nèi)積的計(jì)算可以由兩個(gè)復(fù)數(shù)實(shí)現(xiàn)：，其中 表示共軛復(fù)數(shù)， 表示取復(fù)數(shù)中的實(shí)部。

兩個(gè)二維向量的內(nèi)積，等于把它們當(dāng)復(fù)數(shù)看時(shí)，一個(gè)復(fù)數(shù)與另一個(gè)復(fù)數(shù)的共軛的乘積的實(shí)部。

因此當(dāng)分別對(duì) 和 融入絕對(duì)位置時(shí)，即得到：

RoPE求解過程推導(dǎo)了如何確定獲得每個(gè)位置編碼是 。

可以發(fā)現(xiàn)，當(dāng)乘以絕對(duì)位置 和 時(shí)，等價(jià)于復(fù)數(shù)運(yùn)算中乘以 ，即相當(dāng)于在復(fù)數(shù)空間中是相對(duì)位置 ，這樣就巧妙地通過復(fù)數(shù)運(yùn)算的形式將絕對(duì)位置轉(zhuǎn)換為相對(duì)位置。

復(fù)數(shù)乘法的幾何意義是向量的旋轉(zhuǎn)，假設(shè) 表示向量 在位置的位置編碼，則有：

其中

等價(jià)于

后面兩項(xiàng)的乘積本質(zhì)就是向量 （或）的兩個(gè)二維行向量。

當(dāng)向量維度為 時(shí)（ 為偶數(shù)），則可以擴(kuò)展為：

每?jī)蓚€(gè)為一組二維向量，一共有 個(gè)組合，可以直接拼接作為維度的旋轉(zhuǎn)位置編碼。
二維擴(kuò)展到多維的解讀：Transformer升級(jí)之路：4、二維位置的旋轉(zhuǎn)式位置編碼 - 科學(xué)空間|Scientific SpacesTransformer升級(jí)之路：6、旋轉(zhuǎn)位置編碼的完備性分析 - 科學(xué)空間|Scientific Spaces

2.2 RoPE的實(shí)現(xiàn)

當(dāng)輸入一個(gè)句子“Enhanced Transformer with Rotary Position Embedding”時(shí)，首先獲得其Query和Key向量q、 k，其對(duì)應(yīng)的維度均為d，然后對(duì)于向量中相鄰兩個(gè)元素為一組，可以得到 d/2組（圖中左下部分相同顏色的兩個(gè)元素作為一組，對(duì)于每一組，一個(gè)文本則可以得到兩個(gè)行向量）；

獲得每個(gè)詞的絕對(duì)位置編號(hào)（該句子由6個(gè)詞，位置編號(hào)分別為1，2，3，4，5，6），假設(shè)取“Enhanced”單詞為例，其第一組元素為θ1，位置為 m=1，那么通過旋轉(zhuǎn)位置編碼可以的到新的元素值。

所有單詞的d/2個(gè)組合都按照這種形式進(jìn)行“旋轉(zhuǎn)”，即可得到新的位置編碼（右下角）

RoPE一種線性的實(shí)現(xiàn)如下所示：

RoPE的性質(zhì)

（1）遠(yuǎn)程衰減

從圖中我們可以可以看到隨著相對(duì)距離的變大，內(nèi)積結(jié)果有衰減趨勢(shì)的出現(xiàn)。因此，選擇，確實(shí)能帶來一定的遠(yuǎn)程衰減性。當(dāng)然，能帶來遠(yuǎn)程衰減性的不止這個(gè)選擇，幾乎任意的光滑單調(diào)函數(shù)都可以。如果以為初始化，將視為可訓(xùn)練參數(shù)，然后訓(xùn)練一段時(shí)間后發(fā)現(xiàn)并沒有顯著更新，因此干脆就直接固定了。

（2）優(yōu)勢(shì)

用一個(gè)旋轉(zhuǎn)矩陣rotation matrix來對(duì)絕對(duì)位置進(jìn)行編碼，于此同時(shí)，meanwhile；

在自注意力機(jī)制中導(dǎo)入顯式的位置依賴。

自由的序列長度；

隨著相對(duì)位置的增大，而逐步延緩?fù)嘶?衰減）的inter-token dependency；

用相對(duì)位置編碼來“武裝”線性自注意力。

具體來說，RoPE 使用 旋轉(zhuǎn)矩陣對(duì)絕對(duì)位置進(jìn)行編碼 ，同時(shí)將 顯式的相對(duì)位置依賴性納入自注意公式中 。

【核心的兩個(gè)點(diǎn)，一個(gè)是“旋轉(zhuǎn)矩陣”，一個(gè)是“顯式的相對(duì)位置依賴”】。

三、長文本外推性

外推性 的含義是在長文本表征過程中，如何在訓(xùn)練階段只需要學(xué)習(xí)有限的長度，即可以在推理階段能夠延伸長度至若干倍且依然保持不錯(cuò)的性能和效果。

長度外推性是一個(gè)訓(xùn)練和預(yù)測(cè)的長度不一致的問題，主要體現(xiàn)在兩個(gè)方面：

預(yù)測(cè)的時(shí)候用到了沒訓(xùn)練過的位置編碼（不論是絕對(duì)位置還是相對(duì)位置）；

預(yù)測(cè)的時(shí)候注意力機(jī)制所處理的token數(shù)量遠(yuǎn)超訓(xùn)練時(shí)的數(shù)量。

解決長文本外推性問題的一個(gè)簡(jiǎn)單有效的方法是Attention Mask，如圖所示：

通過類似滑動(dòng)窗口的結(jié)構(gòu)，約束一個(gè)每個(gè)token只能對(duì)局部區(qū)域的token計(jì)算Attention值，因此對(duì)于相對(duì)位置大小不會(huì)超過窗口大小，解決了第一個(gè)問題；

Attention只會(huì)在窗口內(nèi)計(jì)算，避免了對(duì)大量的token的Attention進(jìn)行加權(quán)平均導(dǎo)致最終權(quán)重過度“平滑”現(xiàn)象。

在實(shí)現(xiàn)過程中，本質(zhì)上是在計(jì)算完之后減去一個(gè)矩陣，即 ，其中的形狀如下圖所示：

可以看出，藍(lán)色區(qū)域（即滑動(dòng)窗口內(nèi)的局部區(qū)域）為0，說明保持原始的Attention歸一化前的值；其他區(qū)域則為一個(gè)INT內(nèi)最大的整數(shù)，說明Attention值是一個(gè)非常小的數(shù)（在softmax歸一化后幾乎為0）。

3.1 ALIBI

論文：《Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation》

其與上面的思想一樣，只是改進(jìn)了上面的 M矩陣為 ，即Attention在歸一化前的計(jì)算為：，其中為超參數(shù)，Transformer的多頭注意力中的每個(gè)頭的值可設(shè)置不同。矩陣 的形狀如下所示：

相比于原始的方法，相對(duì)距離越長，值就越大，越遠(yuǎn)的位置Attention值被歸一化后就約小，相對(duì)于“滑動(dòng)窗口”采用的方法是hard（在窗口內(nèi)就計(jì)算attention，不在窗口內(nèi)不計(jì)算），AIBLI是比較soft的（離得近attention就會(huì)比較大，離得遠(yuǎn)就比較小）。

3.2 KERPLE

論文：《KERPLE: Kernelized Relative Positional Embedding for Length Extrapolation》

其對(duì)ALIBI進(jìn)行了一些改進(jìn)，引入了兩個(gè)可學(xué)習(xí)的參數(shù) 和 來“動(dòng)態(tài)”學(xué)習(xí)局部區(qū)域。如下圖所示，左側(cè)為原始的，通過引入?yún)?shù)來動(dòng)態(tài)減去AIBLI中的 矩陣：

定義了兩種模式，分別是power和logarithmic，分別對(duì)應(yīng)沒有對(duì)數(shù)和有對(duì)數(shù)的形式：

在logarithmic模式中， 控制了整體的量，相當(dāng)于ALIBI中的 ， c是一個(gè)常數(shù)。蘇神版簡(jiǎn)化寫作：

3.3 Sandwich

論文：《Receptive Field Alignment Enables Transformer Length Extrapolation》

Sandwich與KEPRLE是同一個(gè)作者提出的，其對(duì)KEPRLE進(jìn)行了少量改進(jìn)，即對(duì)應(yīng)的公式改寫為：，其中 和 可以使用Sinusoidal編碼表示，即：

由于Sinusoidal編碼在單調(diào)性上等價(jià)于 ，都是線性遞增形式，因此Sandwich只是改頭換面了。

3.4 XPOS

論文：《A Length-Extrapolatable Transformer》
參考解讀：Transformer升級(jí)之路：7、長度外推性與局部注意力

其在RoPE的基礎(chǔ)上引入了局部注意力。RoPE的本質(zhì)是：

其中：

在第2大節(jié)中已經(jīng)介紹了RoPE是通過使用復(fù)數(shù)來實(shí)現(xiàn)絕對(duì)位置表示相對(duì)位置的。XPOS通過引入一個(gè)新的標(biāo)量 ，即有：

由于RoPE相對(duì)位置是 ，而不是 ，因此XPOS約束在了單向Transformer，從而避免了負(fù)數(shù)出現(xiàn)。

XPOS又設(shè)計(jì)了一個(gè)局部感知的注意力機(jī)制Blockwise Causal Attention，進(jìn)一步提升了局部注意力的性能，提升了長文本的外推性。

四、外推性的其他探索

（1）混合注意力Mask

在解決長文本位置表征時(shí)，典型的代表有Transformer-XL、BigBird、LongFormer，他們除了局部注意力機(jī)制以外，還引入了隨機(jī)位置的性質(zhì)：

如上圖，第2張圖為局部注意力（滑動(dòng)窗口），第3章圖為有限的全局感知（例如只限制前兩個(gè)token可以看見所有的token）。而第一張圖則是隨機(jī)mask，以緩解過度hard的局部注意力。三者注意力混合起來后得到第四張圖，這也是普遍訓(xùn)練超長文本大模型時(shí)采用的方法。

（2）隨機(jī)位置表征

論文：《Randomized Positional Encodings Boost Length Generalization of Transformers》

絕對(duì)位置表征時(shí)，會(huì)存在位置上的OOV問題，隨機(jī)位置編碼則是通過在訓(xùn)練過程中采用如下策略：

對(duì)應(yīng)的代碼也很簡(jiǎn)單：

def?random_position_ids(N,?L=2048):
????"""從[0,?L)中隨機(jī)不重復(fù)挑N個(gè)整數(shù)，并從小到大排列
????"""
????return?np.sort(np.random.permutation(L)[:N])

蘇神對(duì)隨機(jī)位置編碼的新探索：

對(duì)應(yīng)的代碼為：

def?random_position_ids(N):
????"""先隨機(jī)采樣n，然后從[0,?n]均勻取N個(gè)點(diǎn)
????"""
????n?=?sample_from_xxx()
????return?np.linspace(0,?1,?N)?*?n

（3） Attention Scale

原始的Attention計(jì)算公式為：

只需要簡(jiǎn)單的改成下面即可：

其中 是訓(xùn)練時(shí)的最大長度， 是預(yù)測(cè)時(shí)的位置，通常來說。直覺來看，就是直接在計(jì)算時(shí)，根據(jù)其相對(duì)位置來控制Attention的值。當(dāng) 和距離很遠(yuǎn)時(shí)， 的值會(huì)很大，會(huì)使得整體Attention歸一化后會(huì)比較平緩，有助于解決外推性問題。

（4）全局依賴

滑動(dòng)窗口的方法如果在一層Transformer Layer里看，本質(zhì)上類似長度為 的N-Gram模型，即如下圖所示：

如果Transformer又 層，那么，從輸入層開始，長度為 的窗口內(nèi)的信息，可以在經(jīng)過 層之后傳給一個(gè)更廣的區(qū)域，區(qū)域長度為 ，如下圖所示：

蘇神給出的一種新的想法，就是假設(shè)我有層Transformer，則可以在前 層利用這種擴(kuò)張?zhí)匦?，得到最終 長度的區(qū)域后，在最后一層采用上面提到的 logn Attention Scale方法，講前 層擴(kuò)張的信息快速在最后一層與所有token進(jìn)行交互。引入蘇神的原文為：

這種局部注意力+Attention Scale的結(jié)合也是一種很巧妙的idea。實(shí)驗(yàn)也發(fā)現(xiàn)這種策略的外推性很驚艷。

編輯：黃飛