引言

今天給大家介紹的外文精品博客是來自Google Brain研究員關(guān)于Performer模型?(opens new window)的講解。該博客并未涉及太多的數(shù)學(xué)推導(dǎo)，而是從整體設(shè)計思想來講解Performer，包括廣義注意力的定義以及如何節(jié)省顯存與計算時間。在介紹完P(guān)erformer的設(shè)計思路之后，作者還用實驗證明了該方法的在長序列任務(wù)上的有效性，最后還介紹了Performer能夠應(yīng)用的場景以及研究結(jié)論。

作者介紹

Krzysztof Choromanski Google Brain首席研究員，同時也是Performer的一作。

Lucy Colwell Google Brain研究員。

譯者說

當(dāng)今Transformer模型可以說是在nlp領(lǐng)域處于統(tǒng)治者地位的模型，其超強的性能使得在其出現(xiàn)以后的大多數(shù)研究都以Transformer為骨架進行設(shè)計。但由于Transformer的核心結(jié)構(gòu)是注意力模塊，而原始計算注意力的復(fù)雜度是序列長度的平方（自注意力中Q×K的結(jié)果是一個L×L的注意力矩陣，L為序列長度），因此想要將Transformer模型應(yīng)用到長序列任務(wù)上不是一件容易的事。同時，本文后續(xù)提到的復(fù)雜度的概念均是針對序列長度而言的。

在博客的前半部分，作者講述了前人在使用注意力建模長序列這一問題上的一些嘗試，其中主流方法就是稀疏化注意力矩陣，同時也有研究人員表示能夠?qū)⑾∈韬蟮木仃嚳醋饕粡垷o向圖，即把圖上的節(jié)點看作token，把邊看做attention，兩個節(jié)點有邊則表明這兩個token的注意力操作沒有被稀疏。但作者也提到這些方法都有一些不足：

1. 稀疏算子的加速往往只有部分設(shè)備支持，這就導(dǎo)致該加速手段并不能輕松地應(yīng)用到所有設(shè)備上。

2. 前人的稀疏方法除了Big Bird對稀疏后的矩陣的表示能力有一個理論性的探討與證明，其余工作都沒有保證稀疏后的注意力矩陣的表示能力不受很大影響。

3. 這些工作大多只考慮了以Transformer為基礎(chǔ)的模型和生成式預(yù)訓(xùn)練模型。

4. 稀疏注意力矩陣往往因為不具備與原始注意力矩陣相當(dāng)?shù)谋硎灸芰?，所以需要更多的層?shù)來彌補這一缺陷，從而無法直接使用預(yù)訓(xùn)練好的模型參數(shù)。這就付出了更多的訓(xùn)練代價，包括重新預(yù)訓(xùn)練時間與能量消耗（電費）。

5. 存在著一些稀疏注意力矩陣無法代替原始注意力矩陣的領(lǐng)域，如指針網(wǎng)絡(luò)（Pointer Network?(opens new window)）。

因此Performer的作者提出一種由線性復(fù)雜度計算近似原始注意力矩陣的方法，這種方法能夠?qū)⒆⒁饬τ嬎憬档偷骄€性復(fù)雜度的同時依然保留原始注意力矩陣的各種特性，從而在很大程度上避免了上述不足。具體來說，作者首先使用特殊的指數(shù)變換和高斯投射變換得到Q和K矩陣的近似Q'和K'（先把Q*K從softmax中提取出來），然后重新安排了QKV矩陣乘法的順序，先用K'乘V，再用Q'乘上步計算結(jié)果（即K'乘V的結(jié)果），調(diào)整后的計算復(fù)雜度降低為線性復(fù)雜度（可以參考本文FAVOR+方法的第一張圖）。實驗方面，作者通過幾組實驗證明該方法是具備線性時間與空間復(fù)雜度的。同時還證明使用了Performer的模型速度與不使用attention機制的Transformer相當(dāng)。

原博客精華內(nèi)容概括

1.廣義注意力矩陣

