Revisiting MAE pre-training for 3D medical image segmentation

• Very unclear and unorganized manuscript. I believe that it can be improved substantially. I specified many details in "Suggestions", including the following: previous related works were not described, unclear concepts are not introduced, parts of what should be the conclusion (e.g., that MAEs dominate, SSL pretraining works) are in the "Results and discussion" section, there is no section/subsection where the experiments are clearly described and instead they're mixed with "results and discussion". Another example of mixing: right before the conclusion, in only one paragraph (L508-516), we can find an experiment description, the results, and the discussion, all mixed together.
• Limited novelty.
  • Limited methodological novelty. The framework is based on well-established Masked AutoEncoders "with the recent adaptations introduced by Tian et al. (2023); Woo et al. (2023)".
  • Partially limited application novelty since the pretrained models are not publicly available. Although the code is shared, researchers may not have access to large datasets; L53 reads that there seems to be "a public decrease in the community’s willingness to share data" (I don't agree or disagree with this statement, but this may be only regarding brain MRI).
• In many cases, it is unclear if one approach is better than another because no standard deviations are shown. In other words, it cannot be understood whether a method achieving a 71.66 dice coefficient is actually better than another method achieving 71.35.

• The paper resembles a technical report rather than an academic paper, lacking demonstration of its innovation. It addresses the problem by simply combining methods without discussing the essence of the issue. MAE has already proven its competitiveness in previous work, yet the paper merely applies MAE to the backbone without further exploration.
• The writing quality is poor, especially with the confusing use of symbols (e.g., the confusion of D[1~9], DS[1~9], and dataset names). The excessive use of items and textbf (too much bold text reduces readability) and quotes (all right quotes, which can be displayed correctly in $\LaTeX$ using `word') makes the paper difficult to read.
• The paper lacks a symbolic representation and adequate explanation of the task setup and model, instead focusing extensively on dataset selection and hyperparameter choices.
• The figures are confusing. Figure 1 is hard to understand, appearing to mix development and testing in a workflow without showing the model pipeline. Figure 2 is poorly organized, with excessive whitespace and mixed use of pie and bar charts. Figure 3 seems to be generated directly by a program, lacking good organization and sufficient explanation, with a lot of meaningless whitespace.
• The paper lacks visualizations of some results.
• The experimental section only describes performance improvements and changes in tables without further discussion. The results show that the model does not achieve significant performance gains in many experiments (the large model size yields only slight improvements or none at all), suggesting that simply applying MAE does not produce sufficiently good results, and the authors do not propose better methods.
• From Table 1-a, it can be observed that model performance improves based on some sparsification adaptations, raising doubts about whether the results in Table 3 are achieved by stacking tricks rather than the method itself. Table 1-c shows no performance improvement from scaling, and Table 3 even shows performance degradation due to scaling, without explanation, which is disappointing for the method.

The reviewer rated soundness as 2 and contribution as 2, given the following reasons:

Soundness:

1. Given the current experiment setups, it is insufficient to conclude the optimal pretraining strategy.

• The patch size of MAE is a critical parameter, while Kaiming's MAE paper [1] did not ablate on that, some other studies ablated on that parameter and found significant performance differences [2-4]. This paper utilized a patch size of 5x5x5 in the bottleneck, equivalent to 32x32x32 in image space. It seems too large for the 3D MAE. Both [3] and [4] indicate in the 3D medical image, a high masking ratio and a small patch size are the key ([3] used a patch size of 16x16x16, [4] used 8x8x8).
• Regarding scaling of MAE pretraining, this paper only investigates having 8x batch size, larger learning rate, and 4x training iterations. Those are not keys to evaluating scaling. Scaling more refers to performance gain with an increase in data size and an increase in model parameters. On high pretraining iterations, the impact of larger batch size and learning rate may not be significant. Extending training iterations may also not help as the MAE training tends to saturate after prolonged training ([1] Fig 7 upper, 800 vs. 1600 epochs are very close, 84.9 vs. 85.1). So what will be really interesting to see is to ablate on 1. training on 10%, 25%, 50%, 75%, 100% of 40k pretraining datasets; 2. varying model's depth to see how performance changes with model size. In addition, the naming of S3D-L is very misleading, as -L always indicates a larger model (with more parameters) in the ML naming convention.

The above two reasons lead to a rating of soundness of 2, as without experiments on those two perspectives, it is hard to conclude the current manuscript presents the optimal strategy.

Contribution:

The reason for a rate of 2 in the contribution is that the current manuscript, entitled 'Revisiting MAE pre-training for 3D medical image segmentation', did not include any comparison with previous studies that utilized MAE pretraining for 3D medical image analysis, notably [3, 5]. Instead, it only involves comparisons with Model Genesis, Volume Fusion, and VoCo.

The contribution of the current study will be much higher if compared to the existing 3D MAE pretraining framework developed for medical images (i.e., [3,5]).

Others:

The quality of Fig. 2 can be improved.

Overall, the reviewer recommends rejection because the technical flaws and a lack of comparison with existing 3D MAE frameworks (as presented above) outweigh the benefits brought by large-scale datasets and diverse downstream evaluations.

• [1]: He, Kaiming, et al. "Masked autoencoders are scalable vision learners." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2022.
• [2]: Xie, Zhenda, et al. "Simmim: A simple framework for masked image modeling." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2022.
• [3]: Chen, Zekai, et al. "Masked image modeling advances 3d medical image analysis." Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision. 2023.
• [4]: Zhang, Xuzhe, et al. "MAPSeg: Unified Unsupervised Domain Adaptation for Heterogeneous Medical Image Segmentation Based on 3D Masked Autoencoding and Pseudo-Labeling." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2024.
• [5]: Tang, Yucheng, et al. "Self-supervised pre-training of swin transformers for 3d medical image analysis." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2022.

• (W1) The paper does not provide important details on how the weights are transfered for finetuning. Is finetuning performed in the nnUnet framework? Which augmentations are used when finetuning? Are the learning-rate, augmentations etc. fixed for all evaluation datasets? As noted by the authors, selecting an appropriate configuration for each dataset is important. I assume that the configuration is also dynamic for S3D, however the paper does not contain any mention of how this is achieved with pretrained weights.
• (W2) The authors use a patch size of 160^3 which is significantly larger than most previous works, however does not provide any ablations of the effect of this. The proposed performance gains therefore cannot be ruled out to be mainly from using a larger patch size.
• (W3) The paper lacks references to important related work. Specifically, the authors are suggested to include the following two articles in the related works section:
  • SSL with a convolutional MAE for 3D MRI segmentation on a public dataset of 44K MRI scans, which similarly revisits various design choices for CNN pretraining, yet with inferior evaluation: [1]
  • Implementation of Spark-like sparse masking for large-scale 3D medical image pretraining: [2]
• (W4) The notes on scaling and the S3D-L variant is misleading since it does not use a model of larger size, yet is scaled in other ways. This meaningfully departs from the established literature, and the authors are encouraged to find another way of communicating the different training setup. Scaling the model and data sizes are important ingredients in compound scaling, yet none of these are performed.
• (W5) The pretraining dataset is private and only limited information on the nature of this dataset is included. For reproducibility purposes, it would be beneficial for the community if the authors would release checkpoints trained on Brains-45K (similar size to the used dataset) from [1].
• (W6) The abstract mentions pretraining is on a dataset of 44K 3D MRI volumes, however the actual pretraining dataset is 39K volumes after filtering out low-quality data. This discrepancy is misleading.

References:

[1] Munk, Asbjørn, et al. "AMAES: Augmented Masked Autoencoder Pretraining on Public Brain MRI Data for 3D-Native Segmentation." arXiv preprint arXiv:2408.00640 (2024).

[2] Tang, Fenghe, et al. "Hyspark: Hybrid sparse masking for large scale medical image pre-training." International Conference on Medical Image Computing and Computer-Assisted Intervention. Cham: Springer Nature Switzerland, 2024.