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

AI summary

Overview

Research area: Generative modeling for inverse problems — specifically diffusion-model-based posterior sampling using only forward (black-box) evaluations of the measurement operator.

Technical level: Advanced. The paper builds on stochastic differential equations, diffusion posterior sampling, path-integral stochastic control, statistical linearization, and constrained optimization on hyperspheres.

Scope (one sentence): The paper proposes Constrained Particle Seeking (CPS), a gradient-free method that reformulates diffusion inverse-problem solving as a constrained optimization over all candidate particles, and evaluates it on four image restoration tasks plus black hole imaging and fluid data assimilation.

What This Paper Is About

Many diffusion-based solvers for inverse problems need the gradient of the forward measurement operator H(·) in order to guide sampling toward observations. In many real settings that operator is a black box — a costly numerical simulation, or something non-differentiable such as JPEG compression — so gradients are unavailable. CPS aims to solve such inverse problems using only forward passes through H, while staying as accurate as gradient-based methods and clearly beating existing gradient-free alternatives.

Key Contributions

1. A reformulation of the inverse problem as constrained optimization. Instead of passively picking the best of several sampled candidates, CPS solves arg min Φ(x̂_{0|t}) subject to x_t ~ p(x_t | x_{t+1}), so it can seek a particle beyond the sampled set (Eq. 9).

2. A local linear surrogate of the black-box forward operator. Using statistical linearization, CPS samples n particles from the reverse kernel, evaluates H(x̂_{0|t}^i) for each, and fits H ≈ A x_t + b with a closed-form least-squares solution (Eqs. 11–13).

3. A tractable constraint derived from the diffusion prior. Because the reverse kernel is an isotropic Gaussian, the mass of the χ^d norm distribution concentrates near σ_t √d, so CPS constrains x_t to the hypersphere S^{d-1}(μ_t, σ_t √d). The Lagrange-multiplier system has an asymptotic closed-form solution: x_t* ≈ μ_t + σ_t √d · A^T(y − H̄) / ||A^T(y − H̄)|| (Eq. 16).

4. A Restart strategy for error correction. At each timestep the sample is re-noised and returned to the previous timestep, then re-denoised with CPS, iteratively correcting cumulative sampling error. The paper reports this improves overall performance, stability, and robustness (quantitative Restart ablation is in Appendix C).

5. A public code release: https://github.com/deng-ai-lab/CPS.

Main Findings

• Diagnostic evidence that all particles carry usable information. In a black hole imaging study, not only the top-1 particle but also the top-2 and top-3 particles guided denoising better than unconditioned sampling. Even the five worst-performing particles (bottom-1 to bottom-5), when their noise signs were reversed, still steered denoising in the correct direction.

• Image inverse problems on FFHQ (100 images, 256×256, noise σ_y = 0.05). CPS is the best gradient-free method on all four tasks: 95% inpainting (PSNR 24.90, SSIM 0.794, LPIPS 0.161), 4× super-resolution (27.70, 0.834, 0.115), Gaussian deblurring (27.21, 0.822, 0.128), and JPEG restoration at QF=5 (23.23, 0.784, 0.223). It beats the other gradient-free methods (DPG, EnKG) and competes with gradient-based methods; on super-resolution the best gradient-based result is DAPS at 28.53 PSNR, and on deblurring DAPS at 28.70 PSNR. SCG fails on image problems (e.g., PSNR 6.15 for inpainting). ΠGDM is reported as the only method able to do JPEG restoration via pseudo-inverse construction, but it performs poorly in noisy scenes (18.25 PSNR, 0.593 SSIM, 0.304 LPIPS).

• Black hole imaging (100 resized 64×64 images; VLBI closure phases, log closure amplitudes, total flux; observation time ratios 100%, 10%, 3%). CPS achieves the best PSNR and BPSNR at every ratio — 100%: 25.74 / 31.84; 10%: 25.34 / 31.54; 3%: 25.03 / 31.81. SCG is a close second (25.45 / 30.56; 24.97 / 31.45; 24.65 / 31.28). DPG performs poorly on this highly nonlinear problem (e.g., 12.16 / 13.71 at 3%).

• Fluid data assimilation (Navier-Stokes, T = 1, recovery of 2D initial vorticity, 10 fields of 128×128, noise σ_y = 2.0). Measured by relative L2 error, CPS is best at 2× (0.295), 4× (0.469), and 8× (0.684) downscaling. EnKG is second best (0.320, 0.528, 0.821); SCG fails badly (1.087, 1.095, 1.123).

• Particle efficiency. Tested with 8, 16, 32, 64, and 128 particles, CPS remains stable and competitive and performs well with only 8 particles. EnKG struggles with small particle sets and needs many particles; DPG struggles on highly ill-posed black hole imaging; SCG fails on image problems.

• Local versus global linearization. Computing MAE of the Jensen gap for H(E[x_0 | ·]) on 4× super-resolution with FFHQ, EnKG shows significant errors in early generation stages while CPS consistently maintains small errors — attributed to CPS linearizing locally around p(x_t | x_{t+1}) rather than over the whole global space.

