Free Lunch in Pathology Foundation Model: Task-specific Model Adaptation with Concept-Guided Feature Enhancement

We sincerely appreciate your detailed and constructive comments. We will address your concerns as follows:

1. Aligning patches with tasks

Thanks for your comments. We address your concerns from the following aspects.

• The motivation behind patch-level alignment in this work is to enhance the patch features extracted by the pathology VLM to adapt them to specific downstream tasks and make the CATE more general and able to be incorporated with any MIL method. The alignment is based on the information bottleneck principle and does not explicitly rely on the patch labels. For patches that don't contain task-relevant information, the CIB module will filter out the task-agnostic information, making the alignment more accurate and less noisy.

• We agree that slide-level alignment is more straightforward. However, aligning only at the slide level may be difficult, as the slide-level representation (aggregated from patch features) is noisy without patch-level feature enhancement, hindering the model’s training, as shown in the following table. Incorporating the slide-level alignment with our patch-level alignment can further improve the model’s performance, which demonstrates the importance and rationale of patch-level alignment in the CATE.

2. Unclear improvements

Thanks for your comments. The filtering step alone, without concept alignment, does not provide any benefit, and the increased number of parameters in the CIB and CFI modules does not directly lead to performance improvement.

• The filtering step is used to obtain a subset of the most task-relevant patch features to align with concepts and supervise the training of the information bottleneck. During the training of MIL and the prediction process, all patch features are input to the CIB module to filter task-agnostic information. Without concept alignment, the CIB module degenerates into an encoder of VAE and does not provide any benefit.

• To show this, we conducted experiments on CATE-MIL without concept alignment (discarding PIM loss and SIM loss of the CIB module) and replaced the CIB module with an MLP to investigate the effect of concept alignment and the increased number of parameters. The results are shown in the following table. The performance of CATE-MIL significantly decreases in both cases, demonstrating the importance of concept alignment in the CIB module and that the improvements of CATE are not due to the increased number of parameters. We will add this results in the revised manuscript.

3. Task concept alignment and applicability to more complex MIL problems

• As we pointed out in the original paper, we agree that the performance of CATE relies on the quality of the pathology VLM and it may not be difficult to be applied to other analytical tasks that are difficult to describe with prompts. However, we argue that the main contribution of this work is that it offers a novel and promising idea to adapt the pathology VLM to downstream tasks by leveraging the inherent consistency between image and text in pathology VLM.

• The proposed CATE can benefit more complex tasks beyond subtyping, such as Gleason grading. Experiments on the TCGA-PRAD dataset show that CATE enhances Gleason grading performance, demonstrating that CATE does not affect the aggregation of patch features in the MIL model.

• In the future, as more studies reveal the connection between morphological features and molecular biomarkers and more powerful pathology VLMs are developed, our framework has the potential to benefit more complex tasks.

Thanks to the authors for clarifying the patch-filtering step and the ablations. Agree increased complexity can lead to overfitting for ABMIL, although it isnt super clear if CATE-MIL w/o concept alignment has more params overall as the classifiers and attention layers in ABMIL now operate on bottleneck dimension which is lower than original encoder dimension. The results also arent super consistent as replacing CIB with MLP helps performance in some cases.

I appreciate the authors running these additional experiments and ablations. I think its worth adding these to the paper.

I still feel the the step of aligning patches to tasks is noisy and depends a lot of the quality of the VLM, the nature of the tasks (tasks where there is clear morphological feature descriptions), these being represented in the VLM pre-training and the quality of prompts used for generating task/concept anchors.

However I think this provides a useful way to incorporate domain specific knowledge into WSI classification problems when a well trained pathology VLM is available. Based on this I am changing my decision from 4 to 5.

We deeply appreciate your positive comments and valuable suggestions. We would like to address your concerns below one by one:

1. The quality of concept anchors

Thank you very much for your understanding. As we discuss in the paper, we agree that CATE's performance highly relies on the quality of concept anchors. Fortunately, some operations can be done to improve the quality of the prompts and concepts.

For example, to ensure the quality of the concepts, we can ensemble multiple prompts with various templates and use their average embedding for each class to generate robust and stable class concepts, reducing the risk that the concepts cannot fully capture the characteristics of the class. Additionally, the quality of the concepts will be further improved with the development of pathology research and pathology VLMs.

2. Evaluation metrics

Thank you very much for your valuable comments. In the original paper, we follow the experiment settings of previous papers [1,2] and chose AUC as the primary metric and accuracy as the secondary metric. We agree that the accuracy metric can be misleading in imbalanced datasets. Therefore, we mainly focused on the AUC metric in the original paper. To further improve the reliability of the results, we will supplement the F1-score metric in the revised manuscript.

3. Dataset split strategy

