稀疏模型在深度學(xué)習(xí)領(lǐng)域發(fā)揮著越來(lái)越重要的作用。對(duì)于給定的 token 或樣本，它可以只激活模型的一小部分，從而在擁有很大的參數(shù)量的同時(shí)也能做到計(jì)算友好。但是，如何可靠地訓(xùn)練這類模型依然是一個(gè)需要解決的問(wèn)題。在這篇文章中，來(lái)自谷歌的 Barret Zoph、Irwan Bello、William Fedus 、Jeff Dean 等研究者給出了一份「高效稀疏專家模型設(shè)計(jì)指南」。

稀疏專家神經(jīng)網(wǎng)絡(luò)展示了純規(guī)模的優(yōu)勢(shì)，并為當(dāng)今常用的靜態(tài)神經(jīng)網(wǎng)絡(luò)架構(gòu)提供了一種有效的替代方案。稀疏專家網(wǎng)絡(luò)不是對(duì)所有輸入應(yīng)用相同的參數(shù)，而是為每個(gè)輸入動(dòng)態(tài)選擇使用哪些參數(shù)。這允許網(wǎng)絡(luò)極大地?cái)U(kuò)展參數(shù)的數(shù)量，同時(shí)保持每個(gè) token 的 FLOPs 大致不變。這些方法的采用已經(jīng)帶來(lái)了 SOTA 翻譯模型、4-7 倍的預(yù)訓(xùn)練加速，以及僅使用 1/3 的訓(xùn)練成本就能達(dá)到 GPT-3 級(jí)的 one-shot 性能。盡管參數(shù)數(shù)量驚人，但稀疏模型將訓(xùn)練大型神經(jīng)網(wǎng)絡(luò)的碳足跡降低了一個(gè)數(shù)量級(jí)。然而，困難依然存在。

Fedus et al. (2021)觀察到，與之前的 SOTA 方法（Raffel et al., 2019）相比，稀疏 1.6T 參數(shù)模型實(shí)現(xiàn)了 4 倍的預(yù)訓(xùn)練加速，但在 SuperGLUE 等常用基準(zhǔn)上進(jìn)行微調(diào)時(shí)，卻落后于較小的模型。在 Artetxe et al. (2021)中，研究者在域外數(shù)據(jù)上對(duì) MoE 語(yǔ)言模型進(jìn)行了微調(diào)，并觀察到了相似的差距。

為了解決這一問(wèn)題，Switch-XXL 模型被提出，該模型參數(shù)較少，但計(jì)算占用空間增加到原來(lái)的 8 倍（FLOPs 大約等于最大的 T5 模型），在自然語(yǔ)言理解任務(wù)上的性能有所提高。然而，必要的預(yù)訓(xùn)練受到先前在小規(guī)模研究中未檢測(cè)到的訓(xùn)練不穩(wěn)定性的阻礙。這些不穩(wěn)定性后來(lái)在其他稀疏模型中被識(shí)別出來(lái)。這些結(jié)果揭示了參數(shù)和計(jì)算的必要平衡，但如何可靠地訓(xùn)練這種模型依然是一個(gè)待解決的問(wèn)題。

這篇論文的目的就是提高稀疏模型的實(shí)用性和可靠性。他們研究了這兩個(gè)問(wèn)題，并給出了設(shè)計(jì)指南。最后，他們將稀疏模型的參數(shù)縮放到 269B，其計(jì)算成本與 32B 密集編碼器 - 解碼器 Transformer（穩(wěn)定、可遷移的 Mixture-of-Experts、ST-MoE-32B）相當(dāng)。這是稀疏模型首次在遷移學(xué)習(xí)中實(shí)現(xiàn) SOTA 性能，跨越了一系列不同的任務(wù)，包括推理（SuperGLUE、ARC Easy、ARC Challenge)、摘要（XSum、CNN-DM）、閉卷問(wèn)答（WebQA、Natural Questions）和對(duì)抗式構(gòu)造任務(wù)（Winogrande、ANLI R3）。

本文的貢獻(xiàn)可以概括如下：

