技術(shù)前沿：AICG——混合專家模型 (MoEs)

什么是混合專家模型？

混合專家模型主要由兩個(gè)關(guān)鍵部分組成:

稀疏 MoE 層: 這些層代替了傳統(tǒng) Transformer 模型中的前饋網(wǎng)絡(luò) (FFN) 層。MoE 層包含若干“專家”(例如 8 個(gè))，每個(gè)專家本身是一個(gè)獨(dú)立的神經(jīng)網(wǎng)絡(luò)。在實(shí)際應(yīng)用中，這些專家通常是前饋網(wǎng)絡(luò) (FFN)，但它們也可以是更復(fù)雜的網(wǎng)絡(luò)結(jié)構(gòu)，甚至可以是 MoE 層本身，從而形成層級(jí)式的 MoE 結(jié)構(gòu)。 ? ?

門控網(wǎng)絡(luò)或路由: 這個(gè)部分用于決定哪些令牌 (token) 被發(fā)送到哪個(gè)專家。例如，在下圖中，“More”這個(gè)令牌可能被發(fā)送到第二個(gè)專家，而“Parameters”這個(gè)令牌被發(fā)送到第一個(gè)專家。有時(shí)，一個(gè)令牌甚至可以被發(fā)送到多個(gè)專家。令牌的路由方式是 MoE 使用中的一個(gè)關(guān)鍵點(diǎn)，因?yàn)槁酚善饔蓪W(xué)習(xí)的參數(shù)組成，并且與網(wǎng)絡(luò)的其他部分一同進(jìn)行預(yù)訓(xùn)練。

?Switch Transformers paper 論文中的 MoE layer

總結(jié)來說，在混合專家模型 (MoE) 中，我們將傳統(tǒng) Transformer 模型中的每個(gè)前饋網(wǎng)絡(luò) (FFN) 層替換為 MoE 層，其中 MoE 層由兩個(gè)核心部分組成: 一個(gè)門控網(wǎng)絡(luò)和若干數(shù)量的專家。

盡管混合專家模型 (MoE) 提供了若干顯著優(yōu)勢(shì)，例如更高效的預(yù)訓(xùn)練和與稠密模型相比更快的推理速度，但它們也伴隨著一些挑戰(zhàn):

訓(xùn)練挑戰(zhàn): 雖然 MoE 能夠?qū)崿F(xiàn)更高效的計(jì)算預(yù)訓(xùn)練，但它們?cè)谖⒄{(diào)階段往往面臨泛化能力不足的問題，長(zhǎng)期以來易于引發(fā)過擬合現(xiàn)象。 ? ?

推理挑戰(zhàn): MoE 模型雖然可能擁有大量參數(shù)，但在推理過程中只使用其中的一部分，這使得它們的推理速度快于具有相同數(shù)量參數(shù)的稠密模型。然而，這種模型需要將所有參數(shù)加載到內(nèi)存中，因此對(duì)內(nèi)存的需求非常高。以 Mixtral 8x7B 這樣的 MoE 為例，需要足夠的 VRAM 來容納一個(gè) 47B 參數(shù)的稠密模型。之所以是 47B 而不是 8 x 7B = 56B，是因?yàn)樵?MoE 模型中，只有 FFN 層被視為獨(dú)立的專家，而模型的其他參數(shù)是共享的。此外，假設(shè)每個(gè)令牌只使用兩個(gè)專家，那么推理速度 (以 FLOPs 計(jì)算) 類似于使用 12B 模型 (而不是 14B 模型)，因?yàn)殡m然它進(jìn)行了 2x7B 的矩陣乘法計(jì)算，但某些層是共享的。

混合專家模型簡(jiǎn)史

混合專家模型 (MoE) 的理念起源于 1991 年的論文 Adaptive Mixture of Local Experts。這個(gè)概念與集成學(xué)習(xí)方法相似，旨在為由多個(gè)單獨(dú)網(wǎng)絡(luò)組成的系統(tǒng)建立一個(gè)監(jiān)管機(jī)制。在這種系統(tǒng)中，每個(gè)網(wǎng)絡(luò) (被稱為“專家”) 處理訓(xùn)練樣本的不同子集，專注于輸入空間的特定區(qū)域。那么，如何選擇哪個(gè)專家來處理特定的輸入呢？這就是門控網(wǎng)絡(luò)發(fā)揮作用的地方，它決定了分配給每個(gè)專家的權(quán)重。在訓(xùn)練過程中，這些專家和門控網(wǎng)絡(luò)都同時(shí)接受訓(xùn)練，以優(yōu)化它們的性能和決策能力。

在 2010 至 2015 年間，兩個(gè)獨(dú)立的研究領(lǐng)域?yàn)榛旌蠈＜夷Ｐ?(MoE) 的后續(xù)發(fā)展做出了顯著貢獻(xiàn):

組件專家: 在傳統(tǒng)的 MoE 設(shè)置中，整個(gè)系統(tǒng)由一個(gè)門控網(wǎng)絡(luò)和多個(gè)專家組成。在支持向量機(jī) (SVMs) 、高斯過程和其他方法的研究中，MoE 通常被視為整個(gè)模型的一部分。然而，Eigen、Ranzato 和 Ilya 的研究 探索了將 MoE 作為更深層網(wǎng)絡(luò)的一個(gè)組件。這種方法允許將 MoE 嵌入到多層網(wǎng)絡(luò)中的某一層，使得模型既大又高效。

條件計(jì)算: 傳統(tǒng)的神經(jīng)網(wǎng)絡(luò)通過每一層處理所有輸入數(shù)據(jù)。在這一時(shí)期，Yoshua Bengio 等研究人員開始探索基于輸入令牌動(dòng)態(tài)激活或停用網(wǎng)絡(luò)組件的方法。 ? ?

這些研究的融合促進(jìn)了在自然語言處理 (NLP) 領(lǐng)域?qū)旌蠈＜夷Ｐ偷奶剿鳌Ｌ貏e是在 2017 年，Shazeer 等人 (團(tuán)隊(duì)包括 Geoffrey Hinton 和 Jeff Dean，后者有時(shí)被戲稱為 “谷歌的 Chuck Norris”) 將這一概念應(yīng)用于 137B 的 LSTM (當(dāng)時(shí)被廣泛應(yīng)用于 NLP 的架構(gòu)，由 Schmidhuber 提出)。通過引入稀疏性，這項(xiàng)工作在保持極高規(guī)模的同時(shí)實(shí)現(xiàn)了快速的推理速度。這項(xiàng)工作主要集中在翻譯領(lǐng)域，但面臨著如高通信成本和訓(xùn)練不穩(wěn)定性等多種挑戰(zhàn)。

?Outrageously Large Neural Network 論文中的 MoE layer

混合專家模型 (MoE) 的引入使得訓(xùn)練具有數(shù)千億甚至萬億參數(shù)的模型成為可能，如開源的 1.6 萬億參數(shù)的 Switch Transformers 等。這種技術(shù)不僅在自然語言處理 (NLP) 領(lǐng)域得到了廣泛應(yīng)用，也開始在計(jì)算機(jī)視覺領(lǐng)域進(jìn)行探索。

最近再次大火的原因，評(píng)測(cè)超Llama2，混合專家模型（MoE）會(huì)是大模型新方向嗎？?

