Original: Multi-Task Accelerated MR Reconstruction Schemes for Jointly Training Multiple Contrasts
Revised: Multi-Task Accelerated MR Reconstruction Schemes for Jointly Training Multiple Contrasts

Original: Anonymous Author(s)

Original: AfﬁliationAddress email

Original: Abstract
Revised: Abstract

Original: Model-based accelerated MRI reconstruction methods leverage large datasets toreconstruct diagnostic-quality images from undersampled k-space.
Revised: Model-based accelerated MRI reconstruction methods leverage large datasets to reconstruct diagnostic-quality images from undersampled k-space.

Original: These networksrequire matching training and test time distributions to achieve high quality reconstructions.
Revised: These networks require matching training and test time distributions to achieve high quality reconstructions.

Original: However, there is inherent variability in MR datasets, including differentcontrasts, orientations, anatomies, and institution-speciﬁc protocols.
Revised: However, there is inherent variability in MR datasets, including different contrasts, orientations, anatomies, and institution-specific protocols.

Original: The currentparadigm is to train separate models for each dataset.
Revised: The current paradigm is to train separate models for each dataset.

Original: However, this is a demandingprocess and cannot exploit information that may be shared amongst datasets.
Revised: However, this is a demanding process and cannot exploit information that may be shared amongst datasets.

Original: Toaddress this issue, we propose multi-task learning (MTL) schemes that can jointlyreconstruct multiple datasets.
Revised: To address this issue, we propose multi-task learning (MTL) schemes that can jointly reconstruct multiple datasets.

Original: We test multiple MTL architectures and weightedloss functions against single task learning (STL) baselines.
Revised: We test multiple MTL architectures and weighted loss functions against single task learning (STL) baselines.

Original: Our quantitative andqualitative results suggest that MTL can outperform STL across a range of datasetratios for two knee contrasts.
Revised: Our quantitative and qualitative results suggest that MTL can outperform STL across a range of dataset ratios for two knee contrasts.

Original:  Introduction
Revised:  Introduction

Original: To reduce MRI scan time, various iterative reconstruction schemes have been investigated [1–4].
Revised: To reduce MRI scan time, various iterative reconstruction schemes have been investigated [1–4].

Original: Recently, deep learning approaches, which train a network to estimate the reconstructed image usingretrospectively undersampled k-space, have shown superior efﬁcacy over previous non-network basedmethods [5].
Revised: Recently, deep learning approaches, which train a network to estimate the reconstructed image using retrospectively undersampled k-space have shown superior efficacy over previous non-network based methods [5].

Original: However, these networks require sufﬁcient collection of fully sampled k-space datafrom similar acquisition protocols as the test-time inference data [6].
Revised: However, these networks require sufficient collection of fully sampled k-space data from similar acquisition protocols as the test-time inference data [6].

Original: For example, to train multiplecontrasts, the current paradigm is to collect multiple, fully sampled k-space data for each contrastand train each contrast-speciﬁc network separately to avoid domain shift.
Revised: For example, to train multiple contrasts, the current paradigm is to collect multiple, fully sampled k-space data for each contrast and train each contrast-specific network separately to avoid domain shift.

Original: Considering theexceptionally large variability of MR images (i.e. different contrasts, orientations, anatomies, pulsesequences), separate training requires a large effort and limits the additional information that can begained from multiple datasets.
Revised: Considering the exceptionally large variability of MR images (i e. different contrasts, orientations, anatomies, pulse sequences), separate training requires a large effort and limits the additional information that can be gained from multiple datasets.

Original: To address this barrier, we propose a novel multi-task learning (MTL) scheme that can jointly traina single network on a variety of datasets.
Revised: To address this barrier we propose a novel multi-task learning (MTL) scheme that can jointly train a single network on a variety of datasets.

Original: The scheme jointly trains various fully sampled k-spacedatasets by treating them as different tasks within the same network.
Revised: The scheme jointly trains various fully sampled k-space datasets by treating them as different tasks within the same network.

