在深度學(xué)習(xí)的廣闊天地里，我們時(shí)常追求模型的高效與精準(zhǔn)。精兵簡(jiǎn)政，這一理念在專家混合架構(gòu)MoE中得到了淋漓盡致的體現(xiàn)。MoE，即Mixture of Experts，專家混合架構(gòu)，它的核心思想在于將復(fù)雜的任務(wù)拆解為多個(gè)子任務(wù)，每個(gè)子任務(wù)由專門的“專家”模型負(fù)責(zé)處理。

這種架構(gòu)的優(yōu)勢(shì)在于，它可以根據(jù)輸入數(shù)據(jù)的不同特點(diǎn)，動(dòng)態(tài)地選擇最合適的專家進(jìn)行處理，從而在保證模型性能的同時(shí)，降低了計(jì)算復(fù)雜度和資源消耗。正如古人所言，“多兵不如善用兵”，MoE正是通過(guò)優(yōu)化資源配置，實(shí)現(xiàn)了表達(dá)力的大幅增強(qiáng)。

在接下來(lái)的內(nèi)容中，我們將深入探討MoE的具體實(shí)現(xiàn)原理、應(yīng)用場(chǎng)景以及未來(lái)的發(fā)展趨勢(shì)，以期在這一領(lǐng)域有更深入的理解和探索。

原來(lái)這就是專家混合架構(gòu)

讓我們繼續(xù)借助第2章中的案例，深入了解DeepSeek所帶來(lái)的創(chuàng)新之處。

在春秋末年，吳國(guó)的杰出將領(lǐng)孫武先生負(fù)責(zé)訓(xùn)練和指揮軍隊(duì)。軍營(yíng)中搭建了一座名為“點(diǎn)將臺(tái)”的高臺(tái)。每天黎明時(shí)分，全軍列隊(duì)。隨著號(hào)角的響起，數(shù)萬(wàn)名士兵整齊劃一地排列，靜候?qū)O將軍挑選將領(lǐng)，帶領(lǐng)他們奔赴戰(zhàn)場(chǎng)。

而這，正是后人口中的“稀疏點(diǎn)將制”——如今我們?cè)谌斯ぶ悄苤蟹Q之為MoE。這套制度背后，蘊(yùn)含著孫武對(duì)兵力調(diào)度的極致智慧，也正好道出了MoE模型為何強(qiáng)大又高效的原理。

將多兵不如善用兵：表達(dá)力大幅增強(qiáng)

孫武的軍營(yíng)中，將領(lǐng)如云，人才濟(jì)濟(jì)。營(yíng)外，是經(jīng)歷過(guò)百戰(zhàn)的精兵強(qiáng)將——左側(cè)有擅長(zhǎng)夜襲的田將軍，右側(cè)有精通水戰(zhàn)的呂都尉，更有深諳兵法、洞察敵情的鐘謀士。每位將領(lǐng)都各自統(tǒng)領(lǐng)著萬(wàn)人的軍隊(duì)，他們都是能夠獨(dú)當(dāng)一面的杰出人才。然而，孫武從不一次性派遣所有將領(lǐng)。他深知：兵力眾多并不總是優(yōu)勢(shì)，戰(zhàn)場(chǎng)上的關(guān)鍵在于“用人得當(dāng)”。因此，每當(dāng)敵情發(fā)生突變，他便登上點(diǎn)將臺(tái)，從眾多將軍中挑選出最適合的兩位來(lái)應(yīng)對(duì)敵軍。在出戰(zhàn)之日，雖然看起來(lái)只有兩支隊(duì)伍奔赴前線，但實(shí)際上，這是整個(gè)將軍團(tuán)隊(duì)智慧的體現(xiàn)。

這正體現(xiàn)了MoE模型的運(yùn)作原理：盡管整個(gè)網(wǎng)絡(luò)由眾多“專家”組成，每個(gè)專家的參數(shù)數(shù)量可達(dá)到數(shù)億，但在進(jìn)行推理時(shí)，通常只激活其中兩個(gè)或少數(shù)幾個(gè)專家。以一個(gè)普通的 MLP 層為例，其參數(shù)量為100M；而在 MoE 層中，即便有10個(gè)專家，每個(gè)專家的參數(shù)量也是100M，總計(jì)參數(shù)量達(dá)到 1B；然而，在實(shí)際操作中，每次僅啟用2個(gè)專家，因此實(shí)際的計(jì)算成本保持在200M，大大提高了用兵的效率。

