這篇來(lái)自東京科學(xué)技術(shù)大學(xué)、香港中文大學(xué)和香港理工大學(xué)的聯(lián)合研究發(fā)表于2025年5月12日的arXiv預(yù)印本平臺(tái)（論文編號(hào)：arXiv:2505.07260v1），有興趣深入了解的讀者可以通過(guò)該編號(hào)在arXiv網(wǎng)站上找到完整論文。研究團(tuán)隊(duì)由來(lái)自東京科學(xué)技術(shù)大學(xué)的楊遠(yuǎn)航（Yuanhang Yang）、香港中文大學(xué)的王朝政（Chaozheng Wang）以及香港理工大學(xué)的李晶（Jing Li）組成。

當(dāng)我們談?wù)撊斯ぶ悄埽貏e是大型語(yǔ)言模型時(shí)，可以把它們想象成一個(gè)龐大的公司。在傳統(tǒng)的AI模型中，就像一個(gè)小公司里每個(gè)員工都要處理所有類型的工作——既要會(huì)寫文案，又要懂財(cái)務(wù)，還要處理客戶服務(wù)。但隨著公司規(guī)模越來(lái)越大，這種"全能型員工"的模式就顯得效率低下了。于是，AI研究者們開發(fā)了一種叫做"混合專家"（Mixture of Experts，簡(jiǎn)稱MoE）的技術(shù)，就像在公司里建立專業(yè)部門——有專門的財(cái)務(wù)部、市場(chǎng)部、技術(shù)部等，讓每個(gè)部門專注于自己最擅長(zhǎng)的工作。

然而，在目前的AI"公司"里，存在一個(gè)奇怪的現(xiàn)象：負(fù)責(zé)"理解和分析"的注意力部門和負(fù)責(zé)"信息處理"的前饋網(wǎng)絡(luò)部門使用的是完全不同的專家團(tuán)隊(duì)和管理方式。這就好比同一家公司的兩個(gè)重要部門完全不溝通、不共享資源，各自為政。研究團(tuán)隊(duì)意識(shí)到這種分離可能導(dǎo)致資源浪費(fèi)和效率低下，于是提出了一個(gè)革命性的想法：能否讓這兩個(gè)部門使用同一套專家團(tuán)隊(duì)，實(shí)現(xiàn)真正的資源共享和協(xié)同工作？

要理解這個(gè)問(wèn)題的復(fù)雜性，我們需要先了解AI模型的基本工作原理。在傳統(tǒng)的Transformer架構(gòu)中，注意力機(jī)制就像一個(gè)善于"察言觀色"的接待員，能夠判斷在處理某個(gè)問(wèn)題時(shí)應(yīng)該重點(diǎn)關(guān)注哪些信息。而前饋網(wǎng)絡(luò)則像一個(gè)信息處理器，負(fù)責(zé)對(duì)接收到的信息進(jìn)行深度加工和轉(zhuǎn)換。這兩個(gè)組件雖然在同一個(gè)AI系統(tǒng)中工作，但它們的內(nèi)部結(jié)構(gòu)和專家配置方式卻截然不同，就像兩個(gè)部門使用完全不同的工作流程和團(tuán)隊(duì)配置。

研究團(tuán)隊(duì)面臨的核心挑戰(zhàn)是：如何重新設(shè)計(jì)注意力機(jī)制，使其能夠與前饋網(wǎng)絡(luò)使用相同的專家架構(gòu)？這個(gè)問(wèn)題看似簡(jiǎn)單，實(shí)際上卻極其復(fù)雜。注意力機(jī)制涉及查詢、鍵值計(jì)算、softmax操作等多個(gè)步驟，而前饋網(wǎng)絡(luò)則是相對(duì)簡(jiǎn)單的兩層矩陣乘法。要讓兩者使用相同的專家架構(gòu)，就必須找到一種全新的方式來(lái)重新組織注意力的計(jì)算過(guò)程。

