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