Original: The network can train multipletasks simultaneously and exploit shared, common features to prevent individual tasks from overﬁttingand foster better performance compared to conventional STL counterparts.
Revised: The network can train multiple tasks simultaneously and exploit shared, common features to prevent individual tasks from overfitting and foster better performance compared to conventional STL counterparts.

Original: MTL has recently gained traction in various areas [8–15], but has yet to be applied to MRI reconstruction.
Revised: MTL has recently gained traction in various areas [8–15], but has yet to be applied to MRI reconstruction.

Original: Our study investigates how this scheme can be useful for jointly training diverse datasets.
Revised: Our study investigates how this scheme can be useful for jointly training diverse datasets.

Original: Speciﬁcally, we focus on exploring the performance on two multi-contrast datasets with different datasize ratios.
Revised: Specifically, we focus on exploring the performance on two multi-contrast datasets with different data size ratios.

Original: We test multiple MTL architectures and multi-task loss functions against STL baselines.
Revised: We test multiple MTL architectures and multi-task loss functions against STL baselines.

Original: Our quantitative and qualitative results suggest that the proposed scheme can perform better than
Revised: Our quantitative and qualitative results suggest that the proposed scheme can perform better than

Original: STL across all dataset ratios.
Revised: STL across all dataset ratios.

Original:  Methods
Revised:  Methods

Original: The baseline STL network is a typical unrolled network (i.e., Variational Network) composed of aseries of unrolled blocks as in [16].
Revised: The baseline STL network is a typical unrolled network (i.e., Variational Network) composed of a series of unrolled blocks as in [16].

Original: Each unrolled block solves an unrolled iterative optimizationproblem:
Revised: Each unrolled block solves an unrolled iterative optimization problem:

Original: [Equation]
Revised: [Equation]

Original: where Ψ is a trainable regularization function modelled as a fully convolutional U-Net; A is a linearoperator that consists of sensitivity map projection, Fourier transform, and discrete sampling; ˜ kis a vector of acquired k-space.
Revised: where Ψ is a trainable regularization function modelled as a fully convolutional U-Net; A is a linear operator that consists of sensitivity map projection, Fourier transform, and discrete sampling; ˜ k is a vector of acquired k-space.

Original: The MTL network structure is composed of shared layers andtask-speciﬁc layers, each consisting of multiple unrolled blocks inside the layer.
Revised: The MTL network structure is composed of shared layers and task-specific layers, each consisting of multiple unrolled blocks inside the layer.

Original: In the followingsections, we describe two types of network architectures developed for this study and three weightedloss functions that account for deviations between tasks.
Revised: In the following sections, we describe two types of network architectures developed for this study and three weighted loss functions that account for deviations between tasks.

Original: 2.1 Network architecture
Revised: 2.1 Network architecture

Original: Two architectures are used in the study: split and multi-head, as shown in Figure 1.
Revised: Two architectures are used in the study: split and multi-head, as shown in Figure 1.

Original: Both networksstart with a number of shared blocks before splitting into task-speciﬁc blocks.
Revised: Both networks start with a number of shared blocks before splitting into task-specific blocks.

Original: In the split architecture,the task-speciﬁc blocks do not share further information (1a); we refer to these task-speciﬁc blocksas split blocks.
Revised: In the split architecture, the task-specific blocks do not share further information (1a); we refer to these task-specific blocks as split blocks.

Original: In the multi-head architecture, the U-Net encoder continues to be shared amongsttasks, but the decoder is task-speciﬁc (1b); we refer to these partially split blocks as multi-headblocks.
Revised: In the multi-head architecture, the U-Net encoder continues to be shared amongst tasks, but the decoder is task-specific (1b); we refer to these partially split blocks as multi-head blocks.

Original: During inference, a sample is fed through the appropriate shared and task-speciﬁc layers.
Revised: During inference, a sample is fed through the appropriate shared and task-specific layers.

Original: Note that during inference, the MTL architecture is equivalent to STL, and each task can be inferredindividually.
Revised: Note that during inference, the MTL architecture is equivalent to STL, and each task can be inferred individually.

