Automated Tuning for Diffusion Inverse Problem Solvers without Generative Prior Retraining
This work addresses a critical bottleneck for researchers and practitioners using diffusion models in medical imaging, offering an incremental improvement in adaptive tuning without retraining.
The paper tackles the problem of tuning data fidelity weights for diffusion-based inverse problem solvers in accelerated MRI reconstruction, proposing Zero-shot Adaptive Diffusion Sampling (ZADS) which outperforms existing methods on the fastMRI knee dataset.
Diffusion/score-based models have recently emerged as powerful generative priors for solving inverse problems, including accelerated MRI reconstruction. While their flexibility allows decoupling the measurement model from the learned prior, their performance heavily depends on carefully tuned data fidelity weights, especially under fast sampling schedules with few denoising steps. Existing approaches often rely on heuristics or fixed weights, which fail to generalize across varying measurement conditions and irregular timestep schedules. In this work, we propose Zero-shot Adaptive Diffusion Sampling (ZADS), a test-time optimization method that adaptively tunes fidelity weights across arbitrary noise schedules without requiring retraining of the diffusion prior. ZADS treats the denoising process as a fixed unrolled sampler and optimizes fidelity weights in a self-supervised manner using only undersampled measurements. Experiments on the fastMRI knee dataset demonstrate that ZADS consistently outperforms both traditional compressed sensing and recent diffusion-based methods, showcasing its ability to deliver high-fidelity reconstructions across varying noise schedules and acquisition settings.