Paper: Robustness of Deep Learning for Accelerated MRI: Benefits of Diverse Training Data (ICLR-2024)

Reviewer comment (Novelty):
The scientific contributions of the manuscript are limited: despite the detailed analysis and discussion, the distributional robustness of deep learning MRI reconstruction has been discussed by a series of prior works [1-6]. Despite more detailed experiments (imbalanced data and healthy versus disease images), the major conclusions do not go beyond those of early works [1-3]. The authors failed to make significant theoretical and methodological contributions either (while most of [1-6] proposed either theoretical insights and/or methodological contributions).

Author response:
Thanks for noting that the problem we consider has 'significant real-world relevance' and that 'the experiments are performed on a large range of MRI reconstruction datasets with multiple real-world types of distributional shifts'. 

The main concern of the reviewer is whether our results go beyond those in [1-6]. Here is how our contribution are different from [1-6], and how they go significantly beyond those existing works:

- F. Knoll et al. [1] and J. Huang et al. [6] characterise the severity of specific distribution-shifts, and suggest the use of transfer learning to avoid distribution shifts. Contrary, we investigate the role of training data diversity in influencing robustness and find that diverse training data can significantly improve robustness. 

- C. Ouyang et al. [2] propose an approach that modifies natural images for training MRI reconstruction models. Contrary, we consider models trained in the conventional way, i.e., training on raw k-space data, and analyse systematically how combining different data-distributions affects robustness, and we find that it can significantly improve out-of-distribution robustness. 

- Similar to F. Knoll et al. [1] and C. Ouyang et al. [2], SUH Dar et al. [3] propose transfer-learning to train a reconstruction model on natural images and then fine-tune it on MR images. The goal of their work is to increase training data efficiency. X. Liu et al. [5] propose a special network architecture to improve the performance of training on multiple anatomies simultaneously for. Their results are based on retrospectively under-sampled DICOM images and the method only works for emulated single-coil data. Both works focus on improving in-distribution performance and do not study the robustness issue. Our objective is to study how performance and robustness of models is affected by data diversity of raw k-space data.

- F. Caliva et al. [4] consider the problem of adversarial robustness, whereas we consider the problem of robustness under \emph{natural} distribution-shifts.

Reviewer comment (Experiments):
The choice of using a U-Net is over-simplistic, given that the mainstream reconstruction works are based on unrolled proximal gradients with deep cascade networks, variational networks, as well as probabilistic diffusion models, which may also bring stronger distributional robustness due to better inductive biases compared with a plain U-Net.

Author response:
Additionally, the reviewer is concerned that the choice of studying U-net is over-simplistic as better models exist. We choose the Unet model since it is a well known baseline and our results are agnostic to the architecture. However, we agree it is important to explicitly show this, and have carried out a significant number of additional simulations after submitting the paper. Specifically, we have expanded our results in Section 3 and 7 with results obtained from the End-to-end VarNet (mentioned by the reviewer), a state-of-the-art model for MRI reconstruction, and also with results from ViT to demonstrate that our results continue to hold true even for non-convolutional architectures.

Reviewer comment (Experiments):
The authors are encouraged to take the effect of different reconstruction techniques into consideration: despite diverse implementation details, these methods can be generally categorized as 1. plain feed-forward networks; 2. unrolled cascaded networks; 3. variational networks; and 4. probabilistic diffusion models. The authors may want to consider the effects of inductive biases on distributional robustness.

Author response:
Additionally, the reviewer is concerned that the choice of studying U-net is over-simplistic as better models exist. We choose the Unet model since it is a well known baseline and our results are agnostic to the architecture. However, we agree it is important to explicitly show this, and have carried out a significant number of additional simulations after submitting the paper. Specifically, we have expanded our results in Section 3 and 7 with results obtained from the End-to-end VarNet (mentioned by the reviewer), a state-of-the-art model for MRI reconstruction, and also with results from ViT to demonstrate that our results continue to hold true even for non-convolutional architectures.