經(jīng)過(guò)深入研究，團(tuán)隊(duì)發(fā)現(xiàn)了一個(gè)關(guān)鍵洞察：可以將注意力機(jī)制重新表述為兩個(gè)連續(xù)的操作——先進(jìn)行"信息混合"，然后進(jìn)行"專家處理"。具體來(lái)說(shuō)，他們提出了"預(yù)混合"（pre-mixing）和"后混合"（post-mixing）兩種注意力重構(gòu)方式。在預(yù)混合方法中，系統(tǒng)首先通過(guò)注意力權(quán)重將所有相關(guān)信息融合成一個(gè)綜合表示，然后將這個(gè)綜合表示送給專家進(jìn)行處理。這就像先把所有相關(guān)文件整理成一份綜合報(bào)告，然后交給專家分析。

這種重構(gòu)的妙處在于，它揭示了注意力機(jī)制中價(jià)值投影和輸出投影這兩個(gè)步驟實(shí)際上可以組合成一個(gè)類似前饋網(wǎng)絡(luò)的結(jié)構(gòu)。換句話說(shuō)，研究團(tuán)隊(duì)發(fā)現(xiàn)注意力機(jī)制內(nèi)部隱藏著一個(gè)與前饋網(wǎng)絡(luò)高度相似的"基因"，只是被復(fù)雜的計(jì)算步驟掩蓋了。一旦識(shí)別出這個(gè)共同結(jié)構(gòu)，就可以設(shè)計(jì)出能夠同時(shí)服務(wù)于注意力和前饋網(wǎng)絡(luò)的統(tǒng)一專家架構(gòu)。

基于這一發(fā)現(xiàn)，研究團(tuán)隊(duì)開發(fā)了UMoE（Unified Mixture of Experts）架構(gòu)。在這個(gè)統(tǒng)一架構(gòu)中，整個(gè)AI系統(tǒng)被抽象為三個(gè)基本組件：專家團(tuán)隊(duì)、信息混合操作和路由器。專家團(tuán)隊(duì)負(fù)責(zé)核心的信息處理工作，實(shí)現(xiàn)為標(biāo)準(zhǔn)的兩層前饋網(wǎng)絡(luò)。信息混合操作通過(guò)加權(quán)求和的方式促進(jìn)不同信息之間的交流。路由器則像一個(gè)智能調(diào)度員，負(fù)責(zé)將不同的任務(wù)分配給最合適的專家。

在UMoE架構(gòu)中，注意力層和前饋網(wǎng)絡(luò)層的區(qū)別僅僅在于專家接收的輸入不同：前饋網(wǎng)絡(luò)的專家獨(dú)立處理每個(gè)信息片段，而注意力層的專家處理的是通過(guò)加權(quán)混合后的綜合信息。這種設(shè)計(jì)不僅簡(jiǎn)化了整體架構(gòu)，還實(shí)現(xiàn)了真正的參數(shù)共享——同一組專家可以同時(shí)為注意力和前饋網(wǎng)絡(luò)提供服務(wù)。

為了實(shí)現(xiàn)這種參數(shù)共享，研究團(tuán)隊(duì)還解決了一個(gè)技術(shù)難題：如何為不同的專家生成不同的查詢向量。他們采用了低秩矩陣的技術(shù)，為每個(gè)專家配備了專門的查詢投影矩陣，同時(shí)保持鍵值對(duì)在所有專家之間共享。這就像給每個(gè)專家配備了專門的"問(wèn)題清單"，但他們查閱的是同一套"知識(shí)庫(kù)"。

實(shí)驗(yàn)驗(yàn)證是檢驗(yàn)任何新技術(shù)的關(guān)鍵環(huán)節(jié)。研究團(tuán)隊(duì)在多個(gè)數(shù)據(jù)集上進(jìn)行了廣泛的實(shí)驗(yàn)，包括大規(guī)模的FineWeb-Edu數(shù)據(jù)集（包含1000億個(gè)詞元）和經(jīng)典的Wikitext-103數(shù)據(jù)集。他們比較了UMoE與多種基線模型的性能，包括傳統(tǒng)的密集模型、基于前饋網(wǎng)絡(luò)的MoE模型，以及現(xiàn)有的注意力MoE模型如MoA和SwitchHead。

