Improving Diffusion Inverse Problem Solving
with Decoupled Noise Annealing

Diffusion models [17, 18, 19, 1, 20] generate data by reversing a predefined noising process. Let the data distribution be $p({\mathbf{x}}_{0})$. We can define a series of noisy data distributions $p({\mathbf{x}};\sigma)$ by adding Gaussian noise with a standard deviation of $\sigma$ to the data. These form the time-marginals of a stochastic differential equation (SDE) [20], given by

where $\sigma_{t}$ is the predefined noise schedule with $\sigma_{0}=0$ and $\sigma_{T}=\sigma_{\max}$, $\dot{\sigma}_{t}$ is the time derivative of $\sigma_{t}$, and ${\mathbf{w}}_{t}$ is a standard Wiener process. We use ${\mathbf{x}}_{t}$ interchangeably with ${\mathbf{x}}_{\sigma_{t}}$. For a sufficiently large $\sigma_{\max}$, the distribution $p({\mathbf{x}};\sigma_{\max})$ converges to pure Gaussian noise $\mathcal{N}({\bm{0}},\sigma_{\max}^{2}{\bm{I}})$.

To sample from data distribution $p({\mathbf{x}}_{0})$, we first draw an initial sample from $\mathcal{N}({\bm{0}},\sigma_{\max}^{2}{\bm{I}})$, then solve the reverse SDE

where $\nabla_{{\mathbf{x}}_{t}}\log p({\mathbf{x}}_{t};\sigma_{t})$ is the time-dependent score function [17, 1]. Here $\nabla_{{\mathbf{x}}_{t}}\log p({\mathbf{x}}_{t};\sigma_{t})$ is unknown, but can be approximated by training a diffusion model ${\bm{s}}_{\bm{\theta}}({\mathbf{x}}_{t},\sigma_{t})$ such that ${\bm{s}}_{\bm{\theta}}({\mathbf{x}}_{t},\sigma_{t})\approx\nabla_{{\mathbf{x}}_{t}}\log p({\mathbf{x}}_{t};\sigma_{t})$.

Inverse problems aim to recover data from partial, potentially noisy measurements. Formally, solving an inverse problem involves finding the inversion to a forward model that describes the measurement process. In general, a forward model takes the form of

where $\mathcal{A}$ is the measurement function, ${\mathbf{x}}_{0}$ represents the original data, ${\mathbf{y}}$ is the observed measurement, and ${\mathbf{n}}$ symbolizes the noise in the measurement process, often modeled as ${\mathbf{n}}\sim\mathcal{N}({\bm{0}},\beta_{{\mathbf{y}}}^{2}{\bm{I}})$. In a Bayesian framework, ${\mathbf{x}}_{0}$ comes from the posterior distribution $p({\mathbf{x}}_{0}\mid{\mathbf{y}})\propto p({\mathbf{x}}_{0})p({\mathbf{y}}\mid{\mathbf{x}}_{0})$. Here $p({\mathbf{x}}_{0})$ is a prior distribution that can be estimated from a given dataset, and $p({\mathbf{y}}\mid{\mathbf{x}}_{0})=\mathcal{N}(\mathcal{A}({\mathbf{x}}_{0}),\beta_{{\mathbf{y}}}^{2}{\bm{I}})$ models the noisy measurement process.

When the prior $p({\mathbf{x}}_{0})$ is modeled by a pre-trained diffusion model, we can modify Eq. 2 to approximately sample from the posterior distribution following Bayes’ rule, i.e.,

Here, the noisy likelihood $\nabla_{{\mathbf{x}}_{t}}\log p({\mathbf{y}}\mid{\mathbf{x}}_{t})$ is generally intractable. Multiple methods have been proposed to estimate the noisy likelihood [8, 29, 22, 23, 30, 31, 32, 6]. One predominant approach is the DPS algorithm [3], which estimates $p({\mathbf{y}}\mid{\mathbf{x}}_{t})\approx p({\mathbf{y}}\mid{\mathbf{x}}_{0}=\mathbb{E}[{\mathbf{x}}_{0}\mid{\mathbf{x}}_{t}])$. Another line of work [2, 33, 31] solves linear inverse problems by running the reverse diffusion process in the spectral domain via singular value decomposition (SVD). Other approaches bypass direct computation of this likelihood by interleaving optimization [25, 34, 35, 36, 37] or projection [2, 9, 7, 31] steps with normal diffusion sampling steps.