Reviewer comment (Evaluation):
For the first conclusion, no statistical tests or confidence intervals are used to qualify the statement that P+Q gives similar performance to P alone. I also think that a more rigorous analysis should have been done on a case level to compare the methods, as average SSIM scores can obscure what is going on at the tails. Medicine is inherently risk averse, so the tails are critical for this application.

Author response:
- Regarding concern (i), we changed Figure 3 to include the mean and standard deviations over 5 independent runs and added Figure 14 which also reports the mean and standard deviations over 5 independent runs for the same experiment as in Figure 2 but on smaller datasets. It can be seen from Figure 14 that the confidence intervals are extremely small (almost not visible) for models trained on more than 3k images, which concerns the vast majority of experiments in our paper, and thus we did not add confidence intervals throughout, since this wouldn't be justified given the cost of thousands of GPU hours and the fact that the confidence intervals would be very small.

Reviewer comment (Novelty):
For data diversity, this effect seems to somewhat rehash the results of (Knoll, 2019), but with larger quantities of data. As with (1), a deeper analysis of edge cases could have been useful. However, in general I don't have major issues with this section.

Author response:
Regarding our reason to submit to the 'datasets and benchmark category': One category mentioned regarding the scope of the datasets and benchmarks tracks is often (e.g., at NeurIPS) 'Data-centric AI methods and tools, e.g. studies in data-centric AI that bring important new insight', we believe our work falls into this category.

Reviewer comment (Evaluation):
The analysis of pathology could be particularly problematic, as SSIM average scores obscure what is going on in small regions of the image. For example, in (Radmanesh et al., 2022) the SSIM for T1 images at a 6X acceleration is still quite high, but only 80% of the cases were accepted by radiologists for being clinically acceptable. In the same paper, Figure 1 shows that for very low SSIMs, some large and prominent pathologies can still be seen, whereas more modest pathologies such as MS lesions are erased at low accelerations. All this is to say that the SSIM metric is not indicative of performance for pathology cases, and a deeper analysis is needed to substantiate the claim that models trained on healthy subjects can reconstruct pathology.

Author response:
- As a response to concern (ii), we have now re-evaluated the models only on the region containing the pathology (instead of computing metrics on the entire image), see the revised Figure 5. We further distinguished between small pathologies and relatively large pathologies. 

- We also note that we do not claim that models trained only on healthy images can reconstruct pathologies. Rather, our results indicate that models trained on data without pathologies can reconstruct pathologies as accurately as models trained on data with pathologies.

Reviewer comment (Novelty):
The paper is submitted to the datasets and benchmarks track, but I did not identify what it considered to be its core contribution to that area.

Author response:
Regarding our reason to submit to the 'datasets and benchmark category': One category mentioned regarding the scope of the datasets and benchmarks tracks is often (e.g., at NeurIPS) 'Data-centric AI methods and tools, e.g. studies in data-centric AI that bring important new insight', we believe our work falls into this category.

Reviewer comment (Experiments):
**Limited tuning of hyperparameters.** From Appendix D, it appears that all models are trained with the exact same hyper parameters. A more thorough version of this experiment would tune these hyperparameters for each architecture and dataset. However, I recognize that, given the number of models trained in the paper, this would be exceptionally computationally expensive. The chosen hyperparameters seem to be fairly standard values that a deep learning practitioner might use to initialize a model, and the observed trends are largely consistent across architecture, dataset split, etc, so I am okay with the current experimental setup — but it may be good to include a note about this in the Limitations section.

Author response:
The reviewer emphasized the need for a discussion on model capacity and hyperparameter tuning. To address this, we have conducted experiments akin to those in Section 3 but utilizing a much smaller U-net, see the new Appendix E.1. The results align with those obtained with the larger U-net. Furthermore, we describe the process of determining hyperparameters in the new Appendix D. We thank the reviewer for suggesting this.

