近日,来自加州大学尔湾分校等机构的研究人员,利用延迟掩蔽、MoE、分层扩展等策略,将扩散模型的训练成本降到了1890美元。

训练一个扩散模型要多少钱?

之前最便宜的方法(Wuerstchen)用了28400美元,而像Stable Diffusion这样的模型还要再贵一个数量级。

大模型时代,一般人根本玩不起。想要各种文生小姐姐,还得靠厂商们负重前行

为了降低这庞大的开销,研究者们尝试了各种方案。

比Stable Diffusion便宜118倍!1890美元训出11.6亿参数高质量文生图模型

比如,原始的扩散模型从噪声到图像大约需要1000步,目前已经被减少到20步左右,甚至更少。

当扩散模型中的基础模块逐渐由Unet(CNN)替换为DiT(Transformer)之后,一些根据Transformer特性来做的优化也跟了上来。

比Stable Diffusion便宜118倍!1890美元训出11.6亿参数高质量文生图模型

比如量化,比如跳过Attention中的一些冗余计算,比如pipeline。

而近日,来自加州大学尔湾分校等机构的研究人员,把「省钱」这个目标直接向前推进了一大步:

比Stable Diffusion便宜118倍!1890美元训出11.6亿参数高质量文生图模型

论文地址:https://arxiv.org/abs/2407.15811

——从头开始训练一个11.6亿参数的扩散模型,只需要1890美元!

对比SOTA有了一个数量级的提升,让普通人也看到了能摸一摸预训练的希望。

更重要的是,降低成本的技术并没有影响模型的性能,11.6亿个参数给出了下面这样非常不错的效果。

比Stable Diffusion便宜118倍!1890美元训出11.6亿参数高质量文生图模型

比Stable Diffusion便宜118倍!1890美元训出11.6亿参数高质量文生图模型

除了观感,模型的数据指标也很优秀,比如下表给出的FID分数,非常接近Stable Diffusion 1.5和DALL·E 2。

相比之下,Wuerstchen的降成本方案则导致自己的考试分数不甚理想。

比Stable Diffusion便宜118倍!1890美元训出11.6亿参数高质量文生图模型

省钱的秘诀

抱着「Stretching Each Dollar」的目标,研究人员从扩散模型的基础模块DiT入手。

首先,序列长度是Transformer计算成本的大敌,需要除掉。

对于图像来说,就需要在不影响性能的情况下,尽量减少参加计算的patch数量(同时也减少了内存开销)。

比Stable Diffusion便宜118倍!1890美元训出11.6亿参数高质量文生图模型

减少图像切块数可以有两种方式,一是增大每块的尺寸,二是干掉一部分patch(mask)。

比Stable Diffusion便宜118倍!1890美元训出11.6亿参数高质量文生图模型

因为前者会显著降低模型性能,所以我们考虑进行mask的方式。

最朴素的mask(NAIve token masking)类似于卷积UNet中随机裁剪的训练,但允许对图像的非连续区域进行训练。

比Stable Diffusion便宜118倍!1890美元训出11.6亿参数高质量文生图模型

而之前最先进的方法(MaskDiT),在输出之前增加了一个恢复重建的结构,通过额外的损失函数来训练,希望通过学习弥补丢掉的信息。

比Stable Diffusion便宜118倍!1890美元训出11.6亿参数高质量文生图模型

这两种mask都为了降低计算成本,在一开始就丢弃了大部分patch,信息的损失显著降低了Transformer的整体性能,即使MaskDiT试图弥补,也只是获得了不太多的改进。

——丢掉信息不可取,那么怎样才能减小输入又不丢信息呢?

延迟掩蔽

本文提出了一种延迟掩蔽策略(deferred masking strategy),在mask之前使用混合器(patch-mixer)进行预处理,把被丢弃patch的信息嵌入到幸存的patch中,从而显著减少高mask带来的性能下降。

比Stable Diffusion便宜118倍!1890美元训出11.6亿参数高质量文生图模型

在本架构中,patch-mixer是通过注意力层和前馈层的组合来实现的,使用二进制掩码进行mask,整个模型的损失函数为:

比Stable Diffusion便宜118倍!1890美元训出11.6亿参数高质量文生图模型

与MaskDiT相比,这里不需要额外的损失函数,整体设计和训练更加简单。

而混合器本身是个非常轻量的结构,符合省钱的标准。

微调

由于非常高的掩蔽比(masking ratio)会显著降低扩散模型学习图像中全局结构的能力,并引入训练到测试的分布偏移,所以作者在预训练(mask)后进行了小幅度的微调(unmask)。

另外,微调还可以减轻由于使用mask而产生的任何不良生成伪影。

MoE和分层扩展

MoE能够增加模型的参数和表达能力,而不会显著增加训练成本。

作者使用基于专家选择路由的简化MoE层,每个专家确定路由到它的token,而不需要任何额外的辅助损失函数来平衡专家之间的负载。

比Stable Diffusion便宜118倍!1890美元训出11.6亿参数高质量文生图模型