孫武布兵如 MoE 布網(wǎng)：整軍備戰(zhàn)，局部出擊，智取勝于力拼。

稀而不弱：高效稀疏調(diào)度，輕裝上陣

孫武實(shí)施“點(diǎn)將制”，這不僅體現(xiàn)了他的智慧，更源于他需要迅速調(diào)配兵力、輕裝上陣。他深知兵貴神速，多余的兵力只會(huì)拖慢進(jìn)軍的步伐。每當(dāng)調(diào)動(dòng)軍隊(duì)出征，他僅需向軍令官輕聲下令：“點(diǎn)將謀士、田將軍。”軍令如山，其他人員保持不動(dòng)，唯獨(dú)這兩員大將迅速整頓部隊(duì)。

這與在 MoE 模型中引入的稀疏前向傳播極為相似：通過(guò)門控機(jī)制，僅激活 Top-k 個(gè)專家，而其他專家則保持靜默狀態(tài)。這樣，每次只需計(jì)算一小部分子網(wǎng)絡(luò)，從而節(jié)省了內(nèi)存和計(jì)算資源；同時(shí)避免了重復(fù)學(xué)習(xí)，提升了專家的特化程度。這正是稀疏機(jī)制的核心所在：確保專家僅在必要時(shí)參與計(jì)算。

調(diào)兵無(wú)上限：專家可擴(kuò)展，模型更靈活**

隨著敵軍戰(zhàn)術(shù)的不斷演變，孫武不再滿足于現(xiàn)有的十位將領(lǐng)。他開(kāi)始著手?jǐn)U充軍隊(duì)，吸納更多的專家型將軍。然而，這位智者并未安排所有將領(lǐng)同時(shí)進(jìn)行訓(xùn)練或出征。他深知，在戰(zhàn)斗中每次只需兩位最合適的將軍出戰(zhàn)（參見(jiàn)下圖），即使軍營(yíng)中擁有百將千帥，其目的也是為了在面對(duì)各種不同的戰(zhàn)場(chǎng)時(shí)，能夠有更多的選擇和更精確的調(diào)度。

這正是 MoE 模型可擴(kuò)展性的體現(xiàn)：在模型訓(xùn)練或性能不佳時(shí)，可以通過(guò)新增專家（子網(wǎng)絡(luò)）來(lái)擴(kuò)充“知識(shí)容量”；但每次計(jì)算仍然保持稀疏，僅激活最相關(guān)的 Top-k 專家，不增加實(shí)際計(jì)算負(fù)擔(dān)；隨著專家數(shù)量的增加，模型能適應(yīng)更多任務(wù)場(chǎng)景，遷移性更強(qiáng)。孫武并非盲目擴(kuò)軍，而是精準(zhǔn)備戰(zhàn)；MoE 亦非盲目堆疊參數(shù)，而是理性稀疏計(jì)算的藝術(shù)。

點(diǎn)將有術(shù)：門控網(wǎng)絡(luò)如統(tǒng)帥

盡管將軍眾多，但真正能夠洞悉敵情、善于用人者，正是站在點(diǎn)將臺(tái)上的孫武本人。

在 MoE 模型中，門控網(wǎng)絡(luò)（Gating Network）扮演著至關(guān)重要的角色。每當(dāng)輸入一句話或一個(gè)樣本，它便負(fù)責(zé)判斷當(dāng)前最適合處理該輸入的專家是哪一位，并激活相應(yīng)的子網(wǎng)絡(luò)。例如：輸入一句詩(shī)詞 → 選擇文學(xué)類專家；輸入一段代碼 → 激活程序?qū)＜遥惠斎脶t(yī)學(xué)記錄 → 調(diào)用醫(yī)學(xué)專家進(jìn)行處理。門控網(wǎng)絡(luò)所學(xué)習(xí)的，不僅僅是“誰(shuí)能完成任務(wù)”，更是“誰(shuí)完成得最為出色”。這與孫武多年征戰(zhàn)后練就的“點(diǎn)將直覺(jué)”如出一轍。

MoE（專家混合架構(gòu)）就像孫武兵法中的用兵之道，是一種兼顧力量與智慧的策略工具：

