PaperlayerView exact revision on arXiv

Source-linked AI summary

Scalable Unrolled Chains for MRI Reconstruction

Emiljo Mëhillaj, Ender Konukoglu

arXiv:2610.33695v1cs.CV

TL;DR

Sampled cascade training addresses the optimization and memory demands of scaling unrolled MRI reconstruction while retaining supervision from the final loss. At matched depth, the widest model surpasses E2E-VarNet at 8× and matches it at 4×, with fewer parameters.

  • Problem

    With independent cascade networks, sampled training updates each parameter block less often as chain length grows, making optimization increasingly difficult and potentially suboptimal.

    Source: §2.3 · p. 4

  • Method

    The framework updates one sampled cascade using the final reconstruction loss, propagates its error through an activation-light backward pass, and shares a time-conditioned network across cascades.

    Source: §1 · p. 2, §2.2 · p. 3

  • Results

    At matched depth, the widest model surpasses E2E-VarNet at 8× and matches it at 4×, while outperforming other memory-efficient baselines.

    Source: §4 · p. 8

  • Takeaways & Limitations

    The framework enables substantially longer and wider chains under fixed memory budgets without sacrificing reconstruction quality.

    Source: §4 · p. 8

    The strategy is demonstrated with a plain E2E-VarNet-style backbone on fastMRI knee and brain data; stronger backbones and other raw k-space datasets remain to be evaluated.

    Source: §4 · p. 8

Abstract

from arXiv · show

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.

1. Introduction

Unrolled MRI networks can improve reconstruction, but end-to-end training becomes memory-intensive as cascade width or chain depth grows. The paper trains a sampled cascade using the final reconstruction loss, with activation-light error propagation and shared, time-conditioned weights.

  • End-to-end training stores activations across cascades, so peak memory grows with network width and chain depth and can make scaling impractical under realistic GPU budgets.Source: §1 · p. 1
  • Existing memory-saving approaches trade activation memory for extra computation, architectural constraints, approximate gradients, or altered intermediate supervision.Source: §1 · p. 1
  • Intermediate-target methods such as GLEAM, k-space cold diffusion, and FLAT avoid or modify full-chain backpropagation, but prescribed targets may not align with final reconstruction optimization and may constrain learned solutions.Source: §1 · pp. 1–2
  • The proposed framework samples one cascade for parameter updates while propagating the final reconstruction-loss signal through the frozen suffix using an activation-light backward pass.Source: §1 · p. 2
  • The suffix backward pass is approximate only at normalization layers, where frozen forward statistics are used; the reported exact-gradient control shows closely matching error signals and indistinguishable training behavior.Source: §1 · p. 2
  • Sharing a refinement network across cascade positions and conditioning it on cascade index lets each iteration update reused parameters while retaining stage-dependent behavior.Source: §1 · p. 2
  • The evaluation compares reconstruction accuracy, peak training memory, and training time on fastMRI single-coil knee and multi-coil brain at 4× and 8× acceleration.Source: §1 · p. 2

2. Method

The method represents reconstruction as cascades that alternate data consistency with learned refinement, then trains one sampled cascade using an error signal propagated from the final loss. Activation-light suffix propagation and shared time-conditioned weights reduce the need to retain activations and separate parameters across the chain.

2.1. Problem setup and unrolled reconstruction

  • Each unrolled cascade combines data consistency with a learned refinement operator, while full-chain backpropagation requires storing or recomputing intermediate activations and memory grows with chain length.Source: §2.1 · p. 2

2.2. Memory-efficient backpropagation

  • At each iteration, a sampled cascade receives standard parameter backpropagation; the prefix is evaluated without activation storage and the frozen suffix propagates the final-loss gradient back to it.Source: §2.2 · p. 3
  • In suffix propagation, linear layers use their weights to compute input gradients without cached activations, while LeakyReLU needs only a binary sign mask for its exact input gradient.Source: §2.2 · p. 3
  • Frozen normalization statistics make normalization the only approximate input-gradient map in the suffix, storing O(BC) statistics rather than a full activation tensor.Source: §2.2 · pp. 3–4

