如果你熟悉當(dāng)前LLM的主流架構(gòu)，混合專家（MoE）技術(shù)想必是老朋友之一。有人甚至?xí)f，MoE是使大模型崛起的關(guān)鍵因素之一。

開源的Mixtral、DBRX、Grok等模型都使用了MoE，而且根據(jù)Soumith Chintala等大佬的推測(cè)，GPT-4也是一個(gè)規(guī)模為8×220B的MoE模型。

類似GPT-4，多數(shù)MoE模型都會(huì)將專家數(shù)量限制在較少數(shù)量，一般不會(huì)超過16或32。

然而，DeepMind研究科學(xué)家Xu Owen He最近就在一篇獨(dú)立發(fā)表的論文中，提出了一種全新的方法——PEER（參數(shù)高效專家檢索，Parameter Efficient Expert Retrieval），可以將專家數(shù)量擴(kuò)展到百萬數(shù)量級(jí)。

論文地址：https://arxiv.org/abs/2407.04153

這究竟是如何做到的？參數(shù)量不會(huì)爆炸嗎？不會(huì)造成收益遞減嗎？如何實(shí)現(xiàn)能在百萬個(gè)專家中實(shí)現(xiàn)高效檢索？

背景與介紹

Transformer架構(gòu)中，每個(gè)塊內(nèi)都包含注意力層和前饋層（FFW），注意力層用于計(jì)算序列中token之間的關(guān)系，F(xiàn)FW網(wǎng)絡(luò)則負(fù)責(zé)存儲(chǔ)模型知識(shí)。

我們當(dāng)然希望LLM能在參數(shù)中隱式存儲(chǔ)更多知識(shí)，但FFW的計(jì)算成本和激活內(nèi)存會(huì)隨之線性增加。稠密模型中，F(xiàn)FW層已經(jīng)占據(jù)了總參數(shù)量的2/3，是擴(kuò)展的主要瓶頸之一。

MoE模型雖然參數(shù)量也很大，但每次推理時(shí)不會(huì)動(dòng)用整個(gè)模型的能力，而是將數(shù)據(jù)路由到小型且專門的「專家模塊」，因此能在LLM參數(shù)增加的同時(shí)，讓推理所需的計(jì)算成本基本不變。

那么專家數(shù)量（即MoE模型的「粒度」）是不是越多越好？

這要考慮多個(gè)因素，包括模型參數(shù)總量、訓(xùn)練token數(shù)量和算力的預(yù)算。

2022年的一項(xiàng)研究認(rèn)為，模型總參數(shù)量不變時(shí)，存在一個(gè)能達(dá)到最優(yōu)性能的「最佳粒度」。專家數(shù)量超過這個(gè)與之后，模型性能就會(huì)進(jìn)入「平臺(tái)期」。

Unified Scaling Laws for Routed Language Models

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

然而，今年年初Krajewski等人發(fā)表的一篇論文反駁了這個(gè)觀點(diǎn)。他們發(fā)現(xiàn)，如果同時(shí)增加訓(xùn)練所用的token數(shù)量，那么更高粒度可以提高性能。

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

受到這種細(xì)粒度MoE Scaling Law的啟發(fā)，作者推斷，模型容量的持續(xù)改進(jìn)將帶來具有高粒度的LLM，即包括大量微型專家的模型。

除了能帶來更高效的擴(kuò)展，增加專家數(shù)量還有另外一層好處——終身學(xué)習(xí)。

之前有研究表明，通過簡(jiǎn)單地添加新專家并進(jìn)行適當(dāng)正則化，MoE模型就可以適應(yīng)連續(xù)的數(shù)據(jù)流。凍結(jié)舊專家、僅讓新專家權(quán)重更新，就可以在保持可塑性的同時(shí)防止災(zāi)難性遺忘。

在終身學(xué)習(xí)環(huán)境中，數(shù)據(jù)流可能達(dá)到無限長(zhǎng)度，甚至永無止境，因此論文探索的專家數(shù)量的擴(kuò)大就顯得十分重要。

百萬MoE所系

PEER層設(shè)計(jì)

形式上，PEER可以表達(dá)為函數(shù)，包含如下三部分：

- N個(gè)專家組成的專家池，其中每個(gè)專家e_i和f有相同的函數(shù)簽名

- 一組相應(yīng)的N個(gè)乘積鍵

- 用于將輸入向量x映射到查詢向量q(x)的查詢網(wǎng)絡(luò)

用??_??表示top-k運(yùn)算符，給定輸入x，從N個(gè)專家中檢索到k個(gè)專家的過程可以表示為：

之后使用softmax或sigmoid等激活函數(shù)應(yīng)用于top-k個(gè)專家的查詢-鍵內(nèi)積，獲得路由分?jǐn)?shù)：