• 兵多將廣：模型參數(shù)多，具備強(qiáng)大的表達(dá)能力；
• 點(diǎn)將得當(dāng)：通過(guò)稀疏激活，精準(zhǔn)調(diào)度專家，降低計(jì)算成本；
• 彈性擴(kuò)軍：可根據(jù)任務(wù)靈活擴(kuò)展或裁剪專家，適應(yīng)性強(qiáng)；
• 門控統(tǒng)帥：由門控機(jī)制挑選最合適的專家，高效又精準(zhǔn)。

MoE 的關(guān)鍵不在于“人多”，而在于“選對(duì)人、用對(duì)人”，讓模型從“全員出動(dòng)”進(jìn)化為“精兵出擊”，大幅提升了效率與智能表現(xiàn)。

從撒豆成兵到撒豆成精：DeepSeek MoE的數(shù)學(xué)推演

讓我們通過(guò)結(jié)合Excel表格和DeepSeek MoE的公式，共同探索機(jī)器學(xué)習(xí)領(lǐng)域中一個(gè)非常重要的創(chuàng)新——DeepSeek MoE。下圖是從傳統(tǒng)的MoE架構(gòu)，到Deepseek MoE架構(gòu)的演變。

在人工智能領(lǐng)域，隨著模型規(guī)模的不斷膨脹和復(fù)雜性的增加，管理這些“知識(shí)專家”變得異常棘手。

若每次任務(wù)都激活所有專家，不僅會(huì)消耗大量電力，還會(huì)導(dǎo)致模型運(yùn)行緩慢，造成資源的浪費(fèi)。

幸運(yùn)的是，一些智慧的科學(xué)家們創(chuàng)造了Mixture of Experts（MoE），它宛如一位指揮官，僅挑選出最適宜的少數(shù)專家參與每一次的任務(wù)。

然而，DeepSeek團(tuán)隊(duì)意識(shí)到，傳統(tǒng)的MoE仍有提升空間。因此，他們提出了一個(gè)升級(jí)版——DeepSeekMoE。這個(gè)創(chuàng)新設(shè)計(jì)融合了兩大核心策略：

• 細(xì)粒度專家分割（圖b）
• 共享專家隔離（圖c）

這兩大策略，旨在實(shí)現(xiàn)一個(gè)共同目標(biāo)：

使每位專家更加專業(yè)，整個(gè)系統(tǒng)運(yùn)行更高效、更節(jié)約、更智能。

細(xì)粒度專家分割

想象一下，如果你要組建一個(gè)救援小組。傳統(tǒng)方法是，每個(gè)人都是"全能戰(zhàn)士"：又要能開(kāi)船、又要能攀巖、又要懂醫(yī)療……雖然厲害，但一旦出任務(wù)，總有人做得不好，因?yàn)樾g(shù)業(yè)有專攻。更聰明的方法是：把任務(wù)細(xì)分，組建小而專的隊(duì)伍。 比如有人專攻救火，有人專攻山地搜救，有人專攻醫(yī)療急救。這樣一來(lái)，每個(gè)人只專注在自己最擅長(zhǎng)的領(lǐng)域，整體效率大大提高！

在 DeepSeekMoE 中，科學(xué)家們也采用了類似的辦法。他們做了兩件事：

一是把每個(gè)原本的大型FFN模塊，切成_m_個(gè)更小的專家；二是為了保持整體運(yùn)算量不變，每次選擇的小專家數(shù)量也同步乘_m_。

舉個(gè)具體例子：假設(shè)以前有16個(gè)大專家，每次從中選擇2個(gè)來(lái)工作，組合方式只有120種。

現(xiàn)在每個(gè)專家被切成4個(gè)小專家，一共有64個(gè)小專家，每次選擇8個(gè)，組合方式一下子飆升到驚人的44億種（準(zhǔn)確是4,426,165,368種）。

這種變化帶來(lái)了巨大的好處：

• 輸入的信息可以被更細(xì)致地分配到最適合的小專家；
• 每位小專家負(fù)責(zé)更窄、更專精的領(lǐng)域；
• 整個(gè)模型變得更聰明、響應(yīng)更快。

關(guān)鍵點(diǎn)提醒：雖然專家變小了，數(shù)量變多了，但總的參數(shù)量和計(jì)算量是一樣的哦，不用擔(dān)心模型負(fù)擔(dān)加重！

共享專家隔離

