Beyond Best-of-k: Prior-Consistent DAPS for Ambiguous Nonlinear Inverse Problems

Measurement consistency is not enough. In Fourier phase retrieval, an image, its 180° rotation, its translations and its color-channel sign flips all produce the same Fourier magnitude. They have identical likelihood \(p(y\mid x)\), but very different posterior probability: an upside-down or color-flipped face fits the measurement perfectly, yet a prior concentrated on upright faces should never produce it. We call such reconstructions prior-inconsistent. They show up in classical phase-retrieval methods, deep-learning approaches, and diffusion-based solvers alike.

Best-of-k hides the failure. Results on ambiguous problems are usually reported as best-of-k: run the solver k times and keep the reconstruction closest to the ground truth. This uses information that does not exist at test time, and it throws away the k−1 other samples, including the prior-consistent ones.

Premature mode selection. Diffusion solvers apply the prior along a sampling trajectory, through a local denoising approximation centred on a single clean-image estimate. At high noise several prior modes remain plausible, but a trajectory can commit early to a prior mode that does not explain the measurement. Measurement guidance then pulls it to the nearest measurement-consistent solution, which the prior may not support, and later renoising rarely lets it escape. PC-DAPS instead delays this decision until the prior is informative enough to make it, and makes it with the prior, not with the ground truth.

PC-DAPS runs one DAPS annealing with a bag of \(K\) particles under the same measurement \(y\), each with its own initialization and noise. It changes two things in DAPS: the denoising distribution \(p(x_0\mid x_t)\) becomes a Gaussian mixture estimated from the bag, and a prior-based gate replaces trapped particles by copies of prior-consistent ones.

Gaussian mixture approximation of the denoiser posterior

DAPS approximates \(p(x_0\mid x_t)\) by a single isotropic Gaussian \(\mathcal N(\hat x_0(x_t), r_t^2 I)\) around one point estimate, so its Langevin step only looks for measurement-consistent points close to that estimate. PC-DAPS uses the bag instead. Before pruning it is a uniform mixture centred at each particle's clean estimate with the bag's diagonal variance; after pruning, each surviving cluster \(\mathcal C_m\) contributes one component:

\[ q(x_0\mid x_t)=\sum_{m=1}^{M}\pi_m\,\mathcal N\big(x_0;\mu_m,\Sigma_m\big), \] \[ \pi_m=\frac{|\mathcal C_m|}{K},\qquad \mu_m=\frac{1}{|\mathcal C_m|}\sum_{x^{(k)}\in\mathcal C_m}\tilde x_0^{(k)},\qquad \Sigma_m=\operatorname{diag}\!\big(\operatorname{Var}_{\mathcal C_m}\tilde x_0^{(k)}\big). \]

Its score is available in closed form and drops into the DAPS Langevin update: \(\nabla_{x_0}\log q(x_0\mid x_t)=\sum_m q_m(x_0)\,\Sigma_m^{-1}(\mu_m-x_0)\). The Langevin dynamics are now discouraged from leaving the prior modes found by the bag, rather than from leaving one point, and the covariances follow the geometry of the prior instead of a heuristic radius.

Prior-consistency pruning

When. For a linear operator, Jensen's inequality gives \(\bar r \ge r^\star\), where \(\bar r=\frac1K\sum_k\|\mathcal A(x^{(k)})-y\|^2\) and \(r^\star=\|\mathcal A(\bar x)-y\|^2\) for the bag mean \(\bar x\). When \(\bar r<r^\star\), the bag occupies several wells of the likelihood, which is when mode-level pruning becomes meaningful.

Which. The bag is clustered (a BIC-selected GMM on a PCA projection of downsampled clean estimates), and each cluster is scored by how well its particles survive re-noising and denoising at several noise scales \(\mathcal S\):

\[ s(\mathcal C_m)=-\frac{1}{|\mathcal C_m|\,|\mathcal S|}\sum_{x\in\mathcal C_m}\sum_{\sigma\in\mathcal S}\big\|x-D_\theta(x+\sigma\epsilon;\sigma)\big\|^2 . \]

A sample near the data manifold denoises back to itself at every scale. Clusters scoring below \((1+\gamma)\max_m s(\mathcal C_m)\) are pruned, and their particles are replaced by copies of survivors. The threshold \(\gamma\) is calibrated once per dataset from within-cluster score variation; no ground truth and no extra classifier are involved.

Full algorithm

1. Draw \(x_{t_{N}}^{(k)}\sim\mathcal N(0,\sigma_{t_N}^2 I)\) for \(k=1,\dots,K\); set pruned = false.
2. For each annealing level \(i=N,\dots,1\):
   1. Clean estimates \(\tilde x_0^{(k)}\gets D_\theta(x_{t_i}^{(k)};\sigma_{t_i})\).
   2. Build \(q(x_0\mid x_t)\): uniform over particles before pruning, one component per cluster after.
   3. If \(i>1\): run \(N_L\) Langevin steps on \(q(x_0\mid x_t)\,p(y\mid x_0)\) from \(\tilde x_0^{(k)}\).
   4. If not pruned, \(\bar r<r^\star\) and clustering finds \(M>1\) clusters: score clusters with \(s(\mathcal C_m)\), replace particles of low-scoring clusters by copies of survivors, set pruned = true.
   5. Re-noise \(x_{t_{i-1}}^{(k)}\gets\hat x_0^{(k)}+\sigma_{t_{i-1}}\epsilon^{(k)}\); at the last level, solve the probability-flow ODE to \(\sigma=0\) instead.
3. Return all \(K\) particles \(\{\hat x_0^{(k)}\}\).