首先，為了分析注意力矩陣是如何具備表示能力的，作者對注意力矩陣進行了拆分。但由于原始注意力矩陣是在Q矩陣與K矩陣相乘后使用了一個softmax非線性變換，因此無法直接將原注意力矩陣直接拆分為Q矩陣和K矩陣。但值得注意的是，我們依然可以將原注意力矩陣拆分為經(jīng)過非線性變換的Q矩陣與K矩陣的乘積。如下圖所示，等式的左側(cè)是代表注意力矩陣，等式的右側(cè)則是Q‘矩陣與K’矩陣。其中Q'矩陣與K'矩陣是原始Q矩陣與K矩陣的非線性變換。

如果將注意力矩陣看作是Q矩陣和K矩陣的非線性變換矩陣的乘積，那么原始的注意力矩陣就可以看作是將Q矩陣與K矩陣進行特殊的指數(shù)變換與高斯投射變換。由于Attention矩陣的計算是先計算Q矩陣與K矩陣的乘積，再進行非線性變換，而作者希望先對Q矩陣與K矩陣做非線性變換再做矩陣乘，因此這里Q矩陣與K矩陣的變換是特殊的指數(shù)變換（對標(biāo)原始注意力計算中Softmax操作包含的指數(shù)變換）與高斯投射變換（對標(biāo)Q矩陣與K矩陣在原始注意力計算中包含的線性變換）。而這種變換是非常稠密的，因為他并沒有縮小矩陣的大小，這就導(dǎo)致變換后的矩陣乘法依然具有序列長度平方的復(fù)雜度。從這個角度出發(fā)，作者認(rèn)為可以將這種非線性變換用一些其他近似方法，如核方法（大致思想就是將低維的向量映射到高維后，使用一些核函數(shù)直接使用低維向量來計算高維向量的內(nèi)積而并不直接計算兩個高維向量的內(nèi)積，這樣有助于加速計算。在這里則是將QK矩陣的維度降維后使用核函數(shù)來計算降維后的QK矩陣的乘積來近似未降維的QK矩陣的乘積。關(guān)于核函數(shù)的詳細解析在這里?(opens new window)），替代得到廣義上的注意力矩陣計算方法。雖然對于大多數(shù)核方法來說解析解并不存在，但是作者認(rèn)為，由于在使用近似替代的過程中并不依賴于解析解，因此使用該方法依然是有效的。

在Performer中，作者使用正值隨機特征（即該變換得到的矩陣元素都是正值，且該變換為非線性變換）來對Q、K矩陣近似變換。這種方法能夠在帶來更穩(wěn)定的訓(xùn)練過程的同時還能更好地近似原始注意力矩陣（因為經(jīng)過Softmax之后矩陣的元素都是正的）。關(guān)于這種近似方法的細節(jié)與數(shù)學(xué)推導(dǎo)可以移步原論文或其他論文解讀?(opens new window)了解。

2.FAVOR

在得到Q、K矩陣的近似矩陣后，我們可以使用線性復(fù)雜度計算與存儲原始注意力矩陣的近似矩陣。但即便是如此，顯示的計算注意力矩陣依然會帶來平方復(fù)雜度的顯存消耗，那么另一個優(yōu)化方法則是重新安排矩陣乘的順序，從而達到以線性計算復(fù)雜度和線性空間復(fù)雜度的代價完成Q、K、V的計算。而這就是FAVOR+。

上圖我們可以看到如果先計算與的乘積則需要序列長度平方復(fù)雜度的顯存消耗，而如果先計算與V的乘積就只會有線性復(fù)雜度的顯存消耗（這里我們只關(guān)心序列長度的復(fù)雜度，因為本文假設(shè)序列長度非常長），從而達到節(jié)省顯存的目的。