實(shí)驗(yàn)結(jié)果令人印象深刻。在基礎(chǔ)模型規(guī)模下，UMoE在FineWeb-Edu數(shù)據(jù)集上達(dá)到了20.44的困惑度（困惑度越低表示模型性能越好），明顯優(yōu)于所有對(duì)比方法。即使是UMoE的僅注意力版本（UMoE-Att）也達(dá)到了20.81的困惑度，顯著超越了之前最好的注意力MoE方法。在大規(guī)模模型中，這種優(yōu)勢(shì)依然保持，UMoE達(dá)到了15.95的困惑度，再次證明了統(tǒng)一架構(gòu)的有效性。

更重要的是，這種性能提升并沒(méi)有以計(jì)算效率為代價(jià)。雖然預(yù)混合注意力引入了適度的計(jì)算開銷，但這種開銷隨著模型規(guī)模的增大而變得相對(duì)微不足道。在基礎(chǔ)模型中，UMoE的計(jì)算開銷約為傳統(tǒng)方法的1.17倍，但在大規(guī)模模型中，這個(gè)比例降低到了1.03倍，幾乎可以忽略不計(jì)。

研究團(tuán)隊(duì)還進(jìn)行了詳細(xì)的消融實(shí)驗(yàn)，探索了不同設(shè)計(jì)選擇對(duì)性能的影響。他們發(fā)現(xiàn)，在注意力層和前饋網(wǎng)絡(luò)層之間共享專家參數(shù)確實(shí)能夠帶來(lái)性能提升，而激活函數(shù)在專家內(nèi)部的使用對(duì)于保持模型表達(dá)能力至關(guān)重要。有趣的是，當(dāng)他們嘗試將更多專家分配給注意力層而非前饋網(wǎng)絡(luò)層時(shí)，模型性能進(jìn)一步提升，這支持了他們關(guān)于注意力機(jī)制具有更強(qiáng)表達(dá)能力的理論分析。

在零樣本評(píng)估任務(wù)中，UMoE同樣表現(xiàn)出色。在包括HellaSwag、PIQA、ARC等多個(gè)常用評(píng)估基準(zhǔn)上，UMoE都取得了最高的平均準(zhǔn)確率。這些評(píng)估任務(wù)涵蓋了常識(shí)推理、物理常識(shí)、科學(xué)問(wèn)題等多個(gè)方面，充分驗(yàn)證了UMoE的通用性和實(shí)用性。

專家特化分析揭示了UMoE的另一個(gè)有趣特性。通過(guò)分析不同專家處理的詞元類型，研究團(tuán)隊(duì)發(fā)現(xiàn)某些專家在注意力層和前饋網(wǎng)絡(luò)層中表現(xiàn)出一致的專業(yè)化模式。例如，專家3始終專門處理限定詞，專家46專門處理指示代詞。但也有專家表現(xiàn)出分化的專業(yè)化模式，在不同層中承擔(dān)不同的專業(yè)職能，這暗示了參數(shù)共享可能帶來(lái)更高效的參數(shù)利用。

從技術(shù)實(shí)現(xiàn)的角度來(lái)看，UMoE的優(yōu)勢(shì)不僅體現(xiàn)在性能上，還體現(xiàn)在架構(gòu)的簡(jiǎn)潔性和可擴(kuò)展性上。傳統(tǒng)的注意力MoE方法需要專門的實(shí)現(xiàn)和復(fù)雜的專家設(shè)計(jì)，而UMoE通過(guò)統(tǒng)一的專家架構(gòu)大大簡(jiǎn)化了實(shí)現(xiàn)復(fù)雜度。這種簡(jiǎn)化不僅降低了開發(fā)和維護(hù)成本，還為未來(lái)的創(chuàng)新留下了更大空間。

研究團(tuán)隊(duì)還探索了預(yù)混合和后混合兩種注意力重構(gòu)方式的差異。雖然從數(shù)學(xué)角度看這兩種方式等價(jià)，但在引入非線性激活函數(shù)后，它們會(huì)產(chǎn)生不同的結(jié)果。實(shí)驗(yàn)表明，預(yù)混合方式顯著優(yōu)于后混合方式，這為注意力機(jī)制的改進(jìn)提供了新的視角。預(yù)混合方式通過(guò)先進(jìn)行信息整合再進(jìn)行專家處理，更好地align了注意力機(jī)制的本質(zhì)目標(biāo)。