Despite promising empirical success, we find that this line of approaches faces challenges in solving more difficult inverse problems when the forward model is highly nonlinear. Accurately solving the reverse SDE in Eq. 4 requires the solver to take a very small step size $\Delta t>0$, causing ${\mathbf{x}}_{t}$ and ${\mathbf{x}}_{t+\Delta t}$ to be very close to each other. Consequently, ${\mathbf{x}}_{t}$ can only correct minor errors in ${\mathbf{x}}_{t+\Delta t}$, but oftentimes fails to address larger, global errors that require substantial changes to ${\mathbf{x}}_{t+\Delta t}$. One such failure case is given in Fig. 4.

Instead of solving the reverse-time SDE in Eq. 4, we propose a new noise annealing process that reduces the dependency between samples at consecutive time steps, as illustrated in Figs. 2, 3(a) and 3(b). Unlike previous methods, we ensure ${\mathbf{x}}_{t}$ and ${\mathbf{x}}_{t+\Delta t}$ are conditionally independent given ${\mathbf{x}}_{0}$. To generate sample ${\mathbf{x}}_{t}$ from ${\mathbf{x}}_{t+\Delta t}$, we follow a two-step procedure: (1) sampling ${\mathbf{x}}_{0\mid{\mathbf{y}}}\sim p({\mathbf{x}}_{0}\mid{\mathbf{x}}_{t+\Delta t},{\mathbf{y}})$, and (2) sampling ${\mathbf{x}}_{t}\sim\mathcal{N}({\mathbf{x}}_{0\mid{\mathbf{y}}},\sigma_{t}^{2}{\bm{I}})$. We repeat this process, gradually reducing noise until ${\mathbf{x}}_{0}$ is sampled. We call this process decoupled noise annealing, which is justified by the proposition below.

Remark. The key idea of Proposition 1 is to enable the sampling from $p({\mathbf{x}}_{t_{2}}\mid{\mathbf{y}})$ for any noise level $\sigma_{t_{2}}$ given any sample $x_{t_{1}}$ at another noise level $\sigma_{t_{1}}$. For a sufficiently large $\sigma_{T}$, one can assume $p({\mathbf{x}}_{T}\mid{\mathbf{y}})\approx p({\mathbf{x}}_{T};\sigma_{T})\approx\mathcal{N}(\bm{0},\sigma_{T}^{2}{\bm{I}})$. Starting from ${\mathbf{x}}_{T}$, we can iteratively sample from $p({\mathbf{x}}_{t}\mid{\mathbf{y}})$ with $\sigma_{t}$ annealed down from $\sigma_{T}$ to 0.

The first step of our decoupled noise annealing requires sampling ${\mathbf{x}}_{0\mid{\mathbf{y}}}\sim p({\mathbf{x}}_{0}\mid{\mathbf{x}}_{t},{\mathbf{y}})$ where ${\mathbf{x}}_{t}$ and ${\mathbf{y}}$ are known. Since ${\mathbf{y}}$ is conditionally independent from ${\mathbf{x}}_{t}$ given ${\mathbf{x}}_{0}$, we can deduce from the Bayes’ rule that

To sample ${\mathbf{x}}_{0\mid{\mathbf{y}}}$ from this unnormalized distribution, we propose to run Langevin dynamics[38], given by

where $\eta>0$ is the step size and $\bm{\epsilon}_{j}\sim\mathcal{N}({\bm{0}},{\bm{I}})$. When $\eta\to 0$ and $j\to\infty$, the sample ${\mathbf{x}}_{0}^{(j)}$ will be approximately distributed according to $p({\mathbf{x}}_{0}\mid{\mathbf{x}}_{t},{\mathbf{y}})$.

Following previous works [3, 22, 23], we approximate the conditional distribution $p({\mathbf{x}}_{0}\mid{\mathbf{x}}_{t})$ with a Gaussian distribution

where $\hat{{\mathbf{x}}}_{0}({\mathbf{x}}_{t})$ is an estimator of ${\mathbf{x}}_{0}$ given ${\mathbf{x}}_{t}$, and the variance $r_{t}^{2}$ is specified using heuristics. Given a pre-trained diffusion model ${\bm{s}}_{\bm{\theta}}({\mathbf{x}}_{t},t)$, one way to compute $\hat{{\mathbf{x}}}_{0}({\mathbf{x}}_{t})$ is to solve the (unconditional) probability flow ODE starting at ${\mathbf{x}}_{t}$. We leave details in Section D.2. When the measurement noise is an isotropic Gaussian, i.e., ${\mathbf{n}}\sim\mathcal{N}({\bm{0}},\beta_{\mathbf{y}}^{2}{\bm{I}})$, the update rule simplifies to:

