Rethinking Transformer for Long Contextual Histopathology Whole Slide Image Analysis

Dear Reviewer UesU,

We appreciate your time and valuable feedback. We are glad that you found that our method is well-motivated and intuitive. Below, please find our point-to-point response to your comments:

W1: Novelty: Although the authors tried hard to distance from HIPT, I still consider LongMIL solution to be very similar to HIPT ...

For the comparison to HIPT, please find the A.4 and Figure 6 of paper manuscript for illustration. Here we provide some details:

• HIPT first slice the whole image into regions (4096x4096, ~50 regions per WSI), then each region r_4096 is sliced into patches (256x256, 256 patches per region). Each patch p_256 are extracted by a ViT as feature, the patch-level operation is same as ours and other methods.
• However, their self-attention on extracted patch features only focuses inside on each region r_4096 with a Transformer layer. As pointed out in our paper (line 56-57, 180-183), the adjacent patches may be separated into two regions, and the interactions between them are ignored in HIPT. The features of each region r_4096 is processed by further pooling, and a higher Transformer layer globally focus on slide-level.
• Conversely, our method do not need 4096 region slicing but use 2-d attention mask to treat all patches equally, thus all patches will interact with other adjacent patches.

W2.1: Motivation: .... There exists lots of morphological redundancy in WSI, so the effective number of distinct features might be much lower than actual n. Therefore low-rank might not always be the issue ...

We find this question being very constructive to our paper. Though in patch-level the morphology and feature semantic is quite similar, in region-level, similar patches with different spatial combinations may construct different tumor type (carcinoma in situ V.S. invasive cancer). Some previous methods, e.g. DSMIL [1] [2] try multi-scale information to solve this problem, which indicates the importance. Our method can reach similar goal as multi-scale by adaptively local-global interactions in a uniformed transformer, which we believe to be more elegant.

W2.2 + W4.1: To concretely show n>>d is indeed the issue... evaluate their algorithm on tissue biopsies (not tissue resections), where the number of patches would be way lower ... On these datasets, the gap between LongMIL and other frameworks should decrease, since this is not a low-rank setting. .... the authors could run few more ablation studies to demonstrate the severity of low rank issue, by trying to reduce the gap between n and d ...*

We run some experiments on larger patch size 448 on BRACS (224 originally), where the largest patch num n is less than 2k, near the feature size of UNI (given the limited time, we will further run biopsies and TMAs in next version). The results are shown in the following table, with main findings that:

1. Simple attentions (ABMIL, CLAM without pair wise interactions) gain improvement, and we speculate that the larger image-size can modelling the local context better.
2. DTFD try to split the whole bag into 3 sub-bags, but smaller bag size may result in larger label noise of sub-bags which may answer its performance drop.
3. The gap between LongMIL and TransMIL decreases given closer n and d. Full attention and LongMIL show small drops, since less interactions can be modelled with less patches.
4. LongMIL still out-performs full attention. We speculate that local attention also works better when dealing with the shape-varying WSI even with less n.

• BRACS, 224x224 VS 448x448 | Patch Encoder | UNI-224| | UNI-448| | |:---|:---|:---|:---|:---| | Method\Metric | F1 | AUC | F1 | AUC | | AB-MIL | 0.692±0.03 | 0.875±0.02 | 0.695±0.01 | 0.875±0.01 | | CLAM-SB | 0.640±0.06 | 0.844±0.03 | 0.654±0.03 | 0.851±0.02 | | DTFD-MIL | 0.655±0.03 | 0.878±0.02 | 0.625±0.03 | 0.839±0.01 | | TransMIL | 0.592±0.04 | 0.859±0.02 | 0.646±0.07 | 0.855±0.02 | | Full Attention | 0.715±0.04 | 0.884±0.02 | 0.700±0.04 | 0.874±0.02 | | LongMIL (ours) | 0.728±0.05 | 0.887±0.01 | 0.722±0.03 | 0.883±0.01 |

W3 + Q3: Presentation: ...

We appreciate your advice and will definitely focus on refining the paper's presentation for the next version to enhance clarity.

W4.2: ... consider using latest pathology foundation models ...

We have finished part of this experiments in this stage, which can be inferred in the general response.

Q1: Since n >> d is MIL model-agnostic, would other approaches such as AB-MIL also suffer from low rank nature?

For AB-MIL, CLAM, etc., their attentions are in fact adaptive weighted average ($1 \times n$) over all patches, and there is no pair-wise interaction like self-attention does ($n \times n$ attention matrix). It's hard to say the rank of a $1 \times n$ array, but we find that they are sparse like what doctors do: determine the lesion level of a slide based on some specific lesion regions. However, this process needs quite high semantics align with doctors. Conversely, our method with self-attention can be understood as trying to convert the feature representation (learned by patch-level self-supervised learning) aligning better to doctors' diagnosis level by introducing important context. After layers of attention, cls-token and average pooling on these features will do the same thing as AB-MIL.

Q2: It seems LongMIL supersedes existing positional encodings (Or am I understanding this correctly?) - Can they be combined?

