Original: Greedy Learning for Large-Scale Neural MRI Reconstruction
Revised: Greedy Learning for Large-Scale Neural MRI Reconstruction

Original: Author(s)

Original: Abstract
Revised: Abstract

Original: Model-based deep learning approaches have recently shown state-of-the-art performance for accelerated MRI reconstruction.
Revised: Model-based deep learning approaches have recently shown state-of-the-art performance for accelerated MRI reconstruction.

Original: These methods unroll iterative proximalgradient descent by alternating between data-consistency and a neural-networkbased proximal operation.
Revised: These methods unroll iterative proximal gradient descent by alternating between data-consistency and a neural-network based proximal operation.

Original: However, they demand several unrolled iterations withsufﬁciently expressive proximals for high resolution and multi-dimensional imaging (e.g., 3D MRI).
Revised: However, they demand several unrolled iterations with sufﬁciently expressive proximals for high resolution and multi-dimensional imaging (e.g., 3D MRI).

Original: This impedes traditional training via backpropagation due toprohibitively intensive memory and compute needed to calculate gradients andstore intermediate activations per layer.
Revised: This impedes traditional training via backpropagation due to prohibitively intensive memory and compute needed to calculate gradients and store intermediate activations per layer.

Original: To address this challenge, we advocate analternative training method by greedily relaxing the objective.
Revised: To address this challenge, we advocate an alternative training method by greedily relaxing the objective.

Original: We split the end-toend network into decoupled network modules, and optimize each network moduleseparately, thereby avoiding the need to compute costly end-to-end gradients.
Revised: We split the end-toend network into decoupled network modules, and optimize each network module separately, thereby avoiding the need to compute costly end-to-end gradients.

Original: Weempirically demonstrate that the proposed greedy learning method requires 6x lessmemory with no additional computations, while generalizing slightly better thanbackpropagation.
Revised: We empirically demonstrate that the proposed greedy learning method requires 6x less memory with no additional computations, while generalizing slightly better than backpropagation.

Original:  Introduction
Revised:  Introduction

Original: Magnetic resonance imaging (MRI) is a widely used medical imaging modality that provides highresolution information of the soft tissue anatomy, but is limited by long scan times.
Revised: Magnetic resonance imaging (MRI) is a widely used medical imaging modality that provides high resolution information of the soft tissue anatomy, but is limited by long scan times.

Original: Deep learningmethods have gained popularity for accelerated MRI, improving image quality of reconstructionsfrom undersampled measurements over parallel imaging and compressive sensing (CS) techniques[1, 2].
Revised: Deep learning methods have gained popularity for accelerated MRI, improving image quality of reconstructions from undersampled measurements over parallel imaging and compressive sensing (CS) techniques [1, 2].

Original: A typical approach is model-based deep learning, where an iterative algorithm is unrolled fora ﬁxed number of iterations.
Revised: A typical approach is model-based deep learning, where an iterative algorithm is unrolled for a ﬁxed number of iterations.

Original: The iterative algorithm alternates between enforcing data-consistencywith the measurement model and a neural proximal operation [3, 4].
Revised: The iterative algorithm alternates between enforcing data-consistency with the measurement model and a neural proximal operation [3, 4].

Original: Model-based networks aretrained with backpropagation where the parameters of the end-to-end network are jointly optimized.
Revised: Model-based networks are trained with backpropagation where the parameters of the end-to-end network are jointly optimized.

Original: The graphics processing unit (GPU) memory requirement of training model-based networks withbackpropagation increases linearly with the measurement count and the number of layers in the neuralproximal operation, limiting the number of unrolled iterations of the model-based network which inturn limits expressivity [5].
Revised: The graphics processing unit (GPU) memory requirement of training model-based networks with backpropagation increases linearly with the measurement count and the number of layers in the neural proximal operation, limiting the number of unrolled iterations of the model-based network which in turn limits expressivity [5].

Original: Therefore, improving the memory footprint of training could enable theuse of such networks in high-dimensional large-scale imaging applications such as Dynamic Contrast
Revised: Therefore, improving the memory footprint of training could enable the use of such networks in high-dimensional large-scale imaging applications such as Dynamic Contrast

