Constrained Particle Seeking: Solving Diffusion Inverse Problems with Just Forward Passes

This paper introduces Constrained Particle Seeking (CPS), a novel gradient-free method that reformulates diffusion-based inverse problems as constrained optimization tasks to effectively solve them using only forward passes, even when the forward observation process is unknown.

The Gist

Imagine you are trying to solve a jigsaw puzzle, but you only have a blurry, incomplete photo of the finished picture to guide you. This is what scientists call an "inverse problem." You have the result (the blurry photo), and you need to figure out the original pieces (the clear image).

For a long time, computers have used a powerful tool called a Diffusion Model to help solve these puzzles. Think of a Diffusion Model as a master artist who knows exactly what a "perfect face" or a "perfect galaxy" should look like because they've studied millions of them.

The Old Way: The "Rejection Sampling" Game

Previously, when the computer tried to solve the puzzle, it would play a game of "guess and check" that was incredibly wasteful.

1. The computer would ask the artist to draw 64 different versions of the next step in the puzzle.
2. It would then look at all 64 drawings, pick the one that looked most like the blurry photo, and throw away the other 63.
3. It would repeat this for every single step of the puzzle.

The Problem: This is like throwing away 63 perfectly good clues just because one looked slightly better. It's slow, expensive, and often misses the best solution because it ignores the valuable information hidden in the "rejected" drawings.

The New Way: Constrained Particle Seeking (CPS)

The authors of this paper, Hongkun Dou and his team, came up with a smarter strategy called Constrained Particle Seeking (CPS).

Instead of throwing away the 63 "rejected" drawings, they decided to listen to all of them.

Here is how CPS works, using a simple analogy:

1. The "Group Brain" (Surrogate Model)

Imagine you are trying to find the best route through a foggy forest to reach a specific destination (the blurry photo).

• Old Method: You send out 64 scouts. They all walk a bit. You pick the one scout who is closest to the destination and send them on. You ignore the other 63.
• CPS Method: You send out 64 scouts. Instead of picking just one, you ask all 64 scouts to report back on the terrain.
  • Scout A says, "If I go left, I get closer."
  • Scout B says, "If I go right, I get closer."
  • Even Scout Z (who was walking the wrong way) says, "If I go the opposite of where I was going, I get closer!"

CPS takes all these reports and builds a local map (a "surrogate model") of the forest. It uses the collective wisdom of the whole group to figure out the exact best direction to go, rather than just guessing based on one person.

2. The "Safety Net" (Constraints)

There's a catch. If you just follow the scouts blindly, you might wander off into a swamp because the map is foggy. You need to stay on the "main path" where the forest is dense and safe.

In the computer world, this "main path" is the Prior. It's the rule that says, "Remember, we are trying to generate a human face, not a picture of a toaster."

CPS adds a constraint: "Find the best path to the destination, BUT you must stay within the safe, high-probability zone of the forest." This ensures the computer doesn't get creative in a bad way; it stays realistic.

3. The "Do-Over" Button (Restart Strategy)

Sometimes, the computer starts with a bad guess (like starting a puzzle with the wrong corner piece). The old methods would get stuck.

CPS has a clever trick: The Restart. If the computer realizes it's going down a bad path, it doesn't panic. It gently "re-noises" the image (scrambles it slightly) and tries the step again, using the group wisdom to correct the mistake. It's like hitting "Undo" on a mistake before it becomes permanent.

Why Does This Matter?

• No Math Degree Needed: Many of the best previous methods required knowing the exact mathematical formula (gradient) of how the photo got blurry. But in the real world (like looking at black holes or fluid dynamics), we often don't have that formula. It's a "black box."
• Efficiency: CPS is much faster and uses less computing power because it doesn't waste the "rejected" candidates.
• Results: The paper shows that CPS can solve complex problems (like reconstructing images of black holes or fluid flows) just as well as the heavy, math-heavy methods, but without needing the complex formulas.

In a Nutshell

The Old Way: "Let's try 64 guesses, pick the best one, and throw the rest in the trash."
The New Way (CPS): "Let's listen to all 64 guesses, combine their advice to find the perfect direction, and make sure we stay on the safe path. If we mess up, we gently restart and try again."

It turns a wasteful guessing game into a collaborative, intelligent search, making it possible to solve difficult scientific mysteries even when we don't have all the mathematical tools.

Technical Summary

1. Problem Statement

The paper addresses inverse problems, where the goal is to recover an original signal $x$ from indirect, noisy observations $y$ (modeled as $y = H(x) + \eta$). This is inherently ill-posed ($m < d$), requiring priors to find a meaningful solution.

• Current Limitations: State-of-the-art methods using Diffusion Models (e.g., DPS) rely on gradient-based guidance. They require computing the gradient of the observation operator $\nabla H(x)$ or its pseudo-inverse.
• The Challenge: In many real-world scientific scenarios (e.g., fluid dynamics, black hole imaging), the forward model $H(\cdot)$ is a complex numerical simulation (black-box), highly nonlinear, or computationally expensive, making gradient computation impossible or impractical.
• Existing Gradient-Free Gaps: Current gradient-free alternatives (e.g., SCG, EnKG, DPG) suffer from significant inefficiencies:
  • SCG: Uses rejection sampling, discarding most candidate particles at each step, leading to wasted computation.
  • EnKG: Requires maintaining thousands of particles to approximate the posterior, causing high computational overhead.
  • DPG: Often pushes samples outside the data manifold.

