2025-12 DSMoE

Efficient Training of Diffusion Mixture-of-Experts Models: A Practical Recipe

📌 核心贡献

系统探索了扩散模型 MoE 的架构设计空间，证明架构调优（DeepSeek-style MoE、缩窄中间维度、共享专家、2D RoPE）比路由机制改进更重要。在参数量减少 44% 的情况下，FID 降低 17.6%，训练效率提升 10 倍，适用于 latent 和 pixel-space 两种扩散框架。

📖 摘要

This work explores underexamined architectural configurations for Diffusion MoE systems. Drawing inspiration from LLM design patterns, the authors identify key factors including DeepSeek-style expert modules, alternative intermediate widths, varying expert counts, and enhanced attention positional encodings. Their findings demonstrate that architectural tuning often produces superior results compared to routing mechanism improvements alone. The study presents novel configurations applicable to both latent and pixel-space diffusion frameworks, enabling models to exceed baseline performance while maintaining equal or reduced activated parameters.

🔍 关键信息

字段	内容
机构	Huawei Noah’s Ark Lab
发表	2025-12-01
分类	cs.LG, cs.CV
链接	arXiv

📝 我的笔记

方法

DeepSeek-style MoE FFN： 共享专家 + 路由专家组合，每个 token 输出为：

$h_t=u_t+\sum FF{N}_i^{(s)}(u_t)+\sum g_{i,t}\cdot FF{N}_i^{(r)}(u_t)$

• 共享专家处理所有 token，提供正则化
• 路由专家通过 Sigmoid 亲和度打分 + Top-K 选择激活
• 无辅助损失的负载均衡（偏置项调节）

关键架构创新：

创新点	做法	效果
缩窄 FFN 中间维度	< 4.0× hidden dim	减参数不降性能
增加专家数量	48 experts（非常规 16）	收敛更稳定
共享专家	每层 1 个	正则化 + 训练稳定
2D RoPE	替换绝对位置编码	FID ↓11%
交替 MoE/Dense 层	非全部替换	全替换反而伤害性能

模型家族：

• DSMoE：latent diffusion（S/92M, B/368M, L/1.3B, 3B/2.9B）
• JiTMoE：pixel-space diffusion（B, L）
• 配置记法如 S1E48A5 = 1 共享 + 48 路由 + 每 token 激活 5 个

关键结果

模型	总参数	FID (CFG=1.0)	FID (CFG=1.5)
DiffMoE-L-E16	1.982B	11.16	2.84
DSMoE-L-E48	1.112B	9.19	2.55
DSMoE-3B (700K steps)	2.9B	—	2.38
DiffMoE (7000K steps)	—	—	2.30

• DSMoE-L 参数少 44%，FID 降低 17.6%
• DSMoE-3B 仅用 1/10 训练步数达到可比性能
• 2D RoPE 是最大单点提升（FID 45.13→39.84）
• 去除共享专家收敛慢 15-20%

🔗 相关论文

基于/改进自： DiT 对比基线： DiT 同方向： LDM, DDPM