环球科创网

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

更新时间:2024-08-12 14:07:14

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

【新智元导读】近日,来自加州大学尔湾分校等机构的研究人员,利用延迟掩蔽、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 Diffusion1.5和DALL·E2。

相比之下,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

免责声明:本文由用户上传,如有侵权请联系删除!