近日,来自加州大学尔湾分校等机构的研究人员,利用延迟掩蔽、MoE、分层扩展等策略,将扩散模型的训练成本降到了1890美元。
训练一个扩散模型要多少钱?
之前最便宜的方法(Wuerstchen)用了28400美元,而像Stable Diffusion这样的模型还要再贵一个数量级。
大模型时代,一般人根本玩不起。想要各种文生小姐姐,还得靠厂商们负重前行
为了降低这庞大的开销,研究者们尝试了各种方案。
比如,原始的扩散模型从噪声到图像大约需要1000步,目前已经被减少到20步左右,甚至更少。
当扩散模型中的基础模块逐渐由Unet(CNN)替换为DiT(Transformer)之后,一些根据Transformer特性来做的优化也跟了上来。
比如量化,比如跳过Attention中的一些冗余计算,比如pipeline。
而近日,来自加州大学尔湾分校等机构的研究人员,把「省钱」这个目标直接向前推进了一大步:
论文地址:https://arxiv.org/abs/2407.15811
——从头开始训练一个11.6亿参数的扩散模型,只需要1890美元!
对比SOTA有了一个数量级的提升,让普通人也看到了能摸一摸预训练的希望。
更重要的是,降低成本的技术并没有影响模型的性能,11.6亿个参数给出了下面这样非常不错的效果。
除了观感,模型的数据指标也很优秀,比如下表给出的FID分数,非常接近Stable Diffusion 1.5和DALL·E 2。
相比之下,Wuerstchen的降成本方案则导致自己的考试分数不甚理想。
省钱的秘诀
抱着「Stretching Each Dollar」的目标,研究人员从扩散模型的基础模块DiT入手。
首先,序列长度是Transformer计算成本的大敌,需要除掉。
对于图像来说,就需要在不影响性能的情况下,尽量减少参加计算的patch数量(同时也减少了内存开销)。
减少图像切块数可以有两种方式,一是增大每块的尺寸,二是干掉一部分patch(mask)。
因为前者会显著降低模型性能,所以我们考虑进行mask的方式。
最朴素的mask(NAIve token masking)类似于卷积UNet中随机裁剪的训练,但允许对图像的非连续区域进行训练。
而之前最先进的方法(MaskDiT),在输出之前增加了一个恢复重建的结构,通过额外的损失函数来训练,希望通过学习弥补丢掉的信息。
这两种mask都为了降低计算成本,在一开始就丢弃了大部分patch,信息的损失显著降低了Transformer的整体性能,即使MaskDiT试图弥补,也只是获得了不太多的改进。
——丢掉信息不可取,那么怎样才能减小输入又不丢信息呢?
延迟掩蔽
本文提出了一种延迟掩蔽策略(deferred masking strategy),在mask之前使用混合器(patch-mixer)进行预处理,把被丢弃patch的信息嵌入到幸存的patch中,从而显著减少高mask带来的性能下降。
在本架构中,patch-mixer是通过注意力层和前馈层的组合来实现的,使用二进制掩码进行mask,整个模型的损失函数为:
与MaskDiT相比,这里不需要额外的损失函数,整体设计和训练更加简单。
而混合器本身是个非常轻量的结构,符合省钱的标准。
微调
由于非常高的掩蔽比(masking ratio)会显著降低扩散模型学习图像中全局结构的能力,并引入训练到测试的分布偏移,所以作者在预训练(mask)后进行了小幅度的微调(unmask)。
另外,微调还可以减轻由于使用mask而产生的任何不良生成伪影。
MoE和分层扩展
MoE能够增加模型的参数和表达能力,而不会显著增加训练成本。
作者使用基于专家选择路由的简化MoE层,每个专家确定路由到它的token,而不需要任何额外的辅助损失函数来平衡专家之间的负载。
此外,作者还考虑了分层缩放方法,线性增加Transformer块的宽度(即注意力层和前馈层中的隐藏层尺寸)。
由于视觉模型中的更深层倾向于学习更复杂的特征,因此在更深层中使用更多的参数将带来更好的性能。
实验设置
作者使用两种DiT的变体:DiT-Tiny/2和DiT-Xl/2,patch大小为2。
使用具有余弦学习率衰减和高权重衰减的AdamW优化器训练所有模型。
模型前端使用Stable-Diffusion-XL模型中的四通道变分自动编码器(VAE)来提取图像特征,另外还测试了最新的16通道VAE在大规模训练(省钱版)中的性能。
作者使用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的权重衰减和余弦学习率)来训练两个模型。
上图的结果显示了延迟屏蔽方法在FID、Clip-FID和Clip score三个指标上都获得了提升。
并且,与基线的性能差距随着掩蔽率的增加而扩大。在掩蔽率为75%的情况下,朴素掩蔽会将FID分数降低至 16.5,而本文的方法则达到5.03,更接近于无掩蔽时的FID分数(3.79)。
超参数
沿着训练LLM的一般思路,这里比较两个任务的超参数选择。
首先,在前馈层中,SwiGLU激活函数优于GELU。其次,较高的权重衰减会带来更好的图像生成性能。
另外,与LLM训练不同的是,当对AdamW二阶矩 (β) 使用更高的运行平均系数时,本文的扩散模型可以达到更好的性能。
最后,作者发现使用少量的训练步骤,而将学习率增加到最大可能值(直到训练不稳定)也显著提高了图像生成性能。
混合器的设计
大力出奇迹一般都是对的,作者也观察到使用更大的patch-mixer后,模型性能得到持续改善。
然而,本着省钱的目的,这里还是选择使用小型的混合器。
作者将噪声分布修改为 (−0.6, 1.2),这改善了字幕和生成图像之间的对齐。
如下图所示,在75% masking ratio下,作者还研究了采用不同patch大小所带来的影响。
当连续区域变多(patch变大)时,模型的性能会下降,因此保留随机屏蔽每个patch的原始策略。
分层缩放
这个实验训练了DiT-Tiny架构的两种变体,一种具有恒定宽度,另一种采用分层缩放的结构。
两种方法都使用Naive masking,并调整Transformer的尺寸,保证两种情况下的模型算力相同,同时执行相同的训练步骤和训练时间。
由上表结果可知发现,在所有三个性能指标上,分层缩放方法都优于基线的恒定宽度方法,这表明分层缩放方法更适合DiT的掩蔽训练。
参考资料:
https://arxiv.org/abs/2407.15811
文章来自于微信公众号新智元 作者新智元
发评论,每天都得现金奖励!超多礼品等你来拿
登录 后,在评论区留言并审核通过后,即可获得现金奖励,奖励规则可见: 查看奖励规则