DiffMoE：動態Token選擇助力擴散模型性能飛躍，快手&清華團隊打造視覺生成新標杆！
本文由清華大學和快手可靈團隊共同完成。第一作者是清華大學智能視覺實驗室在讀本科生史明磊。
在生成式 AI 領域，擴散模型（Diffusion Models）已成為圖像生成任務的主流架構。然而，傳統擴散模型在處理不同噪聲水平和條件輸入時採用統一處理方式，未能充分利用擴散過程的異構特性，導致計算效率低下，近期，可靈團隊推出 DiffMoE（Dynamic Token Selection for Scalable Diffusion Transformers），通過創新的動態token選擇機制和全局token池設計，拓展了擴散模型的效率與性能邊界。
-
論文標題：DiffMoE: Dynamic Token Selection for Scalable Diffusion Transformers
-
項目主頁：https://shiml20.github.io/DiffMoE/
-
論文地址：https://arxiv.org/abs/2503.14487
-
代碼：https://github.com/KwaiVGI/DiffMoE
核心突破：動態token選擇與全局上下文感知
DiffMoE 首次在擴散模型中引入批級全局token池（Batch-level Global Token Pool），打破傳統模型對單一樣本內token的限制，使專家網絡能夠跨樣本訪問全局token分佈。這種設計模擬了完整數據集的token分佈，顯著提升了模型對複雜噪聲模式的學習能力。實驗表明，DiffMoE 在訓練損失收斂速度上超越了同等參數量的密集模型（Dense Models），為多任務處理提供了更強大的上下文感知能力。
針對推理階段的計算資源分配問題，DiffMoE 提出動態容量預測器（Capacity Predictor），通過輕量級 MLP 網絡實時調整專家網絡的計算負載。該機制基於訓練時的token路由模式學習，在不同噪聲水平和樣本複雜度間智能分配資源，實現了性能與計算成本的靈活權衡。例如，在生成困難圖片時自動分配更多計算資源，而在處理簡單圖像時降低負載，真正做到 「按需計算」。
性能提升：以少勝多的參數高效模型
在 ImageNet 256×256 分類條件圖像生成基準測試中，其他結構細節保持一致的公平對比情況下，DiffMoE-L-E8 模型僅用 4.58 億參數 (FID50K 2.13), 超越了擁有 6.75 億參數的 Dense-DiT-XL 模型（FID 2.19）。通過進一步擴展實驗，DiffMoE 實現了僅用 1 倍激活參數就實現了 3 倍於密集模型的性能。此外，DiffMoE 在文本到圖像生成任務中同樣展現出卓越的泛化能力，相較於 Dense 模型有明顯效率提升。
多維度驗證：從理論到實踐
研究團隊通過大規模實驗驗證了 DiffMoE 的優越性：
動態計算優勢：DiffMoE 的平均激活參數量較低的情況下實現了性能指標的顯著提升，證明了動態資源分配的高效性；同時，DiffMoE 能夠根據樣本的難度自動分配計算量。本研究可視化了模型認為最困難和最簡單的十類生成。
模型認為的最困難的十類
模型認為的最簡單的十類擴展性測試：從小型（32M）到大型（458M）配置，DiffMoE 均保持性能正向增長，專家數量從 2 擴展到 16 時 FID 持續下降；
跨任務適配：在文本到圖像生成任務中，DiffMoE 模型在對象生成、空間定位等關鍵指標上全面超越 Dense 基線模型。
總結
在這項工作中，研究團隊通過動態token選擇和全局token可訪問性來高效擴展擴散模型。我們的方法利用專門的專家行為和動態資源分配，有效解決了擴散 Transformer 中固定計算量處理的局限性問題。大量的實驗結果表明，DiffMoE 在性能上超越了現有的 TC-MoE 和 EC-MoE 方法，以及激活參數量是其 3 倍的密集型模型。研究團隊不僅驗證了它在類別條件生成任務中的實用性，也驗證了 DiffMoE 在大規模文本到圖像生成任務的有效性。雖然為了進行公平比較，我們未納入現代混合專家（MoE）模型的改進方法，但在未來的工作中，集成諸如細粒度專家和共享專家等先進技術，將可能帶來新的增益。
更多細節請參閱原論文。