Combining Eq. 9 with Proposition 1, we define an algorithm called Decoupled Annealing Posterior Sampling (DAPS), summarized in Algorithm 1. In particular, given a noise schedule $\sigma_{t}$ and a time discretization $\{t_{i},i=0,\dots,N_{A}\}$, we iteratively sample ${\mathbf{x}}_{t_{i}}$ from the measurement-conditioned time-marginal $p({\mathbf{x}}_{t_{i}}\mid{\mathbf{y}})$ for $i=N_{A},N_{A}-1,\dots,0$, following Eq. 5 in Proposition 1. This algorithm provides an approximate sample ${\mathbf{x}}_{0}$ from the posterior distribution $p({\mathbf{x}}_{0}\mid{\mathbf{y}})$.

The computational cost of Langevin dynamics mostly originates from evaluating the measurement function $\mathcal{A}$. In most image restoration tasks, this operation is significantly more efficient compared to evaluating the diffusion model. As demonstrated in Section C.1, Langevin dynamics introduce only a small overhead to the sampling process. Moreover, this framework can be adapted to sampling with pre-trained latent diffusion models by factorizing the probabilistic graphical model for the latent diffusion process, as shown in Fig. 3(b). Further details are provided in Appendix A.

Comparison with other posterior sampling methods. Unlike most previous algorithms, our sampling method does not solve a specific SDE/ODE; instead, we recursively sample from the time-marginal $p({\mathbf{x}}_{t}\mid{\mathbf{y}})$ with noise annealing to zero. We therefore decouple the dependency of ${\mathbf{x}}_{t}$ and ${\mathbf{x}}_{t+\Delta t}$ in the sampling process. We argue that decoupling helps correct the errors accumulated in early diffusion sampling steps by allowing non-local transitions. This is particularly important when the measurement function is nonlinear. As a concrete example, Fig. 4 compares DAPS and DPS when solving a 2D nonlinear inverse problem with a Gaussian mixture prior. We visualize the trajectories of ${\mathbf{x}}_{t}$ from time $T$ to $0$ for both methods. For DAPS, points on the trajectory have significantly larger variations compared to those of DPS. As a result, DPS converges to wrong solutions, but DAPS is able to approximate the true posterior distribution. We provide more discussions in Appendix E.

Comparison with optimization-based methods. Although our method does not directly involve inner-loop optimization, we highlight that it has connections with existing optimization-based solvers for inverse problems. For example, ReSample [25] alternate between denoising, optimizing, and resampling to solve inverse problems using latent diffusion models. Our method also resembles  [25] if we set standard deviation $\beta_{\mathbf{y}}\to 0$ and assume $\eta\beta_{\mathbf{y}}^{2}$ is constant in Eq. 12.

We evaluate our method using both pixel-space and latent diffusion models. For pixel-based diffusion experiments, we leverage the pre-trained diffusion models trained by [3] on the FFHQ dataset and the pre-trained model from [39] on the ImageNet dataset. For latent diffusion models, we use the same pre-trained models as [25]: the unconditional LDM-VQ4 trained on FFHQ and ImageNet by [40]. The autoencoder’s downsample factor is $4$. We use the same time step discretization and noise schedule as EDM [20].

As mentioned in Section 3, we implement a few-step Euler ODE solver to compute $\hat{\mathbf{x}}_{0}({\mathbf{x}}_{t})$, while maintaining the same number of neural function evaluations (NFE) for different noise levels. In our experiments, we use DAPS-1K for all linear tasks and DAPS-4K for all nonlinear tasks. DAPS-1K uses $4$ ODE solver NFE and $250$ annealing scheduling steps, while DAPS-4K uses $10$ ODE solver NFE and $400$ annealing scheduling steps. Ablation studies on their effects are shown in Section 4.3. The corresponding latent version, LatentDAPS, follows the same settings but performs sampling in the latent space of VAEs. We use $100$ Langevin steps per denoising iteration and tune learning rates separately for each task. Further details on model configurations, samplers, and other hyperparameters are provided in Appendix D, along with more discussion on sampling efficiency in Section C.1.

Datasets and metrics. Adopting the previous convention, we test our method on two image datasets, FFHQ $256\times 256$ [41] and ImageNet $256\times 256$ [42]. To evaluate our method, we use $100$ images from the validation set for both FFHQ and ImageNet. We include Learned Perceptual Image Patch Similarity (LPIPS) score [43] and peak signal-to-noise-ratio (PSNR) as our main evaluation metrics. For both our method and baselines, we use the versions implemented in piq [44] with all images normalized to the range $[0,1]$. The replace-pooling option is enabled for LPIPS evaluation.

