Measurement-Aligned Sampling for Inverse Problem

TL;DR

MAS tackles linear inverse problems by marrying a diffusion-prior with a measurement-consistent objective, yielding a closed-form update $x_0^* = Y^{-1}\bigl[m_{0 \mid t} + H^{\mathsf T} W^{-1} y\bigr]$ where $W = \eta_1 HH^{\mathsf T} + \eta_2 I$ and $Y = I + H^{\mathsf T} W^{-1} H$. It provides a probabilistic interpretation via Bayesian linear regression and introduces an overshooting mechanism through negative $\eta_1$ to strengthen alignment with measurements. MAS extends prior work (DDNM, TMPD) and accommodates known Gaussian noise with calibrated corrections, unknown/non-Gaussian noise via adaptive $\eta_1,\eta_2$, and non-differentiable degradations by treating them as implicit noise, all while maintaining competitive computation times. Empirical results across SR, deblurring, inpainting, colorization, and JPEG/quantization restoration show MAS achieving state-of-the-art metrics (PSNR/SSIM/LPIPS/FID) in many cases, with robust performance under noise and degradations. The framework offers practical gains for robust image restoration in realistic noise settings and provides a principled pathway to calibrate measurement-consistency in diffusion-based inverse problem solvers.

Abstract

Diffusion models provide a powerful way to incorporate complex prior information for solving inverse problems. However, existing methods struggle to correctly incorporate guidance from conflicting signals in the prior and measurement, and often failed to maximizing the consistency to the measurement, especially in the challenging setting of non-Gaussian or unknown noise. To address these issues, we propose Measurement-Aligned Sampling (MAS), a novel framework for linear inverse problem solving that flexibly balances prior and measurement information. MAS unifies and extends existing approaches such as DDNM, TMPD, while generalizing to handle both known Gaussian noise and unknown or non-Gaussian noise types. Extensive experiments demonstrate that MAS consistently outperforms state-of-the-art methods across a variety of tasks, while maintaining relatively low computational cost.

Figures (8)

• Figure 1: Solving various inverse problems using unconditional diffusion models. Our model demonstrates better robustness with unknown noise and strong Gaussian noise.
• Figure 2: 2D illustration of the influence of parameters $\eta_1$ and $\eta_2$. Dots represent $x_0^*$, calculated via \ref{['eq:pm']}. (a) Parameter $\eta_1$ controls the trade-off between $m_{0 \mid t}$ and $\tilde{x}_0^{\text{DDNM}}$: as $\eta_1 \rightarrow \infty$, the posterior mean $x_0^*$ approaches $m_{0 \mid t}$; as $\eta_1 \rightarrow 0$, it converges to $\tilde{x}_0^{\text{DDNM}}$. (b) Adjusting $\eta_2$ alters the posterior trajectory differently from varying $\eta_1$.
• Figure 3: The sample process of solving inverse problems with unknown noise, where $\hat{x}_0^{\theta}\approx m_{0 \mid t}$ is the denoising output. Here we set $\eta_1 = 0$ and $\eta_2 = 0.5 a_t/c_t$.
• Figure 4: Ablation study of $\eta_1$ and $\eta_2$ on solving super-resolution and deblurring. We set NFE=20 for all tasks.
• Figure 5: Results on JPEG (QF=5) and quantization restoration.

Theorems & Definitions (11)

• Remark 1: Connection with DDNM wang2022zero
• Remark 2: Connection with optimization methods
• Proposition 3.0: Bayesian Linear Regression
• Remark 3: Connection to TMPD boys2023tweedie
• Remark 4: 'Overshooting' trick
• proof
• Lemma A.1
• proof
• Proposition A.2
• proof