作者?|?王嘉寧?

整理?|?NewBeeNLP

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

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

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

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

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

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

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

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

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

1.1 絕對位置編碼

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

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

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

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

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

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

1.2 相對位置編碼

（1）顯式的相對位置

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

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

這樣一來，只需要有限個(gè)位置編碼，就可以表達(dá)出任意長度的相對位置（因?yàn)檫M(jìn)行了截?cái)啵还苁沁x擇可訓(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的相對位置向量， 表示待學(xué)習(xí)的第 個(gè)token的相對位置向量。通過分解讓相對位置注入在attention的計(jì)算過程中。

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

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

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

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

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

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

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

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

以及：

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

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

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

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

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

假設(shè) 和 是對應(yīng)位置的二維行向量（即每個(gè)位置都有兩個(gè)向量來表征位置），因此這個(gè)二維向量可以用復(fù)數(shù)來代替（包括實(shí)部和虛部），因此他們的內(nèi)積可以作為其對應(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)分別對 和 融入絕對位置時(shí)，即得到：

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

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

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

其中

等價(jià)于

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

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

每兩個(gè)為一組二維向量，一共有 個(gè)組合，可以直接拼接作為維度的旋轉(zhuǎn)位置編碼。
二維擴(kuò)展到多維的解讀：Transformer升級之路：4、二維位置的旋轉(zhuǎn)式位置編碼 - 科學(xué)空間|Scientific SpacesTransformer升級之路：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，其對應(yīng)的維度均為d，然后對于向量中相鄰兩個(gè)元素為一組，可以得到 d/2組（圖中左下部分相同顏色的兩個(gè)元素作為一組，對于每一組，一個(gè)文本則可以得到兩個(gè)行向量）；

獲得每個(gè)詞的絕對位置編號（該句子由6個(gè)詞，位置編號分別為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)程衰減

從圖中我們可以可以看到隨著相對距離的變大，內(nèi)積結(jié)果有衰減趨勢的出現(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)勢

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

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

自由的序列長度；

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

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

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

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

三、長文本外推性

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

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

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

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

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

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

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

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

可以看出，藍(lán)色區(qū)域（即滑動窗口內(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è)置不同。矩陣 的形狀如下所示：

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

3.2 KERPLE

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

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

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

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

3.3 Sandwich

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

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

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

3.4 XPOS

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

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

其中：

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

由于RoPE相對位置是 ，而不是 ，因此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張圖為局部注意力（滑動窗口），第3章圖為有限的全局感知（例如只限制前兩個(gè)token可以看見所有的token）。而第一張圖則是隨機(jī)mask，以緩解過度hard的局部注意力。三者注意力混合起來后得到第四張圖，這也是普遍訓(xùn)練超長文本大模型時(shí)采用的方法。

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

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

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

對應(yīng)的代碼也很簡單：

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

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

對應(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ì)算公式為：

只需要簡單的改成下面即可：

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

（4）全局依賴

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

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

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

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

編輯：黃飛