近日，来自加州大学尔湾分校等机构的研究人员利用延迟掩蔽、MoE、分层扩展等策略，将扩散模型的训练成本降到了1890美元。这比之前最便宜的方法(Wuerstchen)的28400美元还要低一个数量级，而像Stable Diffusion这样的模型的成本则更高。
研究者们尝试了各种方案来降低这庞大的开销。比如，原始的扩散模型从噪声到图像大约需要1000步，目前已经被减少到20步左右，甚至更少。当扩散模型中的基础模块逐渐由Unet(CNN)替换为DiT(Transformer)之后，一些根据Transformer特性来做的优化也跟了上来，如量化、跳过Attention中的一些冗余计算、pipeline等。
本研究从头开始训练一个11.6亿参数的扩散模型，只需要1890美元。这对比SOTA有了一个数量级的提升，让普通人也看到了能摸一摸预训练的希望。更重要的是，降低成本的技术并没有影响模型的性能，11.6亿个参数给出了非常不错的效果。除了观感，模型的数据指标也很优秀，比如FID分数非常接近Stable Diffusion 1.5和DALL·E 2。
省钱的秘诀主要包括:
-
延迟掩蔽策略：在mask之前使用混合器(patch-mixer)进行预处理，把被丢弃patch的信息嵌入到幸存的patch中，从而显著减少高mask带来的性能下降。
-
微调：在预训练(mask)后进行小幅度的微调(unmask)，以减轻由于使用mask而产生的不良生成伪影。
-
MoE和分层扩展：使用基于专家选择路由的简化MoE层，增加模型的参数和表达能力，而不会显著增加训练成本。同时考虑了分层缩放方法，线性增加Transformer块的宽度。
实验设置方面，作者使用两种DiT的变体：DiT-Tiny/2和DiT-Xl/2，patch大小为2。使用具有余弦学习率衰减和高权重衰减的AdamW优化器训练所有模型。模型前端使用Stable-Diffusion-XL模型中的四通道变分自动编码器(VAE)来提取图像特征，另外还测试了最新的16通道VAE在大规模训练(省钱版)中的性能。
作者使用EDM框架作为所有扩散模型的统一训练设置，使用FID以及CLIP分数来衡量图像生成模型的性能。文本编码器选择了最常用的CLIP模型。
训练数据集使用了三个真实图像数据集(Conceptual Captions、Segment Anything、TextCaps)，包含2200万个图像文本对。由于SA1B不提供真实的字幕，这里使用LLaVA模型生成的合成字幕。