還有一個(gè)問(wèn)題：在現(xiàn)實(shí)任務(wù)中，有些基本知識(shí)是大家都會(huì)用的。比如，不管做什么任務(wù)，基本的溝通能力、常識(shí)判斷都必不可少。如果讓每個(gè)小專家都自己去重新學(xué)習(xí)這些基礎(chǔ)知識(shí)，那就太浪費(fèi)了。所以 DeepSeekMoE 引入了第二個(gè)絕招——共享專家隔離。

具體來(lái)說(shuō)：預(yù)留出_Ks_個(gè)專家，專門作為“公共知識(shí)庫(kù)”；所有輸入數(shù)據(jù)，不管走不走路由器，都必須經(jīng)過(guò)這批共享專家；剩下的動(dòng)態(tài)專家，才由路由器根據(jù)需要?jiǎng)討B(tài)選擇。

這樣設(shè)計(jì)帶來(lái)了雙重好處：

• 公共知識(shí)集中處理，每位小專家不用重復(fù)存儲(chǔ)常識(shí)，大大減少了參數(shù)冗余；
• 非共享專家可以專心做細(xì)分領(lǐng)域，比如有人專門研究生物學(xué)，有人專門研究物理學(xué)，各自深耕。

注意：因?yàn)楣蚕韺＜沂潜剡x的，所以為了保持整體計(jì)算量穩(wěn)定，動(dòng)態(tài)選擇的小專家數(shù)量會(huì)相應(yīng)減少 _Ks_個(gè)。

它們的數(shù)學(xué)公式就如下所表示：

我們可以把它們投影到我們的圖中：

接下來(lái)嘗試在Excel中實(shí)現(xiàn)它們，參見(jiàn)下圖。

接下來(lái)再初始化一下路由權(quán)重，參見(jiàn)下圖：

其中，_E_1 代表的是共享專家，_E_2~_E_4代表的是路由專家。

我們計(jì)算路由的公式是這樣的：

想象一下，你經(jīng)營(yíng)著一個(gè)聰明的AI事務(wù)所，里面有許多聰明的“專家”——每位專家擅長(zhǎng)不同的領(lǐng)域。每天，系統(tǒng)會(huì)接收到一些“任務(wù)”，每個(gè)任務(wù)由多個(gè)特征組成。你要做的，是判斷哪些專家對(duì)哪個(gè)任務(wù)最感興趣，從而派出最合適的人來(lái)處理問(wèn)題。
我們首先有一組輸入任務(wù)，它們由一個(gè)包含 6 個(gè) Token（也可以理解為詞或片段）的矩陣組成。每個(gè) Token 擁有 5 個(gè)特征。換句話說(shuō)，你眼前有一個(gè) 5 行 × 6 列的表格，每一列代表一個(gè)任務(wù)點(diǎn)，每一行是它的某個(gè)描述維度。
另一方面，AI事務(wù)所的專家團(tuán)體（比如 E1、E2、E3 和 E4）每位都各有專長(zhǎng)。他們的專長(zhǎng)可以被表達(dá)成一個(gè)“興趣向量”——也就是他們關(guān)注每種特征的程度。

這里用了 Excel 的一個(gè)函數(shù)來(lái)評(píng)估專家與任務(wù)的匹配度：

=MMULT(P152:T154, V141#)

這行公式的含義是：讓系統(tǒng)把每一位專家的興趣向量，與每個(gè)任務(wù)的特征做一次點(diǎn)積運(yùn)算，得出一個(gè)“匹配打分表”。它就像是每位專家在看每個(gè)任務(wù)時(shí)會(huì)說(shuō)：“這個(gè)任務(wù)我多感興趣?！?/span>
在數(shù)學(xué)上，它可以被寫成：

S = E · X

其中：
E 表示專家矩陣，每一行是一位專家對(duì)五個(gè)特征的偏好；
X 表示輸入矩陣，每一列是一項(xiàng)任務(wù)的五個(gè)特征；
S 就是輸出結(jié)果，也就是專家對(duì)每個(gè)任務(wù)的打分表。
接下來(lái)再用SoftMax對(duì)它們進(jìn)行進(jìn)一步的計(jì)算，參見(jiàn)下圖：

Softmax 是深度學(xué)習(xí)和 AI 模型中非常常用的一個(gè)“變魔法”的函數(shù)，它的主要作用可以用一句話概括：

把一堆“分?jǐn)?shù)”變成“比例”，讓它們代表“誰(shuí)說(shuō)話聲音更大”。

用簡(jiǎn)單例子說(shuō)：