Thank you for your suggestion. We would like to provide more relevant details. We define the in-domain (IND) and out-of-domain (OOD) data based on the source sites of the TCGA dataset. Specifically, each dataset in the TCGA project contains samples from different source sites. The different source sites have different staining and imaging characteristics, causing feature domain shifts among different sites. Therefore, MIL models trained on one site may not generalize well to others. To better evaluate the model performance, we report the testing performance on IND data (in-domain, the testing and training data are from the same sites) and OOD data (out-of-domain, the testing and training data are from different sites).

For the BRCA dataset, we randomly selected one or two sites with samples from two categories as IND data and used the remaining sites as OOD data. For NSCLC (2 categories) and RCC (3 categories), each site contains samples from only one category. We randomly selected one or two corresponding sites for each category as IND data and used the other sites as OOD data, resulting in 1 or 2 IND sites for BRCA, 2 or 4 for NSCLC, and 3 or 6 for RCC. For the IND data, we also randomly split it into training, validation, and testing sets; training the models on the IND training set, and evaluating them on both IND testing sets (IND performance) and OOD data (OOD performance).

4. Definition of class-agnostic concepts

In our paper, class-agnostic concepts refer to the common tissues that are agnostic to the classification task. For example, in the cancer subtyping task, different subtypes have subtype-specific attributes directly related to the task, known as class-specific concepts. Meanwhile, there are many pieces of information in the WSI that are agnostic to different subtypes, such as adipose tissue, connective tissue, and normal tissues. We define these common tissues as class-agnostic concepts.

5. The intuition behind maximizing the distance

Thanks for your comments. We understand your concern. However, the model will not become too specialized and prone to overfitting the training data. The class-specific and class-agnostic concepts are general across different domains and are not specific to the training data. By bringing the enhanced features closer to the class-specific concepts and further away from the class-agnostic concepts, the model can avoid distractions unrelated to the classification task and focus on the general task-relevant concepts. This prevents overfitting on the training dataset, and the maximization of the distance between enhanced features and class-agnostic concepts helps the model generalize better to out-of-domain data.

6. Potential applications beyond subtyping

Thank you for your valuable suggestions. We have conducted experiments to show that CATE can also benefit to other more complex tasks beyond subtyping like Gleason grading, as shown in the following table.

Other tasks like mutation prediction are indeed more challenging compared to subtyping. In the future, as more studies reveal the connection between morphological features and molecular biomarkers or patient prognosis, we will continue to explore whether our CATE can be applied to a wider range of pathology analysis tasks.

[1] Transmil: Transformer based correlated multiple instance learning for whole slide image classification. NeurIPS 2021.

[2] Scaling vision transformers to gigapixel images via hierarchical self-supervised learning. CVPR 2022.

We highly appreciate your positive feedback and constructive suggestions. We would like to address your concerns as follows:

1. More diversified downstream tasks

Thank you very much for your valuable comments. As we discussed in the paper, we agree that the performance of CATE relies on the well-presented concepts generated by pathology VLM. For other tasks such as survival prediction, the concepts may be more complex and hard to write, which may limit the applicability of CATE. However, CATE has the potential to benefit more complex tasks beyond subtyping like Gleason grading. We have conducted experiments and shown that CATE can enhance Gleason grading performance. We will add the complete results in the revised manuscript. In the future, as more studies reveal the connection between morphological features and molecular biomarkers, our framework has the potential to benefit more complex tasks.

2. Dependence on VLM encoders

Thank you very much for your kind comments, we believe that with the development of pathology research and pathology VLMs, pathology VLMs will have broader prospects, and more powerful VLMs will be developed, reducing the gap between the pure pathology foundation models.

3. Explanation of OOD

Thank you for your valuable suggestion. In this work, OOD represents out-of-domain. We use the datasets from the TCGA project, which contain samples from different source sites. The diverse source sites may have different imaging technologies and staining protocols, and each site can be regarded as a domain, causing feature domain shifts between different sites. Therefore, we select several sites for training, referred to as in-domain (IND) data. Consequently, data from other sites with different feature distributions are called out-of-domain (OOD) data. We will add this explanation to the revised manuscript.

4. Hybrid approach of different foundation models

Yes, in our preliminary experiments, we have considered a similar hybrid approach that uses a more powerful vision foundation model to extract image features and a more powerful text encoder to extract text features, followed by a bridge layer to project these images and text features into a shared embedding space. This approach could bring more flexibility to the model design and improve applicability to more challenging tasks. However, the end-to-end training of such a hybrid model is more complex and requires much more computational resources, which is beyond the scope of this work. In the future, we will explore more parameter-efficient fine-tuning strategies and continue to explore this direction.

5. Distribution alignment in equitation (8)

We apologize for the misunderstanding. In fact, the primary goal of SIM is to minimize superfluous information, which can be seen as a form of regularization and is task-agnostic. It does not need to align directly with the Bernoulli distribution of ABMIL. Additionally, we greatly appreciate your suggestion; using the Poisson distribution to define the random occurrence of a non-informative patch is indeed a very promising idea. We will explore this in our future work.

