Technical Summary: Consistency Regularised Gradient Flows for Inverse Problems

1. Problem Statement

The paper addresses the challenge of solving linear inverse problems of the form $y = A(x_0) + n$, where $x_0$ is a signal, $A$ is a known forward operator, and $n$ is Gaussian noise. The goal is to perform Bayesian inference to recover $x_0$ from $y$ using the posterior distribution $p(x_0|y) \propto p(y|x_0)p(x_0)$.

While Vision-Language Latent Diffusion Models (LDMs) provide powerful generative priors $p(x_0)$, existing LDM-based solvers face two primary limitations:

1. Computational Cost: They typically require a large number of Neural Function Evaluations (NFEs) and backpropagation through large pretrained components (specifically the VAE decoder), leading to high memory footprints and slow inference.
2. Prompt Sensitivity: State-of-the-art methods often rely on text prompts to condition generation. If the prompt is not perfectly aligned with the scene, reconstruction quality degrades. While prompt optimization exists, current methods (e.g., LATINO-PRO) alternate between sampling and optimization steps, resulting in slow convergence and high NFE counts (often >50).

2. Methodology

The authors propose Consistency-regularised Wasserstein Gradient Flow (CWGF), a unified framework that jointly performs posterior sampling and prompt optimization in the latent space of a Latent Consistency Model (LCM).

2.1 Theoretical Framework

The method formulates the inverse problem as an optimization over the product space of prompt embeddings $C$ and probability measures on the latent space $P_2(Z)$. The objective function $F[c, \mu]$ combines:

• Regularization ($R$): The Kullback-Leibler (KL) divergence $D_{KL}(\mu \| p_c)$, encouraging the latent distribution $\mu$ to align with the LCM prior $p_c$ conditioned on prompt $c$.
• Likelihood ($L$): The negative log-likelihood $-E_\mu[\log p(y|z)]$, ensuring the samples match the observed data.

The optimization is driven by a continuous-time gradient flow:

• Euclidean Gradient: Updates the prompt embedding $c$ to maximize the marginal likelihood (MMLE).
• Wasserstein-2 Gradient: Updates the probability measure $\mu$ to minimize the objective functional.

2.2 Discretization and Implementation

To make the continuous flow computationally feasible, the authors derive a discrete algorithm using Lie-Trotter splitting and specific approximations:

1. Prior Subflow (Latent Space):

   • Instead of computing intractable score functions directly, the method utilizes the Consistency Model (CM) property.
   • It approximates the gradient of the KL divergence using a "flow-map discrepancy." This involves comparing the CM's prediction of the clean latent $g_\theta(z_t, t, c)$ against a kernelized estimate of the current particle distribution's mean.
   • This step includes a "repulsion" term to prevent particle collapse, pushing particles away from their local barycenter.

2. Prompt Flow (Embedding Space):

   • The gradient with respect to the prompt $c$ is derived using Fisher's identity and denoising score matching.
   • Crucially, this gradient is estimated using the LCM's score output as a proxy, avoiding the need to differentiate through the full reverse SDE/ODE or the VAE decoder.

3. Likelihood Subflow (Data Consistency):

   • To avoid backpropagation through the VAE decoder (which is memory-intensive), the authors exploit the VAE encoder $q_\phi(z|x)$.
   • Using a derived approximation (Proposition 2), the likelihood drift is computed using the encoder's output and the Gaussian posterior of the pixel space given the latent and observation. This allows for data consistency updates without differentiating through the decoder.

Algorithm Summary:
The algorithm iterates $K$ times (e.g., 16 steps). In each iteration, it:

1. Samples a diffusion timestep and adds noise to particles.
2. Updates the prompt embedding $c$ via gradient descent on the KL divergence.
3. Updates the latent particles $\mu$ via a split step: first moving toward the prior (using the CM flow map) and then moving toward the likelihood (using the encoder-based drift).

3. Key Contributions

• Unified Gradient Flow: A novel Euclidean-Wasserstein gradient flow framework that jointly optimizes the text prompt and the posterior distribution in a single flow, eliminating the need for alternating sampling and optimization loops.
• Backpropagation-Free Inference: The method avoids backpropagation through the autoencoder (VAE) components during the likelihood update by utilizing the pretrained encoder, significantly reducing memory usage.
• Efficiency: By leveraging Consistency Models (LCMs), the method achieves high-quality reconstructions with very few NFEs (specifically 16), compared to hundreds required by diffusion-based baselines.
• Robustness: The joint optimization allows the method to correct misaligned or "adversarial" prompts during the inference process, recovering semantic alignment that other methods fail to achieve.

4. Experimental Results

The method was evaluated on standard linear inverse problems (Gaussian deblurring, motion deblurring, and $\times8$ super-resolution) using FFHQ-512 and ImageNet-512 datasets.

• Performance: CWGF achieves State-of-the-Art (SOTA) performance in terms of FID, PSNR, and LPIPS.
  • On FFHQ-512 Gaussian deblurring, CWGF achieved an FID of 21.95 and PSNR of 29.42 with only 16 NFEs.
  • This outperforms LATINO-PRO (65 NFEs, FID 32.48) and P2L (400 NFEs, FID 45.12).
• Prompt Robustness: In experiments where the initial prompt was misinitialized (e.g., "A photo of a cat" for a face reconstruction), CWGF successfully optimized the prompt to retrieve the correct semantics, whereas methods without prompt optimization or with fixed prompts failed to generate plausible high-resolution features.
• Computational Efficiency: CWGF requires significantly fewer NFEs (16) compared to competitors (ranging from 65 to 1000) and demonstrates lower memory footprint than methods requiring decoder backpropagation (e.g., LATINO-PRO).

5. Significance and Claims

The paper claims that CWGF represents a significant step forward in solving inverse problems with generative priors by unifying prompt tuning and posterior sampling. Its primary significance lies in:

1. Reducing Computational Cost: Achieving SOTA results with a fraction of the computational resources (NFEs and memory) required by existing LDM-based solvers.
2. Practicality: Enabling high-quality image restoration on consumer-grade GPUs by removing the need for expensive decoder backpropagation.
3. Theoretical Rigor: Providing a provably convergent framework (under strong log-concavity assumptions) that bridges the gap between empirical Bayesian inference and gradient flows in metric spaces.

The authors note limitations, specifically that the Euclidean geometry assumption for prompt embeddings may not fully capture their intrinsic structure, and the method is currently restricted to Consistency Model priors, though extensions to other few-step models are suggested as future work.