此外,作者还考虑了分层缩放方法,线性增加Transformer块的宽度(即注意力层和前馈层中的隐藏层尺寸)。

由于视觉模型中的更深层倾向于学习更复杂的特征,因此在更深层中使用更多的参数将带来更好的性能。

实验设置

作者使用两种DiT的变体:DiT-Tiny/2和DiT-Xl/2,patch大小为2。

使用具有余弦学习率衰减和高权重衰减的AdamW优化器训练所有模型。

比Stable Diffusion便宜118倍!1890美元训出11.6亿参数高质量文生图模型

模型前端使用Stable-Diffusion-XL模型中的四通道变分自动编码器(VAE)来提取图像特征,另外还测试了最新的16通道VAE在大规模训练(省钱版)中的性能。

比Stable Diffusion便宜118倍!1890美元训出11.6亿参数高质量文生图模型

作者使用EDM框架作为所有扩散模型的统一训练设置,使用FID以及CLIP分数来衡量图像生成模型的性能。

文本编码器选择了最常用的CLIP模型,尽管T5-xxl这种较大的模型在文本合成等具有挑战性的任务上表现更好,但为了省钱的目标,这里没有采用。

训练数据集

使用三个真实图像数据集(Conceptual Captions、Segment Anything、TextCaps),包含2200万个图像文本对。

由于SA1B不提供真实的字幕,这里使用LLaVA模型生成的合成字幕。作者还在大规模训练中添加了两个包含1500万个图像文本对的合成图像数据集:JourneyDB和DiffusionDB。

对于小规模消融,研究人员通过从较大的COYO-700M数据集中对10个CIFAR-10类的图像进行二次采样,构建了一个名为cifar-captions的文本到图像数据集。

评估

使用DiT-Tiny/2模型和cifar-captions数据集(256×256分辨率)进行所有评估实验。

对每个模型进行60K优化步骤的训练,并使用AdamW优化器和指数移动平均值(最后10K步平滑系数为0.995)。

延迟掩蔽

实验的基线选择我们上面提到的Naive masking,而本文的延迟掩蔽则加入一个轻量的patch-mixer,参数量小于主干网络的10%。

一般来说,丢掉的patch越多(高masking ratio),模型的性能会越差,比如MaskDiT在超过50%后表现大幅下降。

这里的对比实验采用默认的超参数(学习率1.6×10e-4、0.01的权重衰减和余弦学习率)来训练两个模型。

比Stable Diffusion便宜118倍!1890美元训出11.6亿参数高质量文生图模型

上图的结果显示了延迟屏蔽方法在FID、Clip-FID和Clip score三个指标上都获得了提升。

并且,与基线的性能差距随着掩蔽率的增加而扩大。在掩蔽率为75%的情况下,朴素掩蔽会将FID分数降低至 16.5,而本文的方法则达到5.03,更接近于无掩蔽时的FID分数(3.79)。

超参数

沿着训练LLM的一般思路,这里比较两个任务的超参数选择。

首先,在前馈层中,SwiGLU激活函数优于GELU。其次,较高的权重衰减会带来更好的图像生成性能。

比Stable Diffusion便宜118倍!1890美元训出11.6亿参数高质量文生图模型

另外,与LLM训练不同的是,当对AdamW二阶矩 (β) 使用更高的运行平均系数时,本文的扩散模型可以达到更好的性能。

最后,作者发现使用少量的训练步骤,而将学习率增加到最大可能值(直到训练不稳定)也显著提高了图像生成性能。

混合器的设计

大力出奇迹一般都是对的,作者也观察到使用更大的patch-mixer后,模型性能得到持续改善。

然而,本着省钱的目的,这里还是选择使用小型的混合器。

作者将噪声分布修改为 (−0.6, 1.2),这改善了字幕和生成图像之间的对齐。

如下图所示,在75% masking ratio下,作者还研究了采用不同patch大小所带来的影响。

比Stable Diffusion便宜118倍!1890美元训出11.6亿参数高质量文生图模型

当连续区域变多(patch变大)时,模型的性能会下降,因此保留随机屏蔽每个patch的原始策略。

分层缩放

这个实验训练了DiT-Tiny架构的两种变体,一种具有恒定宽度,另一种采用分层缩放的结构。

两种方法都使用Naive masking,并调整Transformer的尺寸,保证两种情况下的模型算力相同,同时执行相同的训练步骤和训练时间。

比Stable Diffusion便宜118倍!1890美元训出11.6亿参数高质量文生图模型

由上表结果可知发现,在所有三个性能指标上,分层缩放方法都优于基线的恒定宽度方法,这表明分层缩放方法更适合DiT的掩蔽训练。

参考资料:

https://arxiv.org/abs/2407.15811

文章来自于微信公众号新智元 作者新智元

关联网址

关联标签

文章目录

发评论,每天都得现金奖励!超多礼品等你来拿

后,在评论区留言并审核通过后,即可获得现金奖励,奖励规则可见: 查看奖励规则
暂无评论...