2.3. Shared time-conditioned network for depthscalable training

  • Independent cascade networks make sampled training update each parameter block less often as depth increases; a shared U-Net conditioned on cascade index keeps parameter count constant with depth and updates reused parameters.Source: §2.3 · p. 4

3. Experiments

Experiments compare the method with local-training and end-to-end baselines on fastMRI knee and brain data, then examine training-scheme choices and computational scaling. The method improves on other memory-efficient approaches and, with wider shared networks, matches or exceeds E2E-VarNet on key comparisons while using fewer parameters and less memory.

3. Experiments

  • Evaluations used fastMRI single-coil knee and multi-coil brain validation splits at 4× and 8×; knee results were divided into high- and low-SNR-proxy cohorts.The complete evaluation covered 2,514 high-SNR and 4,621 low-SNR knee slices, and all 21,842 brain slices.Source: §3 · pp. 4–5

3.1. Training scheme analysis

  • At T=8 and C=32, the activation-light approximation left PSNR unchanged (29.57 →29.58 dB), while peak memory fell from 6.7 to 1.1 GB.Its suffix error signals had median cosine similarity 0.996 to exact signals; sampling reduced iteration time from 224 to 136 ms, while activation-light propagation reduced it from 136 to 134 ms.Source: §3.1 · p. 5, Table 1 · p. 5

3.2. Main comparison

  • 29.70 dB on high-SNR knee 8× at C=32 exceeded GLEAM (29.33 dB), FLAT-style (25.50 dB), and K-space Cold Diffusion (29.46 dB).The method had 8.1M trainable parameters; increasing shared-network width further improved reconstruction quality.Source: §3.2 · p. 6, Table 2 · p. 6
  • 30.33 dB at high-SNR knee 8× with C=96 surpassed E2E-VarNet’s 29.90 dB; at 4×, 32.44 dB matched its 32.41 dB.At matched C=32, the method used roughly 14× fewer parameters and 1.1 GB versus 11.0 GB peak training memory.Source: §3.2 · p. 6, Table 2 · p. 6
  • At brain C=96, the method attained the best value of every reported metric at both accelerations: 37.71 dB at 4× and 34.91 dB at 8×.At matched C=32, it used 13.6× fewer parameters than E2E-VarNet.Source: §3.2 · p. 7, Table 3 · p. 7

3.3. Computational efficiency

  • At T=50, peak training memory was 1.3 GB for the method versus 36.5 GB for E2E-VarNet, which exceeded 48 GB beyond T≈64.The shared model kept parameters at 8.1M across depths; at T=15 its iteration time was 231 ms versus 403 ms for E2E-VarNet.Source: §3.3 · pp. 7–8, Table 4 · p. 8

4. Conclusion

The framework scales unrolled MRI training while retaining supervision from the final reconstruction loss. On fastMRI knee and brain, it enables longer and wider chains under fixed memory budgets without sacrificing quality, with performance depending on data quality.

4. Conclusion

  • The framework updates one sampled cascade using an activation-light approximation of the final-loss error, while shared time-conditioned weights keep parameter count independent of chain length and peak memory nearly constant as chains grow.It targets scaling in both chain length and network width.Source: §4 · p. 8
  • On fastMRI knee and brain, longer and wider chains fit fixed memory budgets without sacrificing quality; at matched depth, the widest model beats E2E-VarNet at 8× and matches it at 4× with fewer parameters.Computation remains lower than full end-to-end training at matched depth, although it increases with chain length and width.Source: §4 · p. 8
  • The contribution is a training strategy demonstrated with a plain E2E-VarNet-style backbone on fastMRI knee and brain; evaluation on other raw k-space datasets and stronger backbones remains a next step.The authors suggest volumetric and other high-dimensional settings as possible extensions.Source: §4 · p. 8

5. Extended quantitative results