Original: 2.2 Multi-task loss function
Revised: 2.2 Multi-task loss function

Original: In this study, we consider three different loss weighting schemes.
Revised: In this study, we consider three different loss weighting schemes.

Original: Naive weighting addresses the dataimbalance in the loss function by weighting individual task losses in an inverse relation to dataset size,as seen in Equation 2.
Revised: Naive weighting addresses the data imbalance in the loss function by weighting individual task losses in an inverse relation to dataset size, as seen in Equation 2.

Original: Uncertainty weighting [9] treats the multi-task network as a probabilistic modeland incorporates the task-dependent (or homoscedastic) uncertainty in the loss function.
Revised: Uncertainty weighting [9] treats the multi-task network as a probabilistic model and incorporates the task-dependent (or homoscedastic) uncertainty in the loss function.

Original: Finally,dynamic weight averaging (DWA) [14] assigns task losses based on the learning speed of each task.
Revised: Finally, dynamic weight averaging (DWA) [14] assigns task losses based on the learning speed of each task.

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: 2.3 Dataset
Revised: 2.3 Dataset

Original: We use two public knee datasets that are described in [16] and available at mridata.org [17].
Revised: We use two public knee datasets that are described in [16] and available at mridata.org [17].

Original: Thedatasets contain 19 coronal proton density weighted (PDw) and 20 coronal proton density weightedfat suppression (PDw-FS) knee scans acquired with a 15-channel knee coil.
Revised: The datasets contain 19 coronal proton density weighted (PDw) and 20 coronal proton density weighted fat suppression (PDw-FS) knee scans acquired with a 15-channel knee coil.

Original: The sensitivity maps ofeach coil are estimated from a 24 × 24 block at the center of the k-space using ESPIRiT [18].
Revised: The sensitivity maps of each coil are estimated from a 24 × 24 block at the center of the k-space using ESPIRiT [18].

Original: 2Dslices are treated as separate samples and are randomly divided into training, validation, and test sets.
Revised: 2D slices are treated as separate samples and are randomly divided into training, validation, and test sets.

Original: There are 13 PDw volumes (481 slices) and 14 PDw-FS volumes (497 slices) in the training set.
Revised: There are 13 PDw volumes (481 slices) and 14 PDw-FS volumes (497 slices) in the training set.

Original: Inour experiments, PDw simulates the abundant dataset by using all 481 slices, and PDw-FS simulatesthe scarce dataset by using a percentage (497, 253, 107, or 32 slices) of the 497 slices.
Revised: In our experiments, PDw simulates the abundant dataset by using all 481 slices, and PDw-FS simulates the scarce dataset by using a percentage (497, 253, 107, or 32 slices) of the 497 slices.

Original: 2.4 Training and inference
Revised: 2.4 Training and inference

Original: Our models are implemented in PyTorch and trained on NVIDIA Titan Xp GPUs with 12GB ofmemory.
Revised: Our models are implemented in PyTorch and trained on NVIDIA Titan Xp GPUs with 12GB of memory.

Original: For the experiments, networks with 12 unrolled blocks are used to ensure convergence.
Revised: For the experiments, networks with 12 unrolled blocks are used to ensure convergence.

Original: For
Revised: For

Original: MTL, different weighted loss functions – naive, DWA, uncertainty – are used for training the network.
Revised: MTL, different weighted loss functions – naive, DWA, uncertainty – are used for training the network.

Original: To assess image quality, magnitude images are normalized between 0 and 1, and peak signal-to-noiseratio (pSNR), structural similarity index (SSIM), and normalized root mean square error (nRMSE)are used.
Revised: To assess image quality, magnitude images are normalized between 0 and 1, and peak signal-to-noise ratio (pSNR), structural similarity index (SSIM), and normalized root mean square error (nRMSE) are used.

Original: During inference, k-spaces in the test-set are undersampled identically to guarantee faircomparisons.
Revised: During inference, k-spaces in the test-set are undersampled identically to guarantee fair comparisons.

