純MLP在下游任務上欠佳？Meta AI等提出稀疏MLP，超越transformer

作者：機器之心 2022-04-18 15:56:49

來自 Meta AI 和紐約州立大學布法羅分校的研究者分析了 MLP 在表達能力方面的局限性，并提出了在特征和輸入（token）維度上帶有混合專家系統（MoE）的稀疏激活 MLP。

作為基于注意力模型的替代方案，純 MLP 架構吸引了越來越多的關注。在 NLP 中，gMLP 等近期工作表明，純 MLP 在語言建模方面可以達到與 transformer 相當的性能，但在下游任務中弱于 transformer。

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

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

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

方法

稀疏激活的純 MLP

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

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（右）所示。

門控函數

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

圖 3（左）顯示了現有基于 transformer 的 MoE 的門控函數示例（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 使用由參數化的等式（3）中描述的學習門控函數將這 4 個 token 發送給 FFN 層的 3 個專家。與已有的一些 MoE 不同，在稀疏純 MLP 架構中，該研究提出沿隱藏維度對隱藏表示進行分塊，并將分塊向量發送給不同的專家，如圖 3（右）所示。

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

特征空間中的路由

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

實驗及結果

token 操作比較

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

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

稀疏 MLP 的結果

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

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