We have combined the rotary positional embedding, since we still need to determine the position information within the local areas (about a size of $20 \times 20$). We feel sorry to ignore this implementation detail and will add it into the next version.

Ref:

[1]. Dual-stream multiple instance learning network for whole slide image classification with self-supervised contrastive learning.

[2]. Cross-scale multi-instance learning for pathological image diagnosis.

Dear Reviewer Ghwr,

We appreciate your time and constructive feedback. We are glad that you found that our analysis and method are valuable. Below, please find our point-to-point response to your comments:

W1 & W2: Although the paper identifies and analyzes the bottleneck issues in long-sequence attention and proposes the use of local attention based on this analysis, there is a lack of innovation in using local attention. Other methods, such as LongViT, LongNet, and Prov-GigaPath, have also used local attention more elegantly. Additionally, the paper does not discuss the potential for directly transferring numerous attention optimization methods from the fields of computer vision (CV) and natural language processing (NLP). The paper lacks a review of relevant literature, such as the aforementioned works.

We have great respect for the significant contributions made by Prov-GigaPath, LongViT, and LongNet, as well as the innovative dilated-attention method that underpins them. These advancements have brought substantial progress to the field of Computational Pathology. We regret that our initial submission did not adequately acknowledge these important works. The Prov-GigaPath paper, published in Nature on May 22, 2024, coincided with the NeurIPS submission deadline, while the LongViT was available as an arXiv preprint formatted as a technical report and titled as a 'Case Study' thus we missed it. We sincerely apologize for this oversight and have now included extensive discussions on these works in our rebuttal. We are committed to providing a more thorough analysis in the future version of our paper.

For the detailed comparisons of our method and Prov-GigaPath, please infer to the general response, where we make systematical analysis including:

d_1. Method: their receptive field weigh more on x-axis than y-axis, however our method as 2D locality treat x-y equally. Please also check the figure illustration in the rebuttal PDF.
d_2. Contribution: we focus more on analyzing why previous transformers failed then deriving our method, while they empirically scale up to big data based on dilated attention. 
d_3. We find that when their patch feature is not the best in some task cases, their heavily pretrained WSI head with problem in 'd_1' only shows sub-optimal performance. 
We provide some quick experiments to compare their WSI-architecture and our method.

W3: In the experimental results, for instance, the classification results in Table 1 and the prognostic results in Table 2 show only a 0.01 improvement over full attention, which is not significant.

It is not so easy to totally beat the full attention in performance since our main goal is to achieve comparable results as full attention but with less computational complexity. Even in the area of NLP currently, like LLM, the vanilla full attention is the most widely adopted choice [1][2][3] when equipped with the best hardware, although a lot of methods [4][5][6], and also the LongNet, are trying to replace it to solve long sequence problem. However, in the area of digital pathology given less computational resources, the inference speed is quite important. Moreover, we have tested the training and inference speed when dealing with larger resolution WSI or higher magnification like 40x with 0.25 mpp, the latency is unbearable even with flash attention. Whereas our method can highly alleviate the speed problem without performance drop (even improvement) if we need more detailed features in 40x.

W4: Although the paper claims to reduce computational costs, which is evident, it should provide corresponding comparisons to substantiate this claim.

Please check A.6.5 and Figure 8 of our main paper manuscript, where we have tested the speed (time consumption in deployed GPU), and shown the comparisons in paper submission stage.

The theoretical complexity is also included in the line 232-236 of our main paper manuscript.

If you have additional questions, we’re happy to engage further.

Ref:

[1]. Touvron H, Martin L, Stone K, et al. Llama 2: Open foundation and fine-tuned chat models[J]. arXiv preprint arXiv:2307.09288, 2023.

[2]. Jiang A Q, Sablayrolles A, Roux A, et al. Mixtral of experts[J]. arXiv preprint arXiv:2401.04088, 2024.

[3]. Bai J, Bai S, Chu Y, et al. Qwen technical report[J]. arXiv preprint arXiv:2309.16609, 2023.

[4]. Gu A, Dao T. Mamba: Linear-time sequence modeling with selective state spaces[J]. arXiv preprint arXiv:2312.00752, 2023.

[5]. Sun Y, Dong L, Huang S, et al. Retentive network: A successor to transformer for large language models[J]. arXiv preprint arXiv:2307.08621, 2023.

[6]. Peng B, Alcaide E, Anthony Q, et al. Rwkv: Reinventing rnns for the transformer era[J]. arXiv preprint arXiv:2305.13048, 2023.

1. 40x magnification, with about n=40k~70k patches (also shown in Fig. 8), we have test training speed of full attention, which is about 25x to 35x times than ours. We have talked about this in the caption of Fig. 8, and also show some experimental results of performance in Table 5 of A.6.4. Though currently 40x is not the mainstream, we speculate this is caused by 20x captures better context thus performs better using AB-MIL (without context modelling ability) paradigm. Moreover, there is study [1] show better performance via 40x.
2. Overlapped patching, it will be resulting 2~4x patches if the overlap ratio is 0.25~0.5. This helps alleviating the edge effect of image modelling neither by CNN nor ViT.
3. Some slides contain >=2 histology tissues; Survival prediction need multiple slides for one patience.