從更宏觀的視角來(lái)看，UMoE的貢獻(xiàn)不僅在于技術(shù)創(chuàng)新，更在于它為AI架構(gòu)設(shè)計(jì)提供了新的思路。長(zhǎng)期以來(lái)，研究者們習(xí)慣于將注意力機(jī)制和前饋網(wǎng)絡(luò)視為兩個(gè)獨(dú)立的組件，很少考慮它們之間的深層聯(lián)系。UMoE的成功表明，通過(guò)重新審視和重構(gòu)現(xiàn)有組件，我們可能發(fā)現(xiàn)意想不到的改進(jìn)機(jī)會(huì)。

這種統(tǒng)一視角還開啟了進(jìn)一步研究的可能性。既然注意力層和前饋網(wǎng)絡(luò)層可以統(tǒng)一，那么是否可以進(jìn)一步將它們?nèi)诤铣蓡我坏膶咏Y(jié)構(gòu)？研究團(tuán)隊(duì)的發(fā)現(xiàn)表明前饋網(wǎng)絡(luò)實(shí)際上可以視為注意力機(jī)制的一個(gè)特殊情況（注意力矩陣為單位矩陣），這暗示了更深層次的架構(gòu)簡(jiǎn)化可能性。

計(jì)算復(fù)雜度分析顯示，UMoE在不同操作上的開銷分布與傳統(tǒng)方法略有不同。鍵投影的復(fù)雜度降低了，但加權(quán)求和的復(fù)雜度增加了。關(guān)鍵的是，加權(quán)求和的復(fù)雜度隨隱藏維度線性增長(zhǎng)，而專家處理的復(fù)雜度隨隱藏維度平方增長(zhǎng)。這意味著隨著模型規(guī)模的增大，額外的計(jì)算開銷會(huì)變得相對(duì)微不足道。

在實(shí)際應(yīng)用中，UMoE的另一個(gè)優(yōu)勢(shì)是它對(duì)現(xiàn)有訓(xùn)練和推理基礎(chǔ)設(shè)施的兼容性。由于采用了標(biāo)準(zhǔn)的前饋網(wǎng)絡(luò)作為專家，現(xiàn)有的MoE優(yōu)化技術(shù)和加速方法可以直接應(yīng)用到UMoE上。這大大降低了新技術(shù)的部署門檻，使其更容易被研究社區(qū)和工業(yè)界采用。

從資源利用的角度來(lái)看，UMoE實(shí)現(xiàn)了真正的"一專多能"。同一組專家參數(shù)可以同時(shí)服務(wù)于模型的不同組件，這不僅提高了參數(shù)效率，還可能帶來(lái)更好的泛化能力。在參數(shù)數(shù)量相同的情況下，UMoE能夠?qū)崿F(xiàn)更好的性能，這對(duì)于資源受限的應(yīng)用場(chǎng)景具有重要意義。

研究團(tuán)隊(duì)還提供了詳細(xì)的注意力模式分析，揭示了UMoE中專家的工作機(jī)制。通過(guò)可視化分析，他們發(fā)現(xiàn)高排名的專家展現(xiàn)出更加集中和任務(wù)相關(guān)的注意力分布模式。這種現(xiàn)象表明路由機(jī)制能夠有效識(shí)別出具備特定能力的專家，并將相應(yīng)的任務(wù)分配給它們。

值得注意的是，這種專家特化不是預(yù)先設(shè)定的，而是在訓(xùn)練過(guò)程中自然涌現(xiàn)的。這種自組織的特化模式可能比人工設(shè)計(jì)的分工更加高效和靈活。研究結(jié)果顯示，某些專家專門處理特定類型的詞元，如標(biāo)點(diǎn)符號(hào)、程度副詞等，這種細(xì)粒度的專業(yè)化有助于提高整體性能。

從工程實(shí)現(xiàn)的角度來(lái)看，UMoE的偽代碼實(shí)現(xiàn)相對(duì)簡(jiǎn)潔明了。核心思想是在注意力MoE中添加一個(gè)額外的詞元混合操作，這個(gè)操作可以通過(guò)標(biāo)準(zhǔn)的注意力計(jì)算來(lái)實(shí)現(xiàn)。這種簡(jiǎn)潔性不僅降低了實(shí)現(xiàn)難度，還減少了引入錯(cuò)誤的可能性。

