[Paper Review] Adaptive Diffusion Priors for Accelerated MRI Reconstruction

Alper Güngör, Salman Ul Hassan Dar+5 more|arXiv (Cornell University)|Jul 12, 2022

Radiomics and Machine Learning in Medical Imaging21 citations

TL;DR

AdaDiff introduces an adaptive diffusion prior for MRI reconstruction, enabling rapid diffusion with large steps and subject-specific prior adaptation to improve reconstruction under domain shifts.

ABSTRACT

Deep MRI reconstruction is commonly performed with conditional models that de-alias undersampled acquisitions to recover images consistent with fully-sampled data. Since conditional models are trained with knowledge of the imaging operator, they can show poor generalization across variable operators. Unconditional models instead learn generative image priors decoupled from the operator to improve reliability against domain shifts related to the imaging operator. Recent diffusion models are particularly promising given their high sample fidelity. Nevertheless, inference with a static image prior can perform suboptimally. Here we propose the first adaptive diffusion prior for MRI reconstruction, AdaDiff, to improve performance and reliability against domain shifts. AdaDiff leverages an efficient diffusion prior trained via adversarial mapping over large reverse diffusion steps. A two-phase reconstruction is executed following training: a rapid-diffusion phase that produces an initial reconstruction with the trained prior, and an adaptation phase that further refines the result by updating the prior to minimize data-consistency loss. Demonstrations on multi-contrast brain MRI clearly indicate that AdaDiff outperforms competing conditional and unconditional methods under domain shifts, and achieves superior or on par within-domain performance.

Motivation & Objective

• Motivate robust MRI reconstruction that generalizes across varying imaging operators.
• Develop an unconditional diffusion-prior framework decoupled from the operator to improve domain shift resilience.
• Create an efficient diffusion process using large reverse steps via adversarial mapping for faster inference.
• Enable subject-specific prior adaptation during inference to refine reconstructions.
• Demonstrate cross-domain and within-domain performance across multi-contrast brain MRI.

Proposed method

• Propose AdaDiff, an unconditional diffusion prior trained with adversarial mapping for fast, large-step diffusion.
• Use a two-phase reconstruction: rapid diffusion phase for an initial reconstruction, followed by a prior-adaptation phase that minimizes data-consistency loss.
• Implement rapid diffusion with large step size k and a generator G_thetaG and discriminator D_thetaD to model reverse diffusion without assuming normality.
• Define a data-consistency loss comparing synthesized k-space data to acquired measurements and update the generator to minimize this loss (Eq. 45).
• During training, employ a diffusion objective with an adversarial mapper to implicitly capture q(x_t|x_{t+k}) and derive closed-form-like expressions (Eq. 14-18).
• Reconstruct by interleaving data-consistency projections with reverse-diffusion steps in rapid diffusion, then refine via prior adaptation.

Experimental results

Research questions

• RQ1 Can an unconditional, diffusion-based prior provide reliable MRI reconstructions under varying imaging operators?
• RQ2 Does adapting the diffusion prior at test time improve reconstruction quality for domain shifts without retraining?
• RQ3 Can rapid diffusion with large steps, guided by adversarial mapping, match or exceed performance of traditional diffusion with many steps?
• RQ4 How does AdaDiff perform on multi-contrast brain MRI under within-domain and cross-domain conditions?
• RQ5 What is the comparative benefit of rapid diffusion plus prior adaptation versus conditional models in robustness to domain shifts?

Key findings

• AdaDiff outperforms competing conditional and unconditional methods under domain shifts.
• AdaDiff achieves superior or on-par performance within-domain with matched operator and distribution.
• Two-phase approach (rapid diffusion plus prior adaptation) yields improved reconstruction by aligning the prior to test data.
• Adversarial diffusion priors enable efficient sampling with large reverse steps while maintaining reconstruction quality.
• Prior adaptation refines the diffusion prior by minimizing data-consistency loss on subject data.