Paper: Lost in Transformation: Current roadblocks for Transformers in 3D medical image segmentation (ICLR-2024)

Reviewer comment (Novelty):
There is no contribution or novelty. This is more like comparing 9 frameworks.

Author response:
Novelty is a subjective idea and we believe that we highlight an important downside concerning the domain of medical image segmentation. We took the 9 most popular Transformer-based networks in the domain and demonstrated minimal contribution of the Transformer component. We believe that this should guide research in directions that would better utilize Transformers in the medical image segmentation domain. We respectfully disagree with the notion that this is simply “comparing 9 networks” and would have benefited from the reviewer expanding on this a bit further.

Reviewer comment (Experiments):
Experimental design has some flaws. There are so many transformer networks and picking 9 out of those should have a valid assumption, which is missing in the paper.

Author response:
Given the strong comment on our supposed random network selection, we urge the Reviewer  to Section 2 (Experimental Design) where we state  that: “We dissect a number of massively-influential Transformer architectures for medical image segmentation, regularly used as blueprints for designing newer architectures or state-of-the-art baselines. In this work, we focus on 9 such networks with 5500+ citations collectively in the last 3 years (see Table 5)” Our motivation in network selection is clearly aimed at the most popular networks in the domain of medical image segmentation which we clearly state.

Reviewer comment (Evaluation):
Only quantitative analysis is provided in the manuscript. What about qualitative analysis? Showing segmentation masks would help readers to understand which method performs well, especially when it comes to the medical AI domain, qualitative analysis is a must.

Author response:
We unfortunately disagree with the notion of qualitative analysis being superior or even necessary. As informed researchers in the medical image analysis domain, we value quantitative analysis highly, whereas qualitative analysis is usually nice-to-have but not mandatory. This is because it is prone to cherry-picking and does not allow for an intuitive measure of segmentation quality, especially not in the 3D domain where predictions are very hard to visualize. While relevant in a more clinical setting, in the medical image segmentation, we severely doubt the veracity of “qualitative analysis being a must-have” in comparing segmentation techniques.

Reviewer comment (Experiments):
A comparison of intermediate attention maps is missing.

Author response:
This is a process standard to Transformers but we believe not relevant to our scenario. If the removal of an entire Transformer shows no performance degradation, we are unsure as to how visualizing intermediate attention maps help. We believe that our Transformer ablation is a much stricter demonstration of Transformer Utilization by our networks that the one requested by the reviewer.

Reviewer comment (Evaluation):
Representation similarity can be visualized as a heat map comparing network layers. This would give some understanding to the reader of how feature extraction works and whether it’s similar across all (or part of) networks or not.

Author response:
Tying to the previous point, we value quantitative results over qualitative results. With that in mind, we deliberately chose to not use a heatmap visualization of the similarity but the diagonal values as line-plots. This allows readers to read explicit, quantitative values and differences from the plot, whereas in a heatmap the reader would have to interpret changes due to draw conclusions making for bad interpretability and low added value. We urge the Reviewer to consider that there is little loss of information compared to their preferred style of visualizing representation similarity.

Reviewer comment (Novelty):
While the medical image segmentation task, the utility of the transformer, can be brought under scrutiny, this is not true for panoptic/instance segmentation and video segmentation, not only because of the data set size but also because of the fundamental difference in network architectures. This makes the criticism of transformers very specific to U-net-like models popular in the medical imaging community, which makes the paper relatively less appealing to the general image segmentation community.

Author response:
We agree with the reviewer that our findings do not generalize to other domains like e.g. video segmentation, where data is less hard to come by, but we should be absolutely clear that we also do not claim so and constrain ourselves to the domain of 3D medical image segmentation. We believe that 3D medical image segmentation, while very specific as a domain, is also extremely large and that our findings will significantly influence a plurality of of researchers in this area. We hope that the reviewer will consider this and not use this as a determining factor that influences the score of the paper.

Reviewer comment (Novelty):
While the paper quite convincingly points out flaws in the current practice of architectural design in medical image segmentation, the paper did not bring any new ideas to mitigate the issue, which remains a major weakness and is hard to address within the rebuttal period. Hence, despite being a good review and investigative paper, I am not sure whether it is a good fit for ICLR.

Author response:
The reviewer is correct that we do not provide any new ideas to mitigate the issue, yet we hope to highlight the issues in order to shift the focus from currently _proposing new architectures_ to _proposing effective architectures_ - we hope our research demonstrating the weaknesses of Transformers will enable researchers to improve them further. We also address this further in the general comments for the rebuttal which you may find at the top of this page.

