谷歌、DeepMind提出高效Transformer評估基準

來自:機器之心

自誕生以來,Transformer 在不同領域得到了廣泛應用,研究人員也提出了許多高效 Transformer 模型。如何評估這類模型呢?最近,谷歌和 DeepMind 提出了一項系統化的統一基準——Long-Range Arena,重點關注長語境場景下的模型質量評估。

谷歌、DeepMind提出高效Transformer評估基準

基準專案地址:https://github。com/google-research/long-range-arena

論文地址:https://arxiv。org/pdf/2011。04006。pdf

Transformer 在多個模態(語言、影象、蛋白質序列)中獲得了 SOTA 結果,但它存在一個缺點:自注意力機制的平方級複雜度限制了其在長序列領域中的應用。目前,研究人員提出大量高效 Transformer 模型(「xformer」),試圖解決該問題。其中很多展示出了媲美原版 Transformer 的效能,同時還能有效降低自注意力機制的記憶體複雜度。

谷歌和 DeepMind 的研究人員對比了這些論文的評估和實驗設定,得到了以下幾點發現:

首先,高效 Transformer 缺少統一的基準測試,使用的任務型別也多種多樣:每個模型在不同的任務和資料集上進行評估。

其次,評估所用基準通常是隨意選擇的,未充分考慮該任務是否適用於長程建模評估。

第三,很多論文將歸納偏置的效果和預訓練的優點混為一談,這會模糊模型的真正價值:預訓練本身是計算密集型的,將歸納偏置和預訓練分離開來可降低 xformer 研究的門檻。

於是,谷歌和 DeepMind 的研究人員提出了一個新基準 Long-Range Arena (LRA),用來對長語境場景下的序列模型進行基準測試。該基準包括合成任務和現實任務,研究人員在此基準上對比了十個近期提出的高效 Transformer 模型,包括 Sparse Transformers、Reformer、Linformer、Longformer、Sinkhorn Transformer、Performer、Synthesizer、Linear Transformer 和 BigBird 模型。

該基準主要關注模型在長語境場景下的能力,不過研究人員對 xformer 架構在不同資料型別和條件下的能力也很感興趣。因此,該基準選擇了具備特定先驗結構的資料集和任務。例如,這些架構可以建模層級結構長序列或包含某種空間結構形式的長序列嗎?這些任務的序列長度從 1K 到 16K token 不等,還包括大量資料型別和模態,如文字、自然影象、合成影象,以及需要類似度、結構和視覺 - 空間推理的數學表示式。該基準主要面向高效 Transformer,但也可作為長程式列建模的基準。

除了對比模型質量以外,該研究還進行了大量效率和記憶體使用分析。研究者認為,並行效能基準測試對於社群是有益且珍貴的,能夠幫助大家深入瞭解這些方法的實際效率。總之,該研究提出了一個統一框架,既能對高效 Transformer 模型進行簡單的並行對比分析,還能對長程式列模型進行基準測試。該框架使用 JAX/FLAX1 編寫。

高效 Transformer 評估新基準:Long-Range Arena (LRA)

基準需求

在建立 LRA 基準之前,研究者先列舉了一些需求:

1。 通用性:適用於所有高效 Transformer 模型。例如,並非所有 xformer 模型都能執行自迴歸解碼,因此該基準中的任務僅需要編碼。

2。 簡潔性:任務設定應簡單,移除所有令模型對比複雜化的因素,這可以鼓勵簡單模型而不是笨重的 pipeline 方法。

3。 挑戰性:任務應該對目前模型有一定難度,以確保未來該方向的研究有足夠的進步空間。

4。 長輸入:輸入序列長度應該足夠長,因為評估不同模型如何捕獲長程依賴是 LRA 基準的核心關注點。

5。 探索不同方面的能力:任務集合應當評估模型的不同能力,如建模關係和層級 / 空間結構、泛化能力等。

6。 非資源密集、方便使用:基準應該是輕量級的,方便不具備工業級計算資源的研究者使用。

任務

LRA 基準包含多項任務,旨在評估高效 Transformer 模型的不同能力。具體而言,這些任務包括:Long ListOps、位元級文字分類、位元級文件檢索、基於畫素序列的影象分類、Pathfinder(長程空間依賴性)、Pathfinder-X(極端長度下的長程空間依賴性)。

LRA 任務所需的注意力範圍

LRA 基準的主要目標之一是評估高效 Transformer 模型捕獲長程依賴的能力。為了對注意力機制在編碼輸入時需要考慮的空間範圍進行量化估計,該研究提出了「所需注意力範圍」(required attention span)。給出一個注意力模型和輸入 token 序列,注意力模組的所需注意力範圍是 query token 和 attended token 間的平均距離。

谷歌、DeepMind提出高效Transformer評估基準

圖 2 總結了 LRA 基準中每項任務的所需注意力範圍,從圖中可以看出每項任務的所需注意力範圍都很高。這表明,Transformer 模型不僅僅涉及區域性資訊,在很多工和資料集中,注意力機制通常需要結合鄰近位置的資訊。

實驗

量化結果

實驗結果表明,LRA 中的所有任務都具備一定的挑戰性,不同 xformer 模型的效能存在一定程度的差異。具體結果參見下表 1:

谷歌、DeepMind提出高效Transformer評估基準

效率基準

表 2 展示了 xformer 模型的效率基準測試結果:

谷歌、DeepMind提出高效Transformer評估基準

從中可以看出,低秩模型和基於核的模型通常速度最快。整體最快的模型是 Performer,在 4k 序列長度上的速度是 Transformer 的 5。7 倍,Linformer 和 Linear Transformer 緊隨其後。最慢的模型是 Reformer,在 4k 序列長度上的速度是 Transformer 的 80%,在 1k 序列長度上的速度是 Transformer 的一半。

此外,研究者還評估了這些模型的記憶體消耗情況。結果顯示,記憶體佔用最少的模型是 Linformer,在 4k 序列長度上只使用了 0。99GB per TPU,而原版 Transformer 使用了 9。48GB per TPU,記憶體佔用減少了約 90%。

整體結果:不存在萬能模型

根據研究人員的分析,在 LRA 所有任務中整體效能最好(LRA 分數最高)的模型是 BigBird。但是,BigBird 在每項任務中的效能均不是最好,它只是在所有任務上都能取得不錯的效能。Performer 和 Linear Transformer 在一些任務中表現搶眼,但其平均分被 ListOps 任務拖累。

下圖 3 展示了模型效能、速度和記憶體佔用之間的權衡情況。BigBird 效能最好,但速度幾乎與原版 Transformer 相同。而 Local Attention 模型速度很快,但效能較低。在這些模型中,基於核的模型(如 Performer、Linformer 和 Linear Transformer)能夠在速度和效能之間獲得更好的折中效果,同時記憶體佔用也較為合理。

谷歌、DeepMind提出高效Transformer評估基準

谷歌、DeepMind提出高效Transformer評估基準