Scaling Test-Time Compute of Diffusion-Based Posterior Samplers for Physical Inverse Problems
Abstract
For physical inverse problems, pretrained diffusion models commonly provide priors for posterior samplers whose inference cost is often fixed. Motivated by the recent success of scaling test-time compute in generative AI, particularly for large language models (LLMs) and diffusion models, we investigate whether the same recipe can be applied to diffusion-based solvers of inverse problems in physical sciences. Our methodology here is based on the Approximation-Free Diffusion Posterior Sampler (AFDPS) developed in existing work (Chen et al., 2025), where multiple trajectories are sampled in a parallel way under the sequential Monte Carlo (SMC) framework. We tested our method on three problems from InverseBench (Zheng et al., 2025): linear inverse scattering, full waveform inversion (FWI), and the Navier-Stokes equation. The computational cost at test time is operator-dependent: scattering uses cached linear algebra, FWI uses forward and adjoint wave solves, and Navier-Stokes uses differentiation through a fixed-step simulator. Numerical experiments further show that our method can improve the reconstruction quality by increasing the number of particles during inference, especially for ill-posed problems.