這篇來自東京科學(xué)技術(shù)大學(xué)、香港中文大學(xué)和香港理工大學(xué)的聯(lián)合研究發(fā)表于2025年5月12日的arXiv預(yù)印本平臺（論文編號：arXiv:2505.07260v1），有興趣深入了解的讀者可以通過該編號在arXiv網(wǎng)站上找到完整論文。研究團隊由來自東京科學(xué)技術(shù)大學(xué)的楊遠航（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)象：負責(zé)"理解和分析"的注意力部門和負責(zé)"信息處理"的前饋網(wǎng)絡(luò)部門使用的是完全不同的專家團隊和管理方式。這就好比同一家公司的兩個重要部門完全不溝通、不共享資源，各自為政。研究團隊意識到這種分離可能導(dǎo)致資源浪費和效率低下，于是提出了一個革命性的想法：能否讓這兩個部門使用同一套專家團隊，實現(xiàn)真正的資源共享和協(xié)同工作？

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

說到底，UMoE代表的不僅僅是一個新的技術(shù)方案，更是一種新的思維方式。它告訴我們，有時候最大的創(chuàng)新來自于對現(xiàn)有組件關(guān)系的重新審視和理解。通過發(fā)現(xiàn)注意力機制和前饋網(wǎng)絡(luò)之間的深層聯(lián)系，研究團隊不僅改進了模型性能，還為整個領(lǐ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)新的改進機會，推動技術(shù)的持續(xù)進步。對于AI領(lǐng)域來說，這種基于深層理解的創(chuàng)新可能比純粹的工程優(yōu)化更有價值，因為它不僅解決了當(dāng)前的問題，還為未來的發(fā)展指明了方向。對于有興趣深入了解技術(shù)細節(jié)的讀者，完整的論文可以通過arXiv:2505.07260v1獲取，其中包含了詳細的數(shù)學(xué)推導(dǎo)、實驗設(shè)置和結(jié)果分析。

Q&A

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

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

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