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

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

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

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

方法

稀疏激活的純 MLP

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

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

tMoE 模塊和 sMoE 模塊都包含兩個(gè)元素：

• 專家模塊

專家模塊負(fù)責(zé)處理輸入。對(duì)于 tMoE 模塊，每個(gè)專家都包含一個(gè) FFN，如上圖 2 所示。對(duì)于 sMoE 模塊，每個(gè)專家都包含空間門(mén)控單元，如下圖 6（右）所示。

• 門(mén)控函數(shù)

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

圖 3（左）顯示了現(xiàn)有基于 transformer 的 MoE 的門(mén)控函數(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）中描述的學(xué)習(xí)門(mén)控函數(shù)將這 4 個(gè) token 發(fā)送給 FFN 層的 3 個(gè)專家。與已有的一些 MoE 不同，在稀疏純 MLP 架構(gòu)中，該研究提出沿隱藏維度對(duì)隱藏表示進(jìn)行分塊，并將分塊向量發(fā)送給不同的專家，如圖 3（右）所示。

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

特征空間中的路由

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

實(shí)驗(yàn)及結(jié)果

token 操作比較

該研究將 sMLP 模型與兩個(gè)密集模型進(jìn)行比較：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 操作及其各自的頭部機(jī)制：

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

稀疏 MLP 的結(jié)果

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

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

擴(kuò)展

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

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

? ?