《FNet: Mixing Tokens with Fourier Transforms》曾獲得2022年舉辦的計算機(jī)語言學(xué)會北美分會（NAACL）的最高效 NLP 論文獎，NAACL 委員會高度稱贊該團(tuán)隊對大型語言模型效率所做出的貢獻(xiàn)，并稱這項創(chuàng)新還讓模型能夠處理更長的輸入序列，讓未來研究遠(yuǎn)程上下文成為可能。

論文鏈接：https://arxiv.org/abs/2105.03824

摘要

該論文發(fā)現(xiàn)，對于只含Encoder的Transformer（如Bert、GPT等模型），將自注意力層替換成完全不含參數(shù)的傅立葉變換操作（被稱作FNet）就可以達(dá)到原有模型的92-97%的準(zhǔn)確度。而與此同時帶來了在GPU上80%的提速和在TPU上70%的提速。在長序列任務(wù)上，F(xiàn)Net可以在達(dá)到與現(xiàn)有長序列模型的準(zhǔn)確度的同時，比這些模型在任意序列長度下的速度更快。

背景

當(dāng)前Transformer的核心機(jī)制便是Attention機(jī)制，而現(xiàn)有的許多研究表示Attention矩陣能夠獲取到句子中各種不同的語法與語義關(guān)系。但該機(jī)制的一個問題在于其計算代價過高。由于Attention機(jī)制是要計算每個token之間的聯(lián)系，因此該機(jī)制需要序列長度平方的時間復(fù)雜度和空間復(fù)雜度。因此本文作者嘗試尋找更高效的網(wǎng)絡(luò)結(jié)構(gòu)來替代Attention層，同時還能保留Attention機(jī)制所帶來的序列信息捕獲能力。

前人的嘗試中一個具有代表性的工作便是Mixed MLP，而放在NLP語境下便是對的輸入（其中表示序列長度，代表隱藏層緯度），先在維做線性變換，而后在維做線性變換。這種做法可以先在隱藏層緯度將隱藏層的信息聚合起來，而后在序列的維度讓不同的單詞能夠通過線性變換看到其他單詞的信息。

而本文作者在嘗試過這種方法并確認(rèn)有一定效果后思考是否有更快速的信息聚合方式，而在嘗試傅立葉變換之后，作者意外的發(fā)現(xiàn)僅采用傅立葉變換來聚合信息也能達(dá)到不錯的效果。同時由于傅立葉變換目前已經(jīng)有非常成熟的高效實現(xiàn)以及其不含參數(shù)的特點，將Attention機(jī)制完全替換成傅立葉變換的FNet可以在損失不到10%的精度的同時帶來巨大的速度提升。同時作者還在文末提到，他們目前僅對不同的變換方式做了粗淺的調(diào)查，因此可能存在其他更優(yōu)的信息聚合方式。

方法

1.離散傅里葉變換DFT

介紹本文如何使用傅里葉變換之前，需要簡單闡釋什么是傅里葉變換？用來干啥？學(xué)通信的同學(xué)可能對這個概念并不陌生，但計算機(jī)同學(xué)大部分看見這個概念時應(yīng)該和筆者一樣整個人扭曲成了問號。不急，接下來是一個不涉及高深數(shù)學(xué)的直觀理解：

1.1頻域&時域

也許勤于思考的你們曾想過這樣一個問題：你眼中的世界是什么樣的？我們看到的世界都是以時間貫穿，隨時間改變的，如人的身高、汽車的軌跡、股票的走勢等等。我們習(xí)慣以時間作為參照來觀察動態(tài)世界，這種方式稱為時域分析。但是如果拋開時間，用另一種方法來觀察世界，這個世界居然也可以是永恒不變的，而這個靜止的世界，就是頻域。

引入一個非?！耙夂稀钡睦虞o助大家理解。在你的認(rèn)知中，一段音樂是什么呢？

左圖這是我們對于音樂最直觀且普遍的理解，一段隨著時間變化的震動。而對于音樂小能手來說，一段音樂可能是右圖這樣的。他們可以分別被看作一段音樂在時域和頻域的樣子（映射）。所以對大家來說，頻域這一概念并不陌生，只是我們從未意識到而已。

作為一段音樂的兩種表現(xiàn)形式，時域中上下擺動變幻莫測的美妙旋律，在頻域中卻是一個個永恒不變的音符。所以

