↑ 點(diǎn)擊藍(lán)字 關(guān)注極市平臺
如果你熟悉當(dāng)前LLM的主流架構(gòu)，混合專家（MoE）技術(shù)想必是老朋友之一。有人甚至?xí)f，MoE是使大模型崛起的關(guān)鍵因素之一。
開源的Mixtral、DBRX、Grok等模型都使用了MoE，而且根據(jù)Soumith Chintala等大佬的推測，GPT-4也是一個(gè)規(guī)模為8×220B的MoE模型。
類似GPT-4，多數(shù)MoE模型都會將專家數(shù)量限制在較少數(shù)量，一般不會超過16或32。
然而，DeepMind研究科學(xué)家Xu Owen He最近就在一篇獨(dú)立發(fā)表的論文中，提出了一種全新的方法——PEER（參數(shù)高效專家檢索，Parameter Efficient Expert Retrieval），可以將專家數(shù)量擴(kuò)展到百萬數(shù)量級。
論文地址：https://arxiv.org/abs/2407.04153
這究竟是如何做到的？參數(shù)量不會爆炸嗎？不會造成收益遞減嗎？如何實(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é)存儲模型知識。
我們當(dāng)然希望LLM能在參數(shù)中隱式存儲更多知識，但FFW的計(jì)算成本和激活內(nèi)存會隨之線性增加。稠密模型中，F(xiàn)FW層已經(jīng)占據(jù)了總參數(shù)量的2/3，是擴(kuò)展的主要瓶頸之一。
MoE模型雖然參數(shù)量也很大，但每次推理時(shí)不會動(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è)與之后，模型性能就會進(jìn)入「平臺期」。
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í)。
之前有研究表明，通過簡單地添加新專家并進(jìn)行適當(dāng)正則化，MoE模型就可以適應(yīng)連續(xù)的數(shù)據(jù)流。凍結(jié)舊專家、僅讓新專家權(quán)重更新，就可以在保持可塑性的同時(shí)防止災(zāi)難性遺忘。
在終身學(xué)習(xí)環(huán)境中，數(shù)據(jù)流可能達(dá)到無限長度，甚至永無止境，因此論文探索的專家數(shù)量的擴(kuò)大就顯得十分重要。
PEER層設(shè)計(jì)
形式上，PEER可以表達(dá)為函數(shù)，包含如下三部分：
- N個(gè)專家組成的專家池，其中每個(gè)專家ei和f有相同的函數(shù)簽名
- 用于將輸入向量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所示。
乘積鍵檢索
由于要使用非常多的專家（N≥106），直接計(jì)算公式1中的前k個(gè)指標(biāo)可能會非常耗費(fèi)資源。
為此，研究人員提出了乘積鍵檢索技術(shù)——通過連接來自兩個(gè)獨(dú)立d/2維的子鍵C和C′（C, C′ ? Rd/2）的向量來進(jìn)行創(chuàng)建：
這種K的笛卡爾積結(jié)構(gòu)，能夠高效地找到前k位專家。
也就是，先將查詢向量q(x)拆分為兩個(gè)子查詢q1和q2，并分別對子查詢與子鍵之間的內(nèi)積進(jìn)行前k個(gè)操作：
然后得到一組k2候選鍵集合，，在數(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) + k2)d)。
參數(shù)高效專家和多頭檢索
通常的MoE架構(gòu)中會將專家隱藏層設(shè)計(jì)為FFW層相同大小，但PEER中的每個(gè)專家ei則小得多，僅僅是包含一個(gè)神經(jīng)元、一個(gè)隱藏層的單例MLP：
其中vi、ui都不是參數(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ù)效率和知識遷移的能力。
算法1描述了PEER前向計(jì)算的簡化實(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ù)Pactive以及單個(gè)專家的大小Pexpert。
之前提及的，Krajewski論文中提出的「細(xì)粒度MoE Scaling Law」就可以用這三個(gè)參數(shù)表示為如下形式：
其中，L表示損失值，D是訓(xùn)練token總量，G是激活專家的數(shù)量，a、b、g、α、β、γ等字母都表示常量。
為了降低損失值、提高模型性能，我們需要提升P、D、G的值，同時(shí)也需要限制Pactive的大小，因?yàn)橛?jì)算和內(nèi)存成本主要由激活參數(shù)量決定。
值得注意的是，Pactive對應(yīng)的內(nèi)存占用也和batch中的token數(shù)量有關(guān)，但與batch大小或序列長度無關(guān)，因?yàn)樵谔幚韇atch或序列的每個(gè)token時(shí)只需要存儲模型的一份副本。
由于專家大小Pexpert=Pactive/G，因此專家數(shù)量N=P/Pexpert=P·G/Pactive。
這就意味著，如果想要在增加P、G的同時(shí)保持Pactive不變，就應(yīng)該減小專家大小Pexpert，并增加專家數(shù)量N。
這就解釋了，我們?yōu)槭裁葱枰罅康男＜摇?/span>
一般來說，對于只有單個(gè)隱藏層的專家模型，其參數(shù)大小為：Pexpert=(2?dmodel+1)?dexpert，相應(yīng)地，Pactive=(2?dmodel+1)?dactive。
其中dmodel、dexpert、dactive分別表示表示Transformer的隱藏維度、單個(gè)專家的隱藏神經(jīng)元數(shù)量以及每個(gè)token激活的參數(shù)總量。
PEER中使用了盡可能小的專家大小，即dexpert=1，因此最大程度地降低了dactive，大小僅為檢索頭數(shù)量乘以每次檢索的專家數(shù)量（h·k）。
預(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)。
對于密集基線模型，通過改變層數(shù)、注意力頭數(shù)和模型維度來改變它們的大小。
對于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è)儲存器，8個(gè)頭，每個(gè)頭選擇32個(gè)存儲器（top k=32）。
此外，研究者還應(yīng)用了查詢批量歸一化，這是原始PKM論文中推薦的，用于提高記憶使用率。
在PEER中，同樣使用了10242個(gè)專家，8個(gè)頭，每個(gè)頭選擇16個(gè)存儲器（top k=16）。
默認(rèn)情況下，作者也啟用了查詢批歸一化（BatchNorm）來增加專家使用率。與專家選擇MoE基線不同，PEER代表了一種細(xì)粒度方法，使用了大量小型專家。
在所有模型大小和方法中，保持了一致的批大?。?28）和序列長度（2048）。通過將總計(jì)算預(yù)算除以每個(gè)訓(xùn)練步驟的FLOP來計(jì)算訓(xùn)練步驟數(shù)。
與密集FFW基線相比，稀疏替代方案使isoFLOP曲線向下和向右移動(dòng)，因?yàn)槠湟肓烁嗟目倕?shù)P，但使用了更少或相同數(shù)量的活躍參數(shù)Pactive。
另外，在相同的計(jì)算預(yù)算下，PEER模型達(dá)到了最低的計(jì)算優(yōu)化困惑度。
語言建模數(shù)據(jù)集評估
根據(jù)isoFLOP曲線確定每種方法的計(jì)算最優(yōu)模型后，研究人員在幾個(gè)流行的語言建模數(shù)據(jù)集上評估了這些預(yù)訓(xùn)練模型的性能，包括Curation Corpus、Lambada、the Pile、Wikitext和預(yù)訓(xùn)練數(shù)據(jù)集C4。
評估結(jié)果如表1所示，作者根據(jù)模型在訓(xùn)練期間使用的FLOP預(yù)算將其分組。
消融實(shí)驗(yàn)
圖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ù)量，就可以提高模型性能。
另外，作者還對活躍專家hk數(shù)量進(jìn)行了消融研究，hk等同于等式9中的粒度G。
他系統(tǒng)地改變了活躍專家的數(shù)量（hk = 32, 64, 128, 256, 512），同時(shí)保持專家總數(shù)不變（N = 10242）。
此外，對于給定的hk，同時(shí)改變h和k以找到最優(yōu)組合。結(jié)果中的isoFLOP曲線，以頭（h）數(shù)為橫軸繪制，如圖3（b）所示。
結(jié)果表明，在考慮的數(shù)值范圍內(nèi)，更高的h·k值通常會帶來更好的性能。
值得注意的是，隨著h·k的增加，最優(yōu)的h也在增加。然而，性能逐漸趨于飽和，增加活躍專家的數(shù)量也會增加設(shè)備內(nèi)存消耗，可能需要額外的加速器設(shè)備。
因此，在實(shí)際應(yīng)用中，應(yīng)該根據(jù)性能、設(shè)備數(shù)量和計(jì)算資源需求之間的權(quán)衡來選擇適當(dāng)?shù)膆·k值。
由于PEER層中有超過一百萬個(gè)專家，對于許多人來說，一定好奇其在推理過程中實(shí)際選擇了多少專家？它們的使用是否均勻分布？
為了分析這點(diǎn)，研究者在C4驗(yàn)證集的所有token中，為每個(gè)專家ei保留一個(gè)累積的路由分?jǐn)?shù)（router score），表示為：。
其中，gi(x)是指當(dāng)token x作為輸入時(shí)用于聚合專家輸出的路由分?jǐn)?shù)。如果專家ei未被選中，則gi(x)=0。
從這些累積的路由分?jǐn)?shù)中，可以獲得一個(gè)經(jīng)驗(yàn)概率分布向量，表示為。
然后，作者計(jì)算了Lample等人提出的以下指標(biāo)來評估專家的使用情況和分布：
默認(rèn)設(shè)置下，作者在查詢網(wǎng)絡(luò)上添加了一個(gè)批歸一化（BN）層，目的是增加訓(xùn)練過程中的專家使用率。
對此，他研究添加這個(gè)BN層對上述指標(biāo)的影響。
表2展示了不同專家數(shù)量下的專家使用和不均勻性，包括使用和不使用BN的情況。
可以看到，即使對于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í)研討會，目前擔(dān)任Conference on Lifelong Learning Agents的出版主席。
https://arxiv.org/abs/2407.04153
https://venturebeat.com/ai/deepminds-peer-scales-language-models-with-millions-of-tiny-experts/
公眾號后臺回復(fù)“數(shù)據(jù)集”獲取100+深度學(xué)習(xí)各方向資源整理
點(diǎn)擊閱讀原文進(jìn)入CV社區(qū)
收獲更多技術(shù)干貨