Reviewer comment (Evaluation):
The use of volumetric error overlap is confusing in concluding model behavior. Given that all models considered provide point estimates, how can the author assert that the apparent difference in model behavior is not a result of the underlying uncertainty? It will be good to know the volumetric error map between three runs of the same models as a reference to the uncertainty because the authors took the same approach for representational change measurement. And how are the thresholds 0.95, 0.85, etc. chosen? Seems quite arbitrary.

Author response:
The reviewer echoes our sentiments exactly about the underlying uncertainty. The VEO value is indeed a ratio of: 
1) 3 runs of the same model as reference
2) all permutations of these models with the Transformer removed. 

The same approach as representation change measurement was taken here as well. We can make this explicit in the text and apologize for not doing so. The thresholds themselves are formulated by jointly considering VEO and P_sim and their joint degradation across the plots in Table 2. They are selected specifically at points of degradation of one or both scores on the plot and are for illustrative purposes.

Reviewer comment (Experiments):
The study's scope was confined to a narrow selection of segmentation datasets, and the ablation studies conducted were restricted in its current form.

Author response:
We agree that only 2 datasets is a limitation of our work, hence we extended our experiments to verify our core results on TotalSegmentator, which is the largest annotated dataset in the landscape of 3D medical image segmentation. However, we want to emphasize that we want to investigate whether transformers are feasible given the data that is present in 3D medical image segmentation. Hence we chose the next two best large, non-trivial datasets representing both pathology and organ segmentation. We chose these large datasets, as they represent the datasets where transformers should bring the largest benefits over CNNs, and as they enable our dataset-size ablations.

| Method      | Standard     | Whole Block Replaced |
|-------------|--------------|----------------------|
| CoTr        | 89.88 (0.21) | 89.09 (0.28)         |
| SwinUNETR   |  87.08 (0.35)  | 85.15 (0.21)         |
| TransBTS    | 87.74 (0.38) | 86.31 (0.78)         |
| UNETR       | 80.79 (0.41) | 72.86 (0.56)         |
| nnFormer    | 85.30 (0.28) | 80.37 (0.14)         |
| nnUNet (3d) | 89.04 (0.11) |                      |

The missing values are currently being evaluated but we will include them in Table 2 left.

Reviewer comment (Experiments):
The authors' research on medical imaging segmentation with Transformer-based models, focusing on organ and pathology within specific datasets, may not capture the potential benefits of long-range dependencies in all medical imaging scenarios. For example, cardiac video segmentation may reveal different results due to the temporal dynamics involved. The authors could provide insights on whether their findings are applicable to such medical imaging tasks where Transformers might show utility.

Author response:
In our current experiments we constrain ourselves to 3D organ and pathology segmentation as these are the tasks that these architecures originally were proposed for (Fig. 3 - right side) and as these are the tasks that represent the majority of use-cases, as we show in the appendix A (Fig. 5). We agree that there could be slightly different conclusions if the methods were evaluate on temporal datasets, however, to bring the reviewer’s example to our investigative domain, if for example, one decided to segment objects in cardiac videos ignoring the time axis using these Transformer-based network, our results would hold.

Reviewer comment (Reproducibility):
Can the authors clarify whether the transformer blocks were pretrained with natural image datasets.

