Pruning Unrolled Networks (PUN) at Initialization for MRI Reconstruction Improves Generalization

TL;DR

This work tackles the generalization gap in supervised MRI reconstruction under distribution shifts and proposes pruning unrolled networks at initialization (PUN-IT) to improve robustness. By applying PUN to MoDL, the authors learn a sparse subnetwork via a masked initialization and differentiable Bernoulli sampling (via Gumbel-softmax), achieving a final sparsity around 3% of parameters. Empirically, PUN-IT yields better generalization across different sampling patterns, anatomies, and even pathology-containing datasets (fastMRI+), while offering substantial computational savings compared with pruning during or after training. The findings suggest sparsity as a viable, practical strategy to enhance robustness in inverse problems like MRI without sacrificing, and sometimes improving, in-distribution performance.

Abstract

Deep learning methods are highly effective for many image reconstruction tasks. However, the performance of supervised learned models can degrade when applied to distinct experimental settings at test time or in the presence of distribution shifts. In this study, we demonstrate that pruning deep image reconstruction networks at training time can improve their robustness to distribution shifts. In particular, we consider unrolled reconstruction architectures for accelerated magnetic resonance imaging and introduce a method for pruning unrolled networks (PUN) at initialization. Our experiments demonstrate that when compared to traditional dense networks, PUN offers improved generalization across a variety of experimental settings and even slight performance gains on in-distribution data.

Figures (4)

• Figure 1: Comparison of reconstruction methods for 4x accelerated MRI (in-distribution). The model pruned at initialization (PUN-IT) outperforms the dense MoDL.
• Figure 2: Comparison of reconstruction methods trained for 4x accelerated MRI tested at 8x acceleration. PUN-IT generalizes better to this setting than PUN-AT or dense MoDL, both of which show significant aliasing artifacts.
• Figure 3: Box plots for reconstruction PSNR values (in dB) for different methods for the fastMRI brain test set (30 images) at 4x undersampling.
• Figure 4: Quantitative comparison of methods trained for 4x accelerated MRI on the fastMRI brain dataset tested on the fastMRI+ dataset at 4x and 8x acceleration and 4x acceleration with additional noise. PUN-IT offers improved performance in all settings compared to dense MoDL.