Reviewer comment (Experiments):
**Limited discussion of model capacity.** Related to the above point, I am particularly curious whether the observed trends hold for smaller models. This is motivated by some informal experiments I did on FastMRI data with/without fat suppression with much smaller models than those in this paper, where I observed that including both data splits did decrease model performance compared to models trained separately on each split. One hypothesis could be that models with lower capacity are less able to capture the variability in multiple data splits, leading to this effect. The bigger models in this paper definitely achieve better performance overall and thus are the right ones to report results for here, so I don’t expect the current experiments to be expanded to cover this point and think the paper is strong enough as is. But if the authors have already conducted some related experiments, they may be useful to include in the appendix, and/or this may be another point for the discussion.

Author response:
The reviewer emphasized the need for a discussion on model capacity and hyperparameter tuning. To address this, we have conducted experiments akin to those in Section 3 but utilizing a much smaller U-net, see the new Appendix E.1. The results align with those obtained with the larger U-net. Furthermore, we describe the process of determining hyperparameters in the new Appendix D. We thank the reviewer for suggesting this.

Reviewer comment (Experiments):
**Limited utility of the early stopping criteria for distributional overfitting.** Unlike traditional overfitting, where the early stopping epoch can be chosen by looking at a held-out subset of the current dataset, it seems to be much harder to monitor early stopping for the distributional overfitting case. Because we are interested in performance on out-of-distribution data, in a realistic scenario, we likely don’t have the out-of-distribution data to monitor in the first place (if we did, we’d want to include it in the training dataset). I see in section 7 that early stopping was performed by tracking validation loss on the fastMRI knee validation set, which seems like an “in-distribution” loss to me, because the fastMRI knee training set was included in the training data. So it is not clear to me to what extent distributional overfitting is being properly avoided here. I don’t think this is a major limitation; even as is, the learned model outperforms the baselines on the out-of-distribution tests. But this would be a good point to discuss in section 6.

Author response:
We confirm that the reviewer has correctly observed that we track the in-distribution loss as an attempt to mitigate distributional overfitting. This decision is based on the results in Section 6 which suggest that distributional overfitting begins when training is close to convergence.

Reviewer comment (Presentation):
**It is hard to track the same model across some figures.** In Figures 4 and 5, multiple versions of the same architecture are reported within the same figure, trained on different splits of the data. In the current data visualization in Fig 5 for example, it is hard to track which pair of pentagons correspond to the same model configuration. I wonder if it would be useful to have an additional graphic plotting the *difference* in SSIM on Q for each *difference* in SSIM on P, with one datapoint for each model. This would partition the points into four quadrants showing relative improvement/performance decrease on P and Q, and a compelling result would be to show that no datapoint shows a dramatic decrease in performance on Q when trained on just P, as opposed to being trained on P+Q.

Author response:
Concerning the reviewer's remarks on visualization, we appreciate the interesting suggestion. The current choice of visualization method aligns with the conventions established in the effective robustness literature (see R. Taori et al., 2020, and T. Nguyen et al., 2022). This decision was made to maintain consistency and facilitate comprehension for readers who are already familiar with the effective robustness framework.

Reviewer comment (Evaluation):
**Confidence intervals are missing in Figs 2-3.** I understand these would clutter some of the other visualizations, but in Figs 2-3 it would be great to see confidence intervals for the bar plots to contextualize the differences between the bars.

Author response:
The reviewer has raised a valid and important concern regarding the absence of confidence intervals in Section 3, where we demonstrate that training a model on two distributions yields comparable performance to training on individual distributions. In the revised version, we incorporated mean and standard deviations over 5 independent runs over different data and starting from a different initialization in Figure 3 and we introduced a new Figure 14. From Figure 14, it can be seen that if the dataset size is beyond 3k, the confidence interval is vanishingly small (not visible) and thus we did not add confidence intervals throughout, since this wouldn't be justified given the cost of thousands of GPU hours and the fact that the confidence intervals would be very small.