在KV緩存方面，UMoE與傳統(tǒng)方法相比也有所改進(jìn)。預(yù)混合注意力只需要為每個(gè)詞元緩存一個(gè)鍵和一個(gè)詞元嵌入，而不是多個(gè)鍵值對(duì)。這種簡(jiǎn)化在長(zhǎng)序列處理中可能帶來(lái)內(nèi)存使用上的優(yōu)勢(shì)。

研究團(tuán)隊(duì)的工作還揭示了一個(gè)更深層的理論洞察：前饋網(wǎng)絡(luò)可以被視為注意力機(jī)制的一個(gè)特殊情況，其中注意力矩陣被約束為單位矩陣。這種理論統(tǒng)一不僅優(yōu)雅，還為未來(lái)的架構(gòu)創(chuàng)新提供了指導(dǎo)原則。如果兩種看似不同的組件實(shí)際上是同一機(jī)制的不同表現(xiàn)形式，那么我們可能需要重新思考現(xiàn)有架構(gòu)中的其他組件關(guān)系。

在實(shí)驗(yàn)設(shè)計(jì)上，研究團(tuán)隊(duì)采用了嚴(yán)格的對(duì)比方法，確保所有比較方法使用相似的參數(shù)數(shù)量和計(jì)算預(yù)算。這種公平比較是評(píng)估新方法有效性的關(guān)鍵。他們還進(jìn)行了MAC（乘法累加操作）匹配的比較，進(jìn)一步驗(yàn)證了UMoE在相同計(jì)算開銷下的性能優(yōu)勢(shì)。

訓(xùn)練過(guò)程中的收斂分析顯示，UMoE不僅最終性能更好，收斂速度也更快。這意味著在實(shí)際應(yīng)用中，UMoE可能需要更少的訓(xùn)練時(shí)間就能達(dá)到預(yù)期性能，從而降低訓(xùn)練成本。快速收斂通常也是模型架構(gòu)優(yōu)越性的一個(gè)重要指標(biāo)。

參數(shù)共享策略的選擇也經(jīng)過(guò)了仔細(xì)的實(shí)驗(yàn)驗(yàn)證。研究團(tuán)隊(duì)比較了不同的共享方案，包括完全共享、部分共享和不共享等。結(jié)果顯示，在固定專家和路由器之間采用獨(dú)立參數(shù)的配置能夠獲得最佳性能，這為實(shí)際部署提供了明確的指導(dǎo)。

在專家分配的實(shí)驗(yàn)中，研究團(tuán)隊(duì)發(fā)現(xiàn)了一個(gè)有趣現(xiàn)象：當(dāng)將所有專家都分配給注意力層時(shí)，模型達(dá)到了最佳性能。這個(gè)發(fā)現(xiàn)支持了他們的理論分析，即注意力機(jī)制比前饋網(wǎng)絡(luò)具有更強(qiáng)的表達(dá)能力。不過(guò)，考慮到計(jì)算開銷，在實(shí)際應(yīng)用中需要在性能和效率之間找到平衡點(diǎn)。

激活函數(shù)的重要性也得到了實(shí)驗(yàn)驗(yàn)證。雖然移除激活函數(shù)后模型仍然可以訓(xùn)練（得益于詞元混合操作和層歸一化保持的非線性），但性能會(huì)明顯下降。這強(qiáng)調(diào)了非線性變換在深度學(xué)習(xí)模型中的關(guān)鍵作用，即使在統(tǒng)一架構(gòu)中也不能忽視。

從更廣闊的研究前景來(lái)看，UMoE開啟了多個(gè)有趣的研究方向。研究團(tuán)隊(duì)提到了用更高效的詞元混合替代方案來(lái)進(jìn)一步優(yōu)化計(jì)算效率的可能性，比如線性注意力機(jī)制。此外，既然注意力和前饋網(wǎng)絡(luò)可以統(tǒng)一，那么將它們完全融合成單一層結(jié)構(gòu)也成為了一個(gè)值得探索的方向。

在實(shí)際部署考慮方面，UMoE的另一個(gè)優(yōu)勢(shì)是它與現(xiàn)有MoE基礎(chǔ)設(shè)施的兼容性。由于使用標(biāo)準(zhǔn)的前饋網(wǎng)絡(luò)作為專家，現(xiàn)有的負(fù)載均衡、通信優(yōu)化等技術(shù)可以直接應(yīng)用。這降低了新技術(shù)的采用門檻，有助于其在實(shí)際系統(tǒng)中的快速部署。

