Google DeepMind提出DRaFT算法以高效优化扩散模型

要点:

1、Google DeepMind提出了直接奖励微调（DRaFT）的概念，可以高效微调扩散模型以最大化可微的奖励函数。

2、引入DRaFT-K和DRaFT-LV算法，限制反向传播步数，提高效率。实验证明优于全反向传播。

3、在Stable Diffusion1.4上应用DRaFT，结果显示其效率远超基于强化学习的微调基线。DRaFT-LV效率最高。

扩散模型已经革新了各类数据的生成建模。但是在实际应用中，如根据文本描述生成美观图像，仍需要微调模型。目前文本到图像的扩散模型采用无分类器指导和精心设计的数据集如LAION美学数据集来提升图像质量和模型对齐性能。

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

在研究中，Google DeepMind提出了一种直接而高效的基于梯度的奖励微调方法，其核心是将扩散采样过程微分化。他们提出了直接奖励微调（DRaFT）的概念，其本质是反向传播整个采样链，通常表示为具有50步长度的展开计算图。为有效管理内存和计算成本，他们采用梯度检查点技术，优化LoRA权重而不是修改全部模型参数。

此外，Google DeepMind还提出了DRaFT方法的改进，以提升其效率和性能。首先是DRaFT-K，其限制采样时仅反向传播最后K步时的梯度计算。实验结果显示，与全反向传播相比，这种截断梯度方法的性能明显更好，因为全反向传播会导致梯度爆炸问题。

另外，研究人员提出了DRaFT-LV，其平均多个噪声样本来计算更低方差的梯度估计，进一步提高方法效率。

研究人员在Stable Diffusion1.4上应用DRaFT，使用各种奖励函数和提示进行评估。他们的梯度方法相比基于强化学习的微调基线，效率优势明显。例如，在最大化LAION美学分类器分数时，与强化学习算法相比取得了200倍的加速。

他们提出的变体DRaFT-LV展现出卓越的效率，学习速度约为先前梯度微调方法ReFL的两倍。此外，他们证明了DRaFT可以与预训练模型组合或插值，通过混合或缩放LoRA权重来实现。

总之，直接在可微奖励上微调扩散模型是一个提升生成建模技术的有前景的方向。其效率、通用性和有效性使其成为机器学习和生成建模领域研究者和从业者的有价值工具。