Scalable Unrolled Chains for MRI Reconstruction
Organizations: Computer Vision Lab, ETH Zurich, Zurich, Switzerland · LOOP Zurich - Medical Research Center, Zurich, Switzerland
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 and matches it at , while on the complete brain validation set it surpasses E2E-VarNet at and achieves comparable performance at , with fewer parameters, lower peak memory, and faster per-iteration training. Code is available at https://github.com/emehillaj/scalable-unrolled-mri.
Figures & tables
| Training scheme | PSNR | Mem (GB) | Time (ms) |
| E2E-VarNet (distinct, full BP) | 29.54 | 5.9 | 215 |
| Shared-TE, full BP | 29.68 0.02 | 6.7 | 224 |
| Shared-TE, sampled, exact suffix | 29.57 0.02 | 6.7 | 136 |
| Shared-TE, sampled, activation-light (ours) | 29.58 0.02 | 1.1 | 134 |
| Params (M) | |||||
| Method ( ) | High-SNR proxy 35 | Low-SNR proxy 35 | High-SNR proxy 35 | Low-SNR proxy 35 | |
| E2E-VarNet (32) | 0.0094/32.41/0.790 (0.0068/2.71/0.081) | 0.0574/26.83/0.538 (0.0448/3.08/0.132) | 0.0160/29.90/0.716 (0.0091/2.34/0.093) | 0.0691/25.79/0.450 (0.0470/2.88/0.138) | 116.4 |
| E2E + ckpt. (64) | 0.0093 /32.42/ 0.791 (0.0066/2.72/0.081) | 0.0578/26.82/0.538 (0.0459/3.09/0.132) | 0.0154/30.07/0.716 (0.0090/2.39/0.095) | 0.0699/25.74/0.446 (0.0474/2.88/0.138) | 465.4 |
| E2E + ckpt. (96) | 0.0093 / 32.46 / 0.791 (0.0066/2.73/0.081) | 0.0572/26.84/0.539 (0.0447/3.08/0.132) | 0.0153/30.13/0.717 (0.0090/2.40/0.096) | 0.0699/25.75/0.444 (0.0480/2.89/0.138) | 1047.1 |
| GLEAM (32) | 0.0103/31.94/0.777 (0.0068/2.60/0.082) | 0.0597/26.64/0.521 (0.0466/3.05/0.137) | 0.0181/29.33/0.700 (0.0094/2.27/0.094) | 0.0730/25.51/0.435 (0.0484/2.84/0.138) | 116.4 |
| FLAT-style (32) | 0.0208/28.62/0.688 (0.0092/2.16/0.073) | 0.0928/24.88/0.487 (0.0838/3.02/0.109) | 0.0424/25.50/0.576 (0.0165/2.18/0.085) | 0.1320/23.11/0.361 (0.1073/2.82/0.108) | 116.4 |
| Method ( ) | Params (M) | ||
| E2E-VarNet (32) | 0.0058/37.42/0.915 (0.0067/2.96/0.096) | 0.0106/34.38/0.890 (0.0075/2.64/0.102) | 116.8 |
| GLEAM (32) | 0.0081/35.76/0.895 (0.0081/2.72/0.098) | 0.0180/32.11/0.863 (0.0178/2.71/0.102) | 116.8 |
| FLAT-style (32) † | 0.0514/27.62/0.750 (0.0358/2.81/0.080) | 0.1100/24.07/0.631 (0.0565/2.58/0.084) | 116.8 |
| K-space Cold Diff. (64) | 0.0173/32.57/0.851 (0.0242/3.05/0.102) | 0.0284/29.90/0.799 (0.0285/2.85/0.109) | 32.3 |
| Ours (32) | 0.0063/37.03/0.912 (0.0072/2.93/0.099) | 0.0112/34.16/0.887 (0.0081/2.64/0.103) | 8.6 |
| Ours (64) | 0.0056 /37.64/0.917 (0.0067/3.01/0.097) | 0.0105/34.49/0.891 (0.0087/2.72/0.104) | 32.8 |
| Method | / | NMSE | PSNR | SSIM | Mem (GB) | Time (ms) | Params (M) |
| Depth scaling ( , sweep ) | |||||||
| E2E-VarNet | 8 | 0.0173 | 29.54 | 0.708 | 5.9 | 215 | 62.05 |
| E2E-VarNet | 15 | 0.0160 | 29.90 | 0.716 | 11.0 | 403 | 116.35 |
| E2E + ckpt. | 8 | 0.0173 | 29.54 | 0.708 | 1.3 | 293 | 62.05 |
| E2E + ckpt. | 15 | 0.0160 | 29.90 | 0.716 | 1.7 | 551 | 116.35 |
| E2E + ckpt. | 25 | 0.0156 | 30.02 | 0.715 | 2.3 | 901 | 193.9 |
| Method | high-SNR | low-SNR | high-SNR | low-SNR |
| E2E-VarNet (32) | 32.41 | 26.83 | 29.90 | 25.79 |
| GLEAM (32) | 31.94 0% wins | 26.64 0% wins | 29.33 4% wins | 25.51 1% wins |
| FLAT-style (32) | 28.62 0% wins | 24.88 0% wins | 25.50 0% wins | 23.11 0% wins |
| K-space Cold Diff. (64) | 31.91 22% wins | 27.21 85% wins | 29.46 25% wins | 25.89 67% wins |
| Ours (32) | 32.32 4% wins | 26.78 4% wins | 29.70 21% wins | 25.53 2% wins |
| Method | PSNR | SSIM | NMSE |
| E2E-VarNet (32) | 36.75 | 0.9350 | 0.0093 |
| GLEAM (32) | 34.31 0% wins | 0.9122 0% wins | 0.0160 0% wins |
| FLAT-style (32) | 25.91 0% wins | 0.6852 0% wins | 0.1074 0% wins |
| K-space Cold Diff. (64) | 30.78 4% wins | 0.8313 2% wins | 0.0351 2% wins |
| Ours (32) | 36.44 2% wins | 0.9320 0% wins | 0.0099 2% wins |
| Ours (64) | 36.92 81% wins | 0.9365 84% wins | 0.0089 81% wins |
| Method | PSNR | SSIM | NMSE |
| E2E-VarNet (32) | 40.02 | 0.9546 | 0.0044 |
| GLEAM (32) | 38.24 0% wins | 0.9397 0% wins | 0.0065 0% wins |
| FLAT-style (32) | 29.47 0% wins | 0.7974 0% wins | 0.0477 0% wins |
| K-space Cold Diff. (64) | 32.92 2% wins | 0.8731 2% wins | 0.0221 2% wins |
| Ours (32) | 39.50 0% wins | 0.9520 0% wins | 0.0049 0% wins |
| Ours (64) | 40.26 99% wins | 0.9557 99% wins | 0.0042 99% wins |
| Method | PSNR | SSIM | NMSE |
| E2E-VarNet (32) | 31.09 | 0.6770 | 0.0409 |
| E2E + ckpt. (64) | 31.17 | 0.6758 | 0.0408 |
| E2E + ckpt. (96) | 31.18 | 0.6759 | 0.0408 |
| GLEAM (32) | 30.68 0% wins | 0.6651 0% wins | 0.0440 0% wins |
| FLAT-style (32) | 27.75 0% wins | 0.5578 0% wins | 0.0771 0% wins |
| K-space Cold Diff. (64) | 29.45 21% wins | 0.6653 39% wins | 0.0498 8% wins |
| Method | PSNR | SSIM | NMSE |
| E2E-VarNet (32) | 32.75 | 0.7387 | 0.0310 |
| E2E + ckpt. (64) | 32.78 | 0.7396 | 0.0309 |
| E2E + ckpt. (96) | 32.78 | 0.7397 | 0.0309 |
| GLEAM (32) | 32.45 0% wins | 0.7283 0% wins | 0.0324 0% wins |
| FLAT-style (32) | 30.09 0% wins | 0.6632 0% wins | 0.0490 0% wins |
| K-space Cold Diff. (64) | 30.51 14% wins | 0.7340 40% wins | 0.0367 15% wins |
| Configuration | PSNR | SSIM | NMSE |
| Depth ( ) | |||
| Ours, | 30.77 | 0.6629 | 0.0439 |
| Ours, | 30.87 | 0.6656 | 0.0430 |
| Ours, | 30.96 | 0.6677 | 0.0427 |
| Ours, | 31.07 | 0.6698 | 0.0419 |
| Width ( ) | |||
| Method ( ) | Native | ||
| E2E-VarNet (32) | 32.41 | 31.96 | |
| GLEAM (32) | 31.94 | 31.42 | |
| FLAT-style (32) | 28.62 | 28.21 | |
| K-space Cold Diff. (64) | 31.91 | 31.74 | |
| Ours (32) | 32.32 | 31.55 | |
| Ours (64) | 32.28 | 31.91 |
| Method ( ) | Knee | Knee | Params (M) |
| E2E-VarNet (32) | 0.0405/28.79/0.627 (0.0429/3.98/0.168) | 0.0504/27.24/0.543 (0.0458/3.34/0.178) | 116.4 |
| GLEAM (32) | 0.0423/28.51/0.612 (0.0445/3.85/0.172) | 0.0537/26.86/0.528 (0.0473/3.22/0.177) | 116.4 |
| FLAT-style (32) | 0.0674/26.20/0.558 (0.0759/3.28/0.137) | 0.1004/23.95/0.437 (0.0969/2.85/0.144) | 116.4 |
| K-space Cold Diff. (64) | 0.0362 / 28.87 / 0.673 (0.0338/3.59/0.116) | 0.0489 /27.14/ 0.567 (0.0401/3.13/0.145) | 32.3 |
| Ours (32) | 0.0409/28.74/0.623 (0.0434/3.96/0.169) | 0.0536/27.00/0.529 (0.0497/3.35/0.178) | 8.1 |
| Ours (64) | 0.0420/28.67/0.622 (0.0462/3.97/0.169) | 0.0523/27.23/0.534 (0.0505/3.47/0.182) | 32.3 |
| Proxy band | Brain | Knee |
| dB | ||
| – | ||
| – | ||
| – | ||
| overall |
| Central | Edge | |
| Slices | ||
| Median proxy (dB) | ||
| Below dB | ||
| Below dB | ||
| Minimum (dB) |
| Suffix backward component | Peak (GB) | Time (ms) | Gradients |
| L0 standard autograd | 6.61 | 154 | exact |
| L1 + conv/linear adjoint | 4.93 | 134 | exact |
| L2 + LeakyReLU sign masks | 3.10 | 137 | exact |
| L3 + frozen-statistics norm | 1.14 | 134 | approx. |
| Scheme | Peak (GB) | weights grads | activations | Time (ms) |
| full backprop. | ||||
| ckpt. per cascade | ||||
| ckpt. nested | ||||
| Ours | ||||
Accelerated MRI Reconstruction
Accelerated MRI reconstruction recovers diagnostic images from undersampled k-space data to reduce scan time while preserving anatomical and quantitative information.
See all 85 papers