假設(shè)你有三個(gè)專家對(duì)一個(gè)任務(wù)的打分是：

專家 A：5

專家 B：2

專家 C：1

這只是相對(duì)的“分?jǐn)?shù)”，不能直接用作權(quán)重（因?yàn)樗鼈兛赡芴?、太小、甚至為?fù)數(shù)）。

這時(shí)候，Softmax 會(huì)這樣做：

把所有分?jǐn)?shù)指數(shù)化（變正、放大差距）

再除以總和，讓它們變成 0~1 之間的比例

且所有比例加起來(lái)剛好等于 1（像概率）

比如，softmax 后結(jié)果可能是：

A：0.84

B：0.11

C：0.05

接下來(lái)對(duì)專家得分進(jìn)行排序，也就是對(duì)我們上面的值按列進(jìn)行排序，選擇Top 3的專家。下圖呈現(xiàn)了通過(guò)排序（SORT(...,,-1)）得到的每列專家打分的降序排列結(jié)果，這一過(guò)程常用于混合專家模型（MoE）中進(jìn)行Top-K路由選擇。

其中Excle的公式如下：

=SORT(AC151:AC154,,-1)

圖表內(nèi)容揭示了一個(gè)4行×6列的打分矩陣，其中每列代表一個(gè)Token，每行代表一個(gè)專家。該矩陣列出了每個(gè)Token對(duì)應(yīng)所有專家的親和力得分，并且每列得分均按從高到低排序。

為何要進(jìn)行此類排序？在MoE模型中，我們并非讓所有專家同時(shí)“發(fā)言”，而是讓每個(gè)Token僅選擇得分最高的前K個(gè)專家參與處理。這一過(guò)程被稱為：

Top-K Routing：每個(gè)Token路由至最匹配的K個(gè)專家。

此排序矩陣的作用是什么？你使用的是Excel的函數(shù)：=SORT(...,,-1)，意味著對(duì)每列進(jìn)行降序排列。

排序后的矩陣便于快速判斷哪些專家進(jìn)入了Top-K。

舉個(gè)例子：Token5（第 5 列），原始打分（排序前）可能是：

專家 分?jǐn)?shù)

E1 1.000000000

E2 0.867739153

E3 0.521375896

E4 0.407783590

排序結(jié)果是：

1.000000

0.867739

0.521376

0.407784

所以 Top-3 專家就是：E1、E2、E3

這里，我們同時(shí)對(duì)得分用?Softmax?做歸一化，然后才進(jìn)行的排序。接下來(lái)再對(duì)選出來(lái)的專家，進(jìn)行運(yùn)算：

得到最終的?Gate 權(quán)重

對(duì)應(yīng)的Excel公式：

=IF(AC151:AH154 >= AK153:AP153, AC151:AH154, 0)

下圖是……

這張圖非常清晰地展示了 Mixture of Experts 中的 Top-K 路由過(guò)程，并通過(guò) Excel 函數(shù) =IF(...) 實(shí)現(xiàn)了“打分保留 / 剔除”的門控邏輯。下面我來(lái)解釋這段公式和操作背后的含義。

這段公式的意思是：

“對(duì)于每個(gè)專家對(duì)每個(gè) Token 的打分，如果它進(jìn)入了 Top-3（即分?jǐn)?shù) ≥ 該 Token 第 3 高的分?jǐn)?shù)），那么保留原始值；否則設(shè)為 0?！?/span>

對(duì)應(yīng)的數(shù)學(xué)表達(dá)形式，這是我們之前講過(guò)的 “門控函數(shù)” 的選擇形式：

其中：

s：原始親和力得分（專家打分）

g：門控值（控制專家輸出權(quán)重）

紅色區(qū)域（AK153:AP153）：每列 Top-3 分?jǐn)?shù)底線 每個(gè) Token 的第三高分，用于篩選

黃色門控區(qū)域，滿足條件的得分保留，否則為 0。

為什么這么做？這個(gè)過(guò)程實(shí)現(xiàn)了 Top-K 路由的門控過(guò)濾機(jī)制：

它不是直接根據(jù)“誰(shuí)排第幾名”來(lái)選專家，而是對(duì)每列找出 Top-K 的最小分?jǐn)?shù)，然后保留所有大于等于這個(gè)分?jǐn)?shù)的專家，這樣做的好處是便于向下使用 Softmax（因?yàn)楸Ａ舻氖窃?score）。下圖是……