最后，計(jì)算路由分?jǐn)?shù)加權(quán)的專家分?jǐn)?shù)之和，作為PEER層輸出。

上述的PEER層可以插入到Transformer架構(gòu)骨干的中間，或者代替FFW層，整體計(jì)算流程如圖2所示。

PEER層示意圖

乘積鍵檢索

由于要使用非常多的專家（N≥10⁶），直接計(jì)算公式1中的前k個(gè)指標(biāo)可能會(huì)非常耗費(fèi)資源。

為此，研究人員提出了乘積鍵檢索技術(shù)——通過連接來自兩個(gè)獨(dú)立d/2維的子鍵C和C′（C, C′ ? R^d/2）的向量來進(jìn)行創(chuàng)建：

這種K的笛卡爾積結(jié)構(gòu)，能夠高效地找到前k位專家。

也就是，先將查詢向量q(x)拆分為兩個(gè)子查詢q1和q2，并分別對(duì)子查詢與子鍵之間的內(nèi)積進(jìn)行前k個(gè)操作：

然后得到一組k²候選鍵集合，，在數(shù)學(xué)上，這可以保證從K中與q(x)最相似的k個(gè)鍵在這個(gè)候選集合中。

此外，候選鍵與q(x)之間的內(nèi)積只是子鍵與子查詢之間內(nèi)積的總和：。

因此，可以再次將前k個(gè)操作應(yīng)用于這k的平方個(gè)內(nèi)積，以從原始的乘積鍵集合K中獲得前k個(gè)匹配鍵。

最終，公式1中通過窮舉搜索進(jìn)行的top-k專家檢索的復(fù)雜度，從O(N·d)降到了O((√N(yùn) + k²)d)。

參數(shù)高效專家和多頭檢索

通常的MoE架構(gòu)中會(huì)將專家隱藏層設(shè)計(jì)為FFW層相同大小，但PEER中的每個(gè)專家e_i則小得多，僅僅是包含一個(gè)神經(jīng)元、一個(gè)隱藏層的單例MLP：

其中v_i、u_i都不是參數(shù)矩陣，而是與x維度相同的向量，??表示ReLU或GELU等激活函數(shù)。

每個(gè)專家只有如此少的參數(shù)，怎么可能有強(qiáng)大的表達(dá)能力？

此處，作者借鑒了多頭注意力的做法，使用了「多頭檢索」（multi-head retrieval）機(jī)制。

也就是說，查詢過程并不是僅有一次，而是使用h個(gè)獨(dú)立的查詢網(wǎng)絡(luò)，分別計(jì)算自己的查詢向量并檢索出一組top-k專家，但他們共享相同的專家池和每個(gè)專家的乘積鍵。

h個(gè)「檢索頭」的輸出相加后得到最終檢索結(jié)果：

可以證明，當(dāng)每個(gè)頭只檢索一名專家時(shí)（k=1），使用含h個(gè)檢索頭的PEER層，等效于使用含有h個(gè)神經(jīng)元的專家。

這就意味著，這種多頭檢索的過程相當(dāng)于動(dòng)態(tài)組裝出含有h個(gè)神經(jīng)元的專家MLP網(wǎng)絡(luò)。

與現(xiàn)有MoE中多個(gè)專家之間完全隔離的狀態(tài)相比，這種設(shè)計(jì)允許專家間共享隱藏神經(jīng)元，從而提升了參數(shù)效率和知識(shí)遷移的能力。

算法1描述了PEER前向計(jì)算的簡(jiǎn)化實(shí)現(xiàn)。要想達(dá)到高效的實(shí)現(xiàn)，可能還需要專門的硬件內(nèi)核來加速embedding查找，以及與einsum操作的融合。

為什么用大量的小專家

論文的背景介紹中，我們從直覺和經(jīng)驗(yàn)層面推導(dǎo)出使用大量專家的優(yōu)點(diǎn)，但作者也在此處用公式推導(dǎo)的方法證明了這一點(diǎn)。

首先，我們可以用3個(gè)超參數(shù)表征MoE層的設(shè)置：參數(shù)總量P、每個(gè)token激活的參數(shù)P_active以及單個(gè)專家的大小P_expert。

之前提及的，Krajewski論文中提出的「細(xì)粒度MoE Scaling Law」就可以用這三個(gè)參數(shù)表示為如下形式：

其中，L表示損失值，D是訓(xùn)練token總量，G是激活專家的數(shù)量，a、b、g、α、β、γ等字母都表示常量。

為了降低損失值、提高模型性能，我們需要提升P、D、G的值，同時(shí)也需要限制P_active的大小，因?yàn)橛?jì)算和內(nèi)存成本主要由激活參數(shù)量決定。

