作為基于注意力模型的替代方案，純 MLP 架構(gòu)吸引了越來越多的關注。在 NLP 中，gMLP 等近期工作表明，純 MLP 在語言建模方面可以達到與 transformer 相當?shù)男阅埽谙掠稳蝿罩腥跤?transformer。

來自 Meta AI 和紐約州立大學布法羅分校的研究者分析了 MLP 在表達能力方面的局限性，并提出了在特征和輸入（token）維度上帶有混合專家系統(tǒng)（MoE）的稀疏激活 MLP。這種稀疏的純 MLP 顯著提高了模型容量和表達能力，同時保持計算不變。該研究解決了將條件計算與兩種路由策略結(jié)合起來的關鍵挑戰(zhàn)。

論文地址：https://arxiv.org/pdf/2203.06850.pdf

與基于 transformer 的 MoE、密集 Transformer 和純 MLP 相比，該研究提出的稀疏純 MLP（sMLP） 改進了語言建模的困惑度并獲得了高達 2 倍的訓練效率提升。最后，研究者在六個下游任務上評估了稀疏純 MLP 的零樣本上下文學習性能，發(fā)現(xiàn)它超過了基于 transformer 的 MoE 和密集的 transformer。

方法

稀疏激活的純 MLP

sMLP 的整體架構(gòu)如下圖 2 所示，包含 N_1 個密集塊和 N_2 個稀疏塊。N_1 和 N_2 都是超參數(shù)。每個稀疏塊包含兩個模塊：

• tMoE 模塊：該研究采用 Base Layers 的 MoE (Lewis et al., 2021) 來替換密集 transformer 中的 FFN 模塊 (Vaswani et al., 2017b)；
• sMoE 模塊：該研究設計了 sMoE 模塊來替代 transformer 中的自注意力模塊（Vaswani et al., 2017b）和 gMLP 中的空間門控單元（SGU，Liu et al., 2021a）。

tMoE 模塊和 sMoE 模塊都包含兩個元素：

• 專家模塊

專家模塊負責處理輸入。對于 tMoE 模塊，每個專家都包含一個 FFN，如上圖 2 所示。對于 sMoE 模塊，每個專家都包含空間門控單元，如下圖 6（右）所示。

• 門控函數(shù)

該模塊決定哪個專家應該處理輸入的每個部分，因此需要設計一種獨特的路由方法來將 MoE 結(jié)構(gòu)擴展到特征維度。

圖 3（左）顯示了現(xiàn)有基于 transformer 的 MoE 的門控函數(shù)示例（Lepikhin et al., 2020；Fedus et al., 2021；Lewis et al., 2021；Roller et al., 2021）。x_ij 表示 i_th token 中 j_th 隱藏維度的值。 

如下等式 (3) 所示：

tMoE 使用由參數(shù)化的等式（3）中描述的學習門控函數(shù)將這 4 個 token 發(fā)送給 FFN 層的 3 個專家。與已有的一些 MoE 不同，在稀疏純 MLP 架構(gòu)中，該研究提出沿隱藏維度對隱藏表示進行分塊，并將分塊向量發(fā)送給不同的專家，如圖 3（右）所示。

與這些現(xiàn)有的 MoE 不同，在稀疏的全 MLP 架構(gòu)中，該研究提出沿隱藏維度對隱藏表示進行分塊，并將分塊向量發(fā)送給不同的專家，如圖 3（右）所示。

特征空間中的路由

與路由 token 相比，路由隱藏維度在自回歸模型中面臨著一個獨特的挑戰(zhàn)，如果簡單地預測未來的 token，信息會泄漏。此外，與具有 selfattention 的基于 Transformers 的 MoE 不同，此處不能直接應用適當?shù)难诖a來防止信息泄露，因此不能在基于 transformer 的 MoE 中采用現(xiàn)有的路由方法進行語言建模。該研究比較了以下兩種解決方案：確定性路由（deterministic routing）和部分預測（partial prediction）。

實驗及結(jié)果

token 操作比較

該研究將 sMLP 模型與兩個密集模型進行比較：Transformer (Vaswani et al., 2017b) 和 gMLP (Liu et al., 2021a)。基于全 MLP 和基于 transformer 的模型之間的主要區(qū)別在于 token 操作。該研究比較了這三種 token-wise 操作：Transformers 中的 self-attention 模塊、gMLP 中的 Spatial Gating Unit 和 sMLP 模型中的 sMoE 模塊。表 3 比較了三種 token 操作及其各自的頭部機制：

下圖 4 將模型與不同頭數(shù)的密集模型進行了比較。Transformer 模型極大地受益于多頭機制。然而，gMLP 模型雖然增加了參數(shù)量，但并沒有通過多頭機制提高性能。sMLP 模型也可以看作是 gMLP 的一種多頭解決方案，顯著提高了基于 MLP 模型的性能，并且優(yōu)于 transformer 模型。

稀疏 MLP 的結(jié)果

下圖 5 給出了質(zhì)量（有效困惑度）和訓練效率，通過訓練步驟數(shù)（頂部）和訓練時間（底部）來衡量。研究者發(fā)現(xiàn)，具有兩種路由策略變體的 sMLP 優(yōu)于具有大致相同數(shù)量的 FLOP 的最先進的基于 Transformer 的 MoE 模型。

下表 4 總結(jié)了主要實驗中的詳細比較結(jié)果。研究者將所有模型的 FLOPs 控制為約 0.8T。除了模型層數(shù)不同，它們的嵌入維數(shù)為 1024，隱藏維數(shù)為 4096。可以看到，sMLP 模型在 25k 訓練步驟時實現(xiàn)了最好的泛化，同時實現(xiàn)了最高的訓練速度。HASH 層在所有 Transformer 基線中具有最佳性能，并且需要的時間最少。

擴展

為了測試模型的可擴展性，該研究增加了 2.0 TFLOPs 的模型大小訓練。表 4（底部）總結(jié)了結(jié)果。

與表 4（頂部）中的模型相比，該研究擴大了所有模型，將嵌入從 1024 更改為 2048，并將隱藏維度從 4096 調(diào)整為 8192，如表 5 所示。該研究還增加了預訓練數(shù)據(jù)大小，如表 2 所示。

? ?