你眼中看似落葉紛飛變化無常的世界，實際上是躺在上帝懷中一份早已譜好的樂章。（宿命論狂喜？）

而傅里葉變換(Fourier Transformation)正是貫穿時域與頻域的方法之一。接下來我們回歸論文，關(guān)于背景知識更多精彩請自行移步https://zhuanlan.zhihu.com/p/19763358?(opens new window)。

1.2傅里葉變換

傅里葉變換的作用是將時域中非周期的連續(xù)信號轉(zhuǎn)變?yōu)橐粋€頻域中非周期且連續(xù)的信號；

這句話聽上去又讓人云里霧里。如果拿鋼琴譜舉例，音樂是時域中非周期的連續(xù)信號，而琴譜顯然是頻域空間中非周期的離散信號，即一個個離散的音符，或者說頻率（PS：當(dāng)然變成周期的也可以，這個作曲人說了算）。而使用傅里葉變換，音樂會變成頻域空間中的“連續(xù)譜”。

連續(xù)譜應(yīng)該長啥樣我們這里不做進(jìn)一步探討，但是把傅里葉變換拿到Transformer中，如果把連續(xù)的token序列（自然語言）看作時域中非周期的連續(xù)信號，傅里葉變換能夠?qū)⑵滢D(zhuǎn)換為頻域空間的非周期連續(xù)信號（對于token的一種新的表示）。而使用傅里葉變換成熟的計算方法（快速傅里葉變換和矩陣乘）可以比自注意力機(jī)制更快的進(jìn)行特征運算。

2.具體算法

針對一個序列{},∈[0,]，傅立葉變換可以寫成如下方程：

也就是說，針對新序列中每一個位置，傅立葉變換均根據(jù)原有序列的全部信息生成一個新的表示。同時現(xiàn)有的傅立葉變換方法具有兩個高效實現(xiàn)。一個高效實現(xiàn)便是廣為人知的快速傅立葉變換，該實現(xiàn)方法具有的時間復(fù)雜度。

另一個高效實現(xiàn)便是直接將輸入序列乘一個傅立葉變換矩陣即可完成傅立葉變換。其中傅立葉變換矩陣如下所示：

即便快速傅立葉變換方法具有的復(fù)雜度而此方法具有的復(fù)雜度，但矩陣乘操作對TPU來說更友好，且在相對較短的序列上具有更快的運算時間。

3.FNet結(jié)構(gòu)

FNet的結(jié)構(gòu)如下圖所示：

針對網(wǎng)絡(luò)的輸入，網(wǎng)絡(luò)的輸出為：

其中代表傅立葉變換操作，代表取實數(shù)操作。這樣做可以避免后續(xù)網(wǎng)絡(luò)處理因傅立葉變換而帶來的虛數(shù)。同時作者發(fā)現(xiàn)僅僅取傅立葉變換后的實數(shù)部分就已經(jīng)能達(dá)到最優(yōu)的效果。其中這個最優(yōu)是和作者嘗試的其他變換如阿達(dá)馬（Hadamard）變換、哈特萊（Hartley）變換、離散余弦（Discrete Cosine）變換做比較的。

同時作者使用的Embedding層含詞嵌入表示、絕對位置編碼以及該單詞的類型編碼。同時作者也發(fā)現(xiàn)由于傅立葉變換本身就蘊(yùn)含位置信息，因此不含位置編碼的FNet也能有不錯的性能表現(xiàn)。但是具體實驗的時候為了和Bert等模型有相同模型設(shè)置，F(xiàn)Net依然使用了位置編碼。

4.實現(xiàn)細(xì)節(jié)

作者發(fā)現(xiàn)在GPU上使用快速傅立葉變換能夠在實驗范圍內(nèi)的序列長度（512-8196個tokens）上取得更快的性能；而在TPU上，對于序列長度小于4096的序列，使用矩陣乘的方法更快。因此在GPU上的FNet實現(xiàn)完全采用快速傅立葉變換，而在TPU上的FNet僅在序列長度大于4096時才采用快速傅立葉變換，否則采用矩陣乘來實現(xiàn)。帶來不同硬件上性能區(qū)別的原因在于TPU對于矩陣乘有更好的優(yōu)化，且在GPU上具有更高效的快速傅立葉變換實現(xiàn)。

主要實驗結(jié)果及分析