Original: 2.5 Experiments
Revised: 2.5 Experiments

Original: For MTL, there are an exponential number of ways to compose the shared, split, and multi-headunrolled blocks.
Revised: For MTL, there are an exponential number of ways to compose the shared, split, and multi-head unrolled blocks.

Original: After preliminary experiments for two, six, or ten shared blocks in the beginningor middle of the network, we heuristically narrow our selection to two fully shared blocks at thebeginning, followed by ten task-speciﬁc (either split or multi-head) blocks (see Figure 1).
Revised: After preliminary experiments for two, six, or ten shared blocks in the beginning or middle of the network, we heuristically narrow our selection to two fully shared blocks at the beginning, followed by ten task-specific (either split or multi-head) blocks (see Figure 1).

Original: For thisstudy, we mix and match the two aforementioned architectures with the three loss functions (naive,
Revised: For this study, we mix and match the two aforementioned architectures with the three loss functions (naive,

Original: DWA, uncertainty) for a total of six MTL networks (see Appendix).
Revised: DWA, uncertainty) for a total of six MTL networks (see Appendix).

Original: MTL networks are jointly trainedusing 481 PDw slices and a percentage of PDw-FS slices.
Revised: MTL networks are jointly trained using 481 PDw slices and a percentage of PDw-FS slices.

Original: We also provide comparisons with transferlearning by taking the PDw baseline and ﬁne-tuning all layers using PDw-FS data [6].
Revised: We also provide comparisons with transfer learning by taking the PDw baseline and fine-tuning all layers using PDw-FS data [6].

Original: Source code for implementation will be publicly available via GitHub after acceptance of the paper.
Revised: Source code for implementation is publicly available via GitHub.

Original:  Results
Revised:  Results

Original: As seen in Table 1, an MTL network performs better than STL at every dataset ratio.
Revised: As seen in Table 1, an MTL network performs better than STL at every dataset ratio.

Original: At N = 107 and
Revised: At N = 107 and

Original: N = 253, the naive-weighted, split architecture dominates the other MTL architectures (see Appendix).
Revised: N = 253, the naive-weighted, split architecture dominates the other MTL architectures (see Appendix).

Original: Interestingly, the MTL network not only improves metrics for PDw-FS, but also for PDw at certainratios (see Appendix).
Revised: Interestingly, the MTL network not only improves metrics for PDw-FS, but also for PDw at certain ratios (see Appendix).

Original: Qualitative examination also suggests that MTL reduces errors in reconstruction.
Revised: Qualitative examination also suggests that MTL reduces errors in reconstruction.

Original: A comparisonbetween STL and MTL is seen in Figure 2 for two different inference slices.
Revised: A comparison between STL and MTL is seen in Figure 2 for two different inference slices.

Original: At N = 107, the MTLimage is reconstructed by the naive-weighted, split network.
Revised: At N = 107, the MTL image is reconstructed by the naive-weighted, split network.

Original: Although the naive-weighted, splitnetwork gives the best quantitative metrics at N = 253, qualitative observation suggests that the naiveweighted, multi-head network, a close runner-up, reconstructs comparably high quality images for N= 253.
Revised: Although the naive-weighted, split network gives the best quantitative metrics at N = 253, qualitative observation suggests that the naiveweighted, multi-head network, a close runner-up, reconstructs comparably high quality images for N = 253.

Original: As such, we choose to show the reconstruction from the naive-weighted, multi-head network for N = 253.
Revised: As such, we choose to show the reconstruction from the naive-weighted, multi-head network for N = 253.

Original: Interestingly, transfer learning performs better than MTL quantitatively (see Table 1 and
Revised: Interestingly, transfer learning performs better than MTL quantitatively (see Table 1 and

Original: Appendix).
Revised: Appendix).

Original: Because PDw-FS and PDw data are similar, this result suggests that multi-task learningmay be better suited for more dissimilar tasks.
Revised: Because PDw-FS and PDw data are similar, this result suggests that multi-task learning may be better suited for more dissimilar tasks.

