Diffusion models have recently gained popularity for accelerated MRI reconstruction due to their high sample quality. They can effectively serve as rich data priors while incorporating the forward model flexibly at inference time, and they have been shown to be more robust than unrolled methods under distribution shifts. However, diffusion models require careful tuning of inference hyperparameters on a validation set and are still sensitive to distribution shifts during testing. To address these challenges, we introduce SURE-based MRI Reconstruction with Diffusion models (SMRD), a method that performs test-time hyperparameter tuning to enhance robustness during testing. SMRD uses Stein's Unbiased Risk Estimator (SURE) to estimate the mean squared error of the reconstruction during testing. SURE is then used to automatically tune the inference hyperparameters and to set an early stopping criterion without the need for validation tuning. To the best of our knowledge, SMRD is the first to incorporate SURE into the sampling stage of diffusion models for automatic hyperparameter selection. SMRD outperforms diffusion model baselines on various measurement noise levels, acceleration factors, and anatomies, achieving a PSNR improvement of up to 6 dB under measurement noise. The code is publicly available at https://github.com/batuozt/SMRD .
翻译:扩散模型因生成高质量样本而近年来在加速MRI重建中备受关注。该类模型可有效作为丰富的数据先验,同时在推理阶段灵活融入前向模型,且已被证明在分布偏移下比展开方法更稳健。然而,扩散模型需在验证集上精细调整推理超参数,且在测试期间仍对分布偏移敏感。为应对这些挑战,我们提出基于SURE的扩散模型MRI重建方法(SMRD),该方法通过测试时超参数调整增强鲁棒性。SMRD利用斯坦因无偏风险估计(SURE)在测试过程中估计重建均方误差,并据此自动调整推理超参数及设定早停准则,无需验证集调参。据我们所知,SMRD首次将SURE融入扩散模型采样阶段以实现超参数自动选择。在多种测量噪声水平、加速因子和解剖结构条件下,SMRD均优于扩散模型基线,在测量噪声下峰值信噪比(PSNR)提升达6 dB。代码已开源:https://github.com/batuozt/SMRD。