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

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

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

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

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

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

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

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

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

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

實(shí)驗(yàn)驗(yàn)證是檢驗(yàn)任何新技術(shù)的關(guān)鍵環(huán)節(jié)。研究團(tuán)隊在多個數(shù)據(jù)集上進(jìn)行了廣泛的實(shí)驗(yàn)，包括大規(guī)模的FineWeb-Edu數(shù)據(jù)集（包含1000億個詞元）和經(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)于所有對比方法。即使是UMoE的僅注意力版本（UMoE-Att）也達(dá)到了20.81的困惑度，顯著超越了之前最好的注意力MoE方法。在大規(guī)模模型中，這種優(yōu)勢依然保持，UMoE達(dá)到了15.95的困惑度，再次證明了統(tǒng)一架構(gòu)的有效性。

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

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

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

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

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

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

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

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

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

在實(shí)際應(yīng)用中，UMoE的另一個優(yōu)勢是它對現(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è)界采用。

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

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

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

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

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

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

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

訓(xùn)練過程中的收斂分析顯示，UMoE不僅最終性能更好，收斂速度也更快。這意味著在實(shí)際應(yīng)用中，UMoE可能需要更少的訓(xùn)練時間就能達(dá)到預(yù)期性能，從而降低訓(xùn)練成本?？焖偈諗客ǔＲ彩悄Ｐ图軜?gòu)優(yōu)越性的一個重要指標(biāo)。

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

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

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

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

在實(shí)際部署考慮方面，UMoE的另一個優(yōu)勢是它與現(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)隊提供的開源代碼實(shí)現(xiàn)為社區(qū)進(jìn)一步研究和改進(jìn)UMoE提供了基礎(chǔ)。開源不僅促進(jìn)了技術(shù)的傳播和采用，還有助于發(fā)現(xiàn)潛在問題和改進(jìn)機(jī)會。這種開放的研究態(tài)度是推動AI技術(shù)進(jìn)步的重要因素。

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

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

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

Q&A

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

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

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