Ensemble guidance combines pretrained diffusion priors with black-box forward models to solve inverse problems without differentiating through the physical simulator. However, observation coordinates with large predictive spread or extreme residuals can dominate the ensemble correction, degrading reconstruction accuracy. We show that two simple modifications, weighting and clipping, substantially improve this correction. Our method, Robust Ensemble Guidance (REG), uses ensemble predictive spread to balance observation scales and adaptively clips standardized residuals to limit the influence of extreme discrepancies. Both operations reuse existing particles and forward predictions, requiring no additional denoiser or forward-model evaluations. Under a local linear Gaussian model, we derive conditions for reduced one-step estimation risk, bound the influence of individual observation coordinates, and characterize when these benefits persist with finite ensembles. Experiments on Navier-Stokes inversion, black-hole imaging, and acoustic full-waveform inversion demonstrate improved reconstruction over the underlying ensemble solver. In particular, REG increases black-hole reconstruction PSNR by 6.2-8.2 dB across three observation regimes and reduces Navier-Stokes reconstruction error by 26.4% in a matched-budget comparison. These findings highlight the importance of observation heterogeneity and residual influence in designing reliable generative solvers for scientific inverse problems.
Figures & tables
Figure 1: Overview of Robust Ensemble Guidance. Forward predictions from a denoised particle ensemble determine coordinate weights and a shared clipping threshold. The resulting bounded residuals are transported back to the particle space to form the next correction.
Figure 2: Controlled tests of weighting, clipping, and the conditional one-step risk criterion.
Subsampling
×2
×4
×8
Method / σy
0
1
2
0
1
2
0
1
2
EKI
0.577 -0.6pt
0.609 -0.6pt
0.673 -0.6pt
0.579 -0.6pt
0.669 -0.6pt
0.805 -0.6pt
0.852 -0.6pt
0.940 -0.6pt
1.116 -0.6pt
DPS-fGSG
1.687 -0.6pt
1.612 -0.6pt
1.454 -0.6pt
1.203 -0.6pt
1.209 -0.6pt
1.200 -0.6pt
1.246 -0.6pt
1.221 -0.6pt
1.260 -0.6pt
DPS-cGSG
2.203 -0.6pt
2.117 -0.6pt
1.746 -0.6pt
1.175 -0.6pt
1.133 -0.6pt
1.114 -0.6pt
1.186 -0.6pt
1.204 -0.6pt
1.218 -0.6pt
DPG
0.325 -0.6pt
0.408 -0.6pt
0.466 -0.6pt
0.322 -0.6pt
0.361 -0.6pt
0.454 -0.6pt
0.596 -0.6pt
0.591 -0.6pt
0.846 -0.6pt
SCG
0.908 -0.6pt
0.928 -0.6pt
0.966 -0.6pt
0.869 -0.6pt
0.926 -0.6pt
0.929 -0.6pt
1.260 -0.6pt
1.284 -0.6pt
1.347 -0.6pt
Table 1: Navier–Stokes relative ℓ2 error ( ↓ ). Cells show mean above (s.d.); bold indicates the lowest mean in each column.
J
EnKG
REG
Change
512
0.2752±0.1088
0.2690±0.0890
−2.3%
1024
0.2029±0.1250
0.1493±0.0463
−26.4%
2048
0.1156±0.0385
0.1000±0.0314
−13.4%
Table 2: Matched ablation ( ×2 , σy=0 ). Relative ℓ2 error: mean ± s.d.; bold: lowest mean per comparison. Changes are relative to EnKG at the same J .
Figure 3: Navier–Stokes reconstructions (top) and absolute errors (bottom) for the same test case, with shared color scales across methods.
Obs. time
3%
10%
100%
Method
PSNR ↑
Blur ↑
χcp2
χlogca2
PSNR ↑
Blur ↑
χcp2
χlogca2
PSNR ↑
Blur ↑
χcp2
χlogca2
SMILI
18.51 -0.6pt
23.08 -0.6pt
1.478 -0.6pt
4.348 -0.6pt
20.85 -0.6pt
25.24 -0.6pt
1.209 -0.6pt
21.788 -0.6pt
22.67 -0.6pt
27.79 -0.6pt
1.878 -0.6pt
17.612 -0.6pt
EHT-Imaging
21.72 -0.6pt
25.66 -0.6pt
1.507 -0.6pt
1.695 -0.6pt
22.67 -0.6pt
26.66 -0.6pt
1.166 -0.6pt
1.240 -0.6pt
24.28 -0.6pt
28.57 -0.6pt
1.251 -0.6pt
1.259 -0.6pt
DPS
24.20 -0.6pt
30.83 -0.6pt
8.024 -0.6pt
5.007 -0.6pt
24.36 -0.6pt
30.79 -0.6pt
13.052 -0.6pt
6.614 -0.6pt
25.86 -0.6pt
32.94 -0.6pt
8.759 -0.6pt
5.456 -0.6pt
LGD
22.51 -0.6pt
28.50 -0.6pt
15.825 -0.6pt
12.862 -0.6pt
22.08 -0.6pt
27.48 -0.6pt
10.775 -0.6pt
13.375 -0.6pt
21.22 -0.6pt
26.06 -0.6pt
13.239 -0.6pt
13.233 -0.6pt
RED-diff
20.74 -0.6pt
26.10 -0.6pt
6.713 -0.6pt
9.128 -0.6pt
22.53 -0.6pt
27.67 -0.6pt
2.488 -0.6pt
4.916 -0.6pt
23.77 -0.6pt
29.13 -0.6pt
1.853 -0.6pt
2.050 -0.6pt
Table 3: Black-hole imaging over 100 images. Cells show mean above (s.d.); bold and underline mark the best and second-best PSNR and blurred PSNR (dB).
Figure 4: Selected black-hole reconstructions of one case at 3% , 10% , and 100% observation time, with PSNR (dB). EnKG and REG share observations; other baselines retain native realizations.
Method
Relative ℓ2↓
PSNR (dB) ↑
SSIM ↑
Data misfit ↓
Adam
0.333 (0.086)
9.968 (2.083)
0.305 (0.120)
115.14 (52.10)
DPS
0.250 (0.154)
14.111 (6.820)
0.491 (0.161)
155.08 (92.17)
LGD
0.244 (0.024)
12.288 (0.889)
0.341 (0.047)
258.47 (26.40)
DiffPIR
0.204 (0.129)
16.113 (6.962)
0.554 (0.191)
88.53 (56.91)
DAPS
0.201 (0.103)
14.914 (4.184)
0.321 (0.067)
111.13 (71.33)
PnP-DM
0.259 (0.075)
11.983 (2.269)
0.431 (0.073)
308.84 (26.34)
Table 4: Acoustic FWI results. Mean (s.d.). Bold marks the best mean in each column.
Appendix figures & tables5 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 5: NS reconstructions and absolute errors at ×2 , σy=0 . EnKG and REG use J=2048 ; numbers are relative ℓ2 errors of the displayed reconstructions.
Figure 6: REG reconstructions of NS id0 across downsampling factors ×2/4/8 and noise levels σy=0/1/2 , with J=2048 . All panels share a vorticity scale.
Figure 7: Selected black-hole reconstructions: id0 at 100% , id3 at 3% , and id7 at 3% observation time. Numbers are aligned PSNR (dB); evaluation follows Appendix B.3 .
Figure 8: Additional black-hole targets and REG reconstructions at 3% and 100% observation time, with J=256 . Numbers are aligned PSNR (dB).
Figure 9: FWI reconstructions and absolute errors for id2, id5, and id7, with J=64 . Velocity and error panels use shared color scales; numbers are relative ℓ2 errors.
Diffusion models represent the state-of-the-art for solving inverse problems such as image restoration tasks. Diffusion-based inverse solvers incorporate a likelihood term to guide prior sampling, generating data consistent with the posterior distribution. However, due to the intractability of the likelihood, most methods rely on isotropic Gaussian approximations, which can push estimates off the data manifold and produce inconsistent, poor reconstructions. We propose Equivariance Regularized (EquiReg) diffusion, a general plug-in framework that improves posterior sampling by penalizing trajectories that deviate from the data manifold. EquiReg formalizes manifold-preferential equivariant functions that exhibit low equivariance error for on-manifold samples and high error for off-manifold ones, thereby guiding sampling toward symmetry-preserving regions of the solution space. We highlight that such functions naturally emerge when training non-equivariant models with augmentation or on data with symmetries. EquiReg's largest gains are under reduced sampling and measurement consistency steps, where many methods suffer severe quality degradation. By regularizing trajectories toward the manifold, EquiReg implicitly accelerates convergence and enables high-quality reconstructions. EquiReg consistently improves performance in linear and nonlinear image restoration tasks and solving partial differential equations. Our code is available at https://github.com/Anima-Lab/EquiReg
We introduce a new multivariate statistical problem that we refer to as the Ensemble-conditioned Inverse Problem (EIP). The aim of EIP is to invert for an ensemble that is distributed according to the pushforward of a prior under a forward process. In high energy physics (HEP), this is related to a widely known problem called unfolding, which aims to reconstruct the true physics distribution from observations that are distorted by detector effects. The EIP also arises in full waveform inversion (FWI) and inverse imaging with unknown priors. We propose non-iterative inference-time methods that construct posterior samplers based on a new class of conditional generative models, which we call ensemble inverse generative models. For the posterior modeling, these models additionally use the ensemble information contained in the observation set on top of single observations. Unlike existing methods, our proposed methods avoid explicit and iterative use of the forward model at inference time via training across several sets of truth-observation pairs that are consistent with the same forward model, but originate from a wide range of priors. We empirically demonstrate that this training procedure can implicitly encode the likelihood model, enabling direct posterior inference for unseen priors to some degree. We benchmark the proposed method on several synthetic and real datasets in inverse imaging, HEP, and FWI. Our code is available at https://github.com/ZhengyanHuan/EIP.
Zhengyan Huan, Camila Pazos, Martin Klassen +3
Department of Electrical and Computer Engineering, Tufts University, Medford, Massachusetts · The NSF AI Institute for Artificial Intelligence and Fundamental Interactions · Department of Physics and Astronomy, Tufts University, Medford, Massachusetts +1
Generative diffusion models can provide powerful prior probability models for inverse problems in imaging, but existing implementations suffer from two key limitations: (i) the prior density is represented implicitly, and (ii) they rely on likelihood approximations that introduce sampling biases. We address these challenges by introducing a new energy-based model trained for denoising with a covariance-based regularization term that enforces consistency across different measurement conditions. The trained model can compute normalized posterior densities for diverse linear inverse problems, without additional retraining or fine tuning. In addition to preserving the sampling capabilities of diffusion models, this enables previously unavailable capabilities: energy-guided adaptive sampling that adjusts schedules on-the-fly, unbiased Metropolis-Hastings correction steps, and blind estimation of the degradation operator via Bayes rule. We validate the method on multiple datasets (ImageNet, CelebA, AFHQ) and tasks (inpainting, deblurring), demonstrating competitive or superior performance to established baselines.
Nicolas Zilberstein, Santiago Segarra, Eero Simoncelli +1
Rice University, Houston, TX, USA · Flatiron Institute, New York, NY, USA · New York University, New York, NY, USA