Methodology in Plain English

The starting point is the observation that previous gradient-free approaches such as SCG sample several candidate "particles" at each denoising step and throw away all but the one whose one-step clean estimate best matches the measurement. CPS keeps the whole ensemble. At each timestep it (1) samples n candidates from the same reverse kernel SCG uses, (2) evaluates the black-box forward operator on each candidate's one-step denoised estimate x̂_{0|t}, and (3) fits a linear model relating the candidates x_t^i to the forward-evaluated values H(x̂_{0|t}^i). That fitted model — a weighted combination of all candidates — stands in for the true forward operator, so no derivative of H is ever needed.

The second ingredient is a constraint. Diffusion models expect latents whose norm matches σ_t √d, so rather than optimizing freely, CPS restricts its search to the hypersphere of radius σ_t √d centered at the reverse kernel mean μ_t. Solving this sphere-constrained least-squares problem via Lagrange multipliers yields a closed-form approximate optimum that pushes the particle in the direction of the residual y − H̄, scaled to the sphere.

Finally, the Restart loop re-noises the current estimate back to the previous timestep and reruns CPS from there, N_r times per timestep, to gradually undo accumulated sampling error — especially useful when the initial noise is poor. Every gradient-free method compared in the experiments is given a comparable number of forward passes.

Why This Matters

Impact on research: CPS shows that gradient-free diffusion inversion need not be a crude fallback. By treating all candidate particles as information rather than waste and adding a prior-derived constraint, it narrows the long-standing accuracy gap between gradient-free and gradient-based posterior sampling. It also offers an explicit, local alternative to the global linearization used by EnKG, backed by a Jensen-gap diagnostic.

Real-world applications:

• Radio astronomy / black hole imaging: VLBI forward models are nonlinear and gradient-hostile; CPS operates with only forward evaluations of closure phases, log closure amplitudes, and total flux.
• Geophysics and fluid dynamics: assimilating sparse, noisy observations through Navier-Stokes-style solvers that take thousands of discrete steps, where automatic differentiation is impractical.
• Computational imaging and medicine: recovery from indirect, noisy measurements where the measurement process is proprietary or available only as a simulator.
• Legacy image restoration: JPEG decoding (QF=5) is non-differentiable, so it is a natural fit for a gradient-free solver; CPS reaches 23.23 PSNR there while gradient-based ΠGDM reaches 18.25.

Industry relevance: The method is plug-and-play — no retraining or fine-tuning of the diffusion prior is required — and it only needs an executable forward model. That makes it usable when a physical simulator or legacy codec is the only available forward description, which is common in scientific and engineering pipelines.

Future Directions

The paper's conclusion does not list an explicit future-work agenda; the following are open questions raised by the work.

1. Extending beyond the linear surrogate. The current surrogate is explicitly a best-fit linear model of H around the reverse kernel; how CPS behaves for measurement operators that are strongly nonlinear over the local particle spread is not characterized beyond the reported tasks.
2. Choosing the number of restarts and particles in practice. Results are reported for n ∈ {8, 16, 32, 64, 128} and for a restart count N_r, but no general rule for selecting N_r or n as a function of task difficulty or compute budget is given.
3. Broader benchmark coverage. Evaluation covers FFHQ images and two scientific tasks (black hole imaging, fluid data assimilation); behavior on other modalities, imaging datasets, or measurement noise levels beyond σ_y = 0.05 (images) and σ_y = 2.0 (fluid) is not reported.
4. Closing the remaining gap to gradient-based methods. On FFHQ super-resolution and deblurring, DAPS still leads on PSNR (28.53 and 28.70 versus CPS at 27.70 and 27.21), leaving the question of where the remaining accuracy difference comes from.

Target Audience

Researchers and practitioners in generative modeling and inverse problems who need solvers that work when the forward operator is a black box, non-differentiable, or expensive to differentiate — including computational imaging, astronomy, geophysics, and scientific machine learning groups. It is also relevant to engineers deploying diffusion priors over proprietary simulators. Because of the stochastic-control and constrained-optimization content, readers will benefit from prior familiarity with diffusion models and diffusion posterior sampling.

Authors’ abstract

Diffusion models have gained prominence as powerful generative tools for solving inverse problems due to their ability to model complex data distributions. However, existing methods typically rely on complete knowledge of the forward observation process to compute gradients for guided sampling, limiting their applicability in scenarios where such information is unavailable. In this work, we introduce \textbf{\emph{Constrained Particle Seeking (CPS)}}, a novel gradient-free approach that leverages all candidate particle information to actively search for the optimal particle while incorporating constraints aligned with high-density regions of the unconditional prior. Unlike previous methods that passively select promising candidates, CPS reformulates the inverse problem as a constrained optimization task, enabling more flexible and efficient particle seeking. We demonstrate that CPS can effectively solve both image and scientific inverse problems, achieving results comparable to gradient-based methods while significantly outperforming gradient-free alternatives. Code is available at https://github.com/deng-ai-lab/CPS.