先是 Reddit 上一篇關(guān)于 GPT-4 結(jié)構(gòu)的猜測(cè)帖，暗示了 GPT-4 可能是由 16 個(gè)子模塊組成的專家模型（MoE）的混合體。據(jù)說，這 16 個(gè)子模塊中的每一個(gè) MoE 都有 1110 億個(gè)參數(shù)（作為參考，GPT-3 有 1750 億個(gè)參數(shù)）。盡管不能 100% 確定，但 GPT-4 是一個(gè) MoE 組成的集群這個(gè)事很可能是真的。

然后是法國 AI 公司 MistralAI 發(fā)布了全球首個(gè)基于混合專家技術(shù)的大模型 Mistral-8x7B-MoE，是 8 個(gè) 70 億參數(shù)規(guī)模大模型的混合。

主要特點(diǎn)如下： ? ?

它可以非常優(yōu)雅地處理 32K 上下文數(shù)據(jù)；

除了英語外，在法語、德語、意大利語和西班牙語表現(xiàn)也很好；

在代碼能力上表現(xiàn)很強(qiáng)；

指令微調(diào)后 MT-Bench 的得分 8.3 分（GPT-3.5 是 8.32、LLaMA2 70B 是 6.86）；

Mistral-7B×8-MoE 是首個(gè)被證明有效的開源的 MoE LLM，相比于早期的 Switch Transformer、GLaM 等研究，Mistral-7B×8-MoE 證明了 MoE 真的可以落地，且效果遠(yuǎn)好于相同激活值的 Dense 模型。

而在最近的一個(gè)評(píng)測(cè)中，Mistral-8x7B-MoE 經(jīng)過微調(diào)后的表現(xiàn)超過了 Llama2-65B。

到底什么是 MoE，可以參見?Hugging Face 官方對(duì) MoE 的詳細(xì)技術(shù)解讀。?

01、MoE 的前世今生

混合專家模型（MixtureofExperts：MoE）的思想可以追溯到集成學(xué)習(xí)，集成學(xué)習(xí)是通過訓(xùn)練多個(gè)模型（基學(xué)習(xí)器）來解決同一問題，并且將它們的預(yù)測(cè)結(jié)果簡(jiǎn)單組合（例如投票或平均）。集成學(xué)習(xí)的主要目標(biāo)是通過減少過擬合，提高泛化能力，以提高預(yù)測(cè)性能。常見的集成學(xué)習(xí)方法包括 Bagging，Boosting 和 Stacking。

集成學(xué)習(xí)在訓(xùn)練過程中，利用訓(xùn)練數(shù)據(jù)集訓(xùn)練基學(xué)習(xí)器，基學(xué)習(xí)器的算法可以是決策樹、SVM、線性回歸、KNN 等，在推理過程中對(duì)于輸入的 X，在每個(gè)基學(xué)習(xí)器得到相應(yīng)的答案后將所有結(jié)果有機(jī)統(tǒng)一起來，例如通過求均值的方法解決數(shù)值類問題，通過投票方式解決分類問題。

MoE 和集成學(xué)習(xí)的思想異曲同工，都是集成了多個(gè)模型的方法，但它們的實(shí)現(xiàn)方式有很大不同。與 MoE 的最大不同的地方是集成學(xué)習(xí)不需要將任務(wù)分解為子任務(wù)，而是將多個(gè)基礎(chǔ)學(xué)習(xí)器組合起來。這些基礎(chǔ)學(xué)習(xí)器可以使用相同或不同的算法，并且可以使用相同或不同的訓(xùn)練數(shù)據(jù)。 ? ?

MoE 模型本身也并不是一個(gè)全新的概念，它的理論基礎(chǔ)可以追溯到 1991 年由 MichaelJordan 和 GeoffreyHinton 等人提出的論文，距今已經(jīng)有 30 多年的歷史，但至今依然在被廣泛應(yīng)用的技術(shù)。這一理念在被提出來后經(jīng)常被應(yīng)用到各類模型的實(shí)際場(chǎng)景中，在 2017 年得到了更進(jìn)一步的發(fā)展，當(dāng)時(shí)，一個(gè)由 QuocLe,GeoffreyHinton 和 JeffDean 領(lǐng)銜的團(tuán)隊(duì)提出了一種新型的 MoE 層，它通過引入稀疏性來大幅提高模型的規(guī)模和效率。

大模型結(jié)合混合專家模型的方法屬于老樹發(fā)新芽，隨著應(yīng)用場(chǎng)景的復(fù)雜化和細(xì)分化，大模型越來越大，垂直領(lǐng)域應(yīng)用更加碎片化，想要一個(gè)模型既能回答通識(shí)問題，又能解決專業(yè)領(lǐng)域問題，似乎 MoE 是一種性價(jià)比更高的選擇。在多模態(tài)大模型的發(fā)展浪潮之下，MoE 大有可能成為 2024 年大模型研究的新方向之一，而大模型也會(huì)帶著 MoE，讓其再次偉大。

下面是近些年一部分 MoE 的應(yīng)用發(fā)展事件，可以看出早期 MoE 的應(yīng)用和 Transformer 的發(fā)展時(shí)間節(jié)點(diǎn)差不多，都是在 2017 年左右。 ? ?

2017 年，谷歌首次將 MoE 引入自然語言處理領(lǐng)域，通過在 LSTM 層之間增加 MoE 實(shí)現(xiàn)了機(jī)器翻譯方面的性能提升；

2020 年，Gshard 首次將 MoE 技術(shù)引入 Transformer 架構(gòu)中，并提供了高效的分布式并行計(jì)算架構(gòu)，而后谷歌的 Swtich Transformer 和 GLaM 則進(jìn)一步挖掘 MoE 技術(shù)在自然語言處理領(lǐng)域中的應(yīng)用潛力，實(shí)現(xiàn)了優(yōu)秀的性能表現(xiàn)；

2021 年的 V-MoE 將 MoE 架構(gòu)應(yīng)用在計(jì)算機(jī)視覺領(lǐng)域的 Transformer 架構(gòu)模型中，同時(shí)通過路由算法的改進(jìn)在相關(guān)任務(wù)中實(shí)現(xiàn)了更高的訓(xùn)練效率和更優(yōu)秀的性能表現(xiàn)； ? ?

2022 年的 LIMoE 是首個(gè)應(yīng)用了稀疏混合專家模型技術(shù)的多模態(tài)模型，模型性能相較于 CLIP 也有所提升。

近期 Mistral AI 發(fā)布的 Mistral 8x7B 模型是由 70 億參數(shù)的小模型組合起來的 MoE 模型，直接在多個(gè)跑分上超過了多達(dá) 700 億參數(shù)的 Llama 2。

將混合專家模型（Mixture of Experts：MoE）應(yīng)用于大模型中似乎是不一個(gè)不錯(cuò)的想法，Mistral AI 發(fā)布的 Mistral 8x7B 模型在各項(xiàng)性能和參數(shù)上證明了這一點(diǎn)，使用了更少的參數(shù)卻獲得了遠(yuǎn)超于 Llama 2 的效果，這為大模型的發(fā)展提供了一種新的思路。

02、MoE 的核心思想：術(shù)有專攻