Across the extended evaluations, the method performs strongly against E2E-VarNet, with results varying by SNR regime and evaluation protocol. Ablations support shared time-embedded scaling, while cross-acceleration evaluation shows that every method loses quality under the shift.

5.1. Slice-level win fractions (single-coil knee)

  • At 8× on high-SNR knee slices, Ours (C=96) beats E2E-VarNet on 95% of slices and Ours (C=64) on 70%; on low-SNR slices, it wins against the reference on 50%.K-space Cold Diffusion alone beats the reference on low-SNR cohorts: 67% at 8× and 85% at 4×.Source: §5.1 · p. 11, Table 5 · p. 11

5.2. Volumewise, unstratified evaluation with confidence intervals

  • +0.62 dB at brain 8×, +0.32 dB at brain 4×, +0.21 dB at knee 8×, and +0.016 dB at knee 4× are the paired PSNR margins over E2E-VarNet.The respective 95% confidence intervals are [+0.61, +0.63], [+0.31, +0.33], [+0.17, +0.25], and [+0.013, +0.019] dB.Source: §5.2 · p. 11

5.3. Ablations and capacity scaling under volumewise aggregation

  • Quality increases monotonically with depth and width; independent per-cascade networks lose 0.10 dB from T=8 to T=15, while the shared network gains 0.10 dB.The ablation order is preserved under volumewise aggregation, though the depth ladder spans 0.30 dB rather than 0.46 dB per slice.Source: §5.3 · p. 12, Table 10 · p. 12

5.4. Cross-acceleration generalization

  • Every method loses quality when an 8×-trained knee model transfers to 4×; the widest shared model loses 0.32 dB, less than E2E-VarNet’s 0.45 dB and GLEAM’s 0.52 dB.K-space Cold Diffusion loses 0.17 dB, but from a lower absolute level.Source: §5.4 · p. 12, Table 11 · p. 13

6. Noise sensitivity and the SNR proxy

Pooled knee rankings depend on SNR, and the proxy's usefulness varies by anatomy and range. The knee proxy is informative across its range, whereas on brain it provides essentially no information about difficulty above 30 dB; low-SNR knee slices also occur centrally within volumes.

6.1. Noise sensitivity and pooled evaluation (single-coil knee)

  • On pooled knee validation, K-space Cold Diffusion leads in NMSE and SSIM at both accelerations and PSNR at 4×, while the widest shared model leads in PSNR at 8×.The pooled average is dominated by low-SNR slices; on clean references, the capacity-scaled model reconstructs anatomy with highest fidelity, while the sampler is scored higher against noisy references.Source: §6.1 · p. 13, Table 12 · p. 13

6.2. Sensitivity to the SNR proxy threshold

  • Method rankings remain stable as the knee SNR threshold sweeps from 25 to 45 dB in 1 dB steps, with 95% bootstrap intervals for cohort means.The dashed line marks the main-paper cutoff at 35 dB.Source: §6.2 · p. 13, Figure 3 · p. 15, Figure 4 · p. 15

6.3. Informative range of the SNR proxy

  • The proxy correlates with model PSNR at 0.75 overall on knee and remains informative above 30 dB, but its brain correlation above 30 dB is −0.06 overall.On brain, correlations range from 0.56 below 25 dB to −0.38 above 40 dB; across methods above 30 dB, they range from −0.29 to −0.04.Source: §6.3 · pp. 13–14, Table 13 · p. 15

6.4. SNR proxy relation to slice position

  • Of knee slices below 30 dB, 38% are central; 25.8% of central slices fall below 30 dB, and 51% of proxy variance lies between volumes.The proxy's correlation with distance from volume center is only −0.25, so low SNR is not confined to volume ends.Source: §6.4 · p. 15

7. Method analyses

The analyses characterize the compact backward’s memory savings and gradient fidelity, then compare checkpointing and peak-memory scaling across chain lengths. They also examine how nested checkpointing trades iteration time for memory.