2. Methodology: Constrained Particle Seeking (CPS)

The authors propose CPS, a novel gradient-free framework that reframes the inverse problem as a constrained optimization task. Instead of passively selecting the "best" particle from a set, CPS actively searches for an optimal particle by synthesizing information from all candidates.

Core Components:

1. Leveraging All Particles (Surrogate Modeling):

   • Unlike SCG, which discards all but one particle, CPS samples $n$ candidate particles $\{x_t^1, \dots, x_t^n\}$ from the reverse diffusion kernel $p(x_t|x_{t+1})$.
   • It fits a local linear surrogate model ($H(x) \approx Ax + b$) to the forward process using these samples. This is done via statistical linearization (estimating covariance between the samples and their forward projections).
   • This surrogate captures the local behavior of the black-box forward operator without needing explicit gradients.

2. Constrained Optimization:

   • CPS seeks a particle $x_t^*$ that minimizes the discrepancy between the observation $y$ and the surrogate prediction, subject to the constraint that the particle must lie within the high-probability density region of the unconditional prior.
   • Objective: $\min_{x_t} \|y - (Ax_t + b)\|^2$
   • Constraint: $x_t \in S^{d-1}(\mu_t, \sigma_t\sqrt{d})$, where the constraint ensures the solution stays on the "hypersphere" of the diffusion prior (preventing deviation from the data manifold).

3. Analytical Solution:

   • The constrained optimization is solved using Lagrange multipliers.
   • By exploiting the fact that the diffusion noise variance $\sigma_t$ is small in later steps, the authors derive an asymptotic closed-form solution:
     $$x_t^* \approx \mu_t + \sigma_t\sqrt{d} \frac{A^\top(y - \bar{H})}{\|A^\top(y - \bar{H})\|}$$
   • This formula effectively synthesizes all candidate information (via matrix $A$) to determine the optimal direction and step size on the hypersphere.

4. Restart Strategy:

   • To mitigate cumulative errors from sequential sampling and poor initial noise, CPS incorporates a Restart strategy.
   • If the solution quality is suboptimal, the process re-noises the current state and restarts the sampling from a previous timestep, iteratively correcting errors.

3. Key Contributions

• Gradient-Free Efficiency: CPS solves inverse problems using only forward passes of the observation model, making it applicable to black-box, non-differentiable, or expensive simulations.
• Active Particle Seeking: It shifts the paradigm from "passive selection" (rejection sampling) to "active seeking" by utilizing information from the entire ensemble of particles to construct a local surrogate.
• Theoretical Formulation: It provides a principled constrained optimization framework that balances observation fidelity with prior consistency (data manifold alignment).
• Robustness: The integration of the Restart strategy significantly improves stability and accuracy in the presence of cumulative sampling errors.

4. Experimental Results

The authors evaluated CPS on image inverse problems (inpainting, super-resolution, deblurring, JPEG restoration) and scientific inverse problems (Black Hole Imaging, Fluid Data Assimilation).

• Image Inverse Problems:

  • CPS significantly outperforms other gradient-free methods (SCG, DPG, EnKG) in PSNR, SSIM, and LPIPS metrics.
  • It achieves performance comparable to state-of-the-art gradient-based methods (like DPS and RED-diff), despite lacking access to gradients.
  • In JPEG restoration (a non-differentiable task), CPS succeeds where gradient-based methods fail or struggle.

• Scientific Inverse Problems:

  • Black Hole Imaging: CPS produces reconstructions visually and quantitatively superior to baselines, handling the highly nonlinear VLBI forward model effectively.
  • Fluid Data Assimilation: CPS recovers initial vorticity fields from sparse, noisy Navier-Stokes simulations with lower Relative $L_2$ error than EnKG and DPG.

• Efficiency:

  • CPS achieves high performance with a small number of particles (e.g., 64), whereas EnKG requires significantly more particles to function effectively.
  • It avoids the computational overhead of maintaining large ensembles or computing complex Jacobians.

5. Significance

• Bridging the Gap: CPS bridges the performance gap between gradient-based and gradient-free methods, offering a solution that is both efficient and robust for scenarios where gradients are unavailable.
• Scientific Applicability: By removing the dependency on differentiable forward models, CPS opens the door for applying diffusion priors to complex physical simulations (fluid dynamics, astrophysics, geophysics) that were previously inaccessible to diffusion-based inverse solvers.
• Paradigm Shift: The paper challenges the rejection sampling paradigm, demonstrating that "worse" particles contain valuable local information that can be aggregated to guide the search more effectively than simple selection.

In conclusion, Constrained Particle Seeking (CPS) represents a major advancement in diffusion-based inverse problems, providing a versatile, gradient-free tool that matches the performance of gradient-based methods while maintaining the flexibility to handle black-box forward models.