這里我們可以看到因?yàn)閷＜?是共享專家，所以所有Token都參與了運(yùn)算。

其他專家灰色的區(qū)域都沒(méi)有參加計(jì)算，這樣就大大節(jié)約了計(jì)算資源。

混沌初開(kāi)，眾神歸位：DeepSeek MoE架構(gòu)的代碼實(shí)現(xiàn)

DeepSeek MoE架構(gòu)的代碼實(shí)現(xiàn)基于上述的邏輯和機(jī)制，將復(fù)雜的模型運(yùn)算高效地組織起來(lái)。代碼的核心在于如何精準(zhǔn)地控制專家的選擇和計(jì)算資源的分配。在DeepSeek中，專家們被組織成一個(gè)多層的結(jié)構(gòu)，每一層都負(fù)責(zé)不同的任務(wù)，但都遵循著Top-K路由的門控過(guò)濾原則。

代碼實(shí)現(xiàn)的關(guān)鍵步驟包括：首先，定義每個(gè)專家的計(jì)算函數(shù)和參數(shù)；其次，通過(guò)矩陣運(yùn)算確定每個(gè)Token對(duì)應(yīng)的專家得分；然后，根據(jù)Top-K路由機(jī)制，篩選出得分高于或等于每列Top-K最小分?jǐn)?shù)的專家；最后，對(duì)這些篩選出的專家進(jìn)行Softmax運(yùn)算，以確定每個(gè)專家在最終決策中的權(quán)重。

DeepSeek MoE架構(gòu)的代碼實(shí)現(xiàn)不僅優(yōu)化了模型的性能，還大大提高了計(jì)算效率。通過(guò)精準(zhǔn)地控制專家的選擇和計(jì)算資源的分配，DeepSeek能夠在保證模型精度的同時(shí)，顯著降低計(jì)算成本和時(shí)間。這使得DeepSeek在處理大規(guī)模數(shù)據(jù)和復(fù)雜任務(wù)時(shí)具有顯著的優(yōu)勢(shì)。

下面來(lái)看看筆者的實(shí)現(xiàn)代碼：

引入需要的庫(kù)，示例代碼如下：

import?torch

import?torch.nn?as?nn

import?torch.nn.functional?as?F

定義單個(gè)前饋專家，實(shí)現(xiàn)代碼如下：

class?FeedForwardExpert(nn.Module):

def?__init__(self, d_model):

super().__init__()

self.net = nn.Sequential(

nn.Linear(d_model, d_model),

nn.ReLU(),

nn.Linear(d_model, d_model),

)?# 簡(jiǎn)單兩層 MLP，含 ReLU

def?forward(self, x):

return?self.net(x)?# 前向傳播

定義 DeepSeekGroupedMoE 模塊

class?DeepSeekGroupedMoE(nn.Module):

def?__init__(self, d_model, num_groups, experts_per_group, top_k):

super().__init__()

self.d_model = d_model

self.num_groups = num_groups?# 分幾組

self.experts_per_group = experts_per_group?# 每組幾個(gè)專家

self.total_experts = num_groups * experts_per_group

self.top_k = top_k?# 每個(gè) Token 在組內(nèi)選幾個(gè)專家

# 所有專家（按組攤平）

self.experts = nn.ModuleList([

FeedForwardExpert(d_model)?for?_?in?range(self.total_experts)

])

# 一個(gè)共享專家（所有 Token 都經(jīng)過(guò)）

self.shared_expert = FeedForwardExpert(d_model)

# 路由器部分

self.group_router = nn.Linear(d_model, num_groups)?# 負(fù)責(zé)選擇組

self.intra_router = nn.ModuleList([

nn.Linear(d_model, experts_per_group)?for?_?in?range(num_groups)?# 每組內(nèi)部的路由器

])

前向傳播邏輯

def?forward(self, x):?# 輸入 x: [B, T, D]

B, T, D = x.shape

x_flat = x.view(-1, D)?# [B*T, D] 扁平化為 Token 維度

# 步驟1: 分組路由（決定每個(gè) Token 屬于哪個(gè) group）

group_logits = self.group_router(x_flat)?# [B*T, num_groups]

group_idx = group_logits.argmax(dim=-1)?# 每個(gè) Token 屬于哪個(gè)組 [B*T]

# 步驟2: 在組內(nèi)選擇 top-K 個(gè)專家