Inverse problems. We evaluate our method with a series of linear and nonlinear tasks. For linear inverse problems, we consider (1) super-resolution, (2) Gaussian deblurring, (3) motion deblurring, (4) inpainting (with a box mask), and (5) inpainting (with a 70% random mask). For Gaussian and motion deblurring, kernels of size 61$\times$61 with standard deviations of 3.0 and 0.5, respectively, are used. In the super-resolution task, a bicubic resizer downscales images by a factor of 4. The box inpainting task uses a random box of size 128$\times$128 to mask the original images, while random mask inpainting uses a generated random mask where each pixel has a 70% chance of being masked, following the settings in [25].

We consider three nonlinear inverse problems: (1) phase retrieval, (2) high dynamic range (HDR) reconstruction, and (3) nonlinear deblurring. Due to the inherent instability of phase retrieval, we adopt the strategy from DPS [3], using an oversampling rate of 2.0 and reporting the best result out of four independent samples. The goal of HDR reconstruction is to recover a higher dynamic range image (factor of 2) from a low dynamic range image. For nonlinear deblurring, we use the default setting as described in [45]. All linear and nonlinear measurements are subject to white Gaussian noise with a standard deviation of $\beta_{\mathbf{y}}=0.05$. Further details regarding the forward measuring functions for each task and their respective hyperparameters are provided in the Appendix D.

Baselines. We compare our methods with the following baselines: Denoising Diffusion Restoration Models (DDRM) [2], Diffusion Posterior Sampling (DPS) [3], Denoising Diffusion Null-Space Model (DDNM) [31], and Plug-and-Play Alternating Direction Method of Multipliers (PnP-ADMM) [46] for pixel-based diffusion model experiments. We compare the latent diffusion version of our methods with Posterior Sampling with Latent Diffusion Models (PSLD) [27] and ReSample [25] that operates also in the latent space. Note that DDRM, DDNM, and PSLD are not able to handle nonlinear inverse problems. We also include Regularization by denoising diffusion process (RED-diff) [28] for nonlinear experiments.

We show quantitative results for the FFHQ dataset in Tables 1 and 2, and results for the ImageNet dataset in Tables 3 and 4. The qualitative comparisons are provided in Fig. 1. As shown in Table 1, our method achieves comparable or even better performance across all selected linear inverse problems. Moreover, our methods are remarkably stable when handling nonlinear inverse problems. For example, although existing methods such as DPS can recover high-quality images from phase retrieval measurements, they suffer from an extremely high failure rate, resulting in a relatively low PSNR and a high LPIPS. However, DAPS demonstrates superior stability across different samples and achieves a significantly higher success rate, resulting in a substantial improvement in PSNR and LPIPS as evidenced by the results in Table 2. Moreover, DAPS captures and recovers much more fine-grained details in measurements compared to existing baselines. We include more samples and comparisons in Appendix F for further illustration.

To understand the evolution of the noise annealing procedure, we evaluate two crucial trajectories in our method: 1) the estimated mean of $p({\mathbf{x}}_{0}\mid{\mathbf{x}}_{t})$, denoted as $\hat{{\mathbf{x}}}_{0}({\mathbf{x}}_{t})$; and 2) samples from $p({\mathbf{x}}_{0}\mid{\mathbf{x}}_{t},{\mathbf{y}})$ obtained through Langevin dynamics, denoted as ${\mathbf{x}}_{0\mid{\mathbf{y}}}$. We assess image quality and measurement error across these trajectories, as depicted in Fig. 6. The measurement error for ${\mathbf{x}}_{0\mid{\mathbf{y}}}\sim p({\mathbf{x}}_{0}\mid{\mathbf{x}}_{t},{\mathbf{y}})$ remains consistently low throughout the process, while the measurement error for $\hat{{\mathbf{x}}}_{0}({\mathbf{x}}_{t})$ continuously decreases. Consequently, metrics such as PSNR and LPIPS exhibit stable improvement as the noise level is annealed.

In addition, our method can generate diversified samples when the measurements contain less information. To illustrate this, we selected two tasks where the posterior distributions may have multiple modes: (1) super-resolution by a factor of 16, and (2) box inpainting with a box size of 192$\times$192. Some samples from DAPS are shown in Fig. 5, demonstrating DAPS’s ability to produce diverse samples while preserving the measurement information. We show more samples in Appendix F for further demonstration.