6. Minor comments

Thank you very much for your valuable suggestions. We will adjust the figure and correct the equation in the revised manuscript.

We sincerely thank you for your insightful comments and constructive suggestions. We would like to address your concerns below:

1. Dependence on text prompts and principled ways to obtain prompts

In our experiments, in addition to the expert-designed 'class-specific prompts' from the original CONCH article, we also employed GPT-4 to generate 'class-agnostic prompts', asking questions such as 'In addition to tumor tissues, what types of tissue or cells are present in whole slide images of breast cancer?' The quality of GPT-4 generated prompts has been demonstrated in several recent studies [2,3]. In principle, we can use LLM (e.g., GPT-4) to generate reliable expert-designed prompts and further verified by pathologists. This strategy can ensure the scalability and reliability of the prompts.

2. Experimental Settings

• As we discussed in the general response, each dataset in the TCGA project contains samples from different source sites. The different source sites have different staining and imaging characteristics, causing feature domain shifts among different sites. Therefore, MIL models trained on one site may not generalize well to others. To better evaluate the model performance, we report the testing performance on IND data (in-domain, the testing and training data are from the same sites) and OOD data (out-of-domain, the testing and training data are from different sites).

• For the BRCA dataset, we randomly selected one or two sites as IND data and used the remaining sites as OOD data. For NSCLC (2 categories) and RCC (3 categories) datasets, each site contains samples from only one subtype. Therefore, we cannot select only one site as IND data, as it will include one category/subtype in the training data. Instead, we randomly selected one or two corresponding sites for each category as IND data and used the other sites as OOD data, resulting in 1 or 2 IND sites for BRCA, 2 or 4 for NSCLC, and 3 or 6 for RCC.

3. Performance under a more traditional setting and gain of CATE

• As mentioned above, for NSCLC and RCC, different categories correspond to different sites. In a traditional setting (i.e., training and testing data are from the same sites), MIL models tend to use site-specific features (e.g., staining) as shortcuts for high performance, rather than identifying useful class-specific features, making performance less reflective of the models’ actual capability. Therefore, we did not report the results of the traditional setting on NSCLC and RCC. If necessary, we can provide these results in the comments.

• The gain of CATE on OOD performance is closely related to the performance of the baseline MIL model. As the performance of MIL models has been very high (over 0.95 OOD-AUC for NSCLC) with more training data, further enhancement becomes harder, reducing the gain from CATE. However, CATE’s benefit remains positive. As shown in Table 5 in the Appendix, when all sites are used as IND data for the BRCA dataset (under traditional settings), CATE-MIL still outperforms vanilla ABMIL.

4. Attention Maps

We argue that although the attention maps of CATE-MIL and ABMIL are similar, they are not identical. CATE enhances pre-extracted patch features by preserving task-relevant information and eliminating task-agnostic information, which can strengthen the attention scores of cancerous patches and reduce the scores of non-cancerous patches. As shown in the more detailed attention maps in the accompanying PDF, some non-cancerous patches mistakenly assigned high attention by ABMIL are corrected by CATE-MIL (e.g., the upper left area of the second sample), demonstrating CATE-MIL's ability to correct ABMIL's incorrect predictions.

5. Underperformance in NSCLC and RCC

• Site-specific information such as staining style and imaging characteristics is not relevant to the WSI analysis tasks, while it can be used to distinguish different sites. For the NSCLC and RCC datasets, as each site contains samples from only one category/subtype, the model can use these site-specific features as shortcuts for subtyping. In other words, the model learns how to distinguish different sites, not different subtypes. Our CATE-MIL eliminates these site-specific features, so the performance is lower than other MIL models using these spurious site-specific shortcuts. However, our method is more generalizable, as evident by the OOD performance.

• This issue is not very severe in the BRCA dataset, as each site contains samples from two different categories simultaneously. The MIL models need to learn to classify based on task-relevant features rather than site-specific ones. Therefore, CATE's elimination of site-specific patterns does not decrease BRCA's IND performance but improves it by enhancing task-relevant features.

6. Comparison with ref [1]

The comparison results between CATE with ref [1] are shown below. It can be observed that ref [1] underperforms vanilla ABMIL on OOD data for NSCLC and RCC. This shows that ref [1] overfits IND data and focuses on shortcuts rather than task-relevant features. Furthermore, it involves three separate stages, making it time-consuming and resource-intensive.

[1] Task-specific fine-tuning via variational information bottleneck for weakly-supervised pathology whole slide image classification. CVPR 2023.

[2] The rise of ai language pathologists: Exploring two-level prompt learning for few-shot weakly-supervised whole slide image classification. NeruIPS 2023.

[3] Generalizable Whole Slide Image Classification with Fine-Grained Visual-Semantic Interaction. CVPR 2024.