單一作者論文，谷歌提出百萬(wàn)專(zhuān)家Mixture，超越密集前饋、稀疏MoE

作者：機(jī)器之心 2024-07-10 12:32:48

釋放進(jìn)一步擴(kuò)展 Transformer 的潛力，同時(shí)還可以保持計(jì)算效率。

標(biāo)準(zhǔn) Transformer 架構(gòu)中的前饋（FFW）層會(huì)隨著隱藏層寬度的增加而導(dǎo)致計(jì)算成本和激活內(nèi)存的線(xiàn)性增加。在大語(yǔ)言模型（LLM）體量不斷增大的現(xiàn)在，稀疏混合專(zhuān)家（MoE）架構(gòu)已成為解決此問(wèn)題的可行方法，它將模型大小與計(jì)算成本分離開(kāi)來(lái)。很多新興的 MoE 模型都可以實(shí)現(xiàn)相同體量之上，更好的性能與更強(qiáng)大的表現(xiàn)。

最近發(fā)現(xiàn)的細(xì)粒度 MoE 擴(kuò)展定律表明，更高的粒度可帶來(lái)更好的性能。然而由于計(jì)算和優(yōu)化方面的挑戰(zhàn)，現(xiàn)有的 MoE 模型僅限于低數(shù)量專(zhuān)家。

本周二，Google DeepMind 的新研究引入了一種參數(shù)高效的專(zhuān)家檢索機(jī)制，其利用乘積密鑰技術(shù)從一百萬(wàn)個(gè)微型專(zhuān)家中進(jìn)行稀疏檢索。

鏈接：https://arxiv.org/abs/2407.04153

該方法嘗試通過(guò)用于路由的學(xué)習(xí)索引結(jié)構(gòu)有效地串聯(lián)到大量微小專(zhuān)家，從而將計(jì)算成本與參數(shù)計(jì)數(shù)分離。與密集的 FFW、粗粒度 MoE 和產(chǎn)品密鑰存儲(chǔ)器 (PKM) 層相比，表現(xiàn)出卓越的效率。

這項(xiàng)工作引入了參數(shù)高效專(zhuān)家檢索（PEER) 架構(gòu)（parameter efficient expert retrieval），利用產(chǎn)品密鑰（product key）檢索高效地路由到大量專(zhuān)家，將計(jì)算成本與參數(shù)量分離。這種設(shè)計(jì)在實(shí)驗(yàn)中展示了卓越的計(jì)算性能水平，將其定位為用于擴(kuò)展基礎(chǔ)模型的密集 FFW 層的競(jìng)爭(zhēng)性替代方案。這項(xiàng)工作的主要貢獻(xiàn)是：

極端 MoE 設(shè)置的探索：與以前的 MoE 研究中對(duì)少數(shù)大型專(zhuān)家的關(guān)注不同，這項(xiàng)工作研究了眾多小型專(zhuān)家的未充分探索的情況。

用于路由的學(xué)習(xí)索引結(jié)構(gòu)：首次證明學(xué)習(xí)索引結(jié)構(gòu)可以有效地路由到超過(guò)一百萬(wàn)個(gè)專(zhuān)家。

新的層設(shè)計(jì)：將產(chǎn)品密鑰路由與單神經(jīng)元專(zhuān)家相結(jié)合，我們引入了 PEER 層，它可以擴(kuò)展層容量而無(wú)需大量計(jì)算開(kāi)銷(xiāo)。實(shí)證結(jié)果表明，與密集 FFW、粗粒度 MoE 和產(chǎn)品密鑰內(nèi)存 (PKM) 層相比，其效率更高。

綜合消融研究：我們研究了 PEER 的不同設(shè)計(jì)選擇（例如專(zhuān)家數(shù)量、活動(dòng)參數(shù)、頭數(shù)量和查詢(xún)批量規(guī)范化）對(duì)語(yǔ)言建模任務(wù)的影響。

方法介紹

本節(jié)中，研究者詳解了參數(shù)高效專(zhuān)家檢索 (PEER) 層，它一種混合專(zhuān)家架構(gòu)，使用路由中的產(chǎn)品密鑰和單神經(jīng)元 MLP 作為專(zhuān)家。下圖 2 展示了 PEER 層內(nèi)的計(jì)算過(guò)程。

PEER 層概覽。從形式上看，PEER 層是一個(gè)函數(shù) f : R^n → R^m，它由三部分組成：一個(gè)由 N 個(gè)專(zhuān)家組成的池 E := {e_i}^N_i=1，其中每個(gè)專(zhuān)家 e_i : R^n → R^m 與 f 共享相同的簽名；一組相應(yīng)的 N 個(gè)產(chǎn)品密鑰 K := {k_i}^N_i=1 ? R^d ；以及一個(gè)查詢(xún)網(wǎng)絡(luò) q : R^n → R^d，它將輸入向量 x ∈ R^n 映射到查詢(xún)向量 q (x)。

令 T_k 表示 top-k 運(yùn)算符。給定輸入 x，首先檢索 k 個(gè)專(zhuān)家的子集，這些專(zhuān)家的相應(yīng)產(chǎn)品鍵與查詢(xún) q (x) 具有最高的內(nèi)積。

然后將非線(xiàn)性激活（例如 softmax 或 sigmoid）應(yīng)用于前 k 個(gè)專(zhuān)家的查詢(xún)鍵內(nèi)積，以獲得路由分?jǐn)?shù)。

最后通過(guò)對(duì)路由分?jǐn)?shù)加權(quán)的專(zhuān)家輸出進(jìn)行線(xiàn)性組合來(lái)計(jì)算輸出。

