Memory efficient training method improves accelerated mri reconstruction
Scalable Unrolled Chains for MRI Reconstruction
Computer Vision and Pattern Recognition
Summary
MRI scans can be made faster by using smart computer programs, but training these programs can use a lot of computer memory, limiting how complex they can get. The authors created a way to train these programs by focusing on a small part at a time, saving memory without losing accuracy. They use a shared network across parts and a clever method to pass error information backward without storing everything. Their approach works well on real MRI datasets and allows bigger and better models than before.
What this means in practice
- •For medical imaging engineers: Train larger deep learning models for faster MRI reconstruction using less GPU memory, enabling improved image quality at higher acceleration rates.
- •For hardware system designers: Design GPU systems optimized for deep learning methods that use approximate backward passes and cascade sampling to reduce memory consumption in medical imaging tasks.
Authors
Emiljo Mëhillaj, Ender Konukoglu
Abstract
Unrolled variational networks achieve state-of-the-art performance for accelerated MRI reconstruction, but standard end-to-end training becomes prohibitively memory intensive as the number and width of cascades increase. We propose a memory-efficient training framework that updates a randomly sampled cascade while retaining supervision from the final reconstruction loss. The error signal is propagated from the end of the chain to the sampled cascade using an activation-light backward pass that avoids storing full chain activations, thereby decoupling training memory from chain length. A time-conditioned network shared across cascades further keeps the parameter count independent of chain length. The backward pass is approximate only through frozen normalization statistics, and a trained exact gradient control shows the approximation leaves reconstruction quality unchanged. The framework remains trainable at chain lengths and cascade widths where standard end-to-end backpropagation exhausts GPU memory. On fastMRI single-coil knee and multi-coil brain reconstruction, our method consistently outperforms other local training approaches. On high-SNR knee slices, our method surpasses E2E-VarNet at $8\times$ and matches it at $4\times$, while on the complete brain validation set it surpasses E2E-VarNet at $8\times$ and achieves comparable performance at $4\times$, with fewer parameters, lower peak memory, and faster per-iteration training. Code is available at https://github.com/emehillaj/scalable-unrolled-mri.