從 BERT、GPT 和 Flan-T5 等語言模型到 SAM 和 Stable Diffusion 等圖像模型，Transformer 正以銳不可當(dāng)之勢席卷這個(gè)世界，但人們也不禁會問：Transformer 是唯一選擇嗎？

斯坦福大學(xué)和紐約州立大學(xué)布法羅分校的一個(gè)研究團(tuán)隊(duì)不僅為這一問題給出了否定答案，而且還提出了一種新的替代技術(shù)：Monarch Mixer。近日，該團(tuán)隊(duì)在 arXiv 公布了相關(guān)論文和一些檢查點(diǎn)模型及訓(xùn)練代碼。順帶一提，該論文已入選 NeurIPS 2023 并獲得 Oral Presentation 資格。

論文地址：https://arxiv.org/abs/2310.12109

代碼地址：https://github.com/HazyResearch/m2

該方法去掉了 Transformer 中高成本的注意力和 MLP，代之以富有表現(xiàn)力的 Monarch 矩陣，使之在語言和圖像實(shí)驗(yàn)中以更低的成本取得了更優(yōu)的表現(xiàn)。

這并不是斯坦福大學(xué)第一次提出 Transformer 的替代技術(shù)。今年六月該校的另一個(gè)團(tuán)隊(duì)還曾提出過一種名為 Backpack 的技術(shù)，參閱機(jī)器之心文章《斯坦福訓(xùn)練 Transformer 替代模型：1.7 億參數(shù)，能除偏、可控可解釋性強(qiáng)》。當(dāng)然，這些技術(shù)要取得真正的成功，還需要研究社區(qū)的進(jìn)一步檢驗(yàn)并在應(yīng)用開發(fā)者手中變成切實(shí)好用的產(chǎn)品。

下面我們看看這篇論文中對 Monarch Mixer 的介紹以及一些實(shí)驗(yàn)結(jié)果。

論文介紹

在自然語言處理和計(jì)算機(jī)視覺領(lǐng)域，機(jī)器學(xué)習(xí)模型已能處理更長的序列和更高維度的表征，從而支持更長的上下文和更高的質(zhì)量。然而，現(xiàn)有架構(gòu)的時(shí)間和空間復(fù)雜性在序列長度和 / 或模型維度上呈二次增長模式，這會限制上下文長度并提升擴(kuò)展成本。舉個(gè)例子，Transformer 中的注意力和 MLP 會隨序列長度和模型維度呈二次擴(kuò)展模式。

針對這一問題，斯坦福大學(xué)和紐約州立大學(xué)布法羅分校的這個(gè)研究團(tuán)隊(duì)聲稱找到了一種高性能的架構(gòu)，其復(fù)雜度隨序列長度和模型維度的增長是次二次的（sub-quadratic）。

他們的研究靈感來自 MLP-mixer 和 ConvMixer；這兩項(xiàng)研究觀察到：許多機(jī)器學(xué)習(xí)模型的運(yùn)作方式都是沿序列和模型維度軸對信息進(jìn)行混合，并且它們往往對兩個(gè)軸使用了單個(gè)算子。

尋找表現(xiàn)力強(qiáng)、次二次且硬件效率高的混合算子的難度很大。舉個(gè)例子，MLP-mixer 中的 MLP 和 ConvMixer 中的卷積都頗具表現(xiàn)力，但它們都會隨輸入維度二次擴(kuò)展。近期有一些研究提出了一些次二次的序列混合方法，這些方法使用了較長的卷積或狀態(tài)空間模型，而且它們都會用到 FFT，但這些模型的 FLOP 利用率很低并且在模型維度方面依然是二次擴(kuò)展。與此同時(shí)，不損質(zhì)量的稀疏密集 MLP 層方面也有一些頗具潛力的進(jìn)展，但由于硬件利用率較低，某些模型實(shí)際上可能還比密集模型更慢。

基于這些靈感，這個(gè)研究團(tuán)隊(duì)提出了 Monarch Mixer (M2)，其使用到了一類富有表現(xiàn)力的次二次結(jié)構(gòu)化矩陣：Monarch 矩陣。