Author response:
In the original papers of the 3D architectures none of the original authors used pre-trained transformers, hence we follow that notion in order to keep our evaluation as close to the original as possible and subsequently do not use pre-trained transformer blocks throughout our experiments with these. Similarly for the 2D architectures different authors follow different protocols with some using and others not using pre-trained ViT transformers. To be precise [SwinUNet](https://arxiv.org/pdf/2105.05537.pdf), [TransFuse](https://arxiv.org/pdf/2102.08005.pdf) and[TransUNet](https://arxiv.org/pdf/2102.04306.pdf) were pretrained on ImageNet with UTNet being trained from scratch. We admit this is a deviation from the official use, yet we deliberately chose this considering a recent article showing that pre-trained weights do not help for segmentation tasks. (See [Rethinking pre-training on medical imaging](https://doi.org/10.1016/j.jvcir.2021.103145))

Reviewer comment (Experiments):
The study used nnU-Net as a benchmark for evaluating performance. Its status as an industry benchmark relies heavily on advanced data augmentation and meticulous hyperparameter optimization. The nine Transformer-based architectures assessed may not employ these sophisticated techniques. For an equitable comparison, it's crucial that all other variables, such as data augmentation protocols, are standardized across models. Without this uniformity, any observed differences in performance could be attributed to varying methodologies rather than inherent architectural distinctions.

Author response:
It is true that we do leverage nnU-Net as a benchmark for evaluating performance, yet nnU-Net does not rely on extensive hyperparameter tuning. nnU-Net is a framework that determines a single, fixed setting of hyperparameters. It follows a variety of heuristics, as determined by the author, and selects, e.g. an appropriate spacing, patch size and normalization given the statistics of the dataset. As these are dataset specific choices, we use them for all of our experiments. Similarly nnU-Net selects an augmentation schemes, but generally differentiates between two simple, default augmentation schemes, one for 3D data and one for 2D data.
In our Experiments we do apply the identical spacings, patch sizes, normalization schemes that nnU-Net selects for the 3D U-Net to all architectures.The only difference in the settings is based on learning rate, optimizer and weight decay, where we follow the recommendations of the original papers.

Reviewer comment (Experiments):
The authors only tested on two public datasets, which might not be convincing enough for an investigative paper to validate the claims.

Author response:
We agree that only 2 datasets is a limitation of our work, hence we extended our experiments to verify our core results on TotalSegmentator, which is the largest annotated dataset in the landscape of 3D medical image segmentation. However, we want to emphasize that we want to investigate whether transformers are feasible given the data that is present in 3D medical image segmentation. Hence we chose the next two best large, non-trivial datasets representing both pathology and organ segmentation. We chose these large datasets, as they represent the datasets where transformers should bring the largest benefits over CNNs, and as they enable our dataset-size ablations.

To show the promised Total segmentator results:

| Method      | Standard     | Whole Block Replaced |
|-------------|--------------|----------------------|
| CoTr        | 89.88 (0.21) | 89.09 (0.28)         |
| SwinUNETR   | 87.08 (0.35)  | 85.15 (0.21)         |
| TransBTS    | 87.74 (0.38) | 86.31 (0.78)         |
| UNETR       | 80.79 (0.41) | 72.86 (0.56)         |
| nnFormer    | 85.30 (0.28) | 80.37 (0.14)         |
| nnUNet (3d) | 89.04 (0.11) |                      |

The missing values are currently being evaluated but we will include them in Table 2 left.

Reviewer comment (Novelty):
Some of the findings by the authors were already identified in the ViT paper in 2020, such as the Transformer will be better when facing a larger dataset but might be worse when having a small dataset such as in medical imaging.

Author response:
We agree that the findings in section 2) and 3) are closely related to the ViT paper’s findings that the dataset size matters, however this did neither stop any of the authors to apply transformers to the 3D medical domain, nor did this answer the question ‘Do we have enough data in the 3D medical domain to make transformers work?’ – Since there is a domain gap, 3D volumetric data may be worth ‘more’ hence datasets may be large enough. We show this is currently not the case for any of the prominent architectures on the largest public datasets existing.

Reviewer comment (Experiments):
Observation 2 in section 3.1 might not be too meaningful since convolutional blocks are all removed in up/down sampling paths. Replacing transformer blocks with convolutional blocks might be more fair.

Author response:
We assume the Reviewer mistyped transformer blocks instead of convolutional blocks, as we do not touch convolutions but only the transformer parts of the architecture. Nonetheless, this seems to be a misunderstanding. In our experiments we do not cut off any of the low-resolution paths with our replacements. For 7 of 9 architectures we did not have to apply any special rules to keep the information flow of up/down sampled branches intact. For the remaining two, nnFormer and UTNet,  we had to add slightly more complex behaviour. For junctions in the upsampling branch of nnFormer we joined the two incoming inputs via addition (in the Windowattention). 2) In UTNet when upsampling happens  we add a 1x1 Conv to project to the correct (lower) channel dimension of the higher resolution featuremaps and follow it by bi/tri-linear upsampling. When designing these manual interventions we tried to keep the complexity of the modules to a minimum while still enabling information flow of all branches coming in.

Reviewer comment (Presentation):
What does the tick/cross mean in the right table of Figure 3? It would be more clear to add a description directly in the caption.

Author response:
The ticks in the figure highlight which datasets the authors develop/report/evaluate on in their paper. It highlights that the majority of datasets used in these works are small datasets, which are more prone to noise. We add this to the caption of the figure to convey this more clearly.

Reviewer comment (Experiments):
It would be more interesting to discuss how sensitive the transformers are to hyperparameters and summarize the common practice of selecting a reasonable set of hyperparameters.

Author response:
We agree that this is also an interesting notion and may contribute to the issue of transformers in the domain. However, we maintain that this is slightly beyond the scope of this work given the diverse nature of the datasets.