混合專家模型 -MoE- 到底是什麼?看完這篇,你也能講明白!
當我們查看最新發布的大語言模型(LLM)時,常常會看到標題中出現 MoE 這一術語。那麼,這個 MoE 代表什麼?爲什麼這麼多 LLM 都在使用它呢?
如果你還不知道什麼是 MoE ,那就碼住今天的文章吧!我們將圖解 MoE,用 50 張圖細緻地探索這一重要組件——混合專家模型(MoE)。
在這篇文章中,我們將深入探討 MoE 的兩個主要組件——專家和路由器,並瞭解它們在典型 LLM 架構中的應用。
什麼是混合專家(MoE)?
混合專家(Mixture of Experts,簡稱 MoE)模型,是一種利用多個不同的子模型(或 “專家”)來提升大語言模型(LLM)質量的技術。
MoE 架構主要由兩個核心組件構成:
-
專家模塊:每個前饋神經網絡(FFNN)層包含一組 “專家”,每次可從中選擇一個子集。這些“專家” 本身通常是 FFNN。
-
路由或者網關網絡,決定將哪些標記(token)分配給哪些專家。
在採用 MoE 架構的大語言模型的每一層中,我們會發現這些專家模塊往往具有這樣的特點:
要知道,所謂的 “專家” 並不是專門針對某個領域,如 “心理學” 或“生物學”。最多,它只是語法層面的專家:
更具體來說,它們的專長是在特定的上下文中處理特定的標記。
路由器(網關網絡)來選擇最適合給定輸入的專家(們):
每個專家並不是一個完整的 LLM,而是 LLM 架構中的子模型部分。
專家(Experts)
爲了探討專家(MoE 中的 E)的代表意義和它們的工作方式,我們首先需要了解 MoE 原本要替代的是什麼——密集層。
密集層 (Dense Layer)
混合專家(MoE)起始於 LLM 相對基礎的功能,即前饋神經網絡(FFNN)。
記住,標準的僅解碼器 Transformer 架構中,FFNN 是在層歸一化後應用的:
一個 FFNN 允許模型使用由注意力機制創建的上下文信息,並通過進一步轉換這些信息來捕捉數據中更爲複雜的關係。
然而,FFNN 的規模會迅速增長。爲了學習這些複雜的關係,它通常會擴展它接收到的輸入:
稀疏層 (Sparse Layers)
傳統 Transformer 中的前饋神經網絡(FFNN)被稱爲密集模型,因爲所有參數(權重和偏置)都會被激活。沒有任何參數被遺漏,所有參數都參與輸出計算。
如果我們更仔細地觀察這種密集模型,可以發現輸入會在某種程度上激活所有參數:
與此相對,稀疏模型僅激活一部分參數,並且與混合專家(MoE)緊密相關。
爲了說明這一點,我們可以將密集模型切割成多個部分(所謂的專家),重新訓練它,並且在給定時間內只激活一組專家:
其底層思想是每個專家在訓練過程中學習不同的信息。然後,在進行推理時,僅使用與給定任務最相關的特定專家。
當被問到問題時,我們可以選擇最適合執行給定任務的專家:
專家(Experts)學到什麼?
如前所述,專家學習的內容比整個領域更爲細緻。因此,稱它們爲 “專家” 有時被認爲是誤導性的。
ST-MoE 論文中編碼器模型的專家專業化
然而,解碼器模型中的專家似乎並沒有同樣類型的專業化。不過,這並不意味着所有專家都是平等的。
一個很好的例子可以在 Mixtral 8x7B 論文中找到,其中每個標記都以第一個專家的選擇爲顏色。
這個視覺圖展示了專家傾向於關注語法而非特定領域。
因此,儘管解碼器專家似乎沒有特定的專長,但它們確實會被一致地用於某些類型的標記。
專家的架構
儘管將專家視爲密集模型中的隱藏層並將其切分成塊是很形象的做法,但它們通常本身就是完整的 FFNN:
由於大多數 LLM 具有多個解碼器塊,因此給定文本將在多個專家之間傳遞,直到文本生成:
所選的專家很可能在標記之間有所不同,這就導致了不同的 “路徑” 被選擇:
如果我們更新解碼器塊的可視化,它現在會包含更多的 FFNN(每個專家一個):
解碼器塊現在有多個 FFNN(每個都是一個 “專家”),可以在推理時使用。
路由機制
現在我們有了一組專家,那麼模型是如何知道應該使用哪些專家的呢?
在專家之前,添加了一個路由器(也叫網關網絡),它經過訓練來選擇哪個專家適合給定的標記。
路由器
路由器(或網關網絡)也是一種前饋神經網絡(FFNN),它根據特定的輸入來選擇專家。它輸出概率,並利用這些概率來選擇最匹配的專家:
專家層返回所選專家的輸出,並乘以門值(選擇概率)。
路由器與專家(其中只有少數幾個被選中)一起構成了 MoE 層:
給定的 MoE 層有兩種大小,分別是稀疏的混合專家或密集的混合專家模型。
這兩者都使用路由器來選擇專家,但稀疏 MoE 只選擇幾個,而密集 MoE 則選擇所有專家,但可能以不同的分佈進行選擇。
例如,給定一組標記,MoE 會將標記分配給所有專家,而稀疏 MoE 只選擇少數專家。
在當前的 LLM 狀態下,當你看到 “MoE” 時,它通常是一個稀疏 MoE,因爲它允許你使用專家子集。這在計算上更便宜,這是 LLM 的一個重要特性。
專家的選擇
門控網絡可以說是任何 MoE 模型中最重要的組成部分,因爲它不僅決定了在推理過程中選擇哪些專家,還決定了在訓練過程中選擇哪些專家。
在最基本的形式中,我們將輸入(x)與路由器權重矩陣(W)相乘:
然後,我們對輸出應用 SoftMax 操作,爲每個專家創建概率分佈 G(x):
路由器使用這種概率分佈來選擇與給定輸入最匹配的專家。
最後,我們將每個路由器的輸出與每個選定的專家相乘,並對結果進行求和。
讓我們把所有內容放在一起,看看輸入是如何通過路由器和專家流動的:
路由的複雜性
然而,這個簡單的函數常常導致路由器會選擇相同的專家,因爲某些專家可能比其他專家學習得更快:
不僅會導致選擇的專家分佈不均,還有一些專家幾乎不會被訓練到,這導致在訓練和推理過程中都會出現問題。
因此,我們希望在訓練和推理過程中,專家之間具有相等的重要性,這就是我們所說的負載均衡。從某種意義上說,這是爲了防止對同一專家的過擬合。
負載均衡
爲了平衡專家的重要性,我們需要關注路由器,因爲它是決定在特定時刻選擇哪些專家的主要組成部分。
KeepTopK
負載均衡路由器的一種方法是通過一個簡單的擴展,叫做 KeepTopK。通過引入可訓練的(高斯)噪聲,我們可以防止總是選擇相同的專家:
然後,除了你想激活的前 k 個專家(例如 2 個)之外,其餘專家的權重將被設置爲 -∞:
通過將這些權重設置爲 -∞,這些權重上的 SoftMax 輸出將導致概率爲 0:
KeepTopK 策略是許多大型語言模型仍然使用的一種方法,儘管有很多有前景的替代方案。請注意,KeepTopK 也可以在沒有添加額外噪聲的情況下使用。
Token 選擇
KeepTopK 策略是將每個 token 路由發送到少數幾個選定的專家。這種方法稱爲 Token Choice,允許將給定的 token 發送到一個專家(top-1 路由):
或發送到多個專家(top-k 路由):
一個主要的好處是,它允許對專家的貢獻進行加權和整合。
輔助損失
爲了在訓練期間獲得更均勻的專家分佈,輔助損失(也稱爲負載均衡損失)被添加到網絡的常規損失中。
它添加了一個約束,強制要求專傢俱有相等的重要性。
這個輔助損失的第一個組成部分是對整個批次中每個專家的路由器值進行求和:
這給出了每位專家的重要性分數,表示無論如何輸入,某位專家被選擇的可能性。
我們可以使用這個來計算變異係數(CV),它告訴我們不同專家之間的重要性分數差異有多大。
例如,如果重要性分數差異很大,CV 就會很高:
相反,如果所有專家的權重分數相似,CV 就會很低(這是我們希望的):
使用這個 CV 分數,我們可以在訓練過程中更新輔助損失,目的是儘可能降低 CV 分數(從而給每個專家相等的重要性):
最後,輔助損失作爲一個獨立的損失加入到訓練中進行優化。
專家容量
不平衡不僅存在於被選擇的專家中,也存在於分配給專家的 token 分佈中。
例如,如果輸入 token 不成比例地分配給某個專家,而忽視了其他專家,那麼也可能導致訓練不足:
這裏的問題不僅是哪些專家被使用,而是它們被使用的頻率。
解決這個問題的一種方法是限制每個專家可以處理的 token 數量,即專家容量。當一個專家的容量達到上限時,剩餘的 token 會被髮送到下一個專家:
如果兩個專家的容量都已滿,那麼 token 將不會被任何專家處理,而是被髮送到下一層,這就是 token 溢出。
通過 Switch Transformer 簡化 MoE
Switch Transformer 是第一個基於 transformer 的 MoE 模型,它處理了 MoE 中的訓練不穩定性問題(如負載均衡)。它簡化了架構和訓練過程,同時提高了訓練穩定性。
切換層
Switch Transformer 是一個 T5 模型(編碼器 - 解碼器),它用切換層替代了傳統的 FFNN 層。切換層是一個稀疏 MoE 層,爲每個 token 選擇一個專家(Top-1 路由)。
路由器不進行任何特殊計算,它只是計算輸入與專家權重的乘積,並取其 softmax(與之前做的操作相同)。
這種架構(Top-1 路由)假設每次路由只需要一個專家來學習如何路由輸入。這與我們之前看到的有所不同,後者假設 token 應該路由到多個專家(Top-k 路由)來學習路由行爲。
容量因子
容量因子是一個重要值,它決定了每個專家可以處理多少 token。Switch Transformer 擴展了這一點,引入了一個直接影響專家容量的容量因子。
專家容量的組成部分非常直接:
如果我們增加容量因子,每個專家就能處理更多的 token。
然而,如果容量因子過大,我們就浪費計算資源。相反,如果容量因子過小,由於 token 溢出,模型性能將下降。
輔助損失
爲了進一步防止丟失 token,引入了輔助損失的簡化版本。
與其計算變異係數,這個簡化的損失衡量了每個專家的 token 分派比例與路由器概率分佈之間的比率:
由於目標是讓 N 個專家之間的路由統一,我們希望向量 P 和 f 的值爲 1/N。
α 是一個超參數,我們可以用它來微調訓練過程中這個損失函數的權重。過高的值會超越主損失函數,過低的值對負載均衡幾乎沒有作用。
視覺模型中的混合專家(MoE)****
MoE 不僅是語言模型獨有的技術。視覺模型(如 ViT)利用基於 transformer 的架構,因此也有可能使用 MoE。
簡要回顧一下,ViT(視覺 transformer)是一種將圖像分割爲圖像塊並像 token 一樣處理的架構。
這些圖像塊(或 token)隨後被映射爲嵌入(並添加了額外的位置嵌入),然後再輸入到常規編碼器中:
當這些圖像塊進入編碼器時,它們就像 token 一樣被處理,這使得該架構能夠很好地利用 MoE。
Vision-MoE
Vision-MoE(V-MoE)是 MoE 在圖像模型中的首次實現之一。它以我們之前看到的 ViT 爲基礎,並將編碼器中的密集 FFNN 替換爲稀疏 MoE。
這使得通常比語言模型小的 ViT 模型可以通過添加專家來進行大規模擴展。
爲了減少硬件約束,每個專家使用了小的預定義專家容量,因爲圖像通常有很多圖像塊。然而,低容量往往會導致圖像塊丟失(類似於 token 溢出)。
爲了保持容量低,網絡會爲圖像塊分配重要性分數,並優先處理那些重要性較高的圖像塊,這叫做批處理優先路由。
因此,如果 token 的百分比下降,我們仍然應該看到重要的圖像塊被路由。
優先路由使得我們能集中處理最重要的圖像塊,從而減少處理的圖像塊數量。
從稀疏混合專家到軟混合專家
在 V-MoE 中,優先級評分器幫助區分更重要和不太重要的圖像塊。然而,圖像塊被分配給每個專家後,未處理的圖像塊中的信息會丟失。
軟 MoE 旨在通過混合圖像塊,從離散到軟圖像塊(token)分配。
在第一步中,我們將輸入 x(圖像塊嵌入)與可學習的矩陣 Φ 相乘。這爲我們提供了路由信息,告訴我們某個 token 與給定專家的相關性。
然後,通過對路由信息矩陣(列)進行 softmax 操作,我們更新每個圖像塊的嵌入。
更新後的圖像塊嵌入實際上是所有圖像塊嵌入的加權平均值。
從視覺上看,就好像所有圖像塊都被混合了。這些合成的圖像塊然後被髮送到每個專家。生成輸出後,它們再次與路由矩陣相乘。
路由矩陣在 token 層面影響輸入,在專家層面影響輸出。
最終,我們得到的是 “軟” 圖像塊 /token,而不是離散輸入。
Mixtral 8x7B 的稀疏參數與活躍參數
MoE 的一個重要特點是其計算需求。由於在給定時間只使用專家子集,因此我們可以訪問比我們實際使用更多的參數。
儘管一個 MoE 有更多的參數需要加載(稀疏參數),但因爲我們只在推理過程中使用部分專家,所以激活的參數較少。
換句話說,我們仍然需要將整個模型(包括所有專家)加載到設備上(稀疏參數),但在推理時,我們只需要使用其中的一部分(活躍參數)。MoE 模型需要更多的 VRAM 來加載所有專家,但在推理時運行速度更快。
讓我們通過一個例子——Mixtral 8x7B,來探索稀疏參數與活躍參數的數量。
這裏,我們可以看到每個專家的大小是 5.6B,而不是 7B(儘管只有 8 個專家)。
我們不得不加載 8x5.6B(46.7B)個參數(以及所有共享參數),但我們只需要使用 2x5.6B(12.8B)個參數進行推理。
結論
到這裏關於混合專家(MoE)的探討就結束了!希望這篇文章能讓你對這一有趣技術的潛力有更深的理解。如今,幾乎所有的模型集合中都包含了至少一個 MoE 變體,感覺它已經成爲了長期存在的技術。
本文作者 Maarten Grootendorst,他也是 “袋鼠書”_Hands-On Large Language Models _的作者,如果你想查看更多關於大語言模型可視化的內容,相信 “袋鼠書” 能夠幫到你,這本書的中文版《圖解大模型》,將於 **2025 年 4 月上市,**小夥伴們敬請期待~
英文版封面,中文版封面
《圖解大模型》
Jay Alammar,Maarten Grootendorst | 著
李博傑 | 譯
美亞 4.7 星評,暢銷書 Hands-on 系列新作。只要具備 Python 基礎,就可以通過本書學習大語言模型,並將大語言模型的能力應用到真正的 AI 實踐中。
本書將爲 Python 開發人員提供使用大模型的實用工具和概念,幫助大家掌握實際應用場景。你將學習如何利用預訓練的大型語言模型進行文案撰寫、文本摘要、語義搜索等任務,構建超越關鍵詞匹配的智能系統。
初審:邢璐
複審:董衝
終審:鄭夢鴿
本文由 Readfog 進行 AMP 轉碼,版權歸原作者所有。
來源:https://mp.weixin.qq.com/s/5GwzX7LpWanlJwGZPTv-tA