Original:  Conclusion and discussion
Revised:  Conclusion and discussion

Original: This study explores the use of multi-task learning in jointly training diverse MRI data for reconstruction.
Revised: This study explores the use of multi-task learning in jointly training diverse MRI data for reconstruction.

Original: Our framework introduces inductive biases in the network by enforcing the sharing of usefulinformation between tasks.
Revised: Our framework introduces inductive biases in the network by enforcing the sharing of useful information between tasks.

Original: We see that MTL performs better than STL baselines across a rangeof abundant versus scarce ratios, for both PDw and PDw-FS datasets.
Revised: We see that MTL performs better than STL baselines across a range of abundant versus scarce ratios, for both PDw and PDw-FS datasets.

Original: We did not expect MTL tooutperform STL for the abundant data, and this positive sharing of information may guide selectionof tasks in future experiments.
Revised: We did not expect MTL to outperform STL for the abundant data, and this positive sharing of information may guide selection of tasks in future experiments.

Original: Our ﬁnding that transfer learning marginally outperforms MTL suggests that MTL may be bettersuited for more dissimilar tasks such as different orientations, anatomies rather than multi-contrasts.
Revised: Our finding that transfer learning marginally outperforms MTL suggests that MTL may be better suited for more dissimilar tasks such as different orientations, anatomies rather than multi-contrasts.

Original: Moreover, it is possible that once we use larger datasets, STL will dominate transfer learning, and
Revised: Moreover, it is possible that once we use larger datasets, STL will dominate transfer learning, and

Original: MTL’s main competitor will be STL.
Revised: MTL’s main competitor will be STL.

Original: One noted difﬁculty is selecting an appropriate architecture and loss function, as negative transferbetween the dataset can occur (see Appendix).
Revised: One noted difficulty is selecting an appropriate architecture and loss function, as negative transfer between the dataset can occur (see Appendix).

Original: All in all, our study provides a proof of concept that
Revised: All in all, our study provides a proof of concept that

Original: MTL can be successfully used in jointly training multiple contrasts for MRI reconstruction.
Revised: MTL can be successfully used in jointly training multiple contrasts for MRI reconstruction.

Revised: Acknowledgments and Disclosure of Funding

Revised: This project is supported by NIH R01 EB009690 and NIH R01 EB026136, as well as GE Healthcare and the Marcella Bonsall Fellowship.

Original: A Appendix
Revised: A Appendix

Original: Tables 2 and 3 provide benchmarks for how well various MTL networks perform, compared to STLbaselines and transfer learning.
Revised: Tables 2 and 3 provide benchmarks for how well various MTL networks perform, compared to STL baselines and transfer learning.

Original: Transfer learning is not shown in Table 3 because transfer learning isﬁne-tuned for PDw-FS only.
Revised: Transfer learning is not shown in Table 3 because transfer learning is fine-tuned for PDw-FS only.

Original: Experimental setups are described in detail in Section 2.5.
Revised: Experimental setups are described in detail in Section 2.5.

Original: Blank spacesarise from experiments that were not run
Revised: Blank spaces arise from experiments that were not run.

Original: [Table]

Original: [Table]

Original: Figure 3 compares transfer learning, STL
Revised: Figure 3 compares transfer learning, STL

Original: [Table]
Revised: [Table]

Original: and MTL experiments for PDw-FS.
Revised: and MTL experiments for PDw-FS.

Original: These are the same
Revised: These are the same

Original: PDw-FS slices shown in Figure 2, with the addition of the transfer learning results.
Revised: PDw-FS slices shown in Figure 2, with the addition of the transfer learning results.

Original: The improvementfrom MTL to transfer learning is less pronounced than the improvement from STL to MTL.
Revised: The improvement from MTL to transfer learning is less pronounced than the improvement from STL to MTL.

Original: Figure 4compares STL and MTL for PDw.
Revised: Figurecompares STL and MTL for PDw.