1、開展了一項(xiàng)關(guān)于穩(wěn)定性技術(shù)的質(zhì)量 - 穩(wěn)定性權(quán)衡（quality-stability trade-offs）大規(guī)模研究；2、引入了 router z-loss 來(lái)解決穩(wěn)定性問(wèn)題，同時(shí)略微提高了模型質(zhì)量；3、給出了關(guān)于稀疏和密集模型的微調(diào)分析，揭示了二者對(duì)批大小和學(xué)習(xí)率的不同超參數(shù)敏感性；他們發(fā)現(xiàn)，糟糕的超參數(shù)導(dǎo)致密集模型上幾乎沒(méi)有微調(diào)增益，盡管預(yù)訓(xùn)練有很大的加速；4、給出了分布式環(huán)境下設(shè)計(jì) Pareto 高效稀疏模型的架構(gòu)、routing 和模型設(shè)計(jì)原則；5、給出了追蹤跨專家層的 token routing 決策的定性分析；6、訓(xùn)練出了一個(gè) 269B 稀疏模型，在一組不同的自然語(yǔ)言基準(zhǔn)上實(shí)現(xiàn)了 SOTA 性能。

router z-loss

穩(wěn)定神經(jīng)網(wǎng)絡(luò)最成功的方法之一是對(duì)激活的約束和梯度。一種流行的方法是在通過(guò)深度網(wǎng)絡(luò)反向傳播時(shí)，裁剪梯度范數(shù)來(lái)彌補(bǔ)爆炸梯度。

在這篇論文中，研究者使用 Adafactor 優(yōu)化器是因?yàn)樗膬?nèi)存效率（盡管最近推出的 8 位優(yōu)化器 (Dettmers et al., 2021) 可能會(huì)提供更好的 trade-off）。Adafactor 使用更新裁剪（update clipping），而不是梯度裁剪（gradient clipping），其中對(duì)權(quán)重的更改被限制在一定的范數(shù)以下。他們嘗試將更新裁剪收緊到更小的值。

接下來(lái)，他們研究了即將進(jìn)入 router 的 logit 上的約束。router 以 float32 計(jì)算專家的概率分布。然而，研究者發(fā)現(xiàn)，在最大的規(guī)模下，這不足以帶來(lái)可靠的訓(xùn)練結(jié)果。為了解決這個(gè)問(wèn)題，他們引入了 router z-loss，

其中，B 是 token 的數(shù)目，N 是專家數(shù)，x ∈ RB×N 是將要進(jìn)入 router 的 logit。

下表 4 顯示，在三次運(yùn)行中，update clipping 和 router z-loss 都穩(wěn)定了模型，但是 update clipping 嚴(yán)重影響了模型的質(zhì)量。因此，研究者使用 z-loss 方法來(lái)固定模型穩(wěn)定性。

router z-loss 引入了另一個(gè)超參數(shù) (c_z)，這是一個(gè)加權(quán)系數(shù)，作為優(yōu)化的總損失的一部分。總損失是交叉熵?fù)p失 (cross entropy loss, L_CE)、輔助負(fù)載平衡損失 (auxiliary load balance loss, L_B) 和 router z-loss (L_Z) 的線性加權(quán)組合。

基于用超參數(shù)掃描進(jìn)行預(yù)訓(xùn)練后的最佳模型質(zhì)量，研究者選擇 c_z = 0.001 的值。附錄 B 記錄了預(yù)訓(xùn)練過(guò)程中的損失。

稀疏模型的設(shè)計(jì)

密集模型的設(shè)計(jì)受到 Kaplan et al. (2020) 的基礎(chǔ)工作的指導(dǎo)。但是到了稀疏模型這里，還有無(wú)數(shù)的額外問(wèn)題需要解決，比如：(1)使用多少專家？(2)使用哪種 routing 算法？(3)容量因子（capacity factor）的值是多少？(4)硬件如何改變這些決策？在本文中，研究者給出的建議是：

1、在他們的設(shè)置中，他們推薦容量因子為 1.25 的 top-2 routing，每個(gè)核心最多有一個(gè)專家；2、在評(píng)估期間，可以更改容量因子，以適應(yīng)新的內(nèi)存 / 計(jì)算要求；3、密集層疊加和乘法偏置（multiplicative bias）可以提高質(zhì)量。

 更多細(xì)節(jié)請(qǐng)參考原論文。