快手&清華提出DiffMoE：動態Token選擇助力擴散模型效能飛躍！

2025-07-22 18:01 CVer

點選下方卡片，關注“CVer”公眾號

AI/CV重磅乾貨，第一時間送達
點選進入—>【頂會/頂刊】投稿交流群

新增微訊號：CVer2233，小助手會拉你進群！
掃描下方二維碼，加入CVer學術星球！可以獲得最新頂會/頂刊上的論文idea和CV從入門到精通資料，及最前沿應用！發論文/搞科研/漲薪，強烈推薦！

轉載自：機器之心
本文由清華大學和快手可靈團隊共同完成。第一作者是清華大學智慧視覺實驗室在讀本科生史明磊。
在生成式 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）模型的改進方法，但在未來的工作中，整合諸如細粒度專家和共享專家等先進技術，將可能帶來新的增益。
更多細節請參閱原論文。

何愷明在MIT授課的課件PPT下載
在CVer公眾號後臺回覆：何愷明，即可下載本課程的所有566頁課件PPT！趕緊學起來！

CVPR 2025 論文和程式碼下載

在CVer公眾號後臺回覆：CVPR2025，即可下載CVPR 2025論文和程式碼開源的論文合集

ECCV 2024 論文和程式碼下載

在CVer公眾號後臺回覆：ECCV2024，即可下載ECCV 2024論文和程式碼開源的論文合集
CV垂直方向和論文投稿交流群成立
掃描下方二維碼，或者新增微訊號：CVer2233，即可新增CVer小助手微信，便可申請加入CVer-垂直方向和論文投稿微信交流群。另外其他垂直方向已涵蓋：目標檢測、影像分割、目標跟蹤、人臉檢測&識別、OCR、姿態估計、超解析度、SLAM、醫療影像、Re-ID、GAN、NAS、深度估計、自動駕駛、強化學習、車道線檢測、模型剪枝&壓縮、去噪、去霧、去雨、風格遷移、遙感影像、行為識別、影片理解、影像融合、影像檢索、論文投稿&交流、PyTorch、TensorFlow和Transformer、NeRF、3DGS、Mamba等。
一定要備註：研究方向+地點+學校/公司+暱稱（如Mamba、多模態學習或者論文投稿+上海+上交+卡卡），根據格式備註，可更快被透過且邀請進群

▲掃碼或加微訊號: CVer2233，進交流群
CVer計算機視覺（知識星球）人數破萬！如果你想要了解最新最快最好的CV/DL/AI論文、實戰專案、行業前沿、從入門到精通學習教程等資料，一定要掃描下方二維碼，加入CVer知識星球！最強助力你的科研和工作！

▲掃碼加入星球學習
▲點選上方卡片，關注CVer公眾號
整理不易，請點贊和在看


相關文章