電子發(fā)燒友App
技術前沿：AICG——混合專家模型 (MoEs)
什么是混合專家模型？
混合專家模型主要由兩個關鍵部分組成:
稀疏 MoE 層: 這些層代替了傳統(tǒng) Transformer 模型中的前饋網(wǎng)絡 (FFN) 層。MoE 層包含若干“專家”(例如 8 個)，每個專家本身是一個獨立的神經(jīng)網(wǎng)絡。在實際應用中，這些專家通常是前饋網(wǎng)絡 (FFN)，但它們也可以是更復雜的網(wǎng)絡結(jié)構(gòu)，甚至可以是 MoE 層本身，從而形成層級式的 MoE 結(jié)構(gòu)。 ? ?
門控網(wǎng)絡或路由: 這個部分用于決定哪些令牌 (token) 被發(fā)送到哪個專家。例如，在下圖中，“More”這個令牌可能被發(fā)送到第二個專家，而“Parameters”這個令牌被發(fā)送到第一個專家。有時，一個令牌甚至可以被發(fā)送到多個專家。令牌的路由方式是 MoE 使用中的一個關鍵點，因為路由器由學習的參數(shù)組成，并且與網(wǎng)絡的其他部分一同進行預訓練。
?Switch Transformers paper 論文中的 MoE layer
總結(jié)來說，在混合專家模型 (MoE) 中，我們將傳統(tǒng) Transformer 模型中的每個前饋網(wǎng)絡 (FFN) 層替換為 MoE 層，其中 MoE 層由兩個核心部分組成: 一個門控網(wǎng)絡和若干數(shù)量的專家。
盡管混合專家模型 (MoE) 提供了若干顯著優(yōu)勢，例如更高效的預訓練和與稠密模型相比更快的推理速度，但它們也伴隨著一些挑戰(zhàn):
訓練挑戰(zhàn): 雖然 MoE 能夠?qū)崿F(xiàn)更高效的計算預訓練，但它們在微調(diào)階段往往面臨泛化能力不足的問題，長期以來易于引發(fā)過擬合現(xiàn)象。 ? ?
推理挑戰(zhàn): MoE 模型雖然可能擁有大量參數(shù)，但在推理過程中只使用其中的一部分，這使得它們的推理速度快于具有相同數(shù)量參數(shù)的稠密模型。然而，這種模型需要將所有參數(shù)加載到內(nèi)存中，因此對內(nèi)存的需求非常高。以 Mixtral 8x7B 這樣的 MoE 為例，需要足夠的 VRAM 來容納一個 47B 參數(shù)的稠密模型。之所以是 47B 而不是 8 x 7B = 56B，是因為在 MoE 模型中，只有 FFN 層被視為獨立的專家，而模型的其他參數(shù)是共享的。此外，假設每個令牌只使用兩個專家，那么推理速度 (以 FLOPs 計算) 類似于使用 12B 模型 (而不是 14B 模型)，因為雖然它進行了 2x7B 的矩陣乘法計算，但某些層是共享的。
混合專家模型簡史
混合專家模型 (MoE) 的理念起源于 1991 年的論文 Adaptive Mixture of Local Experts。這個概念與集成學習方法相似，旨在為由多個單獨網(wǎng)絡組成的系統(tǒng)建立一個監(jiān)管機制。在這種系統(tǒng)中，每個網(wǎng)絡 (被稱為“專家”) 處理訓練樣本的不同子集，專注于輸入空間的特定區(qū)域。那么，如何選擇哪個專家來處理特定的輸入呢？這就是門控網(wǎng)絡發(fā)揮作用的地方，它決定了分配給每個專家的權(quán)重。在訓練過程中，這些專家和門控網(wǎng)絡都同時接受訓練，以優(yōu)化它們的性能和決策能力。
在 2010 至 2015 年間，兩個獨立的研究領域為混合專家模型 (MoE) 的后續(xù)發(fā)展做出了顯著貢獻:
組件專家: 在傳統(tǒng)的 MoE 設置中，整個系統(tǒng)由一個門控網(wǎng)絡和多個專家組成。在支持向量機 (SVMs) 、高斯過程和其他方法的研究中，MoE 通常被視為整個模型的一部分。然而，Eigen、Ranzato 和 Ilya 的研究 探索了將 MoE 作為更深層網(wǎng)絡的一個組件。這種方法允許將 MoE 嵌入到多層網(wǎng)絡中的某一層，使得模型既大又高效。
條件計算: 傳統(tǒng)的神經(jīng)網(wǎng)絡通過每一層處理所有輸入數(shù)據(jù)。在這一時期，Yoshua Bengio 等研究人員開始探索基于輸入令牌動態(tài)激活或停用網(wǎng)絡組件的方法。 ? ?
這些研究的融合促進了在自然語言處理 (NLP) 領域?qū)旌蠈＜夷Ｐ偷奶剿?。特別是在 2017 年，Shazeer 等人 (團隊包括 Geoffrey Hinton 和 Jeff Dean，后者有時被戲稱為 “谷歌的 Chuck Norris”) 將這一概念應用于 137B 的 LSTM (當時被廣泛應用于 NLP 的架構(gòu)，由 Schmidhuber 提出)。通過引入稀疏性，這項工作在保持極高規(guī)模的同時實現(xiàn)了快速的推理速度。這項工作主要集中在翻譯領域，但面臨著如高通信成本和訓練不穩(wěn)定性等多種挑戰(zhàn)。
?Outrageously Large Neural Network 論文中的 MoE layer
混合專家模型 (MoE) 的引入使得訓練具有數(shù)千億甚至萬億參數(shù)的模型成為可能，如開源的 1.6 萬億參數(shù)的 Switch Transformers 等。這種技術不僅在自然語言處理 (NLP) 領域得到了廣泛應用，也開始在計算機視覺領域進行探索。
最近再次大火的原因，評測超Llama2，混合專家模型（MoE）會是大模型新方向嗎？?
先是 Reddit 上一篇關于 GPT-4 結(jié)構(gòu)的猜測帖，暗示了 GPT-4 可能是由 16 個子模塊組成的專家模型（MoE）的混合體。據(jù)說，這 16 個子模塊中的每一個 MoE 都有 1110 億個參數(shù)（作為參考，GPT-3 有 1750 億個參數(shù)）。盡管不能 100% 確定，但 GPT-4 是一個 MoE 組成的集群這個事很可能是真的。
然后是法國 AI 公司 MistralAI 發(fā)布了全球首個基于混合專家技術的大模型 Mistral-8x7B-MoE，是 8 個 70 億參數(shù)規(guī)模大模型的混合。
主要特點如下： ? ?
它可以非常優(yōu)雅地處理 32K 上下文數(shù)據(jù)；
除了英語外，在法語、德語、意大利語和西班牙語表現(xiàn)也很好；
在代碼能力上表現(xiàn)很強；
指令微調(diào)后 MT-Bench 的得分 8.3 分（GPT-3.5 是 8.32、LLaMA2 70B 是 6.86）；
Mistral-7B×8-MoE 是首個被證明有效的開源的 MoE LLM，相比于早期的 Switch Transformer、GLaM 等研究，Mistral-7B×8-MoE 證明了 MoE 真的可以落地，且效果遠好于相同激活值的 Dense 模型。
而在最近的一個評測中，Mistral-8x7B-MoE 經(jīng)過微調(diào)后的表現(xiàn)超過了 Llama2-65B。
到底什么是 MoE，可以參見?Hugging Face 官方對 MoE 的詳細技術解讀。?
01、MoE 的前世今生
混合專家模型（MixtureofExperts：MoE）的思想可以追溯到集成學習，集成學習是通過訓練多個模型（基學習器）來解決同一問題，并且將它們的預測結(jié)果簡單組合（例如投票或平均）。集成學習的主要目標是通過減少過擬合，提高泛化能力，以提高預測性能。常見的集成學習方法包括 Bagging，Boosting 和 Stacking。
集成學習在訓練過程中，利用訓練數(shù)據(jù)集訓練基學習器，基學習器的算法可以是決策樹、SVM、線性回歸、KNN 等，在推理過程中對于輸入的 X，在每個基學習器得到相應的答案后將所有結(jié)果有機統(tǒng)一起來，例如通過求均值的方法解決數(shù)值類問題，通過投票方式解決分類問題。
MoE 和集成學習的思想異曲同工，都是集成了多個模型的方法，但它們的實現(xiàn)方式有很大不同。與 MoE 的最大不同的地方是集成學習不需要將任務分解為子任務，而是將多個基礎學習器組合起來。這些基礎學習器可以使用相同或不同的算法，并且可以使用相同或不同的訓練數(shù)據(jù)。 ? ?
MoE 模型本身也并不是一個全新的概念，它的理論基礎可以追溯到 1991 年由 MichaelJordan 和 GeoffreyHinton 等人提出的論文，距今已經(jīng)有 30 多年的歷史，但至今依然在被廣泛應用的技術。這一理念在被提出來后經(jīng)常被應用到各類模型的實際場景中，在 2017 年得到了更進一步的發(fā)展，當時，一個由 QuocLe,GeoffreyHinton 和 JeffDean 領銜的團隊提出了一種新型的 MoE 層，它通過引入稀疏性來大幅提高模型的規(guī)模和效率。
大模型結(jié)合混合專家模型的方法屬于老樹發(fā)新芽，隨著應用場景的復雜化和細分化，大模型越來越大，垂直領域應用更加碎片化，想要一個模型既能回答通識問題，又能解決專業(yè)領域問題，似乎 MoE 是一種性價比更高的選擇。在多模態(tài)大模型的發(fā)展浪潮之下，MoE 大有可能成為 2024 年大模型研究的新方向之一，而大模型也會帶著 MoE，讓其再次偉大。
下面是近些年一部分 MoE 的應用發(fā)展事件，可以看出早期 MoE 的應用和 Transformer 的發(fā)展時間節(jié)點差不多，都是在 2017 年左右。 ? ?
2017 年，谷歌首次將 MoE 引入自然語言處理領域，通過在 LSTM 層之間增加 MoE 實現(xiàn)了機器翻譯方面的性能提升；
2020 年，Gshard 首次將 MoE 技術引入 Transformer 架構(gòu)中，并提供了高效的分布式并行計算架構(gòu)，而后谷歌的 Swtich Transformer 和 GLaM 則進一步挖掘 MoE 技術在自然語言處理領域中的應用潛力，實現(xiàn)了優(yōu)秀的性能表現(xiàn)；
2021 年的 V-MoE 將 MoE 架構(gòu)應用在計算機視覺領域的 Transformer 架構(gòu)模型中，同時通過路由算法的改進在相關任務中實現(xiàn)了更高的訓練效率和更優(yōu)秀的性能表現(xiàn)； ? ?
2022 年的 LIMoE 是首個應用了稀疏混合專家模型技術的多模態(tài)模型，模型性能相較于 CLIP 也有所提升。
近期 Mistral AI 發(fā)布的 Mistral 8x7B 模型是由 70 億參數(shù)的小模型組合起來的 MoE 模型，直接在多個跑分上超過了多達 700 億參數(shù)的 Llama 2。
將混合專家模型（Mixture of Experts：MoE）應用于大模型中似乎是不一個不錯的想法，Mistral AI 發(fā)布的 Mistral 8x7B 模型在各項性能和參數(shù)上證明了這一點，使用了更少的參數(shù)卻獲得了遠超于 Llama 2 的效果，這為大模型的發(fā)展提供了一種新的思路。
02、MoE 的核心思想：術有專攻
「學有所長，術有專攻」，古人早已將告訴過我們?nèi)绾螌碗s的事物簡單化處理。大模型從早期只處理文本數(shù)據(jù)，到后來需要同時處理圖像數(shù)據(jù)和語音數(shù)據(jù)的發(fā)展過程中，其參數(shù)量和模型結(jié)構(gòu)設計也越來復雜和龐大。
如果說單模態(tài)大模型是一個「特長生」，那么多模態(tài)大模型就是一個「全能天才」，想要讓這個「全能天才」學習的更好，那么就需要對其學習任務分類，安排不同科目的老師進行學習任務的輔導，這樣才能讓其高效快速的學習到各科的知識，在考試的時候才有可能在各科成績上有優(yōu)異的表現(xiàn)。
混合專家模型（MixtureofExperts：MoE）正是這樣一個培養(yǎng)「全能天才」的方法，其核心思想就是先把任務分門別類，然后分給各個「專家模型」進行解決。混合專家模型（MoE）是一種稀疏門控制的深度學習模型，它主要由一組專家模型和一個門控模型組成。MoE 的基本理念是將輸入數(shù)據(jù)根據(jù)任務類型分割成多個區(qū)域，并將每個區(qū)域的數(shù)據(jù)分配一個或多個專家模型。每個專家模型可以專注于處理輸入這部分數(shù)據(jù)，從而提高模型的整體性能。
MoE 架構(gòu)的基本原理非常簡單明了，它主要包括兩個核心組件：GateNet 和 Experts。GateNet 的作用在于判定輸入樣本應該由哪個專家模型接管處理。而 Experts 則構(gòu)成了一組相對獨立的專家模型，每個專家負責處理特定的輸入子空間。
門控模型（GateNet）：混合專家模型中「門」是一種稀疏門網(wǎng)絡，它接收單個數(shù)據(jù)元素作為輸入，然后輸出一個權(quán)重，這些權(quán)重表示每個專家模型對處理輸入數(shù)據(jù)的貢獻。一般是通過 softmax 門控函數(shù)通過專家或 token 對概率分布進行建模，并選擇前 K 個。例如，如果模型有三個專家，輸出的概率可能為 0.5 和 0.4、0.1，這意味著第一個專家對處理此數(shù)據(jù)的貢獻為 50%，第二個專家為 40%，第二個專家為 10%，這個時候的 K 就可以選擇為 2，我們認為前兩個專家模型的建議會更好，可以用于更加精確的回答中，而第三個專家模型的建議可以用于更加富有創(chuàng)意性的答案中。 ? ?
專家模型（Experts）：在訓練的過程中，輸入的數(shù)據(jù)被門控模型分配到不同的專家模型中進行處理；在推理的過程中，被門控選擇的專家會針對輸入的數(shù)據(jù)，產(chǎn)生相應的輸出。這些輸出最后會和每個專家模型處理該特征的能力分配的權(quán)重進行加權(quán)組合，形成最終的預測結(jié)果。
混合專家模型在訓練過程中通過門控模型實現(xiàn)「因材施教」，進而在推理過程中實現(xiàn)專家模型之間的「博采眾長」。MoE 的專家模型可以是小型的 MLP 或者復雜的 LLM。
在傳統(tǒng)的密集模型中，每個輸入都必須經(jīng)歷完整的計算流程，這導致了在處理大規(guī)模數(shù)據(jù)時的顯著計算成本。然而，在現(xiàn)代深度學習中，稀疏混合專家（MoE）模型的引入為解決這一問題提供了一種新的方法。在這種模型中，輸入數(shù)據(jù)只激活或利用了少數(shù)專家模型，而其他專家模型保持不活躍狀態(tài)，形成了「稀疏」結(jié)構(gòu)。這種稀疏性被認為是混合專家模型的重要優(yōu)點，不僅在減少計算負擔的同時，還能提高模型的效率和性能。 ? ?
MoE模型的優(yōu)勢在于其靈活性和擴展性。由于可以動態(tài)地調(diào)整專家網(wǎng)絡的數(shù)量和類型，MoE 模型可以有效地處理大規(guī)模和復雜的數(shù)據(jù)集。此外，通過并行處理不同的專家網(wǎng)絡，MoE 模型還可以提高計算效率。
在實際應用中，MoE 模型常用于處理需要大量計算資源的任務，如語言模型、圖像識別和復雜的預測問題。通過將大型問題分解為更小、更易管理的子問題，MoE 模型能夠提供更高效和精確的解決方案。
03、MoE 的優(yōu)勢與缺點
混合專家模型的優(yōu)勢顯而易見，通過 MoE 的方式，可以極大的促進大模型的研究和發(fā)展，但也不能忽視其各方面的問題，在實際應用中應該結(jié)合具體的需求對各方面的性能和參數(shù)進行一個權(quán)衡。
混合專家模型（Mixture of Experts，MoE）的優(yōu)勢：
混合專家模型（Mixture of Experts，MoE）具有多方面的優(yōu)勢，使其在深度學習領域得到廣泛應用。以下是一些混合專家模型的優(yōu)勢： ? ?
1. 任務特異性：采用混合專家方法可以有效地充分利用多個專家模型的優(yōu)勢，每個專家都可以專門處理不同的任務或數(shù)據(jù)的不同部分，在處理復雜任務時取得更卓越的性能。各個專家模型能夠針對不同的數(shù)據(jù)分布和模式進行建模，從而顯著提升模型的準確性和泛化能力，因此模型可以更好地適應任務的復雜性。這種任務特異性使得混合專家模型在處理多模態(tài)數(shù)據(jù)和復雜任務時表現(xiàn)出色。
2. 靈活性：混合專家方法展現(xiàn)出卓越的靈活性，能夠根據(jù)任務的需求靈活選擇并組合適宜的專家模型。模型的結(jié)構(gòu)允許根據(jù)任務的需要動態(tài)選擇激活的專家模型，實現(xiàn)對輸入數(shù)據(jù)的靈活處理。這使得模型能夠適應不同的輸入分布和任務場景，提高了模型的靈活性。
3. 高效性：由于只有少數(shù)專家模型被激活，大部分模型處于未激活狀態(tài)，混合專家模型具有很高的稀疏性。這種稀疏性帶來了計算效率的提升，因為只有特定的專家模型對當前輸入進行處理，減少了計算的開銷。
4. 表現(xiàn)能力：每個專家模型可以被設計為更加專業(yè)化，能夠更好地捕捉輸入數(shù)據(jù)中的模式和關系。整體模型通過組合這些專家的輸出，提高了對復雜數(shù)據(jù)結(jié)構(gòu)的建模能力，從而增強了模型的性能。 ? ?
5. 可解釋性：由于每個專家模型相對獨立，因此模型的決策過程更易于解釋和理解，為用戶提供更高的可解釋性，這對于一些對模型決策過程有強解釋要求的應用場景非常重要。
MoE 構(gòu)架還能向 LLM 添加可學習參數(shù)，而不增加推理成本。
6. 適應大規(guī)模數(shù)據(jù)：混合專家方法是處理大規(guī)模數(shù)據(jù)集的理想選擇，能夠有效地應對數(shù)據(jù)量巨大和特征復雜的挑戰(zhàn)，可以利用稀疏矩陣的高效計算，利用 GPU 的并行能力計算所有專家層，能夠有效地應對海量數(shù)據(jù)和復雜特征的挑戰(zhàn)。其并行處理不同子任務的特性，充分發(fā)揮計算資源，幫助有效地擴展模型并減少訓練時間，提高模型在訓練和推理階段的效率，使其在大規(guī)模數(shù)據(jù)下具有較強的可擴展性，以更低的計算成本獲得更好的結(jié)果。這種優(yōu)勢使得混合專家方法成為在大數(shù)據(jù)環(huán)境下進行深度學習的強有力工具。 ? ?
混合專家模型通過充分利用多個專家模型的優(yōu)勢，實現(xiàn)了在任務處理、靈活性、計算效率和可解釋性等方面的平衡，使其成為處理復雜任務和大規(guī)模數(shù)據(jù)的有效工具。
混合專家模型（Mixture of Experts，MoE）的問題：
盡管混合專家模型在許多方面具有優(yōu)勢，但也存在一些問題和挑戰(zhàn)，這些需要在實際應用中謹慎考慮。以下是一些混合專家模型可能面臨的問題：
訓練復雜性：混合專家模型的訓練相對復雜，尤其是涉及到門控網(wǎng)絡的參數(shù)調(diào)整。為了正確地學習專家的權(quán)重和整體模型的參數(shù)，可能需要更多的訓練時間。
超參數(shù)調(diào)整：選擇適當?shù)某瑓?shù)，特別是與門控網(wǎng)絡相關的參數(shù)，以達到最佳性能，是一個復雜的任務。這可能需要通過交叉驗證等技術進行仔細調(diào)整。
專家模型設計：專家模型的設計對模型的性能影響顯著。選擇適當?shù)膶＜夷Ｐ徒Y(jié)構(gòu)，確保其在特定任務上有足夠的表現(xiàn)力，是一個挑戰(zhàn)。
稀疏性失真：在某些情況下，為了實現(xiàn)稀疏性，門控網(wǎng)絡可能會過度地激活或不激活某些專家，導致模型性能下降。需要謹慎設計稀疏性調(diào)整策略，以平衡效率和性能。 ? ?
動態(tài)性問題：在處理動態(tài)或快速變化的數(shù)據(jù)分布時，門控網(wǎng)絡可能需要更加靈活的調(diào)整，以適應輸入數(shù)據(jù)的變化。這需要額外的處理和設計。
對數(shù)據(jù)噪聲的敏感性：混合專家模型對于數(shù)據(jù)中的噪聲相對敏感，可能在一些情況下表現(xiàn)不如其他更簡單的模型。
此外，還有重要的一點是混合專家模型在分布式計算環(huán)境下可能面臨通信寬帶瓶頸的問題。這主要涉及到混合專家模型的分布式部署，其中不同的專家模型或門控網(wǎng)絡可能分布在不同的計算節(jié)點上。在這種情況下，模型參數(shù)的傳輸和同步可能導致通信開銷過大，成為性能的一個瓶頸。
04、MoE 相關論文粗讀
MoE 相關論文 ? ?
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 是第一個將 MoE 的思想拓展到 Transformer 上的工作，但論文亮點是提出了 GShard 這個框架，可以方便的做對 MoE 結(jié)構(gòu)做數(shù)據(jù)并行或者模型并行。 ? ?
我們可以重點看其中提出的MoE結(jié)構(gòu)，論文具體的做法是：把 Transformer 的 encoder 和 decoder 中，每隔一個（every other）的 FFN 層，替換成 position-wise 的 MoE 層，又加了一個分發(fā)器（Gating），使用的是 Top-2 gating network，即不同的 token 每次會發(fā)給至多兩個專家。
文中還提到了很多其他設計：
Expert capacity balancing：強制每個 expert 處理的 tokens 數(shù)量在一定范圍內(nèi)。
Local group dispatching：通過把一個 batch 內(nèi)所有的 tokens 分組，來實現(xiàn)并行化計算。
Auxiliary loss：也是為了緩解「贏者通吃」問題。
Random routing：在 Top-2 gating 的設計下，兩個 expert 如何更高效地進行 routing。
論文 4 Switch Transformer 的亮點在于它簡化了 MoE 的 routing 算法，每個 FFN 層激活的專家個數(shù)從多個變成了一個，提高了計算效率，可以將語言模型的參數(shù)量擴展至 1.6 萬億。 ? ?
論文 5 GLaM 是 Google 在 2021 年推出的一個超大模型，比 GPT-3 大三倍，但是由于使用了 Sparse MoE 的設計，訓練成本卻只有 GPT-3 的 1/3，而且在 29 個 NLP 任務上超越了 GPT-3。 ? ?
以上三篇文章（GShard，Switch-Transformer，GLaM）都是希望通過 MoE 的方式把模型做得盡可能的大，大到普通人玩不起（動輒使用幾百個 experts）。
但也有更親民一點的，論文 6 和 7 是關于如何利用 MoE 去壓縮模型、提高效率。
手把手教你，從零開始實現(xiàn)一個稀疏混合專家架構(gòu)語言模型（MoE）
選自huggingface? ??
本文介紹了實現(xiàn)一個稀疏混合專家語言模型（MoE）的方法，詳細解釋了模型的實施過程，包括采用稀疏混合專家取代傳統(tǒng)的前饋神經(jīng)網(wǎng)絡，實現(xiàn) top-k 門控和帶噪聲的 top-k 門控，以及采用 Kaiming He 初始化技術。作者還說明了從 makemore 架構(gòu)保持不變的元素，比如數(shù)據(jù)集處理、分詞預處理和語言建模任務。最后還提供了一個 GitHub 倉庫鏈接，用于實現(xiàn)模型的整個過程，是一本不可多得的實戰(zhàn)教科書。
內(nèi)容簡介
在混合專家模型 Mixtral 發(fā)布后，混合專家模型（MoE）越來越受到人們的關注。在稀疏化的混合專家語言模型中，大部分組件都與傳統(tǒng)的 transformers 相同。然而，盡管看似簡單，但經(jīng)驗表明，稀疏混合專家語言模型訓練的穩(wěn)定性還存在著一些問題。
像這樣易于修改的小規(guī)模實現(xiàn)可能有助于快速試驗新方法。Hugging Face 上的一篇博客介紹了一種可配置的小規(guī)模稀疏 MoE 實施方法，也許有助于打算在這個方向深耕的研究者們進行快速試驗自己的新方法，并且給出了基于 PyTorch 的詳細代碼：https://github.com/AviSoori1x/makeMoE/tree/main
本文在 makemore 架構(gòu)的基礎上，進行了幾處更改：
使用稀疏混合專家代替單獨的前饋神經(jīng)網(wǎng)絡；
Top-k 門控和有噪聲的 Top-k 門控；
參數(shù)初始化使用了 Kaiming He 初始化方法，但本文的重點是可以對初始化方法進行自定義，這樣就可以在 Xavier/Glorot 等初始化中進行選擇。
同時，以下模塊與 makemore 保持一致：
數(shù)據(jù)集、預處理（分詞）部分以及 Andrej 最初選擇的語言建模任務 - 生成莎士比亞文風的文本內(nèi)容
Casusal 自注意力機制 ? ?
訓練循環(huán)
推理邏輯
接下來逐步介紹實施方案，先從注意力機制開始。
因果縮放點積注意力機制
下面這段代碼展示了自注意力機制的基本概念，并且側(cè)重于使用經(jīng)典的縮放點積自注意力（scaled dot product self-attention.）實現(xiàn)。在這一自注意力變體機制中，查詢矩陣、鍵矩陣和值矩陣都來自相同的輸入序列。同時為了確保自回歸語言生成過程的完整性，特別是在純解碼器模型中，使用了一種掩碼機制。 ? ?
這種掩碼機制非常關鍵，因為它可以掩蓋當前 token 所處位置之后的任何信息，從而引導模型只關注序列的前面部分。這種了遮擋 token 后面內(nèi)容的注意力被稱為因果自注意力。值得注意的是，稀疏混合專家模型并不局限于僅有解碼器的 Transformer 架構(gòu)。事實上，這一領域的許多重要的成果都是圍繞 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])
然后，因果自注意力和多頭因果自注意力的代碼可整理如下。多頭自注意力并行應用多個注意力頭，每個注意力頭單獨關注通道的一個部分（嵌入維度）。多頭自注意力從本質(zhì)上改善了學習過程，并由于其固有的并行能力提高了模型訓練的效率。下面這段代碼使用了 dropout 來進行正則化，來防止過擬合。 ? ?
#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
多頭自注意力的實現(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òu)中，每個 transformer 區(qū)塊內(nèi)的自注意力機制保持不變。不過，每個區(qū)塊的結(jié)構(gòu)發(fā)生了巨大的變化：標準的前饋神經(jīng)網(wǎng)絡被多個稀疏激活的前饋網(wǎng)絡（即專家網(wǎng)絡）所取代。所謂「稀疏激活」，是指序列中的每個 token 只被分配給有限數(shù)量的專家（通常是一個或兩個）。
這有助于提高訓練和推理速度，因為每次前向傳遞都會激活少數(shù)專家。不過，所有專家都必須存在 GPU 內(nèi)存中，因此當參數(shù)總數(shù)達到數(shù)千億甚至數(shù)萬億時，就會產(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 門控的一個例子
門控網(wǎng)絡，也稱為路由，確定哪個專家網(wǎng)絡接收來自多頭注意力的 token 的輸出。舉個例子解釋路由的機制，假設有 4 個專家，token 需要被路由到前 2 個專家中。首先需要通過線性層將 token 輸入到門控網(wǎng)絡中。該層將對應于（Batch size，Tokens，n_embed）的輸入張量從（2，4，32）維度，投影到對應于（Batch size、Tokens，num_expert）的新形狀：（2、4，4）。其中 n_embed 是輸入的通道維度，num_experts 是專家網(wǎng)絡的計數(shù)。
接下來，沿最后一個維度，找出最大的前兩個值及其相應的索引。
#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]]]))
通過僅保留沿最后一個維度進行比較的前 k 大的值，來獲得稀疏門控的輸出。用負無窮值填充其余部分，在使用 softmax 激活函數(shù)。負無窮會被映射至零，而最大的前兩個值會更加突出，且和為 1。要求和為 1 是為了對專家輸出的內(nèi)容進行加權(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 門控以實現(xiàn)負載平衡
# 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
接下來使用下面這段代碼來測試程序：
#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 的論文沒有提到這一點，但本文的作者相信有噪聲的 Top-k 門控機制是訓練 MoE 模型的一個重要工具。從本質(zhì)上講，不會希望所有的 token 都發(fā)送給同一組「受歡迎」的專家網(wǎng)絡。人們需要的是能在開發(fā)和探索之間取得良好平衡。為此，為了負載平衡，從門控的線性層向 logits 激活函數(shù)添加標準正態(tài)噪聲是有幫助的，這使訓練更有效率。
#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)絡的輸出結(jié)果之后，對于給定的 token，將前 k 個值選擇性地與來自相應的前 k 個專家的輸出相乘。這種選擇性乘法的結(jié)果是一個加權(quán)和，該加權(quán)和構(gòu)成 SparseMoe 模塊的輸出。這個過程的關鍵和難點是避免不必要的乘法運算，只為前 k 名專家進行正向轉(zhuǎn)播。為每個專家執(zhí)行前向傳播將破壞使用稀疏 MoE 的目的，因為這個過程將不再是稀疏的。
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
運行以下代碼來用樣本測試上述實現(xiàn)，可以看到確實如此！
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])
需要強調(diào)的是，如上代碼所示，從路由 / 門控網(wǎng)絡輸出的 top_k 本身也很重要。索引確定了被激活的專家是哪些， 對應的值又決定了權(quán)重大小。下圖進一步解釋了加權(quán)求和的概念。 ? ?
模塊整合
將多頭自注意力和稀疏混合專家相結(jié)合，形成稀疏混合專家 transformer 塊。就像在 vanilla transformer 塊中一樣，也要使用殘差以確保訓練穩(wěn)定，并避免梯度消失等問題。此外，要采用層歸一化來進一步穩(wěn)定學習過程。
#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ù)初始化對于深度神經(jīng)網(wǎng)絡的高效訓練非常重要。由于專家中存在 ReLU 激活，因此這里使用了 Kaiming He 初始化。也可以嘗試在 transformer 中更常用的 Glorot 初始化。杰里米 - 霍華德（Jeremy Howard）的《Fastai》第 2 部分有一個從頭開始實現(xiàn)這些功能的精彩講座：https://course.fast.ai/Lessons/lesson17.html ? ?
Glorot 參數(shù)初始化通常被用于 transformer 模型，因此這是一個可能提高模型性能的方法。
def kaiming_init_weights(m):
if?isinstance (m, (nn.Linear)):
init.kaiming_normal_(m.weight)
model = SparseMoELanguageModel()
model.apply(kaiming_init_weights)
本文作者使用 mlflow 跟蹤并記錄重要指標和訓練超參數(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
記錄訓練和驗證損失可以很好地指示訓練的進展情況。該圖顯示，可能應該在 4500 次時停止（當驗證損失稍微增加時）
接下來可以使用這個模型逐字符自回歸地生成文本。
# 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
....
審核編輯：黃飛
?
工商網(wǎng)監(jiān)
評論