Original: Enhanced (DCE) imaging, or blood ﬂow imaging (4D-ﬂow).
Revised: Enhanced (DCE) imaging, or blood ﬂow imaging (4D-ﬂow).

Original: Related Works .
Revised: Related Works .

Original: Several works have aimed to tackle the memory bottleneck of training modelbased deep networks.
Revised: Several works have aimed to tackle the memory bottleneck of training modelbased deep networks.

Original: One line of work uses invertible networks where intermediate activations arecalculated by reversing each layer, hence removing the need to store intermediate activations usingbackpropagation [6, 7, 8].
Revised: One line of work uses invertible networks where intermediate activations are calculated by reversing each layer, hence removing the need to store intermediate activations using backpropagation [6, 7, 8].

Original: [9] applies stochastic approximations to the data-consistency layers ofmodel-based networks to reduce memory complexity with increasing number of measurements.
Revised: [9] applies stochastic approximations to the data-consistency layers of model-based networks to reduce memory complexity with increasing number of measurements.

Original: Itmust be noted that these approaches make modiﬁcations to the network architecture whereas ourmethod proposes a greedy optimization method that is compatible with any model-based network.
Revised: It must be noted that these approaches make modiﬁcations to the network architecture whereas our method proposes a greedy optimization method that is compatible with any model-based network.

Original: Thus, our method can be synergistically combined with any of these approaches to further reducememory.
Revised: Thus, our method can be synergistically combined with any of these approaches to further reduce memory.

Original: Recent works in computer vision propose alternatives to backpropagation for convolutional neuralnetworks (CNNs) in the context of image classiﬁcation.
Revised: Recent works in computer vision propose alternatives to backpropagation for convolutional neural networks (CNNs) in the context of image classiﬁcation [10, 11, 12].

Original: One approach is layer-wise training, wheresingle hidden-layer subnetworks are trained sequentially until convergence while the previous part ofthe network remains frozen [10].
Revised: One approach is layer-wise training, where single hidden-layer subnetworks are trained sequentially until convergence while the previous part of the network remains frozen [11].

Original: Another alternative is to use a greedy learning approach, wherea network is split into several smaller network modules, and gradient updates for each module aredecoupled, which was shown to achieve a similar generalization performance as backpropagation [11].
Revised: Another alternative is to use a greedy learning approach, where a network is split into several smaller network modules, and gradient updates for each module are decoupled, which was shown to achieve a similar generalization performance as backpropagation [12].

Original: Although memory efﬁciency is not the primary focus of those works, greedy learning can improve thememory footprint of training by performing local gradient updates as opposed to end-to-end gradientupdates in backpropagation.
Revised: Although memory efﬁciency is not the primary focus of those works, greedy learning can improve the memory footprint of training by performing local gradient updates as opposed to end-to-end gradient updates in backpropagation.

Original: In this work, we propose a greedy training objective for learning model-based networks.
Revised: In this work, we propose a greedy training objective for learning model-based networks.

Original: We splitthe end-to-end network into decoupled network modules, and perform gradient updates on eachmodule independently.
Revised: We split the end-to-end network into decoupled network modules, and perform gradient updates on each module independently.

Original: Since each module has a signiﬁcantly lower memory footprint, the memoryrequirement of the overall network decreases as well.
Revised: Since each module has a signiﬁcantly lower memory footprint, the memory requirement of the overall network decreases as well.

Original: Hence, our approach reduces memory whenthe number of measurements or number of layers is large.
Revised: Hence, our approach reduces memory when the number of measurements or number of layers is large.

Original: Contributions .
Revised: Contributions .

Original: Our contributions can be summarized as follows:
Revised: Our contributions can be summarized as follows:

Original: • We propose greedy learning for MRI reconstruction, which requires 6 times less GPU memory during training while the compute time remainsalmost the same.
Revised: • We propose greedy learning for MRI reconstruction, which requires 6x less GPU memory during training while the compute time remains the same.

Original: • We empirically show greedy learning can improve generalizationperformance compared to
Revised: • We demonstrate that greedy learning can improve generalization

Original:  backpropagation for MRI reconstruction with model-based networks.
Revised: performance compared to backpropagation for MRI reconstruction with model-based networks.

Original:  Background
Revised:  Background