Monarch 矩陣是一類泛化了快速傅立葉變換（FFT）的結(jié)構(gòu)化矩陣，并且研究表明其涵蓋了范圍廣泛的線性變換，包括哈達(dá)瑪變換、托普利茲矩陣、AFDF 矩陣和卷積。它們可通過分塊對角矩陣的積進(jìn)行參數(shù)化，這些參數(shù)被稱為 Monarch 因子，與排列交織。

它們的計(jì)算是次二次擴(kuò)展的：如果將因子的數(shù)量設(shè)為 p，則當(dāng)輸入長度為 N 時(shí)，計(jì)算復(fù)雜度為 ，從而讓計(jì)算復(fù)雜度可以位于 p = log N 時(shí)的 O (N log N) 與 p = 2 時(shí)的 之間。

M2 使用了 Monarch 矩陣來沿序列和模型維度軸混合信息。這種方法不僅易于實(shí)現(xiàn)，而且硬件效率也很高：使用支持 GEMM（廣義矩陣乘法算法）的現(xiàn)代硬件就能高效地計(jì)算分塊對角 Monarch 因子。

該研究團(tuán)隊(duì)實(shí)現(xiàn)了一個(gè) M2 層來進(jìn)行概念驗(yàn)證 —— 完全使用 PyTorch 編寫，代碼行數(shù)不到 40（包括 import 軟件包），而且其只需依賴矩陣乘法、轉(zhuǎn)置、reshape 和逐元素乘積（見圖 1 中部的偽代碼）；結(jié)果，對于大小為 64k 的輸入，這些代碼在一臺 A100 GPU 上實(shí)現(xiàn)了 25.6% 的 FLOP 利用率。在 RTX 4090 等更新的架構(gòu)上，對于同樣大小的輸入，一個(gè)簡單的 CUDA 實(shí)現(xiàn)就能實(shí)現(xiàn) 41.4% 的 FLOP 利用率。

有關(guān) Monarch Mixer 的更多數(shù)學(xué)描述和理論分析請參看原論文。

實(shí)驗(yàn)

該研究團(tuán)隊(duì)在 Transformer 已占主導(dǎo)地位的三個(gè)任務(wù)上對 Monarch Mixer 和 Transformer 進(jìn)行了比較：BERT 風(fēng)格的非因果掩碼語言建模任務(wù)、ViT 風(fēng)格的圖像分類任務(wù)、GPT 風(fēng)格的因果語言建模任務(wù)。

在每個(gè)任務(wù)上，實(shí)驗(yàn)結(jié)果表明新提出的方法在不使用注意力和 MLP 的前提下均能達(dá)到與 Transformer 相媲美的水平。他們還在 BERT 設(shè)置中評估了新方法相較于強(qiáng)大 Transformer 基準(zhǔn)模型的加速情況。

非因果語言建模

對于非因果語言建模任務(wù)，該團(tuán)隊(duì)構(gòu)建了一種基于 M2 的架構(gòu)：M2-BERT。M2-BERT 可以直接替代 BERT 風(fēng)格的語言模型，而 BERT 是 Transformer 架構(gòu)的一大主力應(yīng)用。對于 M2-BERT 的訓(xùn)練，使用了在 C4 上的掩碼語言建模，token 化器則是 bert-base-uncased。

M2-BERT 基于 Transformer 骨干，但其中的注意力層和 MLP 被 M2 層替換，如圖 3 所示。

在序列混合器中，注意力被帶殘差卷積的雙向門控卷積替代（見圖 3 左側(cè)）。為了恢復(fù)卷積，該團(tuán)隊(duì)將 Monarch 矩陣設(shè)置為 DFT 和逆 DFT 矩陣。他們還在投射步驟之后添加了逐深度的卷積。

在維度混合器中，MLP 中兩個(gè)密集矩陣被替換成了學(xué)習(xí)得到的分塊對角矩陣（1 階 Monarch 矩陣，b = 4）。

研究者預(yù)訓(xùn)練了 4 個(gè) M2-BERT 模型：其中兩個(gè)是大小分別為 80M 和 110M 的 M2-BERT-base 模型，另外兩個(gè)是大小分別為 260M 和 341M 的 M2-BERT-large 模型。它們分別相當(dāng)于 BERT-base 和 BERT-large。