研究團(tuán)隊(duì)提供的開源代碼實(shí)現(xiàn)為社區(qū)進(jìn)一步研究和改進(jìn)UMoE提供了基礎(chǔ)。開源不僅促進(jìn)了技術(shù)的傳播和采用，還有助于發(fā)現(xiàn)潛在問(wèn)題和改進(jìn)機(jī)會(huì)。這種開放的研究態(tài)度是推動(dòng)AI技術(shù)進(jìn)步的重要因素。

說(shuō)到底，UMoE代表的不僅僅是一個(gè)新的技術(shù)方案，更是一種新的思維方式。它告訴我們，有時(shí)候最大的創(chuàng)新來(lái)自于對(duì)現(xiàn)有組件關(guān)系的重新審視和理解。通過(guò)發(fā)現(xiàn)注意力機(jī)制和前饋網(wǎng)絡(luò)之間的深層聯(lián)系，研究團(tuán)隊(duì)不僅改進(jìn)了模型性能，還為整個(gè)領(lǐng)域提供了新的研究視角。

這項(xiàng)工作的影響可能會(huì)超出技術(shù)層面。在AI系統(tǒng)越來(lái)越復(fù)雜的今天，如何設(shè)計(jì)更加統(tǒng)一、簡(jiǎn)潔、高效的架構(gòu)成為了一個(gè)重要挑戰(zhàn)。UMoE提供的統(tǒng)一框架可能會(huì)啟發(fā)更多類似的架構(gòu)創(chuàng)新，推動(dòng)整個(gè)領(lǐng)域向更加優(yōu)雅和高效的方向發(fā)展。

歸根結(jié)底，UMoE的成功證明了基礎(chǔ)研究的價(jià)值。通過(guò)深入理解現(xiàn)有技術(shù)的本質(zhì)，研究者們能夠發(fā)現(xiàn)新的改進(jìn)機(jī)會(huì)，推動(dòng)技術(shù)的持續(xù)進(jìn)步。對(duì)于AI領(lǐng)域來(lái)說(shuō)，這種基于深層理解的創(chuàng)新可能比純粹的工程優(yōu)化更有價(jià)值，因?yàn)樗粌H解決了當(dāng)前的問(wèn)題，還為未來(lái)的發(fā)展指明了方向。對(duì)于有興趣深入了解技術(shù)細(xì)節(jié)的讀者，完整的論文可以通過(guò)arXiv:2505.07260v1獲取，其中包含了詳細(xì)的數(shù)學(xué)推導(dǎo)、實(shí)驗(yàn)設(shè)置和結(jié)果分析。

Q&A

Q1：UMoE到底是什么？它解決了什么問(wèn)題？ A：UMoE是一種新的AI架構(gòu)，將原本分離的注意力機(jī)制和前饋網(wǎng)絡(luò)統(tǒng)一起來(lái)，讓它們共享同一套"專家"團(tuán)隊(duì)。這就像讓公司里的兩個(gè)部門不再各自為政，而是共享專業(yè)人才，從而提高效率、降低成本，同時(shí)獲得更好的性能。

Q2：UMoE會(huì)不會(huì)比傳統(tǒng)方法更慢或更耗資源？ A：雖然UMoE引入了一些額外計(jì)算，但隨著模型規(guī)模增大，這種開銷變得微不足道。在大規(guī)模模型中，UMoE只比傳統(tǒng)方法慢3%左右，但通過(guò)參數(shù)共享實(shí)現(xiàn)了更好的性能，總體上是更高效的。

Q3：普通開發(fā)者能用UMoE嗎？部署難度如何？ A：研究團(tuán)隊(duì)已經(jīng)開源了代碼實(shí)現(xiàn)，UMoE與現(xiàn)有的MoE基礎(chǔ)設(shè)施兼容，可以直接使用現(xiàn)有的優(yōu)化技術(shù)。對(duì)于熟悉Transformer和MoE的開發(fā)者來(lái)說(shuō)，部署UMoE并不復(fù)雜，主要是在注意力層添加詞元混合操作。