[论文解读] Fast Sampling of Diffusion Models via Operator Learning

Hongkai Zheng, Weili Nie+3 more|arXiv (Cornell University)|Nov 24, 2022

Advanced Neuroimaging Techniques and Applications被引用 17

一句话总结

该论文提出DSNO，一种基于神经算子的新方法，通过单次前向传播实现扩散模型的快速、并行采样。通过将傅里叶域时间卷积层集成到扩散模型主干网络中，DSNO学习连续时间反向过程轨迹，在仅进行一次模型推理的情况下，实现了SOTA级别的FID分数：CIFAR-10上为3.78，ImageNet-64上为7.83。

ABSTRACT

Diffusion models have found widespread adoption in various areas. However, their sampling process is slow because it requires hundreds to thousands of network evaluations to emulate a continuous process defined by differential equations. In this work, we use neural operators, an efficient method to solve the probability flow differential equations, to accelerate the sampling process of diffusion models. Compared to other fast sampling methods that have a sequential nature, we are the first to propose a parallel decoding method that generates images with only one model forward pass. We propose diffusion model sampling with neural operator (DSNO) that maps the initial condition, i.e., Gaussian distribution, to the continuous-time solution trajectory of the reverse diffusion process. To model the temporal correlations along the trajectory, we introduce temporal convolution layers that are parameterized in the Fourier space into the given diffusion model backbone. We show our method achieves state-of-the-art FID of 3.78 for CIFAR-10 and 7.83 for ImageNet-64 in the one-model-evaluation setting.

研究动机与目标

• 加速扩散模型采样，以支持对时间敏感应用的实时推理。
• 克服现有快速采样方法存在的顺序瓶颈问题，这些方法需要多次前向传播。
• 利用神经算子开发可微分的、从噪声到数据的连续时间解轨迹映射。
• 在将推理时间大幅缩短至单次前向传播的同时，保持高样本质量。
• 通过参数高效的傅里叶空间时间卷积，将时间建模集成到扩散模型中。

提出的方法

• 提出DSNO，一种神经算子，用于将初始高斯噪声映射到反向扩散过程的连续时间解轨迹。
• 引入在傅里叶域参数化的时序卷积层，以建模轨迹上的时间相关性。
• 通过基于傅里叶的运算同时计算所有时间步的输出，实现并行解码，避免顺序计算。
• 将时序卷积模块集成到现有的U-Net主干网络中，保持与标准扩散模型架构的兼容性。
• 使用L1损失与感知（LPIPS）损失的组合进行模型训练，以提升样本保真度。
• 采用二次时间离散化和多分辨率监督，以提升泛化能力与FID性能。

实验结果

研究问题

• RQ1 神经算子能否有效用于建模扩散模型的连续时间反向过程，以实现快速采样？
• RQ2 傅里叶域中的时序卷积层能否高效捕捉扩散轨迹中的长程时间依赖？
• RQ3 通过基于傅里叶的运算实现并行解码，是否能在不损失样本质量的前提下实现一步生成？
• RQ4 所提方法是否仅通过一次前向传播即可达到SOTA级别的FID分数？
• RQ5 损失函数的选择（如L1与LPIPS）如何影响最终的样本质量与FID分数？

主要发现

• 在单次模型推理设置下，DSNO在CIFAR-10上实现了SOTA级别的FID分数3.78。
• 在仅进行一次模型推理的情况下，DSNO在ImageNet-64上实现了SOTA级别的FID分数7.83。
• 使用基于VGG的LPIPS损失相比标准L1损失，显著提升了FID表现。
• 将时间分辨率从4提升至8可略微改善FID，但边际收益较小，表明4已足够高效。
• 所提出的时序卷积模块仅使模型参数量增加10%，却实现了并行解码与高性能。
• 该方法实现了真正的并行解码，所有时间步输出在单次前向传播中同时生成，优于顺序求解器。