「學(xué)有所長(zhǎng)，術(shù)有專攻」，古人早已將告訴過我們?nèi)绾螌?fù)雜的事物簡(jiǎn)單化處理。大模型從早期只處理文本數(shù)據(jù)，到后來需要同時(shí)處理圖像數(shù)據(jù)和語音數(shù)據(jù)的發(fā)展過程中，其參數(shù)量和模型結(jié)構(gòu)設(shè)計(jì)也越來復(fù)雜和龐大。

如果說單模態(tài)大模型是一個(gè)「特長(zhǎng)生」，那么多模態(tài)大模型就是一個(gè)「全能天才」，想要讓這個(gè)「全能天才」學(xué)習(xí)的更好，那么就需要對(duì)其學(xué)習(xí)任務(wù)分類，安排不同科目的老師進(jìn)行學(xué)習(xí)任務(wù)的輔導(dǎo)，這樣才能讓其高效快速的學(xué)習(xí)到各科的知識(shí)，在考試的時(shí)候才有可能在各科成績(jī)上有優(yōu)異的表現(xiàn)。

混合專家模型（MixtureofExperts：MoE）正是這樣一個(gè)培養(yǎng)「全能天才」的方法，其核心思想就是先把任務(wù)分門別類，然后分給各個(gè)「專家模型」進(jìn)行解決。混合專家模型（MoE）是一種稀疏門控制的深度學(xué)習(xí)模型，它主要由一組專家模型和一個(gè)門控模型組成。MoE 的基本理念是將輸入數(shù)據(jù)根據(jù)任務(wù)類型分割成多個(gè)區(qū)域，并將每個(gè)區(qū)域的數(shù)據(jù)分配一個(gè)或多個(gè)專家模型。每個(gè)專家模型可以專注于處理輸入這部分?jǐn)?shù)據(jù)，從而提高模型的整體性能。

MoE 架構(gòu)的基本原理非常簡(jiǎn)單明了，它主要包括兩個(gè)核心組件：GateNet 和 Experts。GateNet 的作用在于判定輸入樣本應(yīng)該由哪個(gè)專家模型接管處理。而 Experts 則構(gòu)成了一組相對(duì)獨(dú)立的專家模型，每個(gè)專家負(fù)責(zé)處理特定的輸入子空間。

門控模型（GateNet）：混合專家模型中「門」是一種稀疏門網(wǎng)絡(luò)，它接收單個(gè)數(shù)據(jù)元素作為輸入，然后輸出一個(gè)權(quán)重，這些權(quán)重表示每個(gè)專家模型對(duì)處理輸入數(shù)據(jù)的貢獻(xiàn)。一般是通過 softmax 門控函數(shù)通過專家或 token 對(duì)概率分布進(jìn)行建模，并選擇前 K 個(gè)。例如，如果模型有三個(gè)專家，輸出的概率可能為 0.5 和 0.4、0.1，這意味著第一個(gè)專家對(duì)處理此數(shù)據(jù)的貢獻(xiàn)為 50%，第二個(gè)專家為 40%，第二個(gè)專家為 10%，這個(gè)時(shí)候的 K 就可以選擇為 2，我們認(rèn)為前兩個(gè)專家模型的建議會(huì)更好，可以用于更加精確的回答中，而第三個(gè)專家模型的建議可以用于更加富有創(chuàng)意性的答案中。 ? ?

專家模型（Experts）：在訓(xùn)練的過程中，輸入的數(shù)據(jù)被門控模型分配到不同的專家模型中進(jìn)行處理；在推理的過程中，被門控選擇的專家會(huì)針對(duì)輸入的數(shù)據(jù)，產(chǎn)生相應(yīng)的輸出。這些輸出最后會(huì)和每個(gè)專家模型處理該特征的能力分配的權(quán)重進(jìn)行加權(quán)組合，形成最終的預(yù)測(cè)結(jié)果。

混合專家模型在訓(xùn)練過程中通過門控模型實(shí)現(xiàn)「因材施教」，進(jìn)而在推理過程中實(shí)現(xiàn)專家模型之間的「博采眾長(zhǎng)」。MoE 的專家模型可以是小型的 MLP 或者復(fù)雜的 LLM。

在傳統(tǒng)的密集模型中，每個(gè)輸入都必須經(jīng)歷完整的計(jì)算流程，這導(dǎo)致了在處理大規(guī)模數(shù)據(jù)時(shí)的顯著計(jì)算成本。然而，在現(xiàn)代深度學(xué)習(xí)中，稀疏混合專家（MoE）模型的引入為解決這一問題提供了一種新的方法。在這種模型中，輸入數(shù)據(jù)只激活或利用了少數(shù)專家模型，而其他專家模型保持不活躍狀態(tài)，形成了「稀疏」結(jié)構(gòu)。這種稀疏性被認(rèn)為是混合專家模型的重要優(yōu)點(diǎn)，不僅在減少計(jì)算負(fù)擔(dān)的同時(shí)，還能提高模型的效率和性能。 ? ?

MoE模型的優(yōu)勢(shì)在于其靈活性和擴(kuò)展性。由于可以動(dòng)態(tài)地調(diào)整專家網(wǎng)絡(luò)的數(shù)量和類型，MoE 模型可以有效地處理大規(guī)模和復(fù)雜的數(shù)據(jù)集。此外，通過并行處理不同的專家網(wǎng)絡(luò)，MoE 模型還可以提高計(jì)算效率。

在實(shí)際應(yīng)用中，MoE 模型常用于處理需要大量計(jì)算資源的任務(wù)，如語言模型、圖像識(shí)別和復(fù)雜的預(yù)測(cè)問題。通過將大型問題分解為更小、更易管理的子問題，MoE 模型能夠提供更高效和精確的解決方案。

03、MoE 的優(yōu)勢(shì)與缺點(diǎn)

混合專家模型的優(yōu)勢(shì)顯而易見，通過 MoE 的方式，可以極大的促進(jìn)大模型的研究和發(fā)展，但也不能忽視其各方面的問題，在實(shí)際應(yīng)用中應(yīng)該結(jié)合具體的需求對(duì)各方面的性能和參數(shù)進(jìn)行一個(gè)權(quán)衡。

混合專家模型（Mixture of Experts，MoE）的優(yōu)勢(shì)：

混合專家模型（Mixture of Experts，MoE）具有多方面的優(yōu)勢(shì)，使其在深度學(xué)習(xí)領(lǐng)域得到廣泛應(yīng)用。以下是一些混合專家模型的優(yōu)勢(shì)： ? ?

1. 任務(wù)特異性：采用混合專家方法可以有效地充分利用多個(gè)專家模型的優(yōu)勢(shì)，每個(gè)專家都可以專門處理不同的任務(wù)或數(shù)據(jù)的不同部分，在處理復(fù)雜任務(wù)時(shí)取得更卓越的性能。各個(gè)專家模型能夠針對(duì)不同的數(shù)據(jù)分布和模式進(jìn)行建模，從而顯著提升模型的準(zhǔn)確性和泛化能力，因此模型可以更好地適應(yīng)任務(wù)的復(fù)雜性。這種任務(wù)特異性使得混合專家模型在處理多模態(tài)數(shù)據(jù)和復(fù)雜任務(wù)時(shí)表現(xiàn)出色。