7.1. Components of the compact backward

  • The exact convolution and LeakyReLU levels reduce peak memory from 6.6 to 3.1 GB; frozen-statistics normalization reduces it further to 1.1 GB.The three successive memory-saving increments are 1.7, 1.8, and 2.0 GB; the approximation is confined to normalization statistics.Source: §7.1 · p. 16, Table 15 · p. 16

7.2. Broader gradient fidelity distribution

  • At T=8, median gradient cosine similarity is ≈0.996 for both architectures, while 99% of pairs exceed 0.91 for the shared network and 0.82 for the per-cascade variant.The shared time-embedded model has systematically tighter suffix Jacobian fidelity.Source: §7.2 · p. 17, Figure 11 · p. 18
  • At T=25 and T=50, median cosine similarity is 0.998 at both depths, but the minimum falls to 0.42 at T=50 from 0.83 at T=8.At T=25 and T=50, 99% of pairs remain above 0.98 and 0.93, respectively.Source: §7.2 · p. 17, Figure 11 · p. 18

7.3. Iteration time of activation checkpointing

  • Checkpointing takes 1.3 to 1.4× as long as full backpropagation at measured depths and is 2.2 to 2.5× slower per step than our method.The absolute penalty over our method grows from 159 ms at T=8 to 1.09 s at T=50.Source: §7.3 · p. 17

7.4. Checkpointing granularity

  • Nested checkpointing saves ≈0.3 GB at each measured depth, but adds 25 to 31% more iteration time because block interiors are recomputed twice.It lowers peak memory from 1.26 to 0.96 GB at T=8 and from 3.77 to 3.47 GB at T=50.Source: §7.4 · p. 17
  • Checkpointing’s activation memory stays roughly constant with depth, but its weights and gradients grow linearly; sharing one network keeps that growing term at 0.06 GB.From T=8 to T=50, our peak rises by 0.18 GB versus 2.51 GB for nested checkpointing.Source: §7.4 · p. 17

7.5. Peak memory scaling curves

  • Across chain depths and widths, our peak memory stays nearly flat while E2E-VarNet and shared-network full backpropagation grow steeply and exceed 48 GB.The shared-network control tracks E2E-VarNet, locating the saving in activation handling rather than architecture.Source: §7.5 · p. 17, Figure 12 · p. 18

8. Implementation details

The implementation shares one cascade-conditioned U-Net across the chain, rather than assigning a separate network to each cascade. Memory and timing are profiled at batch size 1 on one A6000, while training uses accumulated microbatches.

8.1. U-Net implementation and cascade conditioning

  • A shared U-Net is conditioned on cascade index through sinusoidal embeddings and per-channel additive shifts, keeping the model’s parameter count independent of chain length.At T=8 and C=32, the shared chain has 8.08M parameters versus 62.05M for independent per-cascade U-Nets.Source: §8.1 · pp. 18–19

8.2. Measurement protocol

  • Profiling uses full fp32 at 640×372 resolution on one 48 GB NVIDIA RTX A6000; batch-1 measurements isolate activation footprint, while training accumulates an effective batch of 8.Peak allocated memory is measured across a full training step, excluding optimizer state.Source: §8.2 · p. 19

9. Additional qualitative examples

Additional examples compare reconstruction methods across single-coil knee and multi-coil brain validation slices. At 4×, methods lie closer together, with remaining differences concentrated in finer anatomical structures.

Single-coil knee and multi-coil brain examples

  • At 4×, methods lie closer together than at the better-conditioned 8× comparisons, and remaining differences concentrate in finer anatomical structures.The same slices are shown at both accelerations.Source: §9 · p. 19, Figure 13 · p. 20, Figure 14 · p. 22

Single-coil knee examples

  • Figure 13 presents single-coil knee validation slices in rows and methods plus the fully sampled reference in columns.The two accelerations use the same slices.Source: Figure 13 · pp. 20–21

Multi-coil brain examples

  • Figure 14 presents multi-coil brain validation slices across contrasts in rows, with methods arranged in columns as in Figure 13a.The two accelerations use the same slices.Source: Figure 14 · pp. 22–23

Related papers

Loading 2610.33695v1…