[论文解读] Fourier Diffusion Models: A Method to Control MTF and NPS in Score-Based Stochastic Image Generation

Matthew Tivnan, Jacopo Teneggi+7 more|arXiv (Cornell University)|Mar 23, 2023

Generative Adversarial Networks and Image Synthesis被引用 4

一句话总结

本文提出傅里叶扩散模型（Fourier Diffusion Models），一种新颖的基于得分的生成建模方法，该方法用平移不变卷积和平稳噪声替代标量扩散过程，从而在每个时间步直接控制调制传递函数（MTF）和噪声功率谱（NPS）。该方法通过建模从真实图像到已知系统特性的模糊、噪声图像的连续概率流，实现更少时间步数下的优越图像生成性能——仅用8个时间步即可达到标量模型在128或1024个时间步时的性能水平。

ABSTRACT

Score-based stochastic denoising models have recently been demonstrated as powerful machine learning tools for conditional and unconditional image generation. The existing methods are based on a forward stochastic process wherein the training images are scaled to zero over time and white noise is gradually added such that the final time step is approximately zero-mean identity-covariance Gaussian noise. A neural network is then trained to approximate the time-dependent score function, or the gradient of the logarithm of the probability density, for that time step. Using this score estimator, it is possible to run an approximation of the time-reversed stochastic process to sample new images from the training data distribution. These score-based generative models have been shown to out-perform generative adversarial neural networks using standard benchmarks and metrics. However, one issue with this approach is that it requires a large number of forward passes of the neural network. Additionally, the images at intermediate time steps are not useful, since the signal-to-noise ratio is low. In this work we present a new method called Fourier Diffusion Models which replaces the scalar operations of the forward process with shift-invariant convolutions and the additive white noise with additive stationary noise. This allows for control of MTF and NPS at intermediate time steps. Additionally, the forward process can be crafted to converge to the same MTF and NPS as the measured images. This way, we can model continuous probability flow from true images to measurements. In this way, the sample time can be used to control the tradeoffs between measurement uncertainty and generative uncertainty of posterior estimates. We compare Fourier diffusion models to existing scalar diffusion models and show that they achieve a higher level of performance and allow for a smaller number of time steps.

研究动机与目标

• 解决传统基于得分的扩散模型因退化为纯白噪声而导致时间步数过多的低效问题。
• 在前向和反向扩散过程中实现对调制传递函数（MTF）和噪声功率谱（NPS）的直接控制。
• 通过使前向过程与真实成像系统特性（如空间模糊和纹理噪声）对齐，提升图像生成性能。
• 通过减少反向采样时间步数，降低计算成本，实现在更少步数内生成高质量图像。
• 在低剂量CT图像重建这一具有临床相关性的应用中，验证该方法的有效性。

提出的方法

• 使用循环矩阵在傅里叶域中建模系统模糊（MTF），以平移不变卷积操作替代标量前向扩散过程。
• 将加性噪声建模为具有指定噪声功率谱（NPS）的平稳高斯噪声，实现对噪声特性的频域控制。
• 将前向过程定义为具有加性平稳高斯噪声（ASGN）的线性空间不变（LSI）系统响应序列，使最终状态与实测图像统计特性匹配，而非纯噪声。
• 训练一个得分匹配神经网络，以估计傅里叶域中的时变得分函数，实现迭代去噪与锐化。
• 使用学习到的得分函数进行反向采样，从实测数据生成高质量图像，其中时间步表示从噪声模糊测量到清晰图像的演化过程。
• 采用适用于任意具有线性系统和加性高斯噪声的高斯过程的得分匹配损失和反向SDE公式化方法，不限于平移不变系统。

Figure 1: A probabilistic graphical model for the stochastic process, consisting of linear shift invariant systems and additive stationary Gaussian noise. The measurement-conditioned diffusion model, in light purple, is trained to approximate the reverse process, in dark purple.

实验结果

研究问题

• RQ1 在前向扩散过程中对MTF和NPS进行控制，是否能提升基于得分的图像生成效率与质量？
• RQ2 若将前向过程建模为收敛于实测图像统计特性（而非纯噪声），是否能减少实现高质量生成所需的反向采样步数？
• RQ3 初始实测图像与前向过程最终状态之间的相似性，如何影响反向采样过程的性能？
• RQ4 傅里叶扩散模型是否能在更少时间步数和更低计算成本下，优于标量扩散模型在图像恢复任务中的表现？
• RQ5 性能提升是源于反向过程起始分布与真实数据分布之间更小的分布偏移，还是源于与卷积神经网络架构的更好兼容性？

主要发现

• 傅里叶扩散模型仅用8个时间步即可实现与标量模型在128或1024个时间步时相当甚至更优的图像质量。
• 在所有时间步数下，傅里叶扩散模型的均方误差（MSE）显著低于标量模型，尤其在低时间步数时改善最为明显。
• 均方偏差和均方方差指标也一致显示傅里叶扩散模型性能更优，表明其具有更高的保真度和更低的估计误差。
• 视觉对比显示，傅里叶扩散模型在低时间步数下生成的重建图像质量更高，伪影更少，结构保留更佳。
• 与标量模型相比，傅里叶扩散模型的反向过程随时间退化更小，表明其得分估计更稳定、更准确。
• 结果表明，性能提升源于初始实测图像与前向过程最终状态之间更接近的分布相似性，从而减轻了神经网络对人工退化过程的逆向重构负担。

Figure 2: Diagram of the score-matching neural network. The inputs are the forward process sample image, low-dose CT measured image, and sample time. The output is the estimated score function.