作者主要實驗了以下模型并對比了實驗結(jié)果，上表列舉了這些模型可學(xué)習(xí)的參數(shù)量和前向傳播時混合層(注意力子層及其各種替換)計算操作的數(shù)量級。FNet由于傅里葉子層沒有可學(xué)習(xí)參數(shù)，和FF-only以及Random具有相同數(shù)量的可學(xué)習(xí)參數(shù)。而采用FFT與矩陣乘來優(yōu)化傅里葉變換的方法分別具有和【其中n為序列長度，d為隱層維度】的時間復(fù)雜度：

• BERT-Base(傳統(tǒng)的transformer encoder結(jié)構(gòu))

• FNet encoder(將自注意力子層替換為傅里葉子層)

• Linear encoder(將自注意力子層替換為兩個可學(xué)習(xí)的稠密線性層，分別應(yīng)用到隱層維度和序列維度)

• Random encoder(將自注意力子層替換為兩個隨機(jī)初始化的常數(shù)矩陣，分別應(yīng)用到上述兩個維度)

• FF-only encoder(完全刪除自注意力子層，這種encoder不會混合token)

在GLUE基準(zhǔn)上，作者通過實驗發(fā)現(xiàn)在這些模型中，BERT-base具有最好的性能表現(xiàn)(在8個下游任務(wù)上取平均)，F(xiàn)Net和Linear模型表現(xiàn)十分相近，Random和FF-only在二分類任務(wù)上準(zhǔn)確率接近50%，基本等于啥也沒學(xué)到。

在訓(xùn)練穩(wěn)定性方面，雖然使用稠密參數(shù)化的Linear encoder進(jìn)行mix token會比無參數(shù)的FNet學(xué)習(xí)的更加靈活，但由于梯度爆炸等問題，Linear encoder更難訓(xùn)練。作者還在BERT-Large的基礎(chǔ)上進(jìn)行了試驗，linear encoder的訓(xùn)練非常不穩(wěn)定，導(dǎo)致其在GLUE基準(zhǔn)上嚴(yán)重低于Base模型。相較于FNet，linear encoder還有以下缺陷：1、Linear encoder會占用更大的顯存；2、在常規(guī)序列(長度為512)訓(xùn)練較慢，在長序列上拓展性更差；

在效率方面，base模型中FF-only速度較快，但性能很差。FNet訓(xùn)練速度明顯快于BERT——在gpu上快80%，在tpu上快70%，F(xiàn)LOPS是BERT 的63%。而且FNet-Hybrid模型(最后兩層encoder保留自注意力子層)達(dá)到了與FNet相近的加速效果和與BERT相近的性能。作者實驗表明，傅里葉子層的正向和反向傳遞速度比自注意子層快了一個數(shù)量級(但FNet的整體訓(xùn)練速度會受到所有模型共享的前饋子層的限制)。

在LRA基準(zhǔn)上，作者通過實驗發(fā)現(xiàn)FNet和vanilla transformer（Tay等人（2021a）的LRA評估中第二準(zhǔn)確模型）具有可比的效果，如表4a所示。效率方面，作者統(tǒng)計了在gpu(8個 V100芯片)上進(jìn)行實驗的訓(xùn)練速度和內(nèi)存使用情況。對于不同的序列長度{512、1024、2048、4096、8192}執(zhí)行實驗。在GPU上，所有序列長度上FNet都比其他模型快得多，這可能因為GPU上FFT的實現(xiàn)比較高效。表4b還表明FNet會占用更少的內(nèi)存。這是因為FNet傅里葉子層中沒有可學(xué)習(xí)的參數(shù)，也可能部分由于FFT在處理較長的序列情況下效率較高。

結(jié)語

FNet(hybrid)在犧牲少量性能的基礎(chǔ)上實現(xiàn)了令人吃驚的加速效果，主要歸功于使用傅里葉變換在頻域和時域(線性層可以看作對時域進(jìn)行建模)中交替進(jìn)行特征提取，并且傅里葉子層不需要可學(xué)習(xí)參數(shù)，相比于self attenion可以大大節(jié)省內(nèi)存占用和計算量。沒想到用傅里葉交換替換掉self attention能使其在transformer中發(fā)光發(fā)熱，不禁讓人覺得對于數(shù)學(xué)基礎(chǔ)的掌握至關(guān)重要！

作者 | 劉新宇 張譯 劉曉雯
單位 | 東北大學(xué)自然語言處理實驗室