值得注意的是，P_active對(duì)應(yīng)的內(nèi)存占用也和batch中的token數(shù)量有關(guān)，但與batch大小或序列長(zhǎng)度無關(guān)，因?yàn)樵谔幚韇atch或序列的每個(gè)token時(shí)只需要存儲(chǔ)模型的一份副本。

由于專家大小P_expert=P_active/G，因此專家數(shù)量N=P/P_expert=P·G/P_active。

這就意味著，如果想要在增加P、G的同時(shí)保持P_active不變，就應(yīng)該減小專家大小P_expert，并增加專家數(shù)量N。

這就解釋了，我們?yōu)槭裁葱枰罅康男＜摇?/span>

一般來說，對(duì)于只有單個(gè)隱藏層的專家模型，其參數(shù)大小為：P_expert=(2?d_model+1)?d_expert，相應(yīng)地，P_active=(2?d_model+1)?d_active。

其中d_model、d_expert、d_active分別表示表示Transformer的隱藏維度、單個(gè)專家的隱藏神經(jīng)元數(shù)量以及每個(gè)token激活的參數(shù)總量。

PEER中使用了盡可能小的專家大小，即d_expert=1，因此最大程度地降低了d_active，大小僅為檢索頭數(shù)量乘以每次檢索的專家數(shù)量（h·k）。

實(shí)驗(yàn)

預(yù)訓(xùn)練isoFLOP分析

作者使用isoFLOP分析將PEER與各種基線方法進(jìn)行了比較。

這里，他選擇了固定的FLOP預(yù)算（6e18和2e19），并同時(shí)改變模型大小和來自C4數(shù)據(jù)集的訓(xùn)練token數(shù)，以獲得isoFLOP曲線。

isoFLOP曲線上的每個(gè)點(diǎn)，都具有相同的計(jì)算成本，作者根據(jù)其模型大小和在C4上的最終驗(yàn)證困惑度（perplexity），來繪制這些點(diǎn)。

對(duì)于密集基線模型，通過改變層數(shù)、注意力頭數(shù)和模型維度來改變它們的大小。

對(duì)于MoE、PKM和PEER方法，作者選取了每個(gè)考慮的密集模型，并分別用MoE、PKM和PEER層，替換中間塊中的FFN層（例如，在12塊Transformer中，替換第6塊中的FFN）。

在MoE中，作者使用了專家選擇路由算法。該算法有效地解決了專家負(fù)載不平衡問題，并且通常優(yōu)于token選擇MoE。

每個(gè)專家的大小與相應(yīng)密集模型中原始MLP的大小相同，作者使用了128個(gè)專家來覆蓋與PEER模型相同的模型大小范圍。

這種類型的MoE代表了標(biāo)準(zhǔn)的粗粒度方法，由少量大型專家組成。

在PKM中，作者使用了1024×2個(gè)儲(chǔ)存器，8個(gè)頭，每個(gè)頭選擇32個(gè)存儲(chǔ)器（top k=32）。

此外，研究者還應(yīng)用了查詢批量歸一化，這是原始PKM論文中推薦的，用于提高記憶使用率。

在PEER中，同樣使用了10242個(gè)專家，8個(gè)頭，每個(gè)頭選擇16個(gè)存儲(chǔ)器（top k=16）。

默認(rèn)情況下，作者也啟用了查詢批歸一化（BatchNorm）來增加專家使用率。與專家選擇MoE基線不同，PEER代表了一種細(xì)粒度方法，使用了大量小型專家。

在所有模型大小和方法中，保持了一致的批大小（128）和序列長(zhǎng)度（2048）。通過將總計(jì)算預(yù)算除以每個(gè)訓(xùn)練步驟的FLOP來計(jì)算訓(xùn)練步驟數(shù)。

圖1展示的是isoFLOP曲線。

與密集FFW基線相比，稀疏替代方案使isoFLOP曲線向下和向右移動(dòng)，因?yàn)槠湟肓烁嗟目倕?shù)P，但使用了更少或相同數(shù)量的活躍參數(shù)P_active。

另外，在相同的計(jì)算預(yù)算下，PEER模型達(dá)到了最低的計(jì)算優(yōu)化困惑度。

語言建模數(shù)據(jù)集評(píng)估

根據(jù)isoFLOP曲線確定每種方法的計(jì)算最優(yōu)模型后，研究人員在幾個(gè)流行的語言建模數(shù)據(jù)集上評(píng)估了這些預(yù)訓(xùn)練模型的性能，包括Curation Corpus、Lambada、the Pile、Wikitext和預(yù)訓(xùn)練數(shù)據(jù)集C4。

評(píng)估結(jié)果如表1所示，作者根據(jù)模型在訓(xùn)練期間使用的FLOP預(yù)算將其分組。

消融實(shí)驗(yàn)

改變專家總數(shù)

圖1中isoFLOP圖所示的模型都有超過一百萬（10242）個(gè)專家。