產(chǎn)品密鑰檢索。由于研究者打算使用大量專(zhuān)家（N ≥ 10^6），單純計(jì)算公式 1 中的前 k 個(gè)索引可能成本非常高，因此應(yīng)用了產(chǎn)品密鑰檢索技術(shù)。他們不使用 N 個(gè)獨(dú)立的 d 維向量作為密鑰 k_i，而是通過(guò)連接兩個(gè)獨(dú)立的 d/2 維子密鑰集（即 C, C ′ ? R d/2）中的向量來(lái)創(chuàng)建它們：

參數(shù)高效專(zhuān)家和多頭檢索。與其他 MoE 架構(gòu)不同，這些架構(gòu)通常將每個(gè)專(zhuān)家的隱藏層設(shè)置為與其他 FFW 層相同的大小。而在 PEER 中，每個(gè)專(zhuān)家 e_i 都是一個(gè)單例 MLP，換句話(huà)說(shuō)，它只有一個(gè)帶有單個(gè)神經(jīng)元的隱藏層：

研究者沒(méi)有改變單個(gè)專(zhuān)家的規(guī)模，而是使用了多頭檢索來(lái)調(diào)整 PEER 層的表達(dá)能力，這類(lèi)似于 transformer 中的多頭注意力機(jī)制和 PKM 中的多頭記憶。

具體來(lái)說(shuō)，他們使用 h 個(gè)獨(dú)立的查詢(xún)網(wǎng)絡(luò)，每個(gè)網(wǎng)絡(luò)計(jì)算自己的查詢(xún)并檢索一組單獨(dú)的 k 個(gè)專(zhuān)家。不過(guò)，不同的頭共享同一個(gè)專(zhuān)家池，具有相同的產(chǎn)品密鑰集。這 h 個(gè)頭的輸出簡(jiǎn)單地總結(jié)如下：

為什么要有大量的小專(zhuān)家？給定的 MoE 層可以用三個(gè)超參數(shù)來(lái)表征它：參數(shù)總數(shù) P、每個(gè) token 的活躍參數(shù)數(shù)量 P_active 和單個(gè)專(zhuān)家的大小 P_expert。Krajewski 等人 (2024) 表明，MoE 模型的 scaling law 具有以下形式：

對(duì)于 PEER，研究者通過(guò)設(shè)置 d_expert = 1 來(lái)使用盡可能最小的專(zhuān)家規(guī)模，激活神經(jīng)元的數(shù)量是檢索頭的數(shù)量乘以每個(gè)頭檢索到的專(zhuān)家數(shù)量：d_active = hk。因此，PEER 的粒度始終為 G = P_active/P_expert = d_active/d_expert = hk。

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

我們首先來(lái)看在語(yǔ)言建模數(shù)據(jù)集上的評(píng)估結(jié)果。

在根據(jù) isoFLOP 曲線(xiàn)確定每種方法的計(jì)算最優(yōu)模型后，研究者在以下幾個(gè)流行語(yǔ)言建模數(shù)據(jù)集上評(píng)估了這些預(yù)訓(xùn)練模型的性能：

Curation Corpus
Lambada
Pile
Wikitext
預(yù)訓(xùn)練數(shù)據(jù)集 C4

下表 1 展示了評(píng)估結(jié)果。研究者根據(jù)訓(xùn)練期間使用的 FLOP 預(yù)算對(duì)模型進(jìn)行了分組。可以看到，PEER 在這些語(yǔ)言建模數(shù)據(jù)集上的困惑度最低。

在消融實(shí)驗(yàn)中，研究者改變了專(zhuān)家總數(shù)量。下圖 1 中 isoFLOP 曲線(xiàn)所示的模型都有超過(guò)一百萬(wàn)個(gè)（1024^2 ）專(zhuān)家。

研究者選擇了 isoFLOP 最優(yōu)位置的模型，并改變了 PEER 層中的專(zhuān)家數(shù)量（N = 128^2、256^2、512^2、1024^2），同時(shí)保持活躍專(zhuān)家數(shù)量不變（h = 8，k = 16）。結(jié)果如下圖 3 (a) 所示。

可以看到，isoFLOP 曲線(xiàn)在具有 1024^2 個(gè)專(zhuān)家的 PEER 模型和相應(yīng)的密集主干之間進(jìn)行插值，而無(wú)需用 PEER 層替換中間塊中的 FFW 層。這表明，只需增加專(zhuān)家數(shù)量就可以提高模型性能。

同時(shí)，研究者改變了活躍專(zhuān)家的數(shù)量。他們系統(tǒng)地改變了活躍專(zhuān)家的數(shù)量（hk = 32、64、128、256、512），同時(shí)保持專(zhuān)家總數(shù)不變（N = 1024^2）。對(duì)于給定的 hk，研究者則聯(lián)合改變 h 和 k 以確定最佳組合。下圖 3 (b) 繪制了關(guān)于頭數(shù)量 (h) 的 isoFLOP 曲線(xiàn)。

下表 2 列出了使用和不使用 BN 時(shí)不同數(shù)量專(zhuān)家的專(zhuān)家使用率和不均勻性。可以看到，即使對(duì)于 1M 個(gè)專(zhuān)家，專(zhuān)家使用率也接近 100%，并且使用 BN 可以使專(zhuān)家的利用率更加均衡，困惑度更低。這些發(fā)現(xiàn)證明了 PEER 模型在利用大量專(zhuān)家方面的有效性。