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

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

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

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

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

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

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

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

多維度驗(yàn)證：從理論到實(shí)踐

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

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

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

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

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

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

總結(jié)

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

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