Source-linked AI summary
Adaptive Diffusion Priors for Accelerated MRI Reconstruction
Alper Güngör, Salman UH Dar, Şaban Öztürk, Yilmaz Korkmaz, Gokberk Elmas, Muzaffer Özbey, Tolga Çukur
TL;DR
Operator shifts can weaken conditional MRI reconstruction, while static diffusion priors may also become suboptimal under test-distribution shifts. AdaDiff uses rapid adversarial diffusion followed by inference-time prior adaptation, achieving competitive within-domain and superior cross-domain reconstructions.
Problem
Conditional models can generalize poorly across imaging operators, and static diffusion priors can be suboptimal when test image distributions shift.
Method
AdaDiff uses an adversarial mapper for few-step diffusion sampling, then adapts the trained prior during inference using data-consistency optimization.
Results
AdaDiff performs competitively within domain and achieves superior reconstructions in cross-domain tasks against state-of-the-art baselines.
Takeaways & Limitations
Adaptive diffusion priors can improve MRI reconstruction reliability across shifts in imaging operators and MR image distributions.
Takeaways & Limitations
Prior adaptation increases inference time over single-forward-pass conditional models because it requires iterative sampling and backward passes.
Abstract
from arXiv · showhide
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.
1. Introduction
Deep MRI reconstruction faces operator-related domain shifts that can undermine conditional models. AdaDiff addresses this with a fast, adaptive diffusion prior and shows strong within-domain and cross-domain performance.
- Undersampling shortens MRI exams but creates an ill-posed inverse problem requiring reconstruction from incomplete k-space acquisitions.
- Conditional models learn mappings from undersampled to fully sampled acquisitions, but operator changes can compromise generalization and require retraining.
- Unconditional models decouple generative image priors from the imaging operator, improving reliability against operator-related domain shifts.
- AdaDiff uses an adversarial mapper to generate images through few large reverse-diffusion steps, reducing sampling time relative to vanilla diffusion.
- Its two-phase inference first produces an initial reconstruction through rapid diffusion, then adapts the prior to minimize data-consistency loss.
- Across multi-contrast brain MRI, AdaDiff is generally superior or comparable within domain and outperforms competing methods under cross-domain shifts.
Contributions
AdaDiff combines rapid adversarial diffusion sampling with inference-time prior adaptation. The method is presented as a diffusion-based approach for improving MRI reconstruction under domain shifts.
- AdaDiff is presented as the first diffusion-based prior adaptation method for accelerated MRI reconstruction.
- The method uses a rapid diffusion process with an adversarial mapper for efficient sampling from the diffusion prior.
- AdaDiff adapts the trained diffusion prior during inference to improve reconstruction performance and reliability against domain shifts.
2. Related Work
MRI reconstruction has moved from operator-dependent conditional models toward unconditional generative priors, including diffusion models. AdaDiff extends this direction by adapting its diffusion prior to the test subject during inference.
- Unconditional generative models learn image priors separately from the imaging operator to improve generalization in MRI reconstruction.
- Diffusion models gradually transform Gaussian noise into image samples and offer enhanced representational diversity over GANs.
- Diffusion-based MRI reconstruction has typically reported high image quality.
- Vanilla diffusion methods require hundreds of reverse steps, while separate initialization methods add a second reconstruction model and static priors may become suboptimal under image-distribution shifts.
- AdaDiff performs subject-specific adaptation during inference to increase prior conformity to the test-data distribution.
3. Theory
AdaDiff combines a rapid adversarial diffusion prior with subject-specific prior adaptation to reconstruct accelerated MRI while addressing operator and image-distribution shifts.
- 3.1. MRI Reconstruction: Accelerated MRI reconstructs images x from undersampled k-space data y using the imaging operator A, but the inverse problem is ill-posed.A captures undersampling, coil sensitivities, and Fourier transformation.
- 3.1. MRI Reconstruction: Unconditional priors remain decoupled from the imaging operator, whereas conditional priors are tied to operator-specific reconstruction mappings.This decoupling is intended to improve reliability when the operator changes.
- 3.2. Diffusion Models: Diffusion models generate images by progressively adding Gaussian noise and then removing it through learned reverse mappings, typically requiring hundreds of steps.The reverse process uses network projections to recover clean samples from noisy ones.
- 3.3.1. Training of the Prior: AdaDiff uses large reverse-diffusion steps and an adversarial mapper because the usual Gaussian assumption for q(x_t|x_{t+k}) breaks down at large step sizes.The generator synthesizes reverse samples while the discriminator distinguishes them from samples from the forward denoising distribution.
- 3.3.2. Reconstruction with the Prior: Reconstruction begins with rapid diffusion alternating reverse-diffusion and data-consistency projections, then adapts the generator per subject by minimizing measured-versus-synthesized k-space discrepancy.The adapted generator output is used as the final reconstruction.
4. Methods
AdaDiff is evaluated as a unified multi-contrast MRI reconstruction framework using datasets, competing methods, and within- and cross-domain tests. Its rapid adversarial diffusion configuration uses few reverse steps, followed by subject-specific prior adaptation.
- Datasets: IXI provides single-coil T1-, T2-, and PD-weighted magnitude images, while fastMRI provides multi-coil complex T1-, T2-, and FLAIR-weighted data.FastMRI data are compressed to 5 or 10 virtual coils while preserving over 95% or 98% of original energy, respectively.
- Datasets: Data are retrospectively undersampled with variable-density random masks at acceleration rates R=4x, 8x, or 12x, with coil sensitivities estimated from central calibration data.Volumetric acquisitions are reconstructed as individual 2D cross-sections across the fully-sampled dimension.
- AdaDiff: AdaDiff uses an adversarial mapper with a residual encoder-decoder generator and discriminator to implement rapid reverse diffusion.The generator uses six encoder stages, residual blocks, downsampling, and attention in the last two stages.
- Competing methods: Comparisons include LORAKS, conditional rGAN and MoDL, adaptive GANprior, and static diffusion models DDPM and DiffRecon.Within-domain results are summarized in Tables 1 and 2, while cross-domain results appear in Tables 3 and 4.
- AdaDiff: Training uses k=125 and T/k=8 diffusion steps, followed by 8 rapid-diffusion iterations and 1000 prior-adaptation iterations during inference.Prior adaptation uses a 10^-3 learning rate; the implementation runs on Nvidia RTX 3090 GPUs via PyTorch.
- Training setup: A single unified model is trained on pooled multi-contrast data without being informed of each sample's contrast.Pooling uses T1, T2, and PD in IXI, and T1, T2, and FLAIR in fastMRI.
5. Results
AdaDiff performs strongly within domain and under shifts in the imaging operator or MR image distribution. Ablations further show that prior adaptation, rapid diffusion, and adversarial mapping each contribute to reconstruction quality.
- Within-Domain Reconstruction: AdaDiff generally outperforms competing methods across within-domain contrasts and acceleration rates, with a few method-specific exceptions.Exceptions include MoDL on T1 and FLAIR, and DiffRecon on T2 at R=4x and PSNR at R=8x.
- Within-Domain Reconstruction: 6.8dB PSNR and 15.8% SSIM are AdaDiff’s average gains over LORAKS within domain.
- Cross-Domain Reconstruction: AdaDiff achieves the highest performance across most acceleration-rate shifts, including testing at R=8x and R=12x after conditional training at R=4x.Exceptions include comparable or higher isolated SSIM or PSNR results for DiffRecon and MoDL.
- Cross-Domain Reconstruction: 5.0dB PSNR and 10.0% SSIM are AdaDiff’s average gains over the traditional method when sampling trajectory or coil count shifts.
- Cross-Domain Reconstruction: Under fastMRI-to-IXI image-distribution shifts, AdaDiff generally achieves the highest performance across contrasts and acceleration rates.DiffRecon yields modestly higher SSIM in an exception reported for this setting.
- Ablation Studies: Removing latent variables for adaptive normalization reduces performance across tasks, except for similar T1 SSIM.
- Ablation Studies: AdaDiff outperforms its static, untrained, and non-adversarial variants across reconstruction tasks, while prior adaptation provides the largest ablation improvement.Against the static variant, average improvements are 363.54 FID, 55.43 LPIPS, 10.72dB PSNR, and 20.09% SSIM.
- Ablation Studies: J=1000 offers a trade-off between reconstruction time and image quality, while benefits beyond T/k = 8 are marginal and cross-validated Ne values are near-optimal.
6. Discussion
AdaDiff improves reliability across several tested operator and image-distribution shifts while retaining strong within-domain performance. Its inference optimization increases runtime and memory demands, and broader anatomy, operator, contrast, data, and architecture settings remain open areas.
- Discussion: AdaDiff improves reliability against imaging-operator shifts relative to conditional models and against MR image-distribution shifts relative to all competing methods.
- Discussion: Prior adaptation improves AdaDiff over static diffusion priors in both within-domain and cross-domain scenarios.
- Future Directions: Broader validation is needed across body parts, Cartesian versus non-Cartesian trajectories, coil arrays, architectures, normalization layers, and complex-valued operations.
- Practical Considerations: AdaDiff’s inference optimization causes longer inference and greater memory load than single-forward-pass conditional models because it requires backward passes and stored gradients.
- Practical Considerations: Regular diffusion priors can impose substantial computational burden and leave less data-consistency improvement for subsequent adaptation because they use many projected reverse-diffusion steps.
- Future Directions: The study pools multiple contrasts into unified models; separate contrast-specific training might improve performance but increases computational burden.
- Future Directions: Joint multi-contrast priors or contrast side information could exploit structural correlations or maintain contrast specificity.
- Future Directions: Cycle-consistent learning could reduce dependence on fully-sampled acquisitions, while adapting all diffusion steps would substantially increase memory load and inference time.
7. Conclusion
AdaDiff introduces diffusion-prior adaptation for MRI reconstruction, using efficient few-step generation and test-subject refinement. It performs competitively within domain and better across domain shifts.
- An adversarial mapper enables efficient image generation through few reverse-diffusion steps.
- Inference first produces an initial reconstruction through the trained prior, then adapts the prior to the test subject.
- AdaDiff performs competitively within domain and achieves superior reconstructions across domain shifts compared with state-of-the-art baselines.