Rethinking Transformer for Long Contextual Histopathology Whole Slide Image Analysis

Summary

The authors provide LongMIL, a hierarchical and hybrid of local and global attention mechanism, to address the inherent low-rank bottleneck of MIL problems in computational pathology. Through extensive evaluations across feature encoders and subtyping/survival tasks, the authors indeed demonstrate the superior performance as well as computational efficiency of LongMIL.

Strengths

Despite the recipes for LongMIL being simple (Masking of attention), I find the problem well-motivated and solution intuitive enough to be implemented in future MIL studies. It's been long well-known that TransMIL, although being the first self-attention mechanism in MIL, have not been demonstrating good performance and was in need of some alternative implementations. I also appreciate the fact that the authors performed extensive ablation studies over several different feature encoder choices as well as tasks of different nature (subtyping and survival), to truly show that LongMIL can be a meaningful contribution to the field.

Weaknesses

There are several weaknesses of the studies that I think the authors need to address for this to be a meaningful contribution to the field. **Novelty**: Although the authors tried hard to distance from HIPT, I still consider LongMIL solution to be very similar to HIPT - For HIPT, if the first patch-level stage ViT is replaced with the pretrained feature extractors, wouldn't this be same solution as LongMIL, with the difference being how ROI regions are masked? Can authors expand on this point? **Motivation**: While I agree that the low-rank nature of MIL problem is problematic due to n >> d, I am not entirely convinced that "the representation ability of self-attention is limited by the low-rank bottleneck, thus vanilla Transformer based model in WSI analysis suffers sub-optimal performance" is always the case. 1. There exists lots of morphological redundancy in WSI, so the effective number of distinct features might be much lower than actual n [1], [2]. Therefore low-rank might not always be the issue? Can authors expand on this point? 2. To concretely show n>>d is indeed the issue, the authors should also evaluate their algorithm on tissue biopsies (not tissue resections), where the number of patches would be way lower (few thousand patches) or TMAs (few hundred patches). On these datasets, the gap between LongMIL and other frameworks should decrease, since this is not a low-rank setting. **Presentation**: Although the paper was not hard to follow, there are several items that needs improvement. There are lots of typos in the paper that need to be ironed out (e.g., line 157 "are got", line 172 "is design", line 298 "which may because"). I am not sure if Figure 2 contributes meaningfully to the paper (also it's impossible to read the axial labels). My suggestion would be to make it smaller and use extra space for the experiments section. Same goes for Equation 12 - I think this can be moved to supplemental section. The "pre-training backbones" section (line 263~276) was hard to follow. Perhaps make it as a table? **Experimentation**: To follow up on the Motivation section, I think 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. This could include larger patch size (256->512) or random sampling patches, both of which have been used in literature and results in lower number of patches. I think the authors emphasize the survival experimentation over the subtyping (BRACS), since prognosis is known to depend on context [3], [4]. The authors might also consider using latest pathology foundation models (all of these have not been trained on TCGA) - PLIP, UNI, GigaPath - for future studies (It will be too much to do in the given time) References [1] Song, Andrew H., et al. "Morphological prototyping for unsupervised slide representation learning in computational pathology." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2024. [2] Vu, Quoc Dang, et al. "Handcrafted Histological Transformer (H2T): Unsupervised representation of whole slide images." Medical image analysis 85 (2023): 102743. [3] Lee, Yongju, et al. "Derivation of prognostic contextual histopathological features from whole-slide images of tumours via graph deep learning." Nature Biomedical Engineering (2022): 1-15. [4] Jaume, Guillaume, Andrew H. Song, and Faisal Mahmood. "Integrating context for superior cancer prognosis." Nature Biomedical Engineering 6.12 (2022): 1323-1325.

Questions

- Since n >> d is MIL model-agnostic, would other approaches such as AB-MIL also suffer from low rank nature? - It seems LongMIL supersedes existing positional encodings (Or am I undsteranding this correctly?) - Can they be combined? - Equation 11, p -> p_{i,j}. Otherwise, readers might confuse it as being uniform probability.

Rating

6

Confidence

5

Soundness

3

Presentation

2

Contribution

3

Limitations

Please see Weakness & Questions.

Summary

This paper focuses on the issue of attention computation for long sequences in WSI (Whole Slide Image) images. The authors first analyze how the low-rank nature of the long-sequence attention matrix constrains the representation ability of WSI modeling. They then propose a method using local attention masks to compute attention within local regions, followed by the computation of global attention. Experimental results demonstrate that the combination of local and global attention computations outperforms full attention.

Strengths

1. The computation of attention for ultra-long sequences is a significant challenge in WSI slide-level feature learning, and addressing this issue is highly valuable. 2. The paper provides a detailed analysis of the low-rank and sparsity problems in the attention matrix of long sequences, based on which the use of local attention is proposed. 3. The authors propose the longMIL method, which achieves better results than the baseline in both subtyping and prognostic tasks.

Weaknesses

1. 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). 2. The paper lacks a review of relevant literature, such as the aforementioned works. 3. 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. 4. Although the paper claims to reduce computational costs, which is evident, it should provide corresponding comparisons to substantiate this claim.

Questions

see weakness

Rating

3

Confidence

5

Soundness

2

Presentation

2

Contribution

2

Limitations

no

Summary