表 3 給出了相當(dāng)于 BERT-base 的模型的性能表現(xiàn)，表 4 給出了相當(dāng)于 BERT-large 的模型的性能表現(xiàn)。

從表中可以看到，在 GLUE 基準(zhǔn)上，M2-BERT-base 的表現(xiàn)可以媲美 BERT-base，同時(shí)參數(shù)還少了 27%；而當(dāng)兩者參數(shù)數(shù)量相當(dāng)時(shí)，M2-BERT-base 勝過 BERT-base 1.3 分。類似地，參數(shù)少 24% 的 M2-BERT-large 與 BERT-large 表現(xiàn)相當(dāng)，而參數(shù)數(shù)量一樣時(shí)，M2-BERT-large 有 0.7 分的優(yōu)勢。

表 5 給出了相當(dāng)于 BERT-base 的模型的前向吞吐量情況。其中報(bào)告的是在 A100-40GB GPU 上每毫秒處理的 token 數(shù)，這能反映推理時(shí)間。

可以看到，M2-BERT-base 的吞吐量甚至超過了經(jīng)過高度優(yōu)化的 BERT 模型；相較于在 4k 序列長度上的標(biāo)準(zhǔn) HuggingFace 實(shí)現(xiàn)，M2-BERT-base 的吞吐量可達(dá)其 9.1 倍！

表 6 則報(bào)告了 M2-BERT-base (80M) 和 BERT-base 的 CPU 推理時(shí)間 —— 結(jié)果是直接運(yùn)行這兩個(gè)模型的 PyTorch 實(shí)現(xiàn)得到的。

當(dāng)序列較短時(shí)，數(shù)據(jù)局部性的影響依然主導(dǎo)著 FLOP 的減少情況，而過濾器生成（BERT 中沒有）等操作的成本更高。而當(dāng)序列長度超過 1K 時(shí)，M2-BERT-base 的加速優(yōu)勢就漸漸起來了，當(dāng)序列長度達(dá) 8K 時(shí)，速度優(yōu)勢可達(dá) 6.5 倍。

圖像分類

在非因果建模方面，為了驗(yàn)證新方法在圖像上也有在語言上一樣的優(yōu)勢，該團(tuán)隊(duì)還評估了 M2 在圖像分類任務(wù)上的表現(xiàn)。

表 7 給出了 Monarch Mixer、ViT-b、HyenaViT-b 和 ViT-b-Monarch（用 Monarch 矩陣替換了標(biāo)準(zhǔn) ViT-b 中的 MLP 模塊）在 ImageNet-1k 上的性能表現(xiàn)。

Monarch Mixer 優(yōu)勢非常明顯：只需一半的參數(shù)量，其表現(xiàn)就能勝過原始 ViT-b 模型。而更讓人驚訝的是，參數(shù)更少的 Monarch Mixer 很能勝過 ResNet-152；要知道，ResNet-152 可是專門針對 ImageNet 任務(wù)設(shè)計(jì)的。

因果語言建模

GPT 風(fēng)格的因果語言建模是 Transformer 的一大關(guān)鍵應(yīng)用。該團(tuán)隊(duì)為因果語言建模構(gòu)建了一個(gè)基于 M2 的架構(gòu)：M2-GPT。

對于序列混合器，M2-GPT 組合使用了來自 Hyena 的卷積過濾器、當(dāng)前最佳的無注意力語言模型以及來自 H3 的跨多頭參數(shù)共享。他們使用因果參數(shù)化替換了這些架構(gòu)中的 FFT，并完全移除了 MLP 層。所得到的架構(gòu)完全沒有注意力，也完全沒有 MLP。

他們在因果語言建模的標(biāo)準(zhǔn)數(shù)據(jù)集 PILE 上對 M2-GPT 進(jìn)行了預(yù)訓(xùn)練。結(jié)果見表 8。

可以看到，盡管基于新架構(gòu)的模型完全沒有注意力和 MLP，但其在預(yù)訓(xùn)練的困惑度指標(biāo)上依然勝過 Transformer 和 Hyena。這些結(jié)果表明，與 Transformer 大不相同的模型也可能在因果語言建模取得出色表現(xiàn)。

了解更多內(nèi)容，請參考原論文。