An Unconditional Representation of the Conditional Score in Infinite-Dimensional Linear Inverse Problems

Technical Summary: An Unconditional Representation of the Conditional Score in Linear Inverse Problems in Infinite-Dimensional Spaces

Problem Statement
Score-based Diffusion Models (SDMs) have evolved into a powerful tool for sampling from posterior distributions in Bayesian inverse problems, offering flexible prior distributions that surpass conventional, manually designed regularizers. However, existing methods for posterior sampling face a critical trade-off between computational efficiency and scalability.

1. Correction-based methods (e.g., adding likelihood gradients to the score or enforcing data consistency) require repeated evaluations of the forward operator $A$ during the sampling process. This presents an insurmountable computational barrier for large-scale, high-dimensional problems.
2. Direct conditional training methods learn a score function conditioned on measurements ($s(x, t; \mu_y)$). While this avoids forward evaluations during sampling, it drastically increases the input dimensionality of the neural network (from $x$ to $(x, y)$), necessitating extensive training data and complex architectures to accurately approximate the conditional score, particularly in infinite-dimensional settings.

This work addresses the challenge of performing efficient and scalable posterior sampling in linear inverse problems in infinite-dimensional spaces without incurring the high online costs of forward model evaluations or the high sample complexity of direct conditional training.

Methodology: Unconditional Representation of the Conditional Score (UCoS)
The authors propose UCoS, a framework that shifts the computational burden of conditioning from the online sampling phase to an offline training phase. The core idea is that for linear inverse problems with Gaussian noise, the conditional score function can be exactly derived from an unconditional score function of a modified, task-dependent diffusion process via affine transformations.

• Theoretical Foundation:

  • The method operates in a separable Hilbert space $\mathcal{H}$, ensuring invariance with respect to discretization.
  • Instead of training a network to approximate $s(x, t; \mu_y)$, the authors define a new diffusion process $\tilde{X}_t$, whose distribution is a mixture of the prior distribution and a Gaussian term dependent on the forward operator $A$.
  • They prove (Theorem 3.7) that the true conditional score $s(x, t; \mu_y)$ is linked to the unconditional score $\tilde{s}$ of this new process via the following identity:
    $$s(x, t; \mu_y) = \lambda(t) \left( r(\xi_t(x, y)) - e^{-t/2}x \right)$$
    where $r$ is a transformed version of the unconditional score, $\lambda(t)$ is a time-dependent scalar, and $\xi_t(x, y)$ is an affine transformation incorporating the measurement $y$ and the adjoint operator $A^*$.
  • Crucially, the term $A^*\Gamma^{-1}y$ (where $\Gamma$ is the noise covariance) can be computed once before sampling begins. The subsequent sampling process requires only the evaluation of the trained network $r_\theta$ and simple vector additions, thereby eliminating the need for further applications of $A$ or $A^*$ during the inverse diffusion steps.

• Implementation:

  • Offline Phase: A neural network (specifically a Fourier Neural Operator, FNO) is trained to approximate the function $r$. Training utilizes a denoising score matching objective on the modified process $\tilde{X}_t$. This phase requires access to $A$ and $A^*$, but no specific measurement data $y$.
  • Online Phase: Posterior samples are generated by solving the SDE in reverse time using the transformed score. The process is "matrix-free" and requires only the pre-calculated shift vector $A^*\Gamma^{-1}y$ and the network evaluation.

Main Contributions

1. Exact Identity in Infinite Dimensions: The work establishes a rigorous theoretical identity (Theorem 3.7) linking the conditional score to a task-dependent unconditional score in infinite-dimensional Hilbert spaces via affine transformations. This generalizes the principle beyond specific finite-dimensional discretizations.
2. Convergence Analysis: The authors provide a quantitative error bound (Theorem 4.1) for the generative process, explicitly characterizing error contributions from training loss, numerical discretization, and initialization. This justifies the convergence of the method to the true posterior measure.
3. Scalability and Efficiency: By decoupling conditioning from the sampling loop, UCoS achieves the computational efficiency of unconditional methods (no forward evaluations during sampling) while maintaining the accuracy of conditional methods. It avoids the curse of dimensionality associated with training networks on the joint space $(x, y)$.

Experimental Results
The authors validate UCoS on three inverse problems: image inpainting, sparse-view computed tomography (CT), and image deblurring.

• Inpainting (Gaussian Prior): In a setting where the posterior is Gaussian and scores can be computed in closed form, UCoS performs comparably to the exact conditional method and other baselines, confirming the theoretical identity.
• CT Imaging and Deblurring (Learned Priors): Using the LIDC-IDRI and CelebA datasets, UCoS is compared against SDE ALD, Diffusion Posterior Sampling (DPS), projection-based methods (Proj), and direct conditional training.
  • Accuracy: UCoS generates samples with bias and standard deviation comparable to or better than unconditional methods, and significantly superior to the conditional method when using architectures with low complexity (e.g., 32 nodes per layer).
  • Efficiency: UCoS achieves the sampling speed of the conditional method (both avoid forward evaluations) but significantly outperforms correction-based methods (DPS, Proj, SDE ALD) that require repeated forward passes.
  • Robustness: The performance of the conditional method degrades noticeably with fewer network parameters due to the increased complexity of approximating the joint distribution, whereas UCoS maintains high quality even with simpler architectures.

Significance and Claims
The work claims that UCoS offers a "best-of-both-worlds" solution: it eliminates the computational bottleneck of forward model evaluations during sampling (a limitation of correction-based methods) while simultaneously avoiding the high sample complexity and architectural demands of direct conditional training.

The authors emphasize that their approach is discretization-invariant, making it suitable for infinite-dimensional inverse problems where performance guarantees in high-dimensional finite-dimensional approximations often fail. They acknowledge that the method is task-dependent and requires retraining if the forward operator $A$ changes, and that the exact correspondence relies on the linearity of the forward model and the Gaussian nature of the likelihood function (which enables the "completing the square" step). The work is presented as a proof of concept for a new class of diffusion-based samplers that shift the complexity of conditioning into the training phase without introducing approximation errors in the theoretical formulation.