Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Injecting Measurement Information Yields a Fast and Noise-Robust Diffusion-Based Inverse Problem Solver

Jonathan Patsenker, Henry Li, Myeongseob Ko, Ruoxi Jia, Yuval Kluger

TL;DR

The paper tackles inverse problems solved by diffusion-based priors by injecting measurement information into the denoising process. It replaces the conventional unconditional Tweedie guidance with a conditional posterior mean E[x0|xt,y], computed through a single-parameter maximum-likelihood correction to the forward process score and integrated into standard samplers. The authors prove correctness and sufficiency results for the conditional Tweedie approach and demonstrate substantial practical gains: fast, memory-efficient, and noise-robust reconstruction across multiple tasks and datasets, with Jacobian-free implementations. This approach enables reliable, resource-friendly inverse problem solving suitable for real-world deployment and limited compute budgets.

Abstract

Diffusion models have been firmly established as principled zero-shot solvers for linear and nonlinear inverse problems, owing to their powerful image prior and iterative sampling algorithm. These approaches often rely on Tweedie's formula, which relates the diffusion variate $\mathbf{x}_t$ to the posterior mean $\mathbb{E} [\mathbf{x}_0 | \mathbf{x}_t]$, in order to guide the diffusion trajectory with an estimate of the final denoised sample $\mathbf{x}_0$. However, this does not consider information from the measurement $\mathbf{y}$, which must then be integrated downstream. In this work, we propose to estimate the conditional posterior mean $\mathbb{E} [\mathbf{x}_0 | \mathbf{x}_t, \mathbf{y}]$, which can be formulated as the solution to a lightweight, single-parameter maximum likelihood estimation problem. The resulting prediction can be integrated into any standard sampler, resulting in a fast and memory-efficient inverse solver. Our optimizer is amenable to a noise-aware likelihood-based stopping criteria that is robust to measurement noise in $\mathbf{y}$. We demonstrate comparable or improved performance against a wide selection of contemporary inverse solvers across multiple datasets and tasks.

Injecting Measurement Information Yields a Fast and Noise-Robust Diffusion-Based Inverse Problem Solver

TL;DR

The paper tackles inverse problems solved by diffusion-based priors by injecting measurement information into the denoising process. It replaces the conventional unconditional Tweedie guidance with a conditional posterior mean E[x0|xt,y], computed through a single-parameter maximum-likelihood correction to the forward process score and integrated into standard samplers. The authors prove correctness and sufficiency results for the conditional Tweedie approach and demonstrate substantial practical gains: fast, memory-efficient, and noise-robust reconstruction across multiple tasks and datasets, with Jacobian-free implementations. This approach enables reliable, resource-friendly inverse problem solving suitable for real-world deployment and limited compute budgets.

Abstract

Diffusion models have been firmly established as principled zero-shot solvers for linear and nonlinear inverse problems, owing to their powerful image prior and iterative sampling algorithm. These approaches often rely on Tweedie's formula, which relates the diffusion variate to the posterior mean , in order to guide the diffusion trajectory with an estimate of the final denoised sample . However, this does not consider information from the measurement , which must then be integrated downstream. In this work, we propose to estimate the conditional posterior mean , which can be formulated as the solution to a lightweight, single-parameter maximum likelihood estimation problem. The resulting prediction can be integrated into any standard sampler, resulting in a fast and memory-efficient inverse solver. Our optimizer is amenable to a noise-aware likelihood-based stopping criteria that is robust to measurement noise in . We demonstrate comparable or improved performance against a wide selection of contemporary inverse solvers across multiple datasets and tasks.

Paper Structure

This paper contains 39 sections, 9 theorems, 67 equations, 25 figures, 7 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Revisiting Tweedie's for Inverse Problems
  3. Diffusion Conditional Sampling
  4. Measurement Likelihood Model
  5. Noise-Aware Maximization
  6. Theory
  7. Correctness
  8. Sufficiency
  9. No Expensive $\nabla s_\theta(\mathbf{x}_t, t)$ Evaluations
  10. A Near-Linear Reverse Process
  11. Experiments
  12. Ablation on the Noise-aware Maximization Optimizer
  13. Improvement on Jacobian-Free Implementations
  14. Impact and Limitations
  15. Conclusion
  16. ...and 24 more sections

Key Result

Theorem 2.1

Let $\mathbf{x}_t$ be sampled from a diffusion process (as in Eq. eq:xt_marginal). $\mathbb{E}[\mathbf{x}_0 | \mathbf{x}_t] = \mathbf{x}_0$ if and only if $p(\mathbf{x}_t)$ is a simple isotropic Gaussian with mean $\sqrt{\alpha_t} \mathbf{x}_0$ and variance $\sigma_t \mathbf{I}$.

Figures (25)

  • Figure 1: The posterior mean before and after conditioning on $\mathbf{y}$ for the random inpainting inverse problem.
  • Figure 1: Overview of pixel-based solvers used for comparisons in this work. We list the type (Section \ref{['sec:tweedies']}), whether it requires backpropagation through a neural function evaluation, runtime, and memory footprint.
  • Figure 2: (a) $\mathbb{E}[\mathbf{x}_0|\mathbf{x}_t]$ via the unconditional score versus (b) $\mathbb{E}[\mathbf{x}_0|\mathbf{x}_t, \mathbf{y}]$ via the forward process score estimate obtained by our likelihood maximizer for an image in the FFHQ $256 \times 256$ dataset with motion blur applied. With the unconditional score, $\hat{\mathbf{x}}_0$ estimates the posterior mean of the dataset, rather than a sample $\mathbf{x}$ that satisfies $\mathcal{A}(\mathbf{x}) \approx \mathbf{y}$, especially at $T \gg 0$ (Section \ref{['sec:tweedies_explained']}).
  • Figure 3: Description of latent and Jacobian-free solvers used for comparisons in text. For each solver we list the type (as described in Section \ref{['sec:tweedies']}), optimization space (pixel or latent), whether it requires backpropagation through a neural function evaluation (NFE, i.e., the score network call), as well as runtime and memory footprint.
  • Figure 4: An illustration of our proposed sampling algorithm. An initial noise prediction $\epsilon_\theta$ is corrected by the solution $\epsilon_\mathbf{y}$ of a noise-aware maximization scheme of the measurement likelihood $p(\mathbf{y} | \mathbf{x}_t, \epsilon_y)$. This results in the corrected forward process noise prediction $(\epsilon_\theta + \epsilon_y) \approx -\sigma_t^{-1} \nabla \log p_t(\mathbf{x}_t | \mathbf{x}_0)$. For details see Section \ref{['sec:ours']}.
  • ...and 20 more figures

Theorems & Definitions (18)

  • Theorem 2.1
  • Theorem 3.1
  • Theorem 3.2: ${\epsilon_\mathbf{y}}_*$ is a sufficient statistic
  • Theorem B.1: Tweedie's Formula
  • proof : Proof (of Lemma \ref{['thm:tweedies']})
  • Lemma B.2: Sufficent condition
  • proof : Proof (of Lemma \ref{['thm:tweedie']})
  • Lemma B.3: Necessary condition
  • proof : Proof (of Lemma \ref{['thm:tweedie_sufficiency']})
  • proof : Proof (of Theorem \ref{['thm:tweedie_iff']})
  • ...and 8 more