Original: The forward model of accelerated MRI with parallel imaging and compressive sensing (CS) [1, 2]can be modeled as
Revised: The forward model of accelerated MRI with parallel imaging and compressive sensing (CS) [1, 2] can be modeled as

Original: [Equation]
Revised: [Equation]

Original: where y is the observed measurements in Fourier space (k-space), x is the real image we would liketo reconstruct, S are coil sensitivity maps associated with the receiver coils used in parallel imaging,
Revised: where y is the observed measurements in Fourier space (k-space), x is the real image we would like to reconstruct, S are coil sensitivity maps associated with the receiver coils used in parallel imaging,

Original: F is the Fourier transform matrix, U is the undersampling mask, and (cid:15) is additive noise.
Revised: F is the Fourier transform matrix, U is the undersampling mask, and (cid:15) is additive noise.

Original: Then, the
Revised: Then, the

Original: MRI reconstruction problem can be formulated as

Original: [Equation] Input: 
Revised: Input: Samples or mini-

Original: [Equation]
Revised: [Equation]

Original: or

Original: [Equation]
Revised: [Equation]

Original: Initialize: Parameters { θ m } m ≤ M for k ∈ 1 , ..., N b do for m ∈ 1 , ..., M do 
Revised: Initialize: Parameters { θ m } m ≤ M for k ∈ 1 , ..., N b do for m ∈ 1 , ..., M do x ( m

Original: [Equation]
Revised: [Equation]

Original: Compute ∇ θ m L
Revised: k ← f θ m ( DC ( x ( m

Original: [Equation]
Revised: [Equation]

Original:  x 
Revised: ) k )) Compute ∇ θ m L ( x ( m ) k , x k ) θ m ← Update Parameters ( θ m ) . end for end for MRI reconstruction problem can be formulated as

Original: [Equation]
Revised: [Equation]

Original: θ m ← Update Parameters ( θ m ) .

Original: end for end for where A = UFS denotes the forward model, R is a regularization function and λ is the regularizationstrength.
Revised: where A = UFS denotes the forward model, R is a regularization function and λ is the regularization strength.

Original: This optimization problem can be solved in an alternating manner using proximal gradientdescent [3][Equation]
Revised: This optimization problem can be solved in an alternating manner using proximal gradient descent [3]

Original: [Equation]
Revised: [Equation]

Original: where i is the iteration number, t is the step size, A H is the Hermitian transpose of the forward modeloperator A. DC denotes data-consistency, and Prox R is the proximal operator.
Revised: where i is the iteration number, t is the step size, A H is the Hermitian transpose of the forward model operator A. DC denotes data-consistency, and Prox R is the proximal operator.

Original: In model-basednetworks, the proximal operator is learned with a CNN with trainable parameters, and the iterativeoptimization algorithm described above is unrolled for a ﬁxed number of iterations.
Revised: In model-based networks, the proximal operator is learned with a CNN with trainable parameters, and the iterative optimization algorithm described above is unrolled for a ﬁxed number of iterations.

Original: Then, (4) can berewritten as
Revised: Then, (4) can be rewritten as

Original: [Equation]
Revised: [Equation]

Original: where f θ i is the network, and θ i its parameters at i th iteration.
Revised: where f θ i is the network, and θ i its parameters at i th iteration.

Original: Such a models parameters can bejointly optimized using backpropagation
Revised: Such a models parameters can be jointly optimized using backpropagation

Original: [Equation]
Revised: [Equation]

Original: where f θ 1 ,..,θ N is the end-to-end network, y k is the k th undersampled image, x k is the k th fullysampled reference, L is the loss function, and N is the number of unrolled iterations.
Revised: where f θ 1 ,..,θ N is the end-to-end network, y k is the k th undersampled image, x k is the k th fullysampled reference, L is the loss function, and N is the number of unrolled iterations.

Original: In our framework,we use the unrolled neural network described above as the baseline model for comparison.
Revised: In our framework, we use the unrolled neural network described above as the baseline model for comparison.

Original:  Greedy Learning for MRI Reconstruction
Revised:  Greedy Learning for MRI Reconstruction

Original: Here, we describe our method for optimizing model-based deep networks.
Revised: Here, we describe our method for optimizing model-based deep networks.

Original: We adapt our methodfrom the greedy learning aproach in [11] for training CNNs, where we extend this to model-basednetworks.
Revised: We adapt our method from the greedy learning aproach in [12] for training CNNs, where we extend this to model-based networks.

Original: We consider a model-based network with N number of unrolled iterations, where thenetwork is split into M modules, such that NM is the number of data-consistency steps and proximalsteps in each module.
Revised: We consider a model-based network with N number of unrolled iterations, where the network is split into M modules, such that NM is the number of data-consistency steps and proximal steps in each module.

Original: For simplicity we outline the case where N = M .
Revised: For simplicity we outline the case where N = M .

Original: In particular, let x (0) =
Revised: In particular, let x (0) =

Original: y bethe input, x the fully sampled reference, x ( m ) =
Revised: y be the input, x the fully sampled reference, x ( m ) =

Original: f θ m ( DC ( x ( m − 1) ))
Revised: f θ m ( DC ( x ( m − 1) ))

Original: the output of the m th networkmodule, and N b is the number of mini-batches.
Revised: the output of the m th network module, and N b is the number of mini-batches.

Original: For the m th network module, let the cost function be
Revised: For the m th network module, let the cost function be

Original: L ( f θ m ( DC ( x ( m − 1) )) , x ) , where we select L to be L 1 loss.
Revised: L ( f θ m ( DC ( x ( m − 1) )) , x ) , where we select L to be L 1 loss.

Original: Then, the parameters of module m canbe optimized using the greedy training objective:
Revised: Then, the parameters of module m can be optimized using the greedy training objective:

Original: [Equation]
Revised: [Equation]

Original: (7) where x ( m − 1) =
Revised: where x ( m − 1) =

Original: f θ ∗ m − 1 ( DC ( x ( m − 2) )) for m > 1 , and θ ∗ m − 1 are the optimal weights for the previousmodule.
Revised: f θ ∗ m − 1 ( DC ( x ( m − 2) )) for m > 1 , and θ ∗ m − 1 are the optimal weights for the previous module.

Original: Our method to optimize the greedy training objective in Eq.
Revised: Our method to optimize the greedy training objective in Eq.

Original: 7 for M modules is outlined in
Revised: 7 for M modules is outlined in

Original: Alg 1.
Revised: Alg 1.

Original: When a mini-batch is sampled during training, the ﬁrst module computes the forward passon the mini-batch of images, performs a gradient update based on its own local loss, and passes theresult of the forward pass to the next module.
Revised: When a mini-batch is sampled during training, the ﬁrst module computes the forward pass [Table] on the mini-batch of images, performs a gradient update based on its own local loss, and passes the result of the forward pass to the next module.

Original: This procedure is repeated until the ﬁnal module whichproduces the ﬁnal prediction of the network, after which a new mini-batch is sampled.
Revised: This procedure is repeated until the ﬁnal module which produces the ﬁnal prediction of the network, after which a new mini-batch is sampled.

Original: Therefore,individual updates of each set of parameters are performed independently across different modules,which drastically reduces the memory requirement.
Revised: Therefore, individual updates of each set of parameters are performed independently across different modules, which drastically reduces the memory requirement.

Original: In our approach, we update each module usingbackpropagation.
Revised: In our approach, we update each module using backpropagation.

Original: The model-based network trained with the proposed greedy approach is illustratedin Fig.[Table]
Revised: The model-based network trained with the proposed greedy approach is illustrated in Fig.

Original:  Results
Revised:  Results

Original: [Figure] We evaluated the effectiveness of our approach for two acceleration factors, R = { 12 , 16 } .
Revised: We evaluated the effectiveness of our approach for two acceleration factors, R = { 12 , 16 } .

Original: Weconsidered greedy learning with two different network splits M = { 4 , 8 } .
Revised: We considered greedy learning with two different network splits M = { 4 , 8 } .

Original: We reported structuralsimilarity (SSIM), normalized root-mean-square-error (nRMSE), and peak signal-to-noise ratio(PSNR) calculated on magnitude images.
Revised: We reported structural similarity (SSIM), normalized root-mean-square-error (nRMSE), and peak signal-to-noise ratio (PSNR) calculated on magnitude images.

Original: Additional experimental details are provided in Appendix
Revised: Additional experimental details are provided in Appendix

Original: A.
Revised: A.

Original: Table 1 shows the performance, maximum GPU memory duringtraining, and trainingtime comparison for greedy learning and backpropagation.
Revised: Table 1 shows the performance and maximum GPU memory during training comparison for greedy learning and backpropagation.

Original: Greedy learning with M = 8 performed best across allimage metrics compared with backpropagation, and greedy learning with M = 4 .
Revised: Greedy learning with M = 8 performed best across all image metrics compared with backpropagation, and greedy learning with M = 4 .

Original: Representativeknee images from the test set are illustrated in Fig.
Revised: Representative knee images from the test set are illustrated in Fig.

Original: 2 in Appendix B. Greedy learning successfullyrecovered relevant structural information, preserving details at least as well as backpropagation.
Revised: 2 in Appendix B. Greedy learning successfully recovered relevant structural information, preserving details at least as well as backpropagation.

Original: We benchmarked the training speed and the maximum memory required during training on a 12GB
Revised: We benchmarked the training speed and the maximum memory required during training on a 12GB

Original: NVIDIA Titan Xp graphics card.
Revised: NVIDIA Titan Xp graphics card.

Original: The maximum GPU memory used by greedy learning with M = 8 was 6x less than maximum GPU memory used in end-to-end backpropagation.
Revised: The maximum GPU memory used by greedy learning with M =was 6x less than maximum GPU memory used in end-to-end backpropagation.

Original: Whilebackpropagation was slightly faster per-iteration, the difference in training time is marginal and canfurther be remedied with asynchronous training of each block as shown in [11]. The slight increasein training time  for  greedy learningmight be attributed to increased communication time of sendingactivations from one network module to the next [11]. Furthermore, inference time is identicalbetween the two approaches since the ﬁnal architecture is the same, so real-time latency measures arenot impacted.
Revised: On the same GPU, training time per-iteration was same for backpropagation and greedy learning (Fig. 3). Details on the training time comparison for backpropagation and greedy learning is provided in Appendix C.

Original:  Discussion and Conclusion
Revised:  Discussion and Conclusion

Original: In this work, we presented an alternative to backpropagation for training model-based networksvia greedy learning.
Revised: In this work, we presented an alternative to backpropagation for training model-based networks via greedy learning.

Original: By performing gradient updates on several network modules, greedy learningallows a large reduction in the memory-footprint of training.
Revised: By performing gradient updates on several network modules, greedy learning allows a large reduction in the memory-footprint of training.

Original: Our experiments on a knee MRI datasetshow the effectiveness of our approach in reducing memory during training while preserving, andpotentially improving, generalization performance.
Revised: Our experiments on a knee MRI dataset show the effectiveness of our approach in reducing memory during training while preserving, and potentially improving, generalization performance.

Original: A stated limitation of greedy learning is the decoupling of gradient updates, which prevents information ﬂow between learned features among network modules that are optimized independently,resulting in less generalizable features.
Revised: A stated limitation of greedy learning is the decoupling of gradient updates, which prevents information ﬂow between learned features among network modules that are optimized independently, resulting in less generalizable features.

Original: In our experiments for this application, however, we found thatgreedy learning with M = 8 generalizes better compared to M = 4 and backpropagation, in additionto great improvements to memory consumption.
Revised: In our experiments for this application, however, we found that greedy learning with M = 8 generalizes better compared to M = 4 and backpropagation, in addition to great improvements to memory consumption.

Original: These memory improvements will undoubtedlybe important for enabling larger model-based networks with higher expressivity for reconstructingimages from 3D MRI, DCE imaging, or 4D-ﬂow.
Revised: These memory improvements will undoubtedly be important for enabling larger model-based networks with higher expressivity for reconstructing images from 3D MRI, DCE imaging, or 4D-ﬂow.

Original: The demonstrated beneﬁts of greedy learning thusprovide a vast potential for enhanced clinical outcomes, as can be investigated in future work.
Revised: The demonstrated beneﬁts of greedy learning thus provide a vast potential for enhanced clinical outcomes, as can be investigated in future work.

Revised: Acknowledgments and Disclosure of Funding

Revised: This work was supported by NIH Grants R01 EB009690, R01 EB026136 and NSF Grant DGE1656518.

Original: A Experimental Details 
Revised: A Experimental Details A.1 Dataset

Original: [Figure] [Figure] Our experiments were performed with the fully-sampled 3D fast-spin echo (FSE) multi-coil knee MRIdataset publicly available in mridata.org [12].
Revised: Our experiments were performed with the fully-sampled 3D fast-spin echo (FSE) multi-coil knee MRI dataset publicly available in mridata.org [13].

Original: Each 3D volume had a matrix size of 320 × 320 × 256 ,where training is performed 2D by treating each axial slice with size 320 × 256 as separate examples.
Revised: Each 3D volume had a matrix size of 320 × 320 × 256 , where training is performed 2D by treating each axial slice with size 320 × 256 as separate examples.

Original: The dataset consists of unique scans from 19 subjects, where 14 subjects (4480 slices) were usedfor training, 2 subjects (640 slices) were used for validation, and 3 subjects (960 slices) were usedfor testing.
Revised: The dataset consists of unique scans from 19 subjects, where 14 subjects (4480 slices) were used for training, 2 subjects (640 slices) were used for validation, and 3 subjects (960 slices) were used for testing.

Original: Sensitivity map estimation was performed in Sigpy [13] using JSENSE [14] with kernelwidth of 8 for each volume.
Revised: Sensitivity map estimation was performed in Sigpy [14] using JSENSE [15] with kernel width of 8 for each volume.

Original: Fully-sampled references were retrospectively undersampled using a 2D
Revised: Fully-sampled references were retrospectively undersampled using a 2D

Original: Poisson Disc undersampling mask.
Revised: Poisson Disc undersampling mask.

Original: A.2 Training Details
Revised: A.2 Training Details

Original: Experiments were conducted using Pytorch [15].
Revised: Experiments were conducted using Pytorch [16].

Original: A modiﬁed version of the network in [3] was usedas the baseline network architecture.
Revised: A modiﬁed version of the network in [3] was used as the baseline network architecture.

Original: Our network had N = 8 unrolled blocks, where each proximalblock consisted of 2 residual blocks with 128 channels.
Revised: Our network had N = 8 unrolled blocks, where each proximal block consisted of 2 residual blocks with 128 channels.

Original: Each residual block consisted of 2 ReLUsfollowed by convolutional layers with no normalization.
Revised: Each residual block consisted of 2 ReLUs followed by convolutional layers with no normalization.

Original: A batch size of 4 was used in all experiments,which was the largest possible batch size to ﬁt the baseline network to a 12 GB GPU.
Revised: A batch size of 4 was used in all experiments, which was the largest possible batch size to ﬁt the baseline network to a 12 GB GPU.

Original: The baselinenetwork, as well as each network module in greedy learning were trained with the Adam optimizerwith β 1 = 0 .
Revised: The baseline network, as well as each network module in greedy learning were trained with the Adam optimizer with β 1 = 0 .

Original: 9 , β 2 = 0 .
Revised: 9 , β 2 = 0 .

Original: All networks were trained for 80000 steps.
Revised: All networks were trained for 80000 steps.

Original: At inference time, themodel checkpoint that achieved the lowest validation nRMSE was selected for each method.
Revised: At inference time, the model checkpoint that achieved the lowest validation nRMSE was selected for each method.

Original: B 
Revised: B Sample Reconstructions

Revised: C Training Time In this section, we report the computation time per-iteration across different methods. We compare backpropagation with greedy learning on the same GPU.

Revised: In addition, since individual updates of modules are independent, optimization of each module can be parallelized.

Revised: In parallel greedy learning, the network modules are split across multiple devices.

Original: [Figure]
Revised: When a mini-batch is sampled, the forward pass is computed sequentially through all modules. After the forward pass is computed, each backward pass can be asynchronously executed among all modules in parallel. Since backward time is the primary bottleneck, asynchronous execution of the backward operation can greatly decrease overall computation time. Computation time comparison for backpropagation on 1 GPU, greedy learning on 1 GPU, and parallel greedy learning on 2 GPUs are illustrated in Fig.

Revised: Since no additional computations are needed for greedy learning, the training time is same as backpropagation using the same GPU.

Revised: Split across two

Revised: GPUs, parallel greedy learning results in a speed up in training compared to backpropagation and greedy learning.

Revised: When multiple GPUs are available, parallel greedy learning can enable layer-wise parallelism to reduce overall training time without compromising reconstruction performance.