Reviewer comment (Experiments):
How do different model hyperparameters (especially pertaining to model size/model capacity) affect these trends?

Author response:
The reviewer emphasized the need for a discussion on model capacity and hyperparameter tuning. To address this, we have conducted experiments akin to those in Section 3 but utilizing a much smaller U-net, see the new Appendix E.1. The results align with those obtained with the larger U-net. Furthermore, we describe the process of determining hyperparameters in the new Appendix D. We thank the reviewer for suggesting this.

Reviewer comment (Presentation):
Please see the data visualization suggestions above.

Author response:
Concerning the reviewer's remarks on visualization, we appreciate the interesting suggestion. The current choice of visualization method aligns with the conventions established in the effective robustness literature (see R. Taori et al., 2020, and T. Nguyen et al., 2022). This decision was made to maintain consistency and facilitate comprehension for readers who are already familiar with the effective robustness framework.

Reviewer comment (Experiments):
According to Table 1, the datasets seem diverse in terms of “View” and “Vendor” as well. So, it would be much more interesting to conduct experiments considering the two factors.

Author response:
The reviewer would like to see experiments regarding distribution-shifts with respect to vendors and anatomic views. In Section 7 of the original manuscript we have the mentioned distribution-shifts, however they were hidden since we did not point out that there are those shifts. For example, the distribution shift from fastMRI knee to NYU dataset is a distribution shift in anatomic views, and from fastMRI to Stanford 2D or CC-359 is a distribution shift related to vendors. We have updated the corresponding paragraph in Section 7 of the revised manuscript to emphasize that those shifts are in terms of vendors and anatomic views.

Reviewer comment (Experiments):
Also, it seems they synthesized accelerated MRI by under-sampling the fully sampled k-space from the original MRI images. Then, it would be interesting to compare models trained on k-space data sampled with different frequency (eg. 4-fold vs 8-fold)

Author response:
The reviewer would like to see experiments beyond 4-fold acceleration. While it could be interesting to see what happens at higher acceleration rates, we chose the 4-fold acceleration setup to provide more practical insights, as exceeding 4-fold acceleration tends to diminish clinical relevance (M. Muckley et al., 2021; A. Radmanesh et al., 2022). We added the explanation to our revised manuscript in Section 2.

Reviewer comment (Experiments):
The paper does not explore vendor distribution shifts, raising questions about the model's generalizability across different MRI vendors.

Author response:
- Regarding results for vendor-related distribution shifts: In Section 7 of the original manuscript we consider vendor shifts. Specifically, in Section 7 we introduce distribution-shifts from the fastMRI dataset, which are obtained by Siemens scanners, to the Stanford 2D or CC-359 dataset, which are obtained by GE scanners. In both cases we find that training on a diverse dataset improves robustness. We have updated the corresponding paragraph in Section 7 in the revised manuscript to emphasize this distribution-shift.

Reviewer comment (Experiments):
The paper focuses on input-distribution shift in visual aspects of MRI images but does not delve into how these shifts manifest in k-space data, the raw data format for MRI.

Author response:
- Regarding distribution-shifts related to visual aspects of MRI images vs. how these shifts manifest in k-space data: Since k-space and image are related through the Fourier transform, a distribution shift in the image space is also a distribution shift in the k-space.

Reviewer comment (Experiments):
The study utilizes full datasets from different sources for training, but it doesn't explore whether using a subset of these datasets could be equally effective.

Author response:
- The reviewer asks if mixing data from two distributions, P and Q, in a disproportionate manner like ‘P+0.005Q is sufficient for good model training’: In terms of in-distribution performance, we find that skewed data up to 90%P and 10%Q do not pose a problem (see Figure 3 in the revised manuscript). In case of distribution-shifts from distribution P to distribution Q, we do not have data from distribution Q available. If data from distribution Q was available, then including it into the training set would eliminate the distribution shift.

