[论文解读] Gradient Guidance for Diffusion Models: An Optimization Perspective

Yingqing Guo, Hui Yuan+3 more|arXiv (Cornell University)|Apr 23, 2024

Spacecraft Dynamics and Control被引用 4

一句话总结

本文提出了一种扩散模型的新型梯度引导方法，将引导采样建模为正则化优化问题，通过前向预测损失保留潜在数据结构。在自适应微调下，该方法可保证收敛至全局最优解，对于凹目标函数实现$\mathcal{O}(1/K)$的收敛速率，同时保持样本质量。

ABSTRACT

Diffusion models have demonstrated empirical successes in various applications and can be adapted to task-specific needs via guidance. This paper studies a form of gradient guidance for adapting a pre-trained diffusion model towards optimizing user-specified objectives. We establish a mathematical framework for guided diffusion to systematically study its optimization theory and algorithmic design. Our theoretical analysis spots a strong link between guided diffusion models and optimization: gradient-guided diffusion models are essentially sampling solutions to a regularized optimization problem, where the regularization is imposed by the pre-training data. As for guidance design, directly bringing in the gradient of an external objective function as guidance would jeopardize the structure in generated samples. We investigate a modified form of gradient guidance based on a forward prediction loss, which leverages the information in pre-trained score functions and provably preserves the latent structure. We further consider an iteratively fine-tuned version of gradient-guided diffusion where guidance and score network are both updated with newly generated samples. This process mimics a first-order optimization iteration in expectation, for which we proved O(1/K) convergence rate to the global optimum when the objective function is concave. Our code will be released at https://github.com/yukang123/GGDMOptim.git.

研究动机与目标

• 为解决梯度引导扩散模型中理论理解不足的问题，特别是关于结构保持与收敛性保证的问题。
• 通过引入基于前瞻损失的结构化引导信号，解决朴素梯度引导导致样本质量下降的问题。
• 将引导扩散形式化为一个优化过程，平衡目标函数最大化与预训练中获得的数据分布先验。
• 设计一种自适应微调机制，通过自生成样本迭代改进得分网络。
• 为引导扩散在优化凹目标函数时建立收敛速率与理论保证。

提出的方法

• 提出一种基于前向预测损失的梯度引导信号，定义为$\texttt{G}_{\text{loss}}(x_t,t) := \beta(t)\left(y - g^\top\mathbb{E}[x_0|x_t]\right)g$，利用预训练得分函数以保持潜在结构。
• 采用基于Tweedie公式构造的前瞻估计器$\hat{\mathbb{E}}[x_0|x_t]$，在实际中近似条件期望。
• 提出算法1：一种非自适应的引导采样过程，通过求解正则化优化问题实现，其中预训练得分函数作为先验。
• 提出算法2：一种自适应微调循环，通过在生成样本处查询梯度，并利用新数据更新得分网络。
• 采用基于得分的采样过程，将引导信号在每个时间步注入，以引导扩散轨迹朝向高目标区域。
• 理论分析表明，当目标函数为凹函数且可微时，算法2以$\mathcal{O}(1/K)$的速率收敛至全局最优解。

实验结果

研究问题

• RQ1 为何朴素梯度引导在实践中会失效，导致样本质量下降？
• RQ2 如何设计梯度引导方法以保留预训练数据的潜在低维结构？
• RQ3 在自适应微调下，引导扩散模型能否在理论上保证收敛至最优解？
• RQ4 梯度引导扩散与一阶优化方法之间存在何种关系？
• RQ5 当目标函数为凹函数时，引导采样过程的收敛速率是多少？

主要发现

• 实验表明，基于前向预测损失（$\texttt{G}_{\text{loss}}$）的所提梯度引导方法在保留潜在子空间结构方面显著优于朴素梯度引导，表现为更低的离支持误差。
• 算法1采用非自适应引导，由于预训练得分函数的正则化效应，收敛至次优解，与理论分析一致。
• 算法2通过使用自生成样本进行迭代微调，实现了对凹目标函数的全局最优收敛，且具有可证明的$\mathcal{O}(1/K)$收敛速率。
• 实验结果表明，即使在微调导致分布偏移的情况下，$\texttt{G}_{\text{loss}}$仍能保持优于朴素梯度引导的结构保真度。
• 算法2生成样本的分布已超出原始数据分布范围（由预训练轮廓捕获），表明其有效探索了高目标区域。
• 理论框架表明，梯度引导扩散等价于求解一个正则化优化问题，其中预训练得分函数充当结构先验。