2. 靈活性：混合專家方法展現(xiàn)出卓越的靈活性，能夠根據(jù)任務(wù)的需求靈活選擇并組合適宜的專家模型。模型的結(jié)構(gòu)允許根據(jù)任務(wù)的需要?jiǎng)討B(tài)選擇激活的專家模型，實(shí)現(xiàn)對(duì)輸入數(shù)據(jù)的靈活處理。這使得模型能夠適應(yīng)不同的輸入分布和任務(wù)場(chǎng)景，提高了模型的靈活性。

3. 高效性：由于只有少數(shù)專家模型被激活，大部分模型處于未激活狀態(tài)，混合專家模型具有很高的稀疏性。這種稀疏性帶來了計(jì)算效率的提升，因?yàn)橹挥刑囟ǖ膶＜夷Ｐ蛯?duì)當(dāng)前輸入進(jìn)行處理，減少了計(jì)算的開銷。

4. 表現(xiàn)能力：每個(gè)專家模型可以被設(shè)計(jì)為更加專業(yè)化，能夠更好地捕捉輸入數(shù)據(jù)中的模式和關(guān)系。整體模型通過組合這些專家的輸出，提高了對(duì)復(fù)雜數(shù)據(jù)結(jié)構(gòu)的建模能力，從而增強(qiáng)了模型的性能。 ? ?

5. 可解釋性：由于每個(gè)專家模型相對(duì)獨(dú)立，因此模型的決策過程更易于解釋和理解，為用戶提供更高的可解釋性，這對(duì)于一些對(duì)模型決策過程有強(qiáng)解釋要求的應(yīng)用場(chǎng)景非常重要。

MoE 構(gòu)架還能向 LLM 添加可學(xué)習(xí)參數(shù)，而不增加推理成本。

6. 適應(yīng)大規(guī)模數(shù)據(jù)：混合專家方法是處理大規(guī)模數(shù)據(jù)集的理想選擇，能夠有效地應(yīng)對(duì)數(shù)據(jù)量巨大和特征復(fù)雜的挑戰(zhàn)，可以利用稀疏矩陣的高效計(jì)算，利用 GPU 的并行能力計(jì)算所有專家層，能夠有效地應(yīng)對(duì)海量數(shù)據(jù)和復(fù)雜特征的挑戰(zhàn)。其并行處理不同子任務(wù)的特性，充分發(fā)揮計(jì)算資源，幫助有效地?cái)U(kuò)展模型并減少訓(xùn)練時(shí)間，提高模型在訓(xùn)練和推理階段的效率，使其在大規(guī)模數(shù)據(jù)下具有較強(qiáng)的可擴(kuò)展性，以更低的計(jì)算成本獲得更好的結(jié)果。這種優(yōu)勢(shì)使得混合專家方法成為在大數(shù)據(jù)環(huán)境下進(jìn)行深度學(xué)習(xí)的強(qiáng)有力工具。 ? ?

混合專家模型通過充分利用多個(gè)專家模型的優(yōu)勢(shì)，實(shí)現(xiàn)了在任務(wù)處理、靈活性、計(jì)算效率和可解釋性等方面的平衡，使其成為處理復(fù)雜任務(wù)和大規(guī)模數(shù)據(jù)的有效工具。

混合專家模型（Mixture of Experts，MoE）的問題：

盡管混合專家模型在許多方面具有優(yōu)勢(shì)，但也存在一些問題和挑戰(zhàn)，這些需要在實(shí)際應(yīng)用中謹(jǐn)慎考慮。以下是一些混合專家模型可能面臨的問題：

訓(xùn)練復(fù)雜性：混合專家模型的訓(xùn)練相對(duì)復(fù)雜，尤其是涉及到門控網(wǎng)絡(luò)的參數(shù)調(diào)整。為了正確地學(xué)習(xí)專家的權(quán)重和整體模型的參數(shù)，可能需要更多的訓(xùn)練時(shí)間。

超參數(shù)調(diào)整：選擇適當(dāng)?shù)某瑓?shù)，特別是與門控網(wǎng)絡(luò)相關(guān)的參數(shù)，以達(dá)到最佳性能，是一個(gè)復(fù)雜的任務(wù)。這可能需要通過交叉驗(yàn)證等技術(shù)進(jìn)行仔細(xì)調(diào)整。

專家模型設(shè)計(jì)：專家模型的設(shè)計(jì)對(duì)模型的性能影響顯著。選擇適當(dāng)?shù)膶＜夷Ｐ徒Y(jié)構(gòu)，確保其在特定任務(wù)上有足夠的表現(xiàn)力，是一個(gè)挑戰(zhàn)。

稀疏性失真：在某些情況下，為了實(shí)現(xiàn)稀疏性，門控網(wǎng)絡(luò)可能會(huì)過度地激活或不激活某些專家，導(dǎo)致模型性能下降。需要謹(jǐn)慎設(shè)計(jì)稀疏性調(diào)整策略，以平衡效率和性能。 ? ?

動(dòng)態(tài)性問題：在處理動(dòng)態(tài)或快速變化的數(shù)據(jù)分布時(shí)，門控網(wǎng)絡(luò)可能需要更加靈活的調(diào)整，以適應(yīng)輸入數(shù)據(jù)的變化。這需要額外的處理和設(shè)計(jì)。

對(duì)數(shù)據(jù)噪聲的敏感性：混合專家模型對(duì)于數(shù)據(jù)中的噪聲相對(duì)敏感，可能在一些情況下表現(xiàn)不如其他更簡(jiǎn)單的模型。

此外，還有重要的一點(diǎn)是混合專家模型在分布式計(jì)算環(huán)境下可能面臨通信寬帶瓶頸的問題。這主要涉及到混合專家模型的分布式部署，其中不同的專家模型或門控網(wǎng)絡(luò)可能分布在不同的計(jì)算節(jié)點(diǎn)上。在這種情況下，模型參數(shù)的傳輸和同步可能導(dǎo)致通信開銷過大，成為性能的一個(gè)瓶頸。

04、MoE 相關(guān)論文粗讀

MoE 相關(guān)論文 ? ?

1. Adaptive mixtures of local experts, Neural Computation'1991

2. Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer, ICLR'17

3. GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding, ICLR'21

4. Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity, JMLR'22

5. GLaM: Efficient Scaling of Language Models with Mixture-of-Experts, 2021

6. Go Wider Instead of Deeper, AAAI'22

7. MoEBERT: from BERT to Mixture-of-Experts via Importance-Guided Adaptation, NAACL'22

論文 3 GShard 是第一個(gè)將 MoE 的思想拓展到 Transformer 上的工作，但論文亮點(diǎn)是提出了 GShard 這個(gè)框架，可以方便的做對(duì) MoE 結(jié)構(gòu)做數(shù)據(jù)并行或者模型并行。 ? ?

我們可以重點(diǎn)看其中提出的MoE結(jié)構(gòu)，論文具體的做法是：把 Transformer 的 encoder 和 decoder 中，每隔一個(gè)（every other）的 FFN 層，替換成 position-wise 的 MoE 層，又加了一個(gè)分發(fā)器（Gating），使用的是 Top-2 gating network，即不同的 token 每次會(huì)發(fā)給至多兩個(gè)專家。

文中還提到了很多其他設(shè)計(jì)：

Expert capacity balancing：強(qiáng)制每個(gè) expert 處理的 tokens 數(shù)量在一定范圍內(nèi)。