1. 40x, overlapped patching ... as we mentioned above.
2. pretraining on over 100k slides just like Prov-Gigapath do.
3. slow speed in the clinical or deployed setting, where the GPU hardware is not so good as training.
4. using larger feature embedding d (e.g. carrying both high-level and low-level feature) as we mentioned in the line 207~208 of our paper: 'An intuitive modification to handle the low-rank problem is to set a larger embedding size d, but this makes computational complexity O(n^2 d) more severe...'
5. more transformer layers, as you said.

Dear Reviewer GgQd,

We would like to express our sincere gratitude for your thoughtful review and insightful feedback on our manuscript. We appreciate your recognition of the thoroughness and approach taken in addressing the limitations of traditional transformer-based aggregators. Your positive evaluation of our work's soundness, presentation, and contribution is greatly encouraging.

Below, we provide responses to your comments and questions:

W1: A recent preprint (https://arxiv.org/abs/2407.07841) shows that context aware aggregations functions offer less performance boost over ABMIL when you have a high feature extraction encoder. In this work the encoders used are less robust than the now publicly available encoder (UNI, Gigapath and Virchow). All of these are very recently released so is understandable that they are not part of this submission. Going forward these encoders should be used for any aggregation function assessment. The supplemental table showing that the boost of this method is much stronger with an image net pretrained cnn.

We find this question being very constructive and insightful to the completeness our paper. On the one hand, we have performed some experiments on stronger patch encoder including UNI and GigaPath in this stage (check it in the general response), with main findings:

1. Our method continues to outperform previous work in complicated tasks such as BRACS and survival prediction, demonstrating a notable performance boost in consistency compared to other methods, with the exception of TransMIL.
2. The results in the TCGA-BRCA tumor subtyping task are similar across almost all the methods. We speculate that this task may have reached an upper limit when using a strong patch encoder. In future versions, we plan to validate our method with robust encoders on larger datasets to further explore the potential for improvement.

On the other hand, motivated by this suggestion, we realize that strong patch encoders with high-level semantics may lose important low-level spatial-context or fine-grained details. Thus, in future work we may be going to aggregate more features from different depth of layers in ViT or different pretrained ViT. This may not only include more detailed spatial contexts, but also helps improving the feature size $d$.

Q1: In final version of paper, can you please improve orientation of figure 3. It is hard to zoom in sufficiently to understand second panel. It would also benefit from having sub labels (eg. a, b, c) to improve the legend description of the sub panels.

Thank you for your detailed feedback on Figure 3. We understand that the current orientation and labeling make it difficult to interpret, particularly the second panel. To address this issue, we will take actions to improve the clarity and readability of the figure in the final version of our manuscript.

Dear Reviewer UUFY,

We appreciate your time and valuable feedback. We are glad that you found the formulation and analysis being novel and sound, and the figures and ablations being illustrative. Below, please find our point-to-point response to your comments:

W1: Main limitation of this work is that a comparison with Prov-GigaPath (a concurrent work that appeared at the time of NeurIPS submission) is warranted. Prov-GigaPath also presents overlapping contributions in solving this problem, though I think there is room for more than 1 study examining this problem.

Since another reviewer Ghwr also points out this problem, we post the detailed comparisons in the general response, where we make systematical comparisons to Prov-GigaPath including:

d_1. Method: their receptive field weigh more on x-axis than y-axis, however our method as 2D locality treat x-y equally.
d_2. Contribution: we focus more on analyzing why previous transformers failed then deriving our method, while they empirically scale up to big data based on dilated attention. 
d_3. We find that when their patch feature is not the best in some task cases, their heavily pretrained WSI head with problem in 'd_1' only shows sub-optimal performance. 
We provide some quick experiments to compare their WSI-architecture and our method.

W2: The writing feels a bit rushed and informal. I was often scrolling back and forth to understand the different comparisons being made. Ideally, there should be one table for results that compare MIL architectures and one table for ablating pretrained encoders. Other areas where the writing / presentation of figures and results could be significantly polished ...

We thank a lot for your suggestions on the presentation of our paper, we will polish it for better reader-experience in next version.

W2.3 + Q1: Would the authors be able to address my concern in updating the results of this work to using a different encoder (to simplify the presentation of results)? Choice of pretrained encoder should not matter significantly (as the main focus is in fairly comparing LongMIL) but having 1-2 comparisons would be nice in a Supplemental Figure.

We indeed find that there are more and more stronger patch encoders pretrained without TCGA should be validated, which is also highlighted by almost all the other reviewers. So, we also add it into the general response, including both UNI and GigaPath pretrained patch encoders with main findings:

1. Our method still outperforms previous work for some complicated tasks like BRACS and survival prediction. 
2. The results in TCGA-BRCA tumor-subtyping are similar for almost all the WSI methods. We speculate that this task as binary classification might have reached some sort of upper limit when equipped with strong patch encoder. We will be going to evaluate our method on more difficult tasks in future version.