本文由清華大學(xué)和快手可靈團(tuán)隊共同完成。第一作者是清華大學(xué)智能視覺實驗室在讀本科生史明磊。

在生成式 AI 領(lǐng)域，擴(kuò)散模型（Diffusion Models）已成為圖像生成任務(wù)的主流架構(gòu)。然而，傳統(tǒng)擴(kuò)散模型在處理不同噪聲水平和條件輸入時采用統(tǒng)一處理方式，未能充分利用擴(kuò)散過程的異構(gòu)特性，導(dǎo)致計算效率低下，近期，可靈團(tuán)隊推出 DiffMoE（Dynamic Token Selection for Scalable Diffusion Transformers），通過創(chuàng)新的動態(tài)token選擇機(jī)制和全局token池設(shè)計，拓展了擴(kuò)散模型的效率與性能邊界。

• 論文標(biāo)題：DiffMoE: Dynamic Token Selection for Scalable Diffusion Transformers
• 項目主頁：https://shiml20.github.io/DiffMoE/
• 論文地址：https://arxiv.org/abs/2503.14487
• 代碼：https://github.com/KwaiVGI/DiffMoE

核心突破：動態(tài)token選擇與全局上下文感知 

DiffMoE 首次在擴(kuò)散模型中引入批級全局token池（Batch-level Global Token Pool），打破傳統(tǒng)模型對單一樣本內(nèi)token的限制，使專家網(wǎng)絡(luò)能夠跨樣本訪問全局token分布。這種設(shè)計模擬了完整數(shù)據(jù)集的token分布，顯著提升了模型對復(fù)雜噪聲模式的學(xué)習(xí)能力。實驗表明，DiffMoE 在訓(xùn)練損失收斂速度上超越了同等參數(shù)量的密集模型（Dense Models），為多任務(wù)處理提供了更強(qiáng)大的上下文感知能力。 

針對推理階段的計算資源分配問題，DiffMoE 提出動態(tài)容量預(yù)測器（Capacity Predictor），通過輕量級 MLP 網(wǎng)絡(luò)實時調(diào)整專家網(wǎng)絡(luò)的計算負(fù)載。該機(jī)制基于訓(xùn)練時的token路由模式學(xué)習(xí)，在不同噪聲水平和樣本復(fù)雜度間智能分配資源，實現(xiàn)了性能與計算成本的靈活權(quán)衡。例如，在生成困難圖片時自動分配更多計算資源，而在處理簡單圖像時降低負(fù)載，真正做到 「按需計算」。 

性能提升：以少勝多的參數(shù)高效模型

在 ImageNet 256×256 分類條件圖像生成基準(zhǔn)測試中，其他結(jié)構(gòu)細(xì)節(jié)保持一致的公平對比情況下，DiffMoE-L-E8 模型僅用 4.58 億參數(shù) (FID50K 2.13), 超越了擁有 6.75 億參數(shù)的 Dense-DiT-XL 模型（FID 2.19）。通過進(jìn)一步擴(kuò)展實驗，DiffMoE 實現(xiàn)了僅用 1 倍激活參數(shù)就實現(xiàn)了 3 倍于密集模型的性能。此外，DiffMoE 在文本到圖像生成任務(wù)中同樣展現(xiàn)出卓越的泛化能力，相較于 Dense 模型有明顯效率提升。

多維度驗證：從理論到實踐

研究團(tuán)隊通過大規(guī)模實驗驗證了 DiffMoE 的優(yōu)越性： 

動態(tài)計算優(yōu)勢：DiffMoE 的平均激活參數(shù)量較低的情況下實現(xiàn)了性能指標(biāo)的顯著提升，證明了動態(tài)資源分配的高效性；同時，DiffMoE 能夠根據(jù)樣本的難度自動分配計算量。本研究可視化了模型認(rèn)為最困難和最簡單的十類生成。

模型認(rèn)為的最困難的十類

模型認(rèn)為的最簡單的十類

擴(kuò)展性測試：從小型（32M）到大型（458M）配置，DiffMoE 均保持性能正向增長，專家數(shù)量從 2 擴(kuò)展到 16 時 FID 持續(xù)下降；

跨任務(wù)適配：在文本到圖像生成任務(wù)中，DiffMoE 模型在對象生成、空間定位等關(guān)鍵指標(biāo)上全面超越 Dense 基線模型。 

總結(jié)

在這項工作中，研究團(tuán)隊通過動態(tài)token選擇和全局token可訪問性來高效擴(kuò)展擴(kuò)散模型。我們的方法利用專門的專家行為和動態(tài)資源分配，有效解決了擴(kuò)散 Transformer 中固定計算量處理的局限性問題。大量的實驗結(jié)果表明，DiffMoE 在性能上超越了現(xiàn)有的 TC-MoE 和 EC-MoE 方法，以及激活參數(shù)量是其 3 倍的密集型模型。研究團(tuán)隊不僅驗證了它在類別條件生成任務(wù)中的實用性，也驗證了 DiffMoE 在大規(guī)模文本到圖像生成任務(wù)的有效性。雖然為了進(jìn)行公平比較，我們未納入現(xiàn)代混合專家（MoE）模型的改進(jìn)方法，但在未來的工作中，集成諸如細(xì)粒度專家和共享專家等先進(jìn)技術(shù)，將可能帶來新的增益。

更多細(xì)節(jié)請參閱原論文。