Local group dispatching：通過把一個(gè) batch 內(nèi)所有的 tokens 分組，來實(shí)現(xiàn)并行化計(jì)算。

Auxiliary loss：也是為了緩解「贏者通吃」問題。

Random routing：在 Top-2 gating 的設(shè)計(jì)下，兩個(gè) expert 如何更高效地進(jìn)行 routing。

論文 4 Switch Transformer 的亮點(diǎn)在于它簡(jiǎn)化了 MoE 的 routing 算法，每個(gè) FFN 層激活的專家個(gè)數(shù)從多個(gè)變成了一個(gè)，提高了計(jì)算效率，可以將語言模型的參數(shù)量擴(kuò)展至 1.6 萬億。 ? ?

論文 5 GLaM 是 Google 在 2021 年推出的一個(gè)超大模型，比 GPT-3 大三倍，但是由于使用了 Sparse MoE 的設(shè)計(jì)，訓(xùn)練成本卻只有 GPT-3 的 1/3，而且在 29 個(gè) NLP 任務(wù)上超越了 GPT-3。 ? ?

以上三篇文章（GShard，Switch-Transformer，GLaM）都是希望通過 MoE 的方式把模型做得盡可能的大，大到普通人玩不起（動(dòng)輒使用幾百個(gè) experts）。

但也有更親民一點(diǎn)的，論文 6 和 7 是關(guān)于如何利用 MoE 去壓縮模型、提高效率。

手把手教你，從零開始實(shí)現(xiàn)一個(gè)稀疏混合專家架構(gòu)語言模型（MoE）

選自huggingface? ??

本文介紹了實(shí)現(xiàn)一個(gè)稀疏混合專家語言模型（MoE）的方法，詳細(xì)解釋了模型的實(shí)施過程，包括采用稀疏混合專家取代傳統(tǒng)的前饋神經(jīng)網(wǎng)絡(luò)，實(shí)現(xiàn) top-k 門控和帶噪聲的 top-k 門控，以及采用 Kaiming He 初始化技術(shù)。作者還說明了從 makemore 架構(gòu)保持不變的元素，比如數(shù)據(jù)集處理、分詞預(yù)處理和語言建模任務(wù)。最后還提供了一個(gè) GitHub 倉庫鏈接，用于實(shí)現(xiàn)模型的整個(gè)過程，是一本不可多得的實(shí)戰(zhàn)教科書。

內(nèi)容簡(jiǎn)介

在混合專家模型 Mixtral 發(fā)布后，混合專家模型（MoE）越來越受到人們的關(guān)注。在稀疏化的混合專家語言模型中，大部分組件都與傳統(tǒng)的 transformers 相同。然而，盡管看似簡(jiǎn)單，但經(jīng)驗(yàn)表明，稀疏混合專家語言模型訓(xùn)練的穩(wěn)定性還存在著一些問題。

像這樣易于修改的小規(guī)模實(shí)現(xiàn)可能有助于快速試驗(yàn)新方法。Hugging Face 上的一篇博客介紹了一種可配置的小規(guī)模稀疏 MoE 實(shí)施方法，也許有助于打算在這個(gè)方向深耕的研究者們進(jìn)行快速試驗(yàn)自己的新方法，并且給出了基于 PyTorch 的詳細(xì)代碼：https://github.com/AviSoori1x/makeMoE/tree/main

本文在 makemore 架構(gòu)的基礎(chǔ)上，進(jìn)行了幾處更改：

使用稀疏混合專家代替單獨(dú)的前饋神經(jīng)網(wǎng)絡(luò)；

Top-k 門控和有噪聲的 Top-k 門控；

參數(shù)初始化使用了 Kaiming He 初始化方法，但本文的重點(diǎn)是可以對(duì)初始化方法進(jìn)行自定義，這樣就可以在 Xavier/Glorot 等初始化中進(jìn)行選擇。

同時(shí)，以下模塊與 makemore 保持一致：

數(shù)據(jù)集、預(yù)處理（分詞）部分以及 Andrej 最初選擇的語言建模任務(wù) - 生成莎士比亞文風(fēng)的文本內(nèi)容

Casusal 自注意力機(jī)制 ? ?

訓(xùn)練循環(huán)

推理邏輯

接下來逐步介紹實(shí)施方案，先從注意力機(jī)制開始。

因果縮放點(diǎn)積注意力機(jī)制

下面這段代碼展示了自注意力機(jī)制的基本概念，并且側(cè)重于使用經(jīng)典的縮放點(diǎn)積自注意力（scaled dot product self-attention.）實(shí)現(xiàn)。在這一自注意力變體機(jī)制中，查詢矩陣、鍵矩陣和值矩陣都來自相同的輸入序列。同時(shí)為了確保自回歸語言生成過程的完整性，特別是在純解碼器模型中，使用了一種掩碼機(jī)制。 ? ?

這種掩碼機(jī)制非常關(guān)鍵，因?yàn)樗梢匝谏w當(dāng)前 token 所處位置之后的任何信息，從而引導(dǎo)模型只關(guān)注序列的前面部分。這種了遮擋 token 后面內(nèi)容的注意力被稱為因果自注意力。值得注意的是，稀疏混合專家模型并不局限于僅有解碼器的 Transformer 架構(gòu)。事實(shí)上，這一領(lǐng)域的許多重要的成果都是圍繞 T5 架構(gòu)展開的，T5 架構(gòu)也包含了 Transformer 模型中的編碼器和解碼器組件。

#This code is borrowed from Andrej Karpathy's makemore repository linked in the repo.The self attention layers in Sparse mixture of experts models are the same asin regular transformer models

torch.manual_seed(1337)B,T,C = 4,8,32 # batch, time, channelsx = torch.randn(B,T,C)

# let's see a single Head perform self-attentionhead_size = 16key = nn.Linear(C, head_size, bias=False)query = nn.Linear(C, head_size, bias=False)value = nn.Linear(C, head_size, bias=False)k = key(x)?? # (B, T, 16)q = query(x) # (B, T, 16)wei =? q @ k.transpose(-2, -1) # (B, T, 16) @ (B, 16, T) ---> (B, T, T)

tril = torch.tril(torch.ones(T, T))#wei = torch.zeros((T,T))wei = wei.masked_fill(tril == 0, float('-inf'))wei = F.softmax(wei, dim=-1) #B,T,T

v = value(x) #B,T,Hout = wei @ v # (B,T,T) @ (B,T,H) -> (B,T,H)out.shape

torch.Size([4, 8, 16])

然后，因果自注意力和多頭因果自注意力的代碼可整理如下。多頭自注意力并行應(yīng)用多個(gè)注意力頭，每個(gè)注意力頭單獨(dú)關(guān)注通道的一個(gè)部分（嵌入維度）。多頭自注意力從本質(zhì)上改善了學(xué)習(xí)過程，并由于其固有的并行能力提高了模型訓(xùn)練的效率。下面這段代碼使用了 dropout 來進(jìn)行正則化，來防止過擬合。 ? ?

#Causal scaled dot product self-Attention Head

n_embd = 64

n_head = 4

n_layer = 4

head_size = 16

dropout = 0.1

class Head(nn.Module):

"""?one head of self-attention """

def?__init__(self, head_size):???????

super().__init__()???????