The authors point out that that MIL often has insufficient ability to offer accurate slide level classifications. There is a long (now) history of attempting to better consider sub-slide level context in the aggregation function. TransMIL, GNN-based methods all have provided attempts to this end. The authors argue that transformer based aggregatioins functions have limited ability to consider both local and global context for tile-attention-rankings. They provide theoretical arguments for why a "low-rank bottleneck" exists because the total number of patch embeddings is much greater than the embedding size. The method addresses this limitation but calculating local self attention followed by pooling function before doing a "global" self attention function. In main paper experiments are done with BRACS dataset using two different encoders for feature extraction. The proposed method is benchmarked against cotext aware and non-context aware aggregations functions. Similar experiment is done with survival predictions. Extensive additional supporting experiments are in supplement.

Strengths

The paper is well written and thorough. The authors are very meticulous about addressing the problem proposed with traditional transformer based aggregators. The claims are backed by the results. The supplement provides extensive figures and ablation studies as well as additional details (including memory efficiency studies) on the methodologies.

Weaknesses

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 publically available encoder (UNI, Gigapath and Virchow). All of these are very recently released so is understanable 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.

Questions

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.

Rating

7

Confidence

5

Soundness

3

Presentation

4

Contribution

4

Limitations

I concur with listed limitations. I have pointed out other limitations in the weakness section.

Summary

This work examines the problem of extrapolating Transformer attention to long sequences in WSI representation learning. The main technical contribution is in examining the low-rank bottleneck problem of Transformer attention for WSIs, and proposing LongMIL which introduces modifications via local attention masking + 2D ALIBI in order to improve the rank of the attention matrix and enable extrapolation capabilities.

Strengths

- Core contribution of this work is novel and would have a lot of interest in the CPath community. Extrapolating to long contexts is an exciting problem that has had little investigation (outside of Prov-GigaPath). I believe this work can be re-organized in presenting a more systematic understanding of they key components needed to extrapolating to long contexts. - Supplement includes some interesting ideas and ablation experiments. A.4 discusses similarities and differences between HIPT, adding local mask attention, and 2D alibi. A.6.1 includes a comparison with the ViT-S in HIPT for equivalent comparisons. A.6.3 examines different hyper-parameters during Transformer training. A.6.4 looks at the difference between 20X and 40X magnification. A.6.6. ablates other straightforward extensions of Transformer attention with subquadratic complexity (including Mamba and V-Mamba). - Figures are illustrative (in both main text and Supplement).

Weaknesses

- Main limitation of this work is that a comparison with Prov-GigaPath [1] (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. - 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 omit the ResNet-50 embedding for survival prediction since it get quite low and unacceptable results." Informal and non-scientific. - - Section 4.3 should be able to summarize the findings in the Supplement in a more clear and descriptive manner. - - One of the main issues I see in this work is that the authors are juggling many different encoders for different tasks. A.6.1 compares ViT-S Lunit and HIPT for TCGA-BRCA subtyping (evaluating multiple MIL models across multiple encoders on 1 task), but A.6.2. shows Resnet-50 features for BRACS and TCGA-BRCA subtyping (evaluating multiple MIL models across multiple tasks with the same encoder). Many of these issues can be drastically simplified if the authors were to use an encoder not pretrained on TCGA such as Prov-GigaPath [1], UNI [2], or PLIP [3] - with the main emphasis on comparing LongMIL with other competing works. Ref [4] provides an example on how the findings of this work can be better organized. References 1. Xu, H., Usuyama, N., Bagga, J., Zhang, S., Rao, R., Naumann, T., Wong, C., Gero, Z., González, J., Gu, Y. and Xu, Y., 2024. A whole-slide foundation model for digital pathology from real-world data. Nature, pp.1-8. 2. Chen, R.J., Ding, T., Lu, M.Y., Williamson, D.F., Jaume, G., Song, A.H., Chen, B., Zhang, A., Shao, D., Shaban, M. and Williams, M., 2024. Towards a general-purpose foundation model for computational pathology. Nature Medicine, 30(3), pp.850-862. 3. Huang, Z., Bianchi, F., Yuksekgonul, M., Montine, T.J. and Zou, J., 2023. A visual–language foundation model for pathology image analysis using medical twitter. Nature medicine, 29(9), pp.2307-2316. 4. Park, N. and Kim, S., How Do Vision Transformers Work?. In International Conference on Learning Representations.

Questions

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. Rating at the time of reviewing this work is slightly negative, but am enthusiastic of this work and would raise my rating to borderline / weak accept if my concerns were addressed.

Rating

6

Confidence

5

Soundness

3

Presentation

2

Contribution

3

Limitations

N/A

response

Since the two-stage method is adopted in this paper, the first stage has used the pre-trained model to extract features offline. The second stage feature dimension has been smaller. Usually the number of effective patches (256) for a WSI image is between a few thousand and ten thousand, and the second stage network is smaller. In the experiment, the second stage of the network usually only needs a few simple transformer layers, which can be completed in a few minutes or ten minutes of training. In view of this, if the local attention that cannot be proposed in this paper is not to solve the performance bottleneck and other effects of full attention caused by too many tokens, it cannot exceed full attention (I think too many tokens will lead to inefficient learning, full attention is not necessarily better than local attention performance), and I have doubts about the actual value in terms of solving computational efficiency.

Response

Thank you for addressing my concerns - I trust the authors in the revised version will do their best to address the clarity issue I raised. I increase my score accordingly.