接下來，研究者進(jìn)行了一項(xiàng)消融研究，研究專家數(shù)量N的影響，N決定了等式9中的總參數(shù)數(shù)量P。

作者選擇了isoFLOP最優(yōu)模型，并改變PEER層中的專家數(shù)量（N = 1282, 2562, 5122, 10242），同時(shí)活躍專家的數(shù)量不變（h = 8, k = 16）。

實(shí)驗(yàn)結(jié)果如圖3(a)所示，可以看出，isoFLOP曲線在10242個(gè)專家的PEER模型和相應(yīng)的密集骨干模型（未將中間塊的FFW層替換為PEER層）之間進(jìn)行插值。

這表明，僅僅增加專家數(shù)量，就可以提高模型性能。

改變活躍專家的數(shù)量

另外，作者還對(duì)活躍專家hk數(shù)量進(jìn)行了消融研究，hk等同于等式9中的粒度G。

他系統(tǒng)地改變了活躍專家的數(shù)量（hk = 32, 64, 128, 256, 512），同時(shí)保持專家總數(shù)不變（N = 10242）。

此外，對(duì)于給定的hk，同時(shí)改變h和k以找到最優(yōu)組合。結(jié)果中的isoFLOP曲線，以頭（h）數(shù)為橫軸繪制，如圖3（b）所示。

結(jié)果表明，在考慮的數(shù)值范圍內(nèi)，更高的h·k值通常會(huì)帶來更好的性能。

值得注意的是，隨著h·k的增加，最優(yōu)的h也在增加。然而，性能逐漸趨于飽和，增加活躍專家的數(shù)量也會(huì)增加設(shè)備內(nèi)存消耗，可能需要額外的加速器設(shè)備。

因此，在實(shí)際應(yīng)用中，應(yīng)該根據(jù)性能、設(shè)備數(shù)量和計(jì)算資源需求之間的權(quán)衡來選擇適當(dāng)?shù)膆·k值。

專家使用和查詢批歸一化

由于PEER層中有超過一百萬個(gè)專家，對(duì)于許多人來說，一定好奇其在推理過程中實(shí)際選擇了多少專家？它們的使用是否均勻分布？

為了分析這點(diǎn)，研究者在C4驗(yàn)證集的所有token中，為每個(gè)專家e_i保留一個(gè)累積的路由分?jǐn)?shù)（router score），表示為：。

其中，g_i(x)是指當(dāng)token x作為輸入時(shí)用于聚合專家輸出的路由分?jǐn)?shù)。如果專家e_i未被選中，則g_i(x)=0。

從這些累積的路由分?jǐn)?shù)中，可以獲得一個(gè)經(jīng)驗(yàn)概率分布向量，表示為。

代表C4驗(yàn)證集上所有專家的分布。

然后，作者計(jì)算了Lample等人提出的以下指標(biāo)來評(píng)估專家的使用情況和分布：

- 專家使用：推理過程中檢索到的專家比例：

- 不均勻性：z與均勻分布之間的KL散度

其中，N是專家總數(shù)。

默認(rèn)設(shè)置下，作者在查詢網(wǎng)絡(luò)上添加了一個(gè)批歸一化（BN）層，目的是增加訓(xùn)練過程中的專家使用率。

對(duì)此，他研究添加這個(gè)BN層對(duì)上述指標(biāo)的影響。

表2展示了不同專家數(shù)量下的專家使用和不均勻性，包括使用和不使用BN的情況。

可以看到，即使對(duì)于100萬個(gè)專家，專家使用率也接近100%，而使用BN可以導(dǎo)致專家的使用更加平衡，并降低困惑度。

這些發(fā)現(xiàn)證明了，PEER模型在利用大量專家方面的有效性。

最后，作者還比較了有和沒有批歸一化（BN）的isoFLOP曲線。

如下圖4所示，使用BN的PEER模型通常能達(dá)到更低的困惑度。雖然差異不是很顯著，但在isoFLOP最有區(qū)域附近最為明顯。

作者介紹

Owen He是Google DeepMind的研究科學(xué)家，專注于持續(xù)學(xué)習(xí)和基礎(chǔ)模型。

在加入DeepMind之前，Owen是在格羅寧根大學(xué)和不來梅雅各布大學(xué)（康斯特大學(xué)）攻讀博士學(xué)位，導(dǎo)師是 Herbert Jaeger。

在此期間，Owen訪問過Mila、馬克斯·普朗克研究所、伯爾尼大學(xué)和蘇黎世聯(lián)邦理工學(xué)院，并在Google DeepMind進(jìn)行過實(shí)習(xí)。

他曾組織過ICML 2021持續(xù)學(xué)習(xí)研討會(huì)，目前擔(dān)任Conference on Lifelong Learning Agents的出版主席。