三模态掩码扩散模型的设计空间

离散扩散模型正在变成自回归语言模型之外的另一条路线。以往做法是先训练一个单模态基础模型,再微调成双模态生成。这篇工作换了个方向:来自Apple、MIT、剑桥和Google DeepMind的研究者从头预训练了一个真正的三模态掩码扩散模型(MDM),同时覆盖文本、图像-文本、音频-文本三种数据。

作者对多模态扩散模型的设计空间做了一次系统性扫描,包括多模态缩放定律、模态混合比例、噪声调度、批大小效应,并给出了推理阶段的采样默认值。批大小这块有个新东西:他们提出了一种基于随机微分方程(SDE)的重参数化,直接把「物理批大小」和「逻辑批大小」解耦。物理批大小由GPU饱和、FLOP效率、墙钟时间这类计算因素决定;逻辑批大小则用来平衡随机优化过程中的梯度方差。这样一来,就不需要再照着经验去手调「最优批大小」。

在这个设计下,他们预训练了一个3B参数的初步模型,用了6.4T tokens,在文本生成、T2I和T2S任务上拿到了较强结果。作者称这是迄今最大规模的多模态离散扩散模型系统性开放研究。

整篇的价值在于,它把多模态离散扩散从「单模态微调」推进到「三模态从头训练」,并给出了一批可复用的设计经验。对想用离散扩散做多模态生成的人来说,这些分析比单个模型指标更有参考意义。

The Design Space of Tri-Modal Masked Diffusion Models

查看原文