self.key = nn.Linear(n_embd, head_size, bias=False)???????

self.query = nn.Linear(n_embd, head_size, bias=False)???????

self.value = nn.Linear(n_embd, head_size, bias=False)???????

self.register_buffer('tril', torch.tril(torch.ones(block_size, block_size)))

self.dropout = nn.Dropout(dropout)

def forward(self, x):

B,T,C?= x.shape

k?= self.key(x)?? # (B,T,C)

q?= self.query(x) # (B,T,C) ? ?

#?compute attention scores ("affinities")

wei?= q @ k.transpose(-2,-1) * C**-0.5 # (B, T, C) @ (B, C, T) -> (B, T, T)

wei?= wei.masked_fill(self.tril[:T, :T] == 0, float('-inf')) # (B, T, T)

wei?= F.softmax(wei, dim=-1) # (B, T, T)

wei?= self.dropout(wei)

#?perform the weighted aggregation of the values

v?= self.value(x) # (B,T,C)

out?= wei @ v # (B, T, T) @ (B, T, C) -> (B, T, C)

return?out

多頭自注意力的實(shí)現(xiàn)方式如下：

#Multi-Headed Self Attention

class MultiHeadAttention(nn.Module):

"""?multiple heads of self-attention in parallel """

def?__init__(self, num_heads, head_size):

super().__init__()

self.heads?= nn.ModuleList([Head(head_size) for _ in range(num_heads)])

self.proj?= nn.Linear(n_embd, n_embd)

self.dropout?= nn.Dropout(dropout)

def?forward(self, x):

out?= torch.cat([h(x) for h in self.heads], dim=-1)

out?= self.dropout(self.proj(out)) ? ?

return?out

創(chuàng)建一個(gè)專家模塊

即一個(gè)簡(jiǎn)單的多層感知器

在稀疏混合專家架構(gòu)中，每個(gè) transformer 區(qū)塊內(nèi)的自注意力機(jī)制保持不變。不過，每個(gè)區(qū)塊的結(jié)構(gòu)發(fā)生了巨大的變化：標(biāo)準(zhǔn)的前饋神經(jīng)網(wǎng)絡(luò)被多個(gè)稀疏激活的前饋網(wǎng)絡(luò)（即專家網(wǎng)絡(luò)）所取代。所謂「稀疏激活」，是指序列中的每個(gè) token 只被分配給有限數(shù)量的專家（通常是一個(gè)或兩個(gè)）。

這有助于提高訓(xùn)練和推理速度，因?yàn)槊看吻跋騻鬟f都會(huì)激活少數(shù)專家。不過，所有專家都必須存在 GPU 內(nèi)存中，因此當(dāng)參數(shù)總數(shù)達(dá)到數(shù)千億甚至數(shù)萬億時(shí)，就會(huì)產(chǎn)生部署方面的問題。

#Expert module

class Expert(nn.Module):

"""?An MLP is a simple linear layer followed by a non-linearity i.e. each Expert """

def?__init__(self, n_embd): ? ?

super().__init__()

self.net?= nn.Sequential(

nn.Linear(n_embd,?4 * n_embd),

nn.ReLU(),?

nn.Linear(4?* n_embd, n_embd),

nn.Dropout(dropout),

)

def?forward(self, x):

return?self.net(x)

Top-k 門控的一個(gè)例子

門控網(wǎng)絡(luò)，也稱為路由，確定哪個(gè)專家網(wǎng)絡(luò)接收來自多頭注意力的 token 的輸出。舉個(gè)例子解釋路由的機(jī)制，假設(shè)有 4 個(gè)專家，token 需要被路由到前 2 個(gè)專家中。首先需要通過線性層將 token 輸入到門控網(wǎng)絡(luò)中。該層將對(duì)應(yīng)于（Batch size，Tokens，n_embed）的輸入張量從（2，4，32）維度，投影到對(duì)應(yīng)于（Batch size、Tokens，num_expert）的新形狀：（2、4，4）。其中 n_embed 是輸入的通道維度，num_experts 是專家網(wǎng)絡(luò)的計(jì)數(shù)。

接下來，沿最后一個(gè)維度，找出最大的前兩個(gè)值及其相應(yīng)的索引。

#Understanding how gating works ? ?

num_experts = 4

top_k=2

n_embed=32

#Example multi-head attention output for a simple illustrative example, consider n_embed=32, context_length=4 and batch_size=2

mh_output = torch.randn(2, 4, n_embed)

topkgate_linear = nn.Linear(n_embed, num_experts) # nn.Linear(32, 4)

logits = topkgate_linear(mh_output)

top_k_logits, top_k_indices = logits.topk(top_k, dim=-1)? # Get top-k experts

top_k_logits, top_k_indices

#output:

(tensor([[[ 0.0246, -0.0190],

[?0.1991,? 0.1513],

[?0.9749,? 0.7185],

[?0.4406, -0.8357]],

[[?0.6206, -0.0503],

[?0.8635,? 0.3784], ? ?

[?0.6828,? 0.5972],

[?0.4743,? 0.3420]]], grad_fn=),

tensor([[[2, 3],

[2,?1],

[3,?1],?

[2,?1]],?

[[0,?2],?

[0,?3],??

[3,?2],?

[3,?0]]]))

通過僅保留沿最后一個(gè)維度進(jìn)行比較的前 k 大的值，來獲得稀疏門控的輸出。用負(fù)無窮值填充其余部分，在使用 softmax 激活函數(shù)。負(fù)無窮會(huì)被映射至零，而最大的前兩個(gè)值會(huì)更加突出，且和為 1。要求和為 1 是為了對(duì)專家輸出的內(nèi)容進(jìn)行加權(quán)。

zeros = torch.full_like(logits, float('-inf')) #full_like clones a tensor and fills it with a specified value (like infinity) for masking or calculations.

sparse_logits = zeros.scatter(-1, top_k_indices, top_k_logits)sparse_logits

#output

tensor([[[?? -inf,??? -inf,? 0.0246, -0.0190],

[?? -inf,? 0.1513,? 0.1991,??? -inf],? ? ?

[?? -inf,? 0.7185,??? -inf,? 0.9749],?

[?? -inf,?-0.8357,? 0.4406,??? -inf]],

[[?0.6206,??? -inf, -0.0503,??? -inf],?

[?0.8635,??? -inf,??? -inf,? 0.3784],?

[?? -inf,??? -inf,? 0.5972,? 0.6828],?

[?0.3420,??? -inf,??? -inf,? 0.4743]]], grad_fn=)

gating_output= F.softmax(sparse_logits, dim=-1)

gating_output

#ouput

tensor([[[0.0000, 0.0000, 0.5109, 0.4891],

[0.0000,?0.4881, 0.5119, 0.0000],

[0.0000,?0.4362, 0.0000, 0.5638],

[0.0000,?0.2182, 0.7818, 0.0000]],

[[0.6617,?0.0000, 0.3383, 0.0000],

[0.6190,?0.0000, 0.0000, 0.3810],

[0.0000,?0.0000, 0.4786, 0.5214],

[0.4670,?0.0000, 0.0000, 0.5330]]], grad_fn=)

使用有噪聲的 top-k 門控以實(shí)現(xiàn)負(fù)載平衡

# First define the top k router module ? ?

class TopkRouter(nn.Module):

