能够实现两倍训练速度提升,并且只需微调总模型参数的约0.12% 。超越Diffusion Model全参数微调
paper:https://arxiv.org/pdf/2304.06648.pdf
大型预训练Diffusion Models微调到新领域在实际应用非常关键,但微调参数量大,训练时间长。最近来自华为的DiffFit方法通过仅微调特定layer中的bias和新增的缩放因子,即可实现快速迁移新领域。相比于全参数微调,DiffFit能够实现两倍训练速度提升,并且只需微调总模型参数的约0.12% 。DiffFit建立了一个简单而强大的基线方法,实现图像生成中的参数高效微调。DiffFit效果如图1所示。
要实现的效果:figure 2中的第一行小猫的source domain向最下面行女人油画domain的快速迁移:
figure 3:实现figure 2迁移效果,达到的指标及需要微调的参数量,从figure 3可以看出,DiffFit不仅参数最少,且效果最佳。
DiffFit技术原理
基础模型:DiT
DiffFit整体框架如图4所示,DiffFit基于Dit模型,DiT(Diffusion Transformers)是最近使用Transformer设计Diffusion Model的代表性方法。DiT遵循潜在扩散模型(latent diffusion models,LDMs)的设计思路,其包含两个部分:(1)一个由编码器E和解码器D组成的自编码器,其中潜在编码为z=E(x),重构数据为x_hat=D(z);(2)一个具有分块、顺序的DiT块和解码化操作的Latent Diffusion Transformers。在每个块Bi中,本文有zi=Bi(x, t, c),其中t和c分别是时间嵌入和类别嵌入。每个块Bi包含一个自注意力模块和一个前馈模块。分块化/解码化操作用于将潜在编码z编码/解码为图像token序列。
Dit模型基础模块如下图:
整体架构对于Dit结构与DiffFit的结合,见图4,DiffFit可以推广到其他扩散模型(例如 Stable Diffusion model)。如图4,主要是将大部分参数在潜在扩散模型中冻结,仅训练偏置项、归一化和类别条件模块。此外,在扩散模型的几个块中插入可学习的缩放因子γ,其中γ的初始值为1.0,并在每个块的相应层上进行相乘。每个块包含多个组件,如多头自注意力、前馈网络和层归一化,并且该块可以堆叠N次。详细信息参考 Algorithm 1。
微调:
在微调过程中,Diffusion Model的参数最初被冻结,然后仅有与偏置、类别嵌入、归一化和缩放因子相关的特定参数被选择性解冻。具体方法如Algorithm 2所述,能够在最小程度干扰预训练权重的情况下实现快速微调。DiT-XL/2仅需要更新其参数的0.12%,训练时间约快于完全微调2倍。DiffFit避免了灾难性遗忘,同时加强了预训练模型的知识,并使其能够适应特定任务。
推理和存储
在对K个数据集进行微调后,只需存储原始模型的完整参数副本和K×数据集特定的可训练参数(通常小于1M)。将这些权重组合到扩散模型中,可以实现对多个领域的类别条件图像生成的适应。
实验
实现细节
基准模型DiT1,在ImageNet 256×256数据集上进行700万次迭代预训练。在遵循DiT基础上,DiffFit将完全微调的学习率lr设置为1e-4,将分类器自由引导参数设置为1.5用于评估,4.0用于可视化。此外,DiffFit还重新实现了几种参数高效的微调方法,如Adaptor、BitFit、Visual Prompt Tuning (VPT)和LoRA。最终发现VPT对深度和token数量非常敏感,训练过程异常不稳定。因此,DiffFit寻求了一种更稳定的配置,将深度设置为5,token数量设置为1,并对所有任务保持最后几层参数解冻。DiffFit将参数高效微调设置的学习率放大10倍,以获得更好的结果。
下游任务迁移实验
实验设置
针对256×256分辨率的下游小型数据集的微调,使用了8个V100 GPU,总batchsize大小256,并进行24,000次迭代训练。选择8个常用的细粒度数据集:Food101、SUN397、DF-20M mini、Caltech101、CUB-200-2011、ArtBench-10、Oxford Flowers和Stanford Cars。对所有任务使用50个采样步骤来计算FID。其中大多数数据集是从CLIP下游任务中选择,除了ArtBench-10。
结果
表1中列出了不同参数高效微调方法的性能。从表中可以看出,通过仅微调0.12%的参数,DiffFit在8个下游任务上平均获得了最低的FID。虽然完全微调是一个强大的基准方法,并且在3个数据集上略微优于其他方法,但需要微调100%的参数。
从低分辨率到高分辨率
实验设置
考虑到以不同分辨率生成图像是一种特殊类型的分布转移,DiffFit可以轻松地将预训练的低分辨率扩散模型适应于生成高分辨率图像。为了证明DiffFit的有效性,加载预训练ImageNet 256×256的DiT-XL/2权重,并在ImageNet 512×512上对模型进行微调。并采用Positional Encoding Trick来加速微调过程。使用32个V100 GPU、1024批量大小和30,000次迭代来在ImageNet 512×512上微调DiT-XL/2。使用250个采样步骤来计算FID。请注意,在这里不需要微调标签嵌入,因为标签不会改变。
Positional Encoding Trick。
DiT 采用了静态正弦曲线2D位置编码方案。为了更好地利用预训练模型中编码的位置信息,DiffFit开发了一种正弦插值方法,将512×512分辨率的位置编码与256×256分辨率的位置编码对齐。通过将位置编码公式中的每个像素坐标(i, j)替换为其一半值(i/2, j/2)来实现,这种方法简单且没有额外计算成本。
结果
如表2所示,DiffFit在ImageNet 512×512的基准测试中获得3.02FID,超过ADM的3.84和官方DiT的3.04。DiffFit在基于扩散的方法中创造了新的最优性能。重要的是,DiffFit比以前的方法要高效得多,因为微调过程只需要额外的51个GPU天。即使考虑到预训练的950个GPU天,本文的方法仍然优于DiT和ADM所需的1500多个GPU天。
另外,本文观察到使用位置编码技巧可以加快训练收敛速度,如图5所示
微调收敛性分析
为便于收敛速度的分析,在Food-101、ArtBench-10、Flowers-102和CUB-200数据集中每隔15,000次迭代报告几种方法的FID分数,如图6所示。观察结果表明,完全微调、BitFit和DiffFit表现出类似的收敛速度,超过了AdaptFormer和VPT的初始性能。虽然AdaptFormer初始性能较差,但在训练中间阶段迅速提升。相比之下,VPT的收敛速度最慢。与完全微调相比,DiffFit冻结了大部分参数,因此最大程度地保留了预训练期间学到的信息,从而实现了更好的微调性能。与BitFit相比,DiffFit通过调整特征的缩放因子来加快收敛速度并获得更好的结果。以上观察和分析验证了DiffFit在适应目标领域和生成图像方面表现出快速的适应能力和出色的性能。
总结
DiffFit可以快速适应大型预训练扩散模型到各种下游领域,包括不同的数据集或不同的分辨率。通过仅微调偏置项和缩放因子,DiffFit提供了一种成本-效益的解决方案,可以减少存储需求并加速微调,工程应用实际中有较大用处。但一个局限性是实验主要集中在类别条件的图像生成上。目前还不清楚这种策略在更复杂的任务,如文本到图像生成或视频/三维生成方面是否能够表现同样出色。