output = torch.zeros_like(x_flat)?# 初始化輸出 [B*T, D]

for?g?in?range(self.num_groups):?# 遍歷每個(gè)組

mask = (group_idx == g)?# 選擇屬于第 g 組的 Token

if?mask.sum() ==?0:

continue?# 當(dāng)前組沒(méi)有 Token，跳過(guò)

x_g = x_flat[mask]?# 提取該組 Token，形狀 [N_g, D]

intra_scores = self.intra_router[g](x_g)?# 對(duì)每個(gè) Token 給組內(nèi)專家打分 [N_g, experts_per_group]

# 選擇 top-k 個(gè)專家

topk_vals, topk_idx = torch.topk(intra_scores, self.top_k, dim=-1)

out_g = torch.zeros_like(x_g)?# 初始化當(dāng)前組輸出

for?i?in?range(self.top_k):?# 遍歷每個(gè) top-k 專家

expert_ids = g * self.experts_per_group + topk_idx[:, i]?# 計(jì)算全局 expert ID

gate = topk_vals[:, i].unsqueeze(-1)?# 獲得門控分?jǐn)?shù) [N_g, 1]

# 對(duì)每個(gè) Token 單獨(dú)調(diào)用對(duì)應(yīng)專家

expert_out = torch.stack([

self.experts[expert_id](x_g[j])?# 每個(gè) Token 分別走自己的專家

for?j, expert_id?in?enumerate(expert_ids)

])

out_g += gate * expert_out?# 按照 gate 加權(quán)相加

output[mask] = out_g?# 將該組 Token 的輸出插回原位置

#步驟 3: 共享專家輸出（所有 Token 都過(guò)一遍）

shared_out = self.shared_expert(x_flat)?# 所有 Token 共享專家 [B*T, D]

final = output + shared_out?# 最終輸出為兩者相加 [B*T, D]

return?final.view(B, T, D)?# reshape 回原始形狀

在上面的代碼清單中，我們對(duì)代碼進(jìn)行了盡量詳細(xì)的注釋，相信大家應(yīng)該可以看得懂。

不過(guò)，為了更深入地理解這個(gè)模塊，我們還是簡(jiǎn)單總結(jié)一下它的工作流程：

首先，輸入的序列 x 會(huì)被展平成二維的形式 x_flat，這樣可以對(duì)每個(gè) Token 進(jìn)行獨(dú)立的處理。然后，這些 Token 會(huì)根據(jù)一定的規(guī)則（例如通過(guò)路由網(wǎng)絡(luò)）被分配到不同的專家組中。每個(gè)專家組會(huì)對(duì)分配給自己的 Token 進(jìn)行處理，并輸出一個(gè)專家輸出 expert_out。

接著，這些專家輸出會(huì)根據(jù) gate（門控機(jī)制）進(jìn)行加權(quán)相加，得到加權(quán)后的輸出 out_g。這個(gè)加權(quán)的過(guò)程實(shí)際上是在選擇性地融合不同專家的知識(shí)，讓模型能夠更好地處理輸入序列。

然后，我們將加權(quán)后的輸出 out_g 插回到原始序列的對(duì)應(yīng)位置上，得到完整的輸出 output。這個(gè)步驟確保了即使某些 Token 被分配到了不同的專家組，它們的輸出仍然能夠按照原始序列的順序進(jìn)行組合。

此外，為了讓所有 Token 都能夠從專家網(wǎng)絡(luò)中獲益，我們還設(shè)計(jì)了一個(gè)共享的專家網(wǎng)絡(luò) self.shared_expert。這個(gè)網(wǎng)絡(luò)會(huì)對(duì)所有 Token 進(jìn)行處理，并輸出一個(gè)共享的專家輸出 shared_out。這個(gè)輸出會(huì)與前面的 output 相加，得到最終的輸出 final。

因此，DeepSeekMoE 模塊實(shí)際上是一個(gè)結(jié)合了條件計(jì)算和知識(shí)蒸餾思想的混合專家系統(tǒng)。它通過(guò)對(duì)輸入序列進(jìn)行細(xì)粒度的分組和處理，以及通過(guò)門控機(jī)制和共享專家網(wǎng)絡(luò)進(jìn)行知識(shí)的融合和提煉，從而實(shí)現(xiàn)了對(duì)復(fù)雜任務(wù)的高效和準(zhǔn)確的建模。