def?__init__(self, n_embed, num_experts, top_k):

super(TopkRouter,?self).__init__()

self.top_k?= top_k

self.linear?=nn.Linear(n_embed, num_experts)

def?forward(self, mh_ouput):

#?mh_ouput is the output tensor from multihead self attention block

logits?= self.linear(mh_output)

top_k_logits,?indices = logits.topk(self.top_k, dim=-1)

zeros?= torch.full_like(logits, float('-inf'))

sparse_logits?= zeros.scatter(-1, indices, top_k_logits)

router_output?= F.softmax(sparse_logits, dim=-1)

return?router_output, indices

接下來使用下面這段代碼來測(cè)試程序：

#Testing this out:

num_experts = 4

top_k = 2

n_embd = 32

mh_output = torch.randn(2, 4, n_embd)? # Example input

top_k_gate = TopkRouter(n_embd, num_experts, top_k) ? ?

gating_output, indices = top_k_gate(mh_output)

gating_output.shape, gating_output, indices#

And it works!!

#output

(torch.Size([2, 4, 4]),

tensor([[[0.5284, 0.0000, 0.4716, 0.0000],

[0.0000,?0.4592, 0.0000, 0.5408],

[0.0000,?0.3529, 0.0000, 0.6471],

[0.3948,?0.0000, 0.0000, 0.6052]],

[[0.0000,?0.5950, 0.4050, 0.0000],

[0.4456,?0.0000, 0.5544, 0.0000],

[0.7208,?0.0000, 0.0000, 0.2792],

[0.0000,?0.0000, 0.5659, 0.4341]]], grad_fn=),

tensor([[[0, 2],

[3,?1],

[3,?1],

[3,?0]],

[[1,?2],

[2,?0],?

[0,?3],

[2,?3]]])) ? ?

盡管最近發(fā)布的 mixtral 的論文沒有提到這一點(diǎn)，但本文的作者相信有噪聲的 Top-k 門控機(jī)制是訓(xùn)練 MoE 模型的一個(gè)重要工具。從本質(zhì)上講，不會(huì)希望所有的 token 都發(fā)送給同一組「受歡迎」的專家網(wǎng)絡(luò)。人們需要的是能在開發(fā)和探索之間取得良好平衡。為此，為了負(fù)載平衡，從門控的線性層向 logits 激活函數(shù)添加標(biāo)準(zhǔn)正態(tài)噪聲是有幫助的，這使訓(xùn)練更有效率。

#Changing the above to accomodate noisy top-k gating

class NoisyTopkRouter(nn.Module):

def?__init__(self, n_embed, num_experts, top_k):

super(NoisyTopkRouter,?self).__init__()

self.top_k?= top_k

#layer?for router logits

self.topkroute_linear?= nn.Linear(n_embed, num_experts)

self.noise_linear?=nn.Linear(n_embed, num_experts)

def?forward(self, mh_output): ? ?

#?mh_ouput is the output tensor from multihead self attention block

logits?= self.topkroute_linear(mh_output)

#Noise?logits???????

noise_logits = self.noise_linear(mh_output)

#Adding?scaled unit gaussian noise to the logits

noise?= torch.randn_like(logits)*F.softplus(noise_logits)

noisy_logits?= logits + noise

top_k_logits,?indices = noisy_logits.topk(self.top_k, dim=-1)

zeros?= torch.full_like(noisy_logits, float('-inf'))

sparse_logits?= zeros.scatter(-1, indices, top_k_logits)

router_output?= F.softmax(sparse_logits, dim=-1)

return?router_output, indices

再次嘗試代碼：

#Testing this out, again:

num_experts = 8

top_k = 2

n_embd = 16

mh_output = torch.randn(2, 4, n_embd)? # Example input

noisy_top_k_gate = NoisyTopkRouter(n_embd, num_experts, top_k)

gating_output, indices = noisy_top_k_gate(mh_output)

gating_output.shape, gating_output, indices

#It works!!

#output

(torch.Size([2, 4, 8]),

tensor([[[0.4181, 0.0000, 0.5819, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],?

[0.4693,?0.5307, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],

[0.0000,?0.4985, 0.5015, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],

[0.0000,?0.0000, 0.0000, 0.2641, 0.0000, 0.7359, 0.0000, 0.0000]],

[[0.0000,?0.0000, 0.0000, 0.6301, 0.0000, 0.3699, 0.0000, 0.0000],

[0.0000,?0.0000, 0.0000, 0.4766, 0.0000, 0.0000, 0.0000, 0.5234],

[0.0000,?0.0000, 0.0000, 0.6815, 0.0000, 0.0000, 0.3185, 0.0000],

[0.4482,?0.5518, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000]]],

grad_fn=),

tensor([[[2, 0],?

[1,?0],

[2,?1],

[5,?3]],

[[3,?5], ? ?

[7,?3],

[3,?6],

[1,?0]]]))

創(chuàng)建稀疏化的混合專家模塊

在獲得門控網(wǎng)絡(luò)的輸出結(jié)果之后，對(duì)于給定的 token，將前 k 個(gè)值選擇性地與來自相應(yīng)的前 k 個(gè)專家的輸出相乘。這種選擇性乘法的結(jié)果是一個(gè)加權(quán)和，該加權(quán)和構(gòu)成 SparseMoe 模塊的輸出。這個(gè)過程的關(guān)鍵和難點(diǎn)是避免不必要的乘法運(yùn)算，只為前 k 名專家進(jìn)行正向轉(zhuǎn)播。為每個(gè)專家執(zhí)行前向傳播將破壞使用稀疏 MoE 的目的，因?yàn)檫@個(gè)過程將不再是稀疏的。

class SparseMoE(nn.Module):

def?__init__(self, n_embed, num_experts, top_k):

super(SparseMoE,?self).__init__()

self.router?= NoisyTopkRouter(n_embed, num_experts, top_k)

self.experts?= nn.ModuleList([Expert(n_embed) for _ in range(num_experts)])?

self.top_k?= top_k

def?forward(self, x):

gating_output,?indices = self.router(x)

final_output?= torch.zeros_like(x)

#?Reshape inputs for batch processing ? ?

flat_x?= x.view(-1, x.size(-1))

flat_gating_output?= gating_output.view(-1, gating_output.size(-1))

#?Process each expert in parallel

for?i, expert in enumerate(self.experts):

#?Create a mask for the inputs where the current expert is in top-k

expert_mask?= (indices == i).any(dim=-1)

flat_mask?= expert_mask.view(-1)

if?flat_mask.any():

expert_input?= flat_x[flat_mask]

expert_output?= expert(expert_input)

#?Extract and apply gating scores

gating_scores?= flat_gating_output[flat_mask, i].unsqueeze(1)?

weighted_output?= expert_output * gating_scores

#?Update final output additively by indexing and adding

final_output[expert_mask]?+= weighted_output.squeeze(1)

return?final_output

運(yùn)行以下代碼來用樣本測(cè)試上述實(shí)現(xiàn)，可以看到確實(shí)如此！

import torch

import torch.nn as nn

#Let's test this outnum_experts = 8

top_k = 2

n_embd = 16

dropout=0.1

mh_output = torch.randn(4, 8, n_embd)? # Example multi-head attention output

sparse_moe = SparseMoE(n_embd, num_experts, top_k)

