Table of Contents
Fetching ...

Coupled Data and Measurement Space Dynamics for Enhanced Diffusion Posterior Sampling

Shayan Mohajer Hamidi, En-Hui Yang, Ben Liang

TL;DR

This paper addresses inverse problems by showing that diffusion priors can be more effectively applied when measurement information is embedded directly into the diffusion dynamics. The authors propose C-DPS, which runs a forward diffusion in the measurement space in parallel with the data-space diffusion, deriving a closed-form posterior p(x_{t-1} | x_t, y_{t-1}) and enabling exact recursive updates without likelihood approximations or constraint tuning. A scalable sampling algorithm based on pre-whitened conjugate gradients is introduced, and a latent variant (LC-DPS) extends the framework to latent spaces. Empirical results on FFHQ and ImageNet across inpainting, deblurring, and super-resolution demonstrate state-of-the-art performance and improved measurement fidelity, with ablations on posterior recovery and stability.

Abstract

Inverse problems, where the goal is to recover an unknown signal from noisy or incomplete measurements, are central to applications in medical imaging, remote sensing, and computational biology. Diffusion models have recently emerged as powerful priors for solving such problems. However, existing methods either rely on projection-based techniques that enforce measurement consistency through heuristic updates, or they approximate the likelihood $p(\boldsymbol{y} \mid \boldsymbol{x})$, often resulting in artifacts and instability under complex or high-noise conditions. To address these limitations, we propose a novel framework called \emph{coupled data and measurement space diffusion posterior sampling} (C-DPS), which eliminates the need for constraint tuning or likelihood approximation. C-DPS introduces a forward stochastic process in the measurement space $\{\boldsymbol{y}_t\}$, evolving in parallel with the data-space diffusion $\{\boldsymbol{x}_t\}$, which enables the derivation of a closed-form posterior $p(\boldsymbol{x}_{t-1} \mid \boldsymbol{x}_t, \boldsymbol{y}_{t-1})$. This coupling allows for accurate and recursive sampling based on a well-defined posterior distribution. Empirical results demonstrate that C-DPS consistently outperforms existing baselines, both qualitatively and quantitatively, across multiple inverse problem benchmarks.

Coupled Data and Measurement Space Dynamics for Enhanced Diffusion Posterior Sampling

TL;DR

This paper addresses inverse problems by showing that diffusion priors can be more effectively applied when measurement information is embedded directly into the diffusion dynamics. The authors propose C-DPS, which runs a forward diffusion in the measurement space in parallel with the data-space diffusion, deriving a closed-form posterior p(x_{t-1} | x_t, y_{t-1}) and enabling exact recursive updates without likelihood approximations or constraint tuning. A scalable sampling algorithm based on pre-whitened conjugate gradients is introduced, and a latent variant (LC-DPS) extends the framework to latent spaces. Empirical results on FFHQ and ImageNet across inpainting, deblurring, and super-resolution demonstrate state-of-the-art performance and improved measurement fidelity, with ablations on posterior recovery and stability.

Abstract

Inverse problems, where the goal is to recover an unknown signal from noisy or incomplete measurements, are central to applications in medical imaging, remote sensing, and computational biology. Diffusion models have recently emerged as powerful priors for solving such problems. However, existing methods either rely on projection-based techniques that enforce measurement consistency through heuristic updates, or they approximate the likelihood , often resulting in artifacts and instability under complex or high-noise conditions. To address these limitations, we propose a novel framework called \emph{coupled data and measurement space diffusion posterior sampling} (C-DPS), which eliminates the need for constraint tuning or likelihood approximation. C-DPS introduces a forward stochastic process in the measurement space , evolving in parallel with the data-space diffusion , which enables the derivation of a closed-form posterior . This coupling allows for accurate and recursive sampling based on a well-defined posterior distribution. Empirical results demonstrate that C-DPS consistently outperforms existing baselines, both qualitatively and quantitatively, across multiple inverse problem benchmarks.

Paper Structure

This paper contains 58 sections, 1 theorem, 64 equations, 11 figures, 11 tables, 4 algorithms.

Key Result

Proposition 1

For the sequence $\{\boldsymbol{y}_t\}_{t=0}^T$ defined in eq:y_eq, the distribution of $\boldsymbol{y}_t$ is a Gaussian whose mean and covariance at step $t$ are given by

Figures (11)

  • Figure 1: Visualization of reconstructed samples for some linear inverse problems using our proposed C-DPS method, compared with two leading baselines: DPS DPS and Resample resample. C-DPS exhibits less visual artifacts and improved measurement consistency.
  • Figure 2: Qualitative results on FFHQ dataset.
  • Figure 3: Visualizing the first two dimensions of the estimated posterior distributions for the configuration ($d=80$, $m=1$, $\sigma = 10^{-1}$) for a random $\bold{A}$.
  • Figure 4: Evolution of the measurement error $\|\boldsymbol{y} - \bold{A} \boldsymbol{x}_t\|_2^2$ over the sampling trajectory for DPS, ReSample, and C-DPS on the FFHQ (left) and ImageNet (right) datasets.
  • Figure 5: Cosine similarity and MSE between score vectors. Cosine similarity remains high and MSE stays low across diffusion steps, supporting the approximation $\boldsymbol{s}_\theta(\boldsymbol{x}_{t-1}, t-1) \approx \boldsymbol{s}_\theta(\boldsymbol{x}_t,\,t)$.
  • ...and 6 more figures

Theorems & Definitions (3)

  • Remark 1: Choice of noise schedule
  • Proposition 1: Distribution of $\{\boldsymbol{y}_t\}$
  • Remark 2