Reviewer comment (Evaluation):
The paper does not adequately discuss its limitations.

Author response:
- The reviewer states that we do not adequately discuss the limitations of our results and further suggests that we discuss the training times of our models: In the revised manuscript we have expanded the discussion on limitations by adding that in practice it might be difficult or expensive to collect diverse and large datasets. We believe that the training time is not a significant concern, given that the current bottleneck in medical imaging is data availability rather than computational resources. Nevertheless, we agree that stating the training times is important, and we have included this information in Appendix D of the revised manuscript.

Reviewer comment (Experiments):
Is there a reason vendor distribution shift was not studied? Will a model trained using data for GE vendors work for Siemens?

Author response:
- Regarding results for vendor-related distribution shifts: In Section 7 of the original manuscript we consider vendor shifts. Specifically, in Section 7 we introduce distribution-shifts from the fastMRI dataset, which are obtained by Siemens scanners, to the Stanford 2D or CC-359 dataset, which are obtained by GE scanners. In both cases we find that training on a diverse dataset improves robustness. We have updated the corresponding paragraph in Section 7 in the revised manuscript to emphasize this distribution-shift.

Reviewer comment (Experiments):
The premise of the paper is based on input-distribution shift. While visually different contrast images, anatomy, etc have distribution shift in the data. However, the input to MRI reconstruction networks is k-space data (subsampled MRI). How much of the data shift is actually present in k-space representation? Is the k-space representation for different contrast/anatomy/Magnetic-field images actually different distributions not clear in the paper? Studying histograms of amplitude and phase information might give an answer to this question. The distribution shift in the visual domain vs in the Fourier(k-space) domain might be completely different.

Author response:
- Regarding distribution-shifts related to visual aspects of MRI images vs. how these shifts manifest in k-space data: Since k-space and image are related through the Fourier transform, a distribution shift in the image space is also a distribution shift in the k-space.

Reviewer comment (Experiments):
Comparing results from Figures 1 and 2, it looks like if the in-distribution data size increases (1.3k vs 20k, 400 vs 6k), it also shows an increase in out-of-distribution performance (SSIM on Q is always higher for a larger dataset in Figure 1). A study of the impact of dataset size increase vs out-of-distribution performance is also warranted. This will support the idea that adding datasets from diverse resources is better than increasing the dataset size of a single source.

Author response:
- Regarding that we do not 'explicitly study the impact of dataset size on model performance with out-of-distribution data': Section 4 and 5 contain experiments with different dataset sizes, but we indeed do not consider different training set sizes throughout for the following reason: Training on larger datasets leads to better models, but not necessarily more robust models. Specifically, in and out-of-distribution performance is correlated for the distribution shifts we consider, see [3] and Figure 4 in our paper, and see [1,2] for correlations of in-and out-of-distribution relations in classification and beyond. In Figure 4, increasing the training set size improves a model in that it improves both in and out-of-distribution performance (it moves on the line) but it is not more robust than expected of a model with a given in-distribution-performance. Contrary, training on more diverse data, leads to an \emph{effective} robustness gain in that it gives better out-of-distribution performance relative to the in-distribution performance.

Reviewer comment (Evaluation):
The paper lacks a discussion on the limitations of the results. For example, this is only applicable in cases where data from different resources are available at training time. Other limitations might also include the amount of training time required to train the model with a large dataset size. Comparison of training times for different models might also be a relevant metric to the study.

Author response:
- The reviewer states that we do not adequately discuss the limitations of our results and further suggests that we discuss the training times of our models: In the revised manuscript we have expanded the discussion on limitations by adding that in practice it might be difficult or expensive to collect diverse and large datasets. We believe that the training time is not a significant concern, given that the current bottleneck in medical imaging is data availability rather than computational resources. Nevertheless, we agree that stating the training times is important, and we have included this information in Appendix D of the revised manuscript.