[论文解读] Deep Equilibrium Approaches to Diffusion Models

Ashwini Pokle, Zhengyang Geng+1 more|arXiv (Cornell University)|Oct 23, 2022

Generative Adversarial Networks and Image Synthesis被引用 7

一句话总结

本文提出了一种用于去噪扩散隐式模型（DDIM）的深度均衡（DEQ）公式，将传统的顺序采样过程替换为所有去噪步骤的并行、联合固定点优化。该方法实现了更快的推理速度，并通过隐式微分显著加速了模型反演，在CIFAR-10和CelebA数据集上实现了最高达2倍的加速，同时保持了较高的感知质量。

ABSTRACT

Diffusion-based generative models are extremely effective in generating high-quality images, with generated samples often surpassing the quality of those produced by other models under several metrics. One distinguishing feature of these models, however, is that they typically require long sampling chains to produce high-fidelity images. This presents a challenge not only from the lenses of sampling time, but also from the inherent difficulty in backpropagating through these chains in order to accomplish tasks such as model inversion, i.e. approximately finding latent states that generate known images. In this paper, we look at diffusion models through a different perspective, that of a (deep) equilibrium (DEQ) fixed point model. Specifically, we extend the recent denoising diffusion implicit model (DDIM; Song et al. 2020), and model the entire sampling chain as a joint, multivariate fixed point system. This setup provides an elegant unification of diffusion and equilibrium models, and shows benefits in 1) single image sampling, as it replaces the fully-serial typical sampling process with a parallel one; and 2) model inversion, where we can leverage fast gradients in the DEQ setting to much more quickly find the noise that generates a given image. The approach is also orthogonal and thus complementary to other methods used to reduce the sampling time, or improve model inversion. We demonstrate our method's strong performance across several datasets, including CIFAR10, CelebA, and LSUN Bedrooms and Churches.

研究动机与目标

• 为解决扩散模型采样速度缓慢的问题，其根源在于固有的顺序去噪过程。
• 实现通过采样链的高效反向传播，以支持模型反演等任务，而在标准扩散模型中此类计算在计算上不可行。
• 通过将整个DDIM采样链建模为多变量固定点系统，统一扩散模型与深度均衡（DEQ）框架。
• 证明DEQ公式化可实现跨多个GPU的并行化，即使在单图生成任务中亦可实现。
• 表明在DEQ框架中使用隐式微分可显著加速扩散模型中的模型反演，相比顺序采样。

提出的方法

• 将DDIM采样过程公式化为一个联合多变量固定点系统，使所有T个去噪步骤作为深度均衡问题被同时求解。
• 使用黑箱根求解器（如Anderson加速）并行计算整个系统的固定点，从而实现多GPU推理。
• 通过DEQ框架利用隐式微分，以O(1)内存复杂度计算通过固定点的梯度，避免存储完整计算图的需要。
• 通过将噪声项引入固定点系统，将DEQ公式化扩展至处理DDIM的随机变体。
• 将DEQ框架应用于单图生成与模型反演，两者均使用相同的固定点求解器。
• 利用DEQ框架优化潜在代码以实现图像重建，从而在模型反演任务中实现更快收敛。

实验结果

研究问题

• RQ1 整个DDIM采样链能否被重新公式化为一个深度均衡系统，以实现并行推理？
• RQ2 与顺序DDIM相比，DEQ公式化是否能带来更快的收敛速度和更高的采样效率？
• RQ3 在DEQ框架中使用隐式微分能否显著加速扩散模型中的模型反演？
• RQ4 在不同数据集上，DEQ-DDIM与标准DDIM在FID分数和感知质量方面相比如何？
• RQ5 DEQ公式化能否推广至如DDPM这样的随机扩散模型？

主要发现

• 与标准顺序DDIM相比，DEQ-DDIM在CIFAR-10和CelebA数据集上的采样时间最高可实现2倍加速，同时FID分数保持相当。
• 由于高效的隐式梯度计算，使用DEQ-DDIM进行模型反演的收敛速度显著快于顺序DDIM。
• DEQ公式化实现了跨多个GPU的并行采样，即使在单图生成任务中也克服了标准扩散模型中串行处理的瓶颈。
• 通过DEQ反演恢复的初始潜在代码能更准确地重构原始图像，相比顺序DDIM反演更完整地保留了细节。
• DEQ框架支持通过采样链反向传播的O(1)内存复杂度，使得对长扩散过程进行微分成为可能。
• 该方法可扩展至随机扩散模型如DDPM，仅需对固定点系统进行简单调整以包含噪声项。