pDAPS
A diffusion-prior sampler for accelerated MRI.
pDAPS uses a preconditioned solve for observed k-space and a diffusion prior for the missing frequencies. It produces multiple reconstructions so uncertainty can be estimated from their spread.
Incomplete measurements
For multi-coil MRI, the forward operator is \(A=PFS\): coil sensitivities \(S\), Fourier transform \(F\), and an undersampling mask \(P\). The observation model is \(y=Ax+\varepsilon\).
The target is the posterior \(p(x\mid y)\), not only a point estimate. Its samples should agree on measured k-space while varying where the acquisition leaves the image underdetermined.
\(A^HA\) inherits the mask, coil geometry, and a wide spectrum of curvature.
The likelihood constrains observed directions; the diffusion model supplies the image prior.
pDAPS update
At noise level \(\sigma_t\), the denoiser produces an anchor \(\hat{x}_0\). Data correction is a regularized least-squares problem around that anchor.
The first term enforces the measurement; the second controls how far the correction may move from the current prior estimate.
This is the curvature that makes an isotropic inner Langevin chain stiff.
The inverse is never formed. A truncated conjugate-gradient solve evaluates its action.
Denoise
Integrate a short probability-flow path from \(x_t\) to obtain a fresh clean-image anchor \(\hat{x}_0\).
Data correction
Solve \(B_tx_0=A^Hy+\nu_t\hat{x}_0\) with 30 truncated CG iterations.
Re-noise
Draw \(z\) and set \(x_{t'}=x_0+\sigma_{t'}z\). This outer step carries the sampler's stochasticity.
Conditioning
First-order inner Langevin updates require \(\mathcal{O}(\kappa(B_t))\) iterations in the standard quadratic analysis. CG reduces the dependence to \(\mathcal{O}(\sqrt{\kappa(B_t)})\).
Temperature
The implemented inner correction is deterministic. Injecting inner noise degraded validation performance; independent posterior samples instead come from initialization and repeated outer re-noising.
For an exact null projector, \(A\Pi_\perp=0\), so changing \(w\) moves only the underdetermined component. In multi-coil MRI the Fourier-mask split is approximate because sensitivity maps mix frequencies; this equality is therefore an organizing model, not an exact implementation guarantee.
A sampler needs more than PSNR
SSIM, PSNR, and HFEN score the posterior mean against a reference. They do not detect whether repeated runs return meaningful alternatives or the same point estimate.
\(D\to0\) indicates collapse, but a large \(D\) is useful only while samples remain data-consistent.
\(z_{\mathrm{rms}}\approx1\) is calibrated; values above one mean the cross-seed spread is too small for the observed error.
- Validation
- 16 slices · 3 seeds
- Test
- 48 slices · 8 seeds
- Fidelity
- SSIM ↑ · PSNR ↑ · HFEN ↓
- Sampling
- \(D\) as spread descriptor · \(z_{\mathrm{rms}}\to1\)
64-D posterior
A four-mode Gaussian-mixture posterior was tested while the condition number of the forward operator increased. Every method was tuned independently. Lower mode-mass error is better.
pDAPS had lower mode-mass error than pULA at all three settings.
| Sampler | κ=10 | 30 | 100 | 300 | 1000 |
|---|---|---|---|---|---|
| DPS | 1.41 | 1.44 | 1.44 | 1.44 | 1.39 |
| DAPS | 0.85 | 1.15 | 1.26 | 1.31 | 1.29 |
| pULA | 0.00 | 0.13 | 0.51 | 0.64 | 0.65 |
| pDAPS | 0.01 | 0.15 | 0.30 | 0.50 | 0.35 |
fastMRI knee and brain
DPS, DAPS, pULA, and pDAPS were evaluated on knee and brain slices at 4× and 8× acceleration. pDAPS was reported at \(w=0.5\) and \(w=1.0\), exposing the null-fill operating point rather than hiding it in one aggregate rank.
Observed
- Sharper reconstructions than pULA.
- No inner-step divergence on the brain experiments.
- Cross-sample spread concentrated in weakly measured regions.
- About half the DAPS runtime at equal NFE.
Limitations
- Not best on every distortion or calibration metric.
- Less diverse and less calibrated than pULA.
- Raw spread underestimated reconstruction error.
- No explicit acquisition-noise model.
Compute per reconstructed slice
NFE counts score-network evaluations. Seconds were measured on one shared GPU and are comparable within this experiment, not as universal runtimes.
| Method | NFE | Knee · seconds | Brain · seconds | ||
|---|---|---|---|---|---|
| R=4 | R=8 | R=4 | R=8 | ||
| DPS | 1000 | 16.6 | 14.6 | 13.9 | 13.9 |
| DAPS | 1000 | 16.4 | 16.3 | 20.2 | 18.7 |
| pULA | 400 / 640 | 4.0 | 3.9 | 6.3 | 6.3 |
| pDAPS | 1000 | 7.1 | 7.7 | 7.1 | 8.9 |
What the uncertainty means here
The benchmark uses \(\sigma_{\mathrm{noise}}=0\). Reported spread comes from random initialization and outer re-noising, and lies mostly in the unmeasured subspace. The \(z_{\mathrm{rms}}\) results compare samplers inside this setup; they are not clinical coverage guarantees under a realistic acquisition-noise model.
Result
pDAPS replaced DAPS's inner Langevin loop with a regularized CG solve. On the stiff synthetic problem it settles at 1.5–2× lower mode-mass error than the next-best sampler once κ ≥ 100. On MRI it takes SSIM on knee at both accelerations, halves DAPS's wall-clock at equal NFE, and removes the brain step-size divergence observed for the larger DAPS setting.
It is a tradeoff, not a sweep: sharper and less diverse than pULA, and not uniformly better than DPS or DAPS on every fidelity metric. The next technical problem is calibrating the magnitude of uncertainty in the underdetermined component.
Implementation
The sampler was implemented from scratch in PyTorch: the truncated-CG data correction, the measured/unmeasured split, and the decoupled outer schedule, built on InverseBench's interface and its pretrained diffusion priors.
Experimental scale
All results come from batch job fleets on a shared GPU cluster: two anatomies at two accelerations with 8-seed posterior sampling, independent validation sweeps for every baseline, and 2-D and 64-D synthetic studies.