正文
核心突破:动态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 模型有明显效率提升。