final_output = sparse_moe(mh_output)

print("Shape of the final output:", final_output.shape)

Shape of the final output: torch.Size([4, 8, 16])

需要強(qiáng)調(diào)的是，如上代碼所示，從路由 / 門控網(wǎng)絡(luò)輸出的 top_k 本身也很重要。索引確定了被激活的專家是哪些， 對(duì)應(yīng)的值又決定了權(quán)重大小。下圖進(jìn)一步解釋了加權(quán)求和的概念。 ? ?

模塊整合

將多頭自注意力和稀疏混合專家相結(jié)合，形成稀疏混合專家 transformer 塊。就像在 vanilla transformer 塊中一樣，也要使用殘差以確保訓(xùn)練穩(wěn)定，并避免梯度消失等問題。此外，要采用層歸一化來進(jìn)一步穩(wěn)定學(xué)習(xí)過程。

#Create a self attention + mixture of experts block, that may be repeated several number of times

class Block(nn.Module):

"""?Mixture of Experts Transformer block: communication followed by computation (multi-head self attention + SparseMoE) """

def __init__(self, n_embed, n_head, num_experts, top_k):

#?n_embed: embedding dimension, n_head: the number of heads we'd like

super().__init__()

head_size?= n_embed // n_head

self.sa?= MultiHeadAttention(n_head, head_size) ? ?

self.smoe?= SparseMoE(n_embed, num_experts, top_k)

self.ln1?= nn.LayerNorm(n_embed)

self.ln2?= nn.LayerNorm(n_embed)

def?forward(self, x):

x?= x + self.sa(self.ln1(x))

x?= x + self.smoe(self.ln2(x))

return?x

最后，將所有內(nèi)容整合在一起，形成稀疏混合專家語言模型。

class SparseMoELanguageModel(nn.Module):

def?__init__(self):

super().__init__()?

#?each token directly reads off the logits for the next token from a lookup table???????

self.token_embedding_table = nn.Embedding(vocab_size, n_embed)

self.position_embedding_table?= nn.Embedding(block_size, n_embed)

self.blocks?= nn.Sequential(*[Block(n_embed, n_head=n_head, num_experts=num_experts,top_k=top_k) for _ in range(n_layer)])

self.ln_f?= nn.LayerNorm(n_embed) # final layer norm

self.lm_head?= nn.Linear(n_embed, vocab_size) ? ?

def?forward(self, idx, targets=None):

B,?T = idx.shape

#?idx and targets are both (B,T) tensor of integers

tok_emb?= self.token_embedding_table(idx) # (B,T,C)

pos_emb?= self.position_embedding_table(torch.arange(T, device=device)) # (T,C)

x?= tok_emb + pos_emb # (B,T,C)

x?= self.blocks(x) # (B,T,C)

x?= self.ln_f(x) # (B,T,C)

logits?= self.lm_head(x) # (B,T,vocab_size)

if?targets is None:

loss?= None???

else:

B,?T, C = logits.shape

logits?= logits.view(B*T, C)

targets?= targets.view(B*T)

loss?= F.cross_entropy(logits, targets)

return?logits, loss

def?generate(self, idx, max_new_tokens):

#?idx is (B, T) array of indices in the current context

for?_ in range(max_new_tokens):

#?crop idx to the last block_size tokens

idx_cond?= idx[:, -block_size:]

#?get the predictions

logits,?loss = self(idx_cond)

#?focus only on the last time step

logits?= logits[:, -1, :] # becomes (B, C)

#?apply softmax to get probabilities

probs?= F.softmax(logits, dim=-1) # (B, C)

#?sample from the distribution

idx_next?= torch.multinomial(probs, num_samples=1) # (B, 1)

#?append sampled index to the running sequence

idx?= torch.cat((idx, idx_next), dim=1) # (B, T+1)

return?idx

參數(shù)初始化對(duì)于深度神經(jīng)網(wǎng)絡(luò)的高效訓(xùn)練非常重要。由于專家中存在 ReLU 激活，因此這里使用了 Kaiming He 初始化。也可以嘗試在 transformer 中更常用的 Glorot 初始化。杰里米 - 霍華德（Jeremy Howard）的《Fastai》第 2 部分有一個(gè)從頭開始實(shí)現(xiàn)這些功能的精彩講座：https://course.fast.ai/Lessons/lesson17.html ? ?

Glorot 參數(shù)初始化通常被用于 transformer 模型，因此這是一個(gè)可能提高模型性能的方法。

def kaiming_init_weights(m):

if?isinstance (m, (nn.Linear)):

init.kaiming_normal_(m.weight)

model = SparseMoELanguageModel()

model.apply(kaiming_init_weights)

本文作者使用 mlflow 跟蹤并記錄重要指標(biāo)和訓(xùn)練超參數(shù)。

#Using MLFlow

m = model.to(device)

# print the number of parameters in the modelprint(sum(p.numel() for p in m.parameters())/1e6, 'M parameters')

# create a PyTorch optimizer

optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)

#mlflow.set_experiment("makeMoE")with mlflow.start_run():

#If?you use mlflow.autolog() this will be automatically logged. I chose to explicitly log here for completeness???

params = {"batch_size": batch_size , "block_size" : block_size, "max_iters": max_iters, "eval_interval": eval_interval,????????????? "learning_rate": learning_rate, "device": device, "eval_iters": eval_iters, "dropout" : dropout, "num_experts": num_experts, "top_k": top_k }??? mlflow.log_params(params)??? for iter in range(max_iters): ? ?

#?every once in a while evaluate the loss on train and val sets

if?iter % eval_interval == 0 or iter == max_iters - 1:

losses?= estimate_loss()

print(f"step?{iter}: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}")

metrics?= {"train_loss": losses['train'], "val_loss": losses['val']}

mlflow.log_metrics(metrics,?step=iter)

#?sample a batch of data

xb,?yb = get_batch('train')

#?evaluate the loss

logits,?loss = model(xb, yb)

optimizer.zero_grad(set_to_none=True)

loss.backward()

optimizer.step()

8.996545 M parameters ? ?

step 0: train loss 5.3223, val loss 5.3166

step 100: train loss 2.7351, val loss 2.7429

step 200: train loss 2.5125, val loss 2.5233...

.

.

step 4999: train loss 1.5712, val loss 1.7508

記錄訓(xùn)練和驗(yàn)證損失可以很好地指示訓(xùn)練的進(jìn)展情況。該圖顯示，可能應(yīng)該在 4500 次時(shí)停止（當(dāng)驗(yàn)證損失稍微增加時(shí)）

接下來可以使用這個(gè)模型逐字符自回歸地生成文本。

# generate from the model. Not great. Not too bad either

context = torch.zeros((1, 1), dtype=torch.long, device=device)

print(decode(m.generate(context, max_new_tokens=2000)[0].tolist()))

DUKE VINCENVENTIO:

If it ever fecond he town sue kigh now,

That thou wold'st is steen 't. ? ?

SIMNA:

Angent her; no, my a born Yorthort,

Romeoos soun and lawf to your sawe with ch a woft ttastly defy,

To declay the soul art; and meart smad.

CORPIOLLANUS:

Which I cannot shall do from by born und ot cold warrike,

What king we best anone wrave's going of heard and good

Thus playvage; you have wold the grace

....

審核編輯：黃飛

?