比Transformer好，無Attention、MLPs的BERT、GPT反而更強了
- WBOY轉載
- 2023-10-30 14:33:04669瀏覽
從BERT、GPT和Flan-T5等語言模型到SAM和Stable Diffusion等影像模型，Transformer正以迅猛之勢席捲全球，但人們不禁會問：Transformer是唯一的選擇嗎？
史丹佛大學和紐約州立大學布法羅分校的一個研究團隊不僅為這個問題給出了否定答案，而且還提出了一種新的替代技術：Monarch Mixer 。近日，團隊在 arXiv 公佈了相關論文和一些檢查點模型及訓練程式碼。順帶一提，該論文已入選 NeurIPS 2023 並獲得 Oral Presentation 資格。
論文連結：https://arxiv.org/abs/2310.12109
GitHub上的程式碼位址為：https://github.com/HazyResearch/m2
該方法去掉了Transformer 中高成本的注意力和MLP，代之以富有表現力的Monarch 矩陣，使之在語言和圖像實驗中以更低的成本取得了更優的表現。
這並不是史丹佛大學第一次提出Transformer的替代技術。今年六月該校的另一個團隊也曾提出過一種名為Backpack的技術，參考機器之心文章《史丹佛訓練Transformer替代模型：1.7億參數，能除偏、可控可解釋性強》。當然，這些技術要取得真正的成功，還需要研究社群的進一步檢驗並在應用開發者手中變成切實好用的產品
下面我們看看這篇論文中對Monarch Mixer 的介紹以及一些實驗結果。
論文介紹
在自然語言處理和電腦視覺領域，機器學習模型已能處理更長的序列和更高維度的表徵，從而支援更長的上下文和更高的品質。然而，現有架構的時間和空間複雜度在序列長度和 / 或模型維度上呈現二次成長模式，這會限制上下文長度並提升擴展成本。舉個例子，Transformer 中的注意力和 MLP 會隨序列長度和模型維度呈現二次擴展模式。
針對這個問題，史丹佛大學和紐約州立大學布法羅分校的這個研究團隊聲稱找到了一個高效能的架構，其複雜度隨序列長度和模型維度的成長是次二次的（sub-quadratic）。
他們的研究靈感來自MLP-mixer和ConvMixer。這兩項研究觀察到：許多機器學習模型以序列和模型維度為軸對資訊進行混合，並且通常使用單一算子對這兩個軸進行操作
尋找表現力強、次二次且硬體效率高的混合算子的難度很高。舉個例子，MLP-mixer 中的 MLP 和 ConvMixer 中的捲積都頗具表現力，但它們都會隨輸入維度二次擴展。近期有一些研究提出了一些次二次的序列混合方法，這些方法使用了較長的捲積或狀態空間模型，而且它們都會用到FFT，但這些模型的FLOP 利用率很低並且在模型維度方面依然是二次擴展。同時，不損品質的稀疏密集 MLP 層方面也有一些頗具潛力的進展，但由於硬體利用率較低，某些模型實際上可能還比密集模型更慢。
根據這些靈感，研究團隊提出了Monarch Mixer (M2)，它使用了一個富有表現力的次二次結構化矩陣：Monarch 矩陣
Monarch 矩陣是一種廣義的快速傅立葉變換（FFT）結構矩陣，研究顯示它包含了多種線性變換，如哈達瑪變換、托普利茲矩陣、AFDF 矩陣和卷積等。這些矩陣可以透過分塊對角矩陣的乘積來參數化，這些參數被稱為Monarch 因子，並與排列交織相關
它們的計算是次二次擴展的：如果將因數的數量設為p，則當輸入長度為N 時，計算複雜度為 
，從而讓計算複雜度可以位於p = log N 時的O (N log N)與p = 2 時的 之間。
M2 使用了 Monarch 矩陣來沿著序列和模型維度軸混合資訊。這種方法不僅易於實現，而且硬體效率也很高：使用支援 GEMM（廣義矩陣乘法演算法）的現代硬體就能有效地計算分塊對角 Monarch 因子。
該研究團隊透過使用PyTorch編寫程式碼，在不到40行的情況下實現了一個M2層，並且只依賴矩陣乘法、轉置、reshape和逐元素乘積（見圖1中部的偽代碼）。對於大小為64k的輸入，在一台A100 GPU上，這些程式碼實現了25.6%的FLOP利用率。在更新的架構如RTX 4090上，對於相同大小的輸入，一個簡單的CUDA實作能夠實現41.4%的FLOP利用率
##有關Monarch Mixer 的更多數學描述和理論分析請參考原文。
實驗
研究團隊對Monarch Mixer 和Transformer 這兩個模型進行了比較，主要針對Transformer 在三個主要任務中佔據主導地位的情況進行了研究。這三個任務分別是：BERT 風格的非因果掩碼語言建模任務、ViT 風格的圖像分類任務以及GPT 風格的因果語言建模任務
在每個任務上，實驗結果顯示新提出的方法在不使用注意力和MLP 的前提下均能達到與Transformer 相媲美的水平。他們還在BERT 設定中評估了新方法相較於強大Transformer 基準模型的加速情況
非因果語言建模的需要進行重寫
對於非因果語言建模的需要進行重寫任務，該團隊建立了基於M2 的架構：M2-BERT。 M2-BERT 可以直接取代 BERT 風格的語言模型，而 BERT 是 Transformer 架構的一大主力應用。對於 M2-BERT 的訓練，使用了在 C4 上的遮罩語言建模，token 化器則是 bert-base-uncased。
M2-BERT 是基於Transformer 骨幹，但M2 層取代了其中的注意力層和MLP，如圖3 所示
在序列混合器中，注意力被帶有殘差卷積的雙向門控卷積取代（見圖3 左側）。為了恢復卷積，團隊將 Monarch 矩陣設定為 DFT 和逆 DFT 矩陣。他們還在投射步驟之後添加了逐深度的捲積。
在維度混合器中，MLP的兩個密集矩陣被替換為學習得到的分塊對角矩陣（Monarch矩陣的階數為1，b=4）
研究者進行了預訓練，共得到了4個M2-BERT模型：其中兩個是大小分別為80M和110M的M2-BERT-base模型，另外兩個是大小分別為260M和341M的M2-BERT-large模型。這些模型分別相當於BERT-base和BERT-large
表3 給出了相當於BERT-base 的模型的效能表現，表4 給出了相當於BERT-large的模型的性能表現。
從表中可以看到，在GLUE 基準上，M2-BERT-base 的表現可以媲美BERT-base，同時參數還少了27%；而當兩者參數數量相當時，M2-BERT-base 勝過BERT-base 1.3 分。類似地，參數少 24% 的 M2-BERT-large 與 BERT-large 表現相當，而參數數量一樣時，M2-BERT-large 有 0.7 分的優勢。
表格5展示了與BERT-base模型相當的模型的前向吞吐量狀況。報告的是在A100-40GB GPU上每毫秒處理的token數量，這可以反映推理時間
當序列較短時，資料局部性的影響依然主導著 FLOP 的減少情況，而過濾器產生（BERT 中沒有）等操作的成本更高。而當序列長度超過 1K 時，M2-BERT-base 的加速優勢就漸漸起來了，當序列長度達 8K 時，速度優勢可達 6.5 倍。
影像分類
#為驗證新方法在影像領域中的優勢是否與語言領域中的優勢一樣，該團隊還評估了M2 在圖像分類任務方面的表現，這是在非因果建模方面進行的
表7 給出了Monarch Mixer、ViT-b、HyenaViT- b 和ViT-b-Monarch（以Monarch 矩陣取代了標準ViT-b 中的MLP 模組）在ImageNet-1k 上的效能表現。
Monarch Mixer的優點非常明顯：它只需要一半的參數量，就能超過原始的ViT-b模型。令人驚訝的是，參數較少的Monarch Mixer甚至能夠超過ResNet-152，而ResNet-152是專門為ImageNet任務設計的
GPT 風格的因果語言建模是Transformer 的重要應用。團隊為因果語言建模開發了一個基於M2 的架構，稱為M2-GPT
對於序列混合器，M2-GPT 組合使用了來自Hyena 的捲積過濾器、目前最佳的無注意力語言模型以及來自H3 的跨多頭參數共享。他們使用因果參數化來取代了這些架構中的 FFT，並完全移除了 MLP 層。所得到的架構完全沒有註意力，也完全沒有 MLP。
他們在因果語言建模的標準資料集 PILE 上對 M2-GPT 進行了預訓練。結果見表 8。
可以看到，儘管基於新架構的模型完全沒有註意力和MLP，但在預訓練的困惑度指標上依然勝過Transformer 和Hyena。這些結果表明，與 Transformer 大不相同的模型也可能在因果語言建模方面取得出色表現。
請參考原始論文以取得更詳細的內容
以上是比Transformer好，無Attention、MLPs的BERT、GPT反而更強了的詳細內容。更多資訊請關注PHP中文網其他相關文章！