但同時我們注意到，由于沒有顯示地構(gòu)造注意力矩陣，這種計算方法無法直接使用原始注意力計算中的future mask矩陣，而mask矩陣卻是非常重要。例如在自回歸機器翻譯系統(tǒng)的訓(xùn)練中，我們不能讓單詞在注意力機制中看到它后面的單詞，而只能讓它看到它自己與它前面的單詞（因為自回歸機器翻譯是逐個地完成譯文單詞的生成，因此在模型實際生成譯文的時候，正在生成的單詞 之后的單詞還沒有被模型生成出來，所以在訓(xùn)練階段也不能讓模型見到未來的單詞）。為了解決這一問題，作者采用了前綴和計算來達到future mask的效果，并且還能節(jié)省常規(guī)注意力矩陣計算時future mask所需的顯存。（詳細計算過程見論文附錄B.1）

3.實驗結(jié)果

實驗方面，作者通過幾組實驗證明該方法是具備線性時間與空間復(fù)雜度的。同時還證明使用了Performer的模型速度基本上是最優(yōu)的（即與完全不使用注意力機制的模型有相似的性能與復(fù)雜度）。其中下圖中畫X的線代表去除注意力計算的Transformer模型，可以看到Performer與這條線具有相當(dāng)?shù)膹?fù)雜度量級（即隨著序列長度增長，模型計算時間的增長趨勢與無注意力計算的模型一致），而僅在常數(shù)項有些區(qū)別。

之后作者還設(shè)置了一組實驗表示Performer與原有的注意力計算方法是兼容的。這里作者將訓(xùn)練好的網(wǎng)絡(luò)的參數(shù)遷移到Performer上，實驗發(fā)現(xiàn)稍稍經(jīng)過微調(diào)的Performer就能夠達到Transformer模型收斂后的效果了。這一實驗表明，Performer可以減少重新訓(xùn)練網(wǎng)絡(luò)的訓(xùn)練代價。

4.潛在應(yīng)用方面：蛋白質(zhì)分析

在生物醫(yī)學(xué)領(lǐng)域，由于蛋白質(zhì)本身可以看作是由20個氨基酸為單詞組成的句子，因此蛋白質(zhì)分析任務(wù)與NLP任務(wù)具有強相關(guān)性?；谶@個特性，也有許多人嘗試使用基于Transformer為骨架的模型來完成這一任務(wù)。但是他們往往會遇到蛋白質(zhì)序列較長導(dǎo)致模型輸入序列過長的問題。這時，Performer的線性時空復(fù)雜度特性就完美地契合這一任務(wù)。作者通過實驗發(fā)現(xiàn)，盡管理論上分析使用Softmax作為非線性變換的Performer能夠更好地擬合原始注意力矩陣，但是使用ReLU作為非線性變換的Performer具有最好的實驗效果。

同時，作者還分析了Performer的注意力矩陣的可視圖，發(fā)現(xiàn)Performer訓(xùn)練得到的注意力權(quán)重與Transformer的注意力權(quán)重非常類似，同時Performer也能在完全沒有除了數(shù)據(jù)集以外的其他先驗知識的情況下注意到相似的氨基酸，從而賦予更高的注意力權(quán)重（下圖的D-E與F-Y）。而Performer所具備的線性時空復(fù)雜度允許研究人員可以在相同設(shè)備上訓(xùn)練更大規(guī)模的模型從而帶來更好的模型性能。

下圖則展示了在同等設(shè)備下，使用Performer方法，從而能夠訓(xùn)練更大的模型帶來的比普通Transformer更好的性能。

結(jié)論

根據(jù)實驗與理論分析，基于核方法設(shè)計的Performer具有非稀疏特性，且能夠與現(xiàn)有模型兼容（能夠直接將原有參數(shù)移植到Performer上）。同時該方法也能與其它基于注意力的方法兼容，如已經(jīng)開源的FAVOR與Reformer結(jié)合的工作?(opens new window)。同時作者希望Performer能夠引發(fā)人們關(guān)于注意力矩陣壓縮、Transformer模型本身以及核方法的其他思考。

來源 | Google AI Blog
譯者 | 張譯 穆永譽
單位| 東北大學(xué)自然語言處理實驗室