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

Authors: yanyan  huang
weiqin  zhao
yihang  chen
yu  fu
lequan  yu
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 10:20:20

Abstract: Whole slide image (WSI) analysis is gaining prominence within the medical imaging field. Recent advances in pathology foundation models have shown the potential to extract powerful feature representations from WSIs for downstream tasks. However, these foundation models are usually designed for general-purpose pathology image analysis and may not be optimal for specific downstream tasks or cancer types. In this work, we present Concept Anchor-guided Task-specific Feature Enhancement (CATE), an adaptable paradigm that can boost the expressivity and discriminativeness of pathology foundation models for specific downstream tasks. Based on a set of task-specific concepts derived from the pathology vision-language model with expert-designed prompts, we introduce two interconnected modules to dynamically calibrate the generic image features extracted by foundation models for certain tasks or cancer types. Specifically, we design a Concept-guided Information Bottleneck module to enhance task-relevant characteristics by maximizing the mutual information between image features and concept anchors while suppressing superfluous information. Moreover, a Concept-Feature Interference module is proposed to utilize the similarity between calibrated features and concept anchors to further generate discriminative task-specific features. The extensive experiments on public WSI datasets demonstrate that CATE significantly enhances the performance and generalizability of MIL models. Additionally, heatmap and umap visualization results also reveal the effectiveness and interpretability of CATE.

Main paper: 
 introduction: 
 Multiple Instance Learning (MIL) [26, 34, 23, 2] is widely adopted for weakly supervised analysis in computational pathology, where the input of MIL is typically a set of patch features generated by a pre-trained feature extractor (i.e., image encoder). Although promising progress has been achieved, the effectiveness of MIL models heavily relies on the quality of the extracted features. A robust feature extractor can discern more distinctive pathological features, thereby improving the predictive capabilities of MIL models. Recently, several studies have explored using pretrained foundation models on large-scale pathology datasets with self-supervised learning as the feature extractors for WSI analysis [37, 7, 3, 36] . Additionally, drawing inspiration from the success of Contrastive Language-Image Pretraining (CLIP) [30, 20] in bridging visual and linguistic modalities, some works have aimed to develop a pathology vision-language foundation model (VLM) to simultaneously learn representations of pathology images and their corresponding captions [15, 25] . The intrinsic consistency between the image feature space and caption embedding space in the pathology VLM enables the image encoder to extract more meaningful and discriminative features for downstream WSI analysis applications [25] .Although the development of these pathology foundation models has significantly advanced computational pathology, these models are designed for general-purpose pathology image analysis and may not be optimal for specific downstream tasks or cancer types, as the features extracted by the image encoder may contain generic yet task-irrelevant information that will harm the performance of specific downstream tasks. For example, as illustrated in Figure 1 (a), the features extracted by the image encoder of a pathology VLM can include both task-relevant information (e.g., arrangement or morphology of tumor cells) and task-irrelevant elements(such as background information, stain styles, etc.). The latter information may act as "noise", distracting the learning process of MIL models tailored to specific tasks, and potentially impairing the generalization performance of these models across different data sources. Consequently, it is crucial to undertake task-specific adaptation to enhance feature extraction of generic foundation models and enable MIL models to concentrate on task-relevant information and thus improve analysis performance and generalization [32, 38] .In this paper, we propose a novel paradigm, named Concept Anchor-guided Task-specific Feature Enhancement (CATE), to enhance the generic features extracted by the pathology VLM for specific downstream tasks (e.g., cancer subtyping). Without requiring additional supervision or significant computational resources, CATE offers an approximately "free lunch" in the context of pathology VLM. Specifically, we first derive a set of task-specific concept anchors from the pathology VLM with task-specific prompts, and these prompts rely on human expert design or are generated through querying large language models (LLMs), necessitating a certain level of pathological background knowledge. Based on these concept anchors, we design two concept-driven modules, i.e., the Conceptguided Information Bottleneck (CIB) module and the Concept-Feature Interference (CFI) module, to calibrate and generate task-specific features for downstream analysis. Particularly, with the taskspecific concepts as the guidance, the CIB module enhances task-relevant features by maximizing the mutual information between the image features and the concept anchors and also eliminates task-irrelevant information by minimizing the superfluous information, as shown in Figure 1(a) . Moreover, the CFI module further generates discriminative task-specific features by utilizing the similarities between the calibrated image features and concept anchors (i.e., concept scores). By incorporating the CATE into existing MIL frameworks, we not only obtain more discriminative features but also improve generalization regarding domain shift by eliminating task-irrelevant features and concentrating on pertinent information, as shown in Figure 1(b) .In summary, the main contributions of this work are threefold:• We introduce a novel method, named CATE, for model adaptation in computational pathology. To the best of our knowledge, this is the first initiative to conduct task-specific feature enhancement based on the pathology foundation model for MIL tasks.• We design a new CIB module to enhance the task-relevant information and discard irrelevant information with the guidance of task-specific concepts, and a new CFI module to generate task-specific features by exploiting the similarities between image features and concept anchors.• Extensive experiments on Whole Slide Image (WSI) analysis tasks demonstrate that CATE significantly enhances the performance and generalization capabilities of MIL models. 
related work: 
 Multiple Instance Learning (MIL) for WSI Analysis. MIL is the predominant paradigm for WSI analysis, treating each WSI as a bag of patch instances and classifying the entire WSI based on aggregated patch-level features. Attention-based methods [16, 26, 17, 39] are highly regarded for their ability to determine the significance of each instance within the bag. For instance, Ilse et al. [16] introduced an attention-based MIL model, while Lu et al. [26] proposed clustering-constrainedattention to refine this mechanism further. To model the relationships among instances, graph-based and Transformer-based methods have been developed [10, 2, 13] . For example, Chen et al. [2] introduced a Transformer-based hierarchical network to capitalize on the inherent hierarchical structure of WSIs.Pathology Foundation Model. With the advancement of foundation models in computer vision, several pathology foundation models have been developed to serve as robust image encoders for WSI analysis. Riasatian et al. [31] proposed fine-tuning the DenseNet [12] on the TCGA dataset, while Filiot et al. [7] utilized iBOT [42] to pretrain a vision Transformer using the Masked Image Modeling framework. Recently, Chen et al. [3] pre-trained a general-purpose foundation model on large-scale pathology datasets using DINOv2 [28] , which has demonstrated strong and readily usable representations for WSI analysis. Inspired by CLIP [30] , Ikezogwo et al. [15] , Huang et al. [14] , and Lu et al. [25] developed vision-language foundation models by training on large-scale pathology datasets with image-caption pairs. These foundation models have demonstrated superior performance in downstream tasks due to their ability to extract more discriminative features for WSI analysis.Feature Enhancement in Computational Pathology. Several methods have been developed to obtain more discriminative features for WSI analysis by adapting pathology foundation models [41, 24] or designing new plug-and-play modules [35] . For instance, Zhang et al. [41] suggested aligning the image features with text features extracted from a pre-trained natural language model to enhance the feature representation of WSI patch images, while it operates solely at the patch level, without considering the informational relationship between image and text features. Recently, Tang et al. [35] introduced Re-embedded Regional Transformer for feature re-embedding, aimed at enhancing WSI analysis when integrated with existing MIL methods. However, while this method considers the spatial information of WSIs and adds flexibility to MIL models, it falls short in extracting task-specific discriminative information for WSI analysis.3 Method 
overview: 
 The proposed CATE can be seamlessly integrated with any MIL framework to adapt the existing pathology foundation model (Pathology VLM) for performance-improved WSI analysis via taskspecific enhancement. Specifically, consider a training set D = {(x, y)} of WSI-label pairs, where x = {x 1 , x 2 , ..., x N } is a set of patch features with dimension of C (i.e., x i ∈ R C ) extracted by the image encoder of pathology VLM, N denotes the number of patches, and y is the corresponding label. The objective of CATE is to obtain the corresponding enhanced task-specific feature set z from the original feature x with the guidance of pre-extracted concepts anchors c (see description below) for downstream usage:EQUATIONAs illustrated in Figure 2 , we design two different modules to enhance the extracted features from foundation models: (1) Concept-guided Information Bottleneck (CIB) module calibrates original image features with the guidance of concept anchors with information bottleneck principle; and(2) Concept-Feature Interference (CFI) module generates discriminative task-specific features by leveraging the similarities between the calibrated image features and concept anchors. Specifically, the enhanced patch features can be represented as z = {z 1 , z 2 , ..., z N }, where z i is the concatenation of the calibrated feature α i and the interference feature β i generated by CIB and CFI module: Concept Extraction. We extract two kinds of task-specific concept anchors, c = {c cs , c ca }, comprising class-specific concepts c cs = {c cs i } m i=1 (e.g., subtyping classes) and class-agnostic concepts c ca = {c ca i } n i=1 (e.g., adipose, connective, and normal tissues), with m and n representing the numbers of class-specific and class-agnostic concepts, respectively. These concepts are generated by the text encoder of pathology VLM with prompt p. Each prompt consists of a class name (e.g., "invasive ductal carcinoma") and a template (e.g., "An image of <CLASSNAME>"). To obtain more robust concepts, we use multiple prompts for each class and the final concept anchor is the average of the embeddings generated by different prompts. Details of class names and templates for various tasks are provided in Appendix G. Note that due to the inherent consistency between the image and text embedding space in VLM, these extracted concepts can also be regarded as image concept vectors.EQUATIONThe objective of this module is to find a distribution p(α|x) that maps the original image feature x into a representation α, which contains enhanced task-discriminative characteristics and suppressed task-irrelevant information. WSIs typically contain various cell types or tissues (e.g., tumor cells, normal cells, adipose tissue, connective tissue), while only a subset of patches (e.g., with tumor cells) is crucial for certain tasks such as tumor subtyping. We thus define x = {x i } k i=1 ⊆ x as the representative subset of the original feature set (e.g., tumor tissue patches), where k denotes the number of representative patches. Note that this selection can be conducted with a simple comparison of image features with class-specific concepts (see discussion below). To this end, the corresponding enhanced feature set is α = { αi } k i=1 ⊆ α and we want to find the conditional distribution p( αi |x i ) to map the selected patch feature xi ∈ R C into a more discriminative enhanced feature αi ∈ R C , which is discriminative enough to identify the label y.Sufficiency and Consistency Requirements. To quantify the informativeness requirement of the calibrated feature α, we consider the sufficiency of α for y. As defined in Appendix E.1, the encoded feature α derived from the original feature x is sufficient for determining the label y if and only if the amount of task-specific information remains unchanged after calibration, i.e., I(x; y) = I( α; y). However, the label y pertains to the slide level and specific labels cannot be assigned to each instance due to the absence of patch-level annotations.To address this challenge, we propose using the task-specific concept anchor as the guidance for each single α. Specifically, we posit that the concept anchor c is distinguishable for the task and contains task-relevant information for label y. Given the consistency between image and text features in pathology VLM, any representation α containing all information accessible from both image feature x and concept c will also encapsulate the discriminative information required for the label. This consistency requirement is detailed in Appendix E.1. Thus, if α is sufficient for c (i.e., I(x; c| α) = 0), then α is as predictive for label y as the joint of original feature x and concept anchor c. Applying the chain rule of mutual information, we derive:EQUATIONAccording to the consistency between the concept anchor and original feature, the mutual information term I( α; c) represents the predictive information for the task, while the conditional information term I(x; α|c) denotes task-irrelevant information (i.e., superfluous information) in original patch feature x, which can be minimized to enhance the robustness and generalization ability of downstream MIL models. As a result, the main objective of the feature calibration in this module can be formalized as maximize predictive information I( α; c) while minimize the superfluous information I(x; α|c).Predictive Information Maximization (PIM). The predictive information in Equ (3) equals to the mutual information between the calibrated feature and concept anchors. To maximize this, we choose the InfoNCE [27] to estimate the lower bound of the mutual information, which can be obtained by comparing positive pairs sampled from the joint distribution α, c cs pos ∼ p( α, c cs ) to pairs α, c cs j and α, c ca j built using a set of negative class concepts c cs j ∼ p (c cs ) and class-agnostic concepts c ca j ∼ p (c ca ):EQUATIONWe set f (c, αi ) = exp αT i c/τ with τ > 0 in practice following [27] . By maximizing this mutual information lower bound, f (c, αi ) will be proportional to the density ratio p (c, αi )/p (c) p ( αi ) as proved in [27] . Hence, f (c, αi ) preserves the mutual information between the calibrated feature and concept anchor. The detailed derivation can be found in Appendix E.2. The loss function for PIM can be denoted as: L P IM = ETo compress task-irrelevant information, we aim to minimize the superfluous information term as defined in Equ (3). This objective can be achieved by minimizing the mutual information I(x; α). In practice, we conduct SIM for all patches in the subset x, as each patch may contain task-irrelevant information. Following [1] , it can be represented as:EQUATIONAfter that, we let the distribution of α: r (α) (e.g., Gaussian distribution in this work), be a variational approximation to the marginal distribution p (α), and we can obtain the upper bound for I(x; α):EQUATIONFurthermore, we use a variational distribution q θ (α|x) with parameter θ to approximate p (α|x) and we implement the parameterization of the variational distribution with MLP by predicting the mean and variance of the Gaussian distribution and sample the calibrated feature α from this distribution:EQUATIONIn practice, we implement this by utilizing the reparameterization trick [18] to obtain an unbiased estimate of the gradient and further optimize the variational distribution. The detailed derivation can be found in Appendix E.3. The minimization of the upper bound of I(x; α) equals to the minimization of the Kullback-Leibler divergence between q θ (α|x) and r (α). Therefore, the loss function can be represented as:EQUATIONDiscussion. We further provide explanation of CIB module with the information plane [8, 6] in Appendix F. It should be noted that the PIM supervises only the representative subset x containing task-relevant information (selected by the similarity between image features and corresponding class-specific concepts). Meanwhile, the SIM is applied to all patches in x, as any patch may carry information irrelevant to the task (e.g., background information and stain styles). Besides, SIM cannot be directly optimized without the guidance of concept anchors (i.e., PIM) due to the absence of patch-level labels. As demonstrated in the ablation study in Section 4.4, the absence of concept anchor guidance leads to the collapse of discriminative information in the calibrated feature, adversely affecting downstream task performance. By maximizing predictive information and minimizing superfluous details, the CIB module effectively enhances the discriminative capacity of the original features and aligns them with the task-specific concept anchors for improved prediction.We also propose the Concept-Feature Interference (CFI) module to utilize the similarity characteristic between calibrated features and concept anchors to further obtain robust and discriminative information for the downstream tasks. Our primary focus is on the class-specific concept anchors c cs . Specifically, for each CIB encoded feature α i , we calculate the cosine similarity between α i and each class-specific concept c cs i . It is important to note that the number of class-specific concepts m is larger than the number of classes, as we use multiple <CLASSNAME> and templates to generate the concept anchor for each class, as shown in the Appendix G. Thus, we can obtain the similarity vector by concatenating the similarity scores between α i and each class-specific concept c cs i . To integrate the interference information (similarity relationship) into the enhanced feature, we align the similarity vector with the calibrated feature α i using a Self-Normalizing Network (SNN) layer [19] . This allows us to obtain the final interference vector β i of CFI:EQUATIONThe interference vector contains superficial information that indicates the similarity between the calibrated feature and concept anchor directly. This is completely different from the calibrated feature of the CIB module, which contains discriminative latent information for the downstream tasks. Therefore, integrating the interference feature can further provide robust and discriminative information for the downstream tasks.Discussion. The CFI module is designed to utilize the similarity characteristic between calibrated feature and concept anchor as a discriminative feature, which can be further integrated into the calibrated feature for downstream tasks. This is different from other studies that directly compare the similarity between visual features and textual concept features of different classes to perform zero-shot classification [25] .The overall training objective of the CATE framework can be represented as the combination of the cross entropy loss L CE for the downstream tasks, the predictive information maximization loss L P IM , and the superfluous information minimization loss L SIM :EQUATIONwhere λ P and λ S are hyperparameters and influence of them is discussed in Appendix C. 
experimental settings: 
 Tasks and Datasets. We conducted cancer subtyping tasks on three public WSI datasets from The Cancer Genome Atlas (TCGA) project: Invasive Breast Carcinoma (BRCA), Non-Small Cell Lung Cancer (NSCLC), and Renal Cell Carcinoma (RCC). Detailed dataset information is available in Appendix D. For the BRCA dataset, we randomly selected one or two sites as IND data and used the remaining sites as OOD data. However, 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 for NSCLC and RCC, and used the other sites as OOD data. Finally, we obtained 1 or 2 IND sites for BRCA, 2 or 4 for NSCLC, and 3 or 6 for RCC.Evaluation. We report the area under the receiver operating characteristic curve (AUC) and accuracy for the OOD and IND test sets, respectively, with means and standard deviations over 10 runs of Monte-Carlo Cross Validation. Notably, the OOD performance is emphasized for NSCLC and RCC, where each site contains samples from only one cancer subtype. Traditional MIL models tend to recognize site-specific patterns (e.g., staining) as shortcuts and excel in in-domain evaluations, rather than identifying useful class-specific features, making performance less reflective of the models' actual capability. Therefore, OOD performance more accurately reflects the models' discriminative and generalization capabilities.Comparisons. Given that CATE is an adaptable method, we evaluated the performance variations across various MIL models both with and without the integration of CATE. We specifically focused on the following state-of-the-art MIL models: the original ABMIL [16] , CLAM [26] , DSMIL [21] , Implementation Details. This study begins the image feature extraction process by segmenting the foreground tissue and then splitting the WSI into 512×512 pixels patches at 20× magnification. Subsequently, these patches are processed through a pre-trained image encoder from CONCH [25] to extract image features. For concept anchors, we utilize CONCH's text encoder to derive task-relevant concepts from predefined text prompts, with detailed prompt information available in Appendix G. Model parameters are optimized using the Adam optimizer with a learning rate of 10 -5 . The batch size is set to 1, and all the experiments are conducted on a single NVIDIA RTX 3090 GPU.Quantitative Results on BRCA Dataset. To fully evaluate the effectiveness of CATE, we assessed its impact on several state-of-the-art MIL models using the BRCA dataset. Qualitative Analysis. To qualitatively investigate the effectiveness of CATE, we visualized attention heatmaps, UMAP, and the similarities between original features and corresponding class concept features, as well as calibrated features and class concept features, as shown in Figure 3 . Additional visualization results are provided in Appendix H. As shown in Figure 3 (a&b), attention heatmap comparisons reveal that CATE-MIL focuses more intensely on cancerous regions, with a clearer delineation between high and low attention areas. By comparing the similarities of original and calibrated features to class concept features in Figure 3 (c&d), it is evident that the enhanced similarity in cancerous regions is significantly higher than in original features. Moreover, the disparity between cancerous and non-cancerous regions' similarities is also expanded, which further verifies the ability of CATE to enhance task-relevant information and suppress irrelevant information. We further performed a UMAP visualization of class concept features, original features, and calibrated features. As depicted in Figure 3 For clarity and to highlight the superiority of ABMIL when enhanced with CATE, we developed CATE-MIL by incorporating CATE into ABMIL and compared it against other leading MIL models on NSCLC and RCC datasets. The comparative results in Table 2 confirm that CATE-MIL consistently outperforms other models in both OOD and IND performance. However, it is noted that CATE-MIL performs poorly on the in-domain testing data for NSCLC and RCC. This underperformance may be attributed to the elimination of task-irrelevant information, including site-specific patterns, by CATE, potentially degrading performance on in-domain data for these datasets. Consequently, OOD performance more accurately reflects the discriminative and generalization capabilities of the models.We conduct ablation studies to assess the effectiveness of each component within CATE, and the results are shown in Table 3 . Initially, incorporating Predictive Information Maximization (PIM) enables ABMIL to achieve improved performance in most experiments, which demonstrates PIM's efficacy in extracting task-relevant information. However, using Superfluous Information Minimization (SIM) alone results in performance degradation across most experiments, which suggests that SIM may discard some task-relevant information without guidance from a task-relevant concept anchor. Incorporating both PIM and SIM consistently enhances ABMIL's performance in all experiments, which further verifies that their combination effectively boosts the generalization capabilities of MIL models. We also conduct experiments by only using the interference features in CFI as the input of ABMIL, and the results show that the interference features are also informative for WSI classification tasks. More ablation analysis about the weights of PIM and SIM in CIB module the number of representative patches are in Appendix C. 
conclusion and discussion: 
 In this paper, we introduce CATE, a new approach that offers a "free lunch" for task-specific adaptation of pathology VLM by leveraging the inherent consistency between image and text modalities. CATE shows the potential to enhance the generic features extracted by pathology VLM for specific downstream tasks, using task-specific concept anchors as guidance. The proposed CIB module calibrates the image features by enhancing task-relevant information while suppressing task-irrelevant information, while the CFI module obtains the interference vector for each patch to generate discriminative task-specific features. Extensive experiments on WSI datasets demonstrate the effectiveness of CATE in improving the performance and generalizability of state-of-the-art MIL methods.Limitations and Social Impact. The proposed CATE offers a promising solution to customize the pathology VLM for specific tasks, significantly improving the performance and applicability of MIL methods in WSI analysis. However, the performance of CATE heavily depends on the quality of the concept anchors, which, in turn, relies on domain knowledge and the quality of the pre-trained pathology VLM. Additionally, while CATE is optimized for classification tasks such as cancer subtyping, it may not be readily applicable to other analytical tasks, such as survival prediction. However, there might be a potential solution to address this challenge. For instance, we could leverage LLMs or retrieval-based LLMs to generate descriptive prompts about the general morphological appearance of WSIs for specific cancer types. By asking targeted questions, we can summarize reliable and general morphological descriptions associated with different survival outcomes or biomarker expressions and further verify these prompts with pathologists. Moreover, since medical data may contain sensitive information, ensuring the privacy and security of such data is crucial. 
experimental setting/details: 
 Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes]Justification: The paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [Yes]Justification: The paper report error bars suitably and correctly defined.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.• The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [Yes] Justification: The paper discuss both potential positive societal impacts and negative societal impacts of the work performed.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: The paper poses no such risks.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 
new assets: 
 Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?Answer: [Yes]Justification: New assets introduced in the paper well documented and is the documentation provided alongside the assets.Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file. 
crowdsourcing and research with human subjects: 
 Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?Answer: [NA]Justification: The paper does not involve crowdsourcing nor research with human subjects.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector. 
institutional review board (irb) approvals or equivalent for research with human subjects: 
 Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?Answer: [NA]Justification: The paper does not involve crowdsourcing nor research with human subjects.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.• Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: The structure of this supplementary material as shown below,• Appendix B presents additional experimental results, including changes in in-domain and out-of-domain site splitting, further comparisons on the BRCA dataset under traditional settings, additional results on the PRAD dataset, and results under site-preserved crossvalidation.• Appendix C discusses additional ablation study results concerning hyperparameters, including the impact of loss weights for PIM and SIM, the effect of the number of representative patches k, and further ablation studies on the CIB module.• Appendix D presents detailed descriptions of datasets and experimental settings.• Appendix E provides the detailed definition and formula derivation.• Appendix F elaborates on the Concept Information Bottleneck (CIB) module using the information plane.• Appendix G details the prompts used in our experiments.• Appendix H presents additional visualization results.To -B6: 39 slides of IDC and 6 slides of ILC.The additional results are presented in Table 4 . It is evident that the proposed CATE-MIL consistently delivers superior performance in both in-domain and out-of-domain settings, demonstrating its effectiveness in enhancing task-specific information and minimizing the impact of irrelevant data.In the main paper, we divided the dataset into in-domain and out-of-domain sites to evaluate the generalization ability of the proposed CATE method. To further evaluate its effectiveness, we conducted additional experiments under traditional settings, randomly splitting the dataset into training, validation, and testing sets. We report the performance metrics, including means and standard deviations over 10 Monte-Carlo Cross-Validation runs, in Table 5 . Additionally, we provide comparisons of training times and parameter sizes for various methods.It is evident that the proposed CATE-MIL consistently outperforms others in both AUC and ACC metrics, underscoring its superiority. Furthermore, CATE-MIL benefits from shorter training times and smaller parameter sizes compared to the R 2 T-MIL method.CATE is a general framework that can be applied to more complex tasks beyond cancer subtyping, such as Gleason grading in prostate cancer. We have conducted conducted experiments on the TCGA-PRAD dataset to evaluate the performance of CATE-MIL in Gleason grading. Specifically,  and 5 , and the task is to classify the samples into these three categories. The results are shown in Table 6 . It is evident that CATE-MIL consistently outperforms the base model ABMIL in both in-domain and out-of-domain settings, demonstrating its effectiveness in enhancing task-specific information and minimizing the impact of irrelevant data. 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.To provide a more comprehensive evaluation of the proposed CATE-MIL, we conduct additional experiments using site-preserved cross-validation [11] , where the samples from the same site are preserved in the same fold. For each fold, we split the data into training and testing sets, and these testing sets are regarded as in-domain testing data. And the other sites are used as out-of-domain testing data. The results are shown in Table 7 . It is evident that CATE-MIL consistently outperforms ABMIL in both in-domain and out-of-domain settings.The overall objective of the proposed CATE-MIL is a weighted sum of the PIM and SIM losses, along with the classification loss:EQUATION)We note that the hyperparameters for the weights of the PIM and SIM losses were not specifically tuned in the main paper. From Figure 4 , it is evident that CATE enhances the performance of the base model ABMIL in most scenarios, particularly in out-of-domain settings. When the weight of PIM loss λ P is too low, model performance suffers, underscoring the significant role of PIM loss in enhancing task-specific information. Conversely, an excessively high λ P also diminishes performance, as the model overly prioritizes maximizing mutual information between the original and calibrated features at the expense of optimizing classification loss. Regarding the weight of SIM loss λ S , optimal performance is achieved when λ S is approximately 30. If the λ S is too low, the model fails to effectively eliminate irrelevant information, thereby impairing performance. Conversely, if λ S is too high, the model risks discarding task-relevant information, leading to performance degradation. In this section, we conduct an ablation study to investigate the impact of the number of representative patches, k, on model performance, as discussed in Section 3.2. In practice, the representative patches are selected based on the similarities between the original image feature and the corresponding class-specific concept anchor.As shown in Figure 5 , it is evident that CATE generally enhances the performance of the base model, ABMIL. When the number of representative patches k is too small, model performance degrades due to insufficient capture of task-specific information. Conversely, an excessively large k also leads to performance degradation in out-of-domain scenarios, as it introduces noise from irrelevant information. Consequently, we have set the number of representative patches k to 10 in the main paper.We further 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 Table 8 .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. In this section, we provide detailed descriptions of the datasets used in the experiments and the detailed experimental settings that we used in the experiments in Section 4.2.• TCGA-BRCA: The dataset contains nine disease subtypes, and this study focuses on the classification of Invasive ductal carcinoma (IDC, 726 slides from 694 cases) and invasive lobular carcinoma (ILC, 149 slides from 143 cases). The dataset is collected from 36 sites with 20 of them having both IDC and ILC slides, and the other 16 sites only have IDC slides or ILC slides. • TCGA-NSCLC: The dataset contains two disease subtypes, including lung adenocarcinoma (LUAD, 492 slides from 430 cases) and lung squamous cell carcinoma (LUSC, 466 slides from 432 cases). The dataset is collected from 66 sites. Different from TCGA-BRCA, each site only contains one disease subtype.• TCGA-RCC: The dataset contains three disease subtypes, including clear cell renal cell carcinoma (CCRCC, 498 slides from 492 cases), papillary renal cell carcinoma (PRCC, 289 slides from 267 cases), and chromophobe renal cell carcinoma (CHRCC, 118 slides from 107 cases). The dataset is collected from 55 sites. Similar to TCGA-NSCLC, each site only contains one disease subtype.To validate that the proposed CATE effectively improves MIL model performance by enhancing the task-specific information and eliminating the disturbance of task-irrelevant information, we conduct extensive experiments across three datasets under various in-domain and out-of-domain settings.Specifically, for TCGA-BRCA dataset, we conduct the following experiments:• For TCGA-NSCLC dataset, we conduct the following experiments:• NSCLC (N IND =2): Two sites as the in-domain data, and the other sites as the out-of-domain data. The details of the in-domain sites are shown below:-44: 43 slides of LUAD and 0 slides of LUSC.-22: 0 slides of LUAD and 36 slides of LUSC.• NSCLC (N IND =4): Four sites as the in-domain data, and the other sites as the out-of-domain data. The details of the in-domain sites are shown below:-44: 43 slides of LUAD and 0 slides of LUSC.-50: 20 slides of LUAD and 0 slides of LUSC.-22: 0 slides of LUAD and 36 slides of LUSC.-56: 0 slides of LUAD and 35 slides of LUSC.For TCGA-RCC dataset, we conduct the following experiments:• RCC (N IND =3): Three sites as the in-domain data, and the other sites as the out-of-domain data. The details of the in-domain sites are shown below:-A3: 48 slides of CCRCC, 0 slides of CHRCC, and 0 slides of PRCC.-KL: 0 slides of CCRCC, 24 slides of CHRCC, and 0 slides of PRCC.-2Z: 0 slides of CCRCC, 0 slides of CHRCC, and 23 slides of PRCC.• RCC (N IND =6): Six sites as the in-domain data, and the other sites as the out-of-domain data. The details of the in-domain sites are shown below:-A3: 48 slides of CCRCC, 0 slides of CHRCC, and 0 slides of PRCC.-AK: 15 slides of CCRCC, 0 slides of CHRCC, and 0 slides of PRCC.-KL: 0 slides of CCRCC, 24 slides of CHRCC, and 0 slides of PRCC.-KM: 0 slides of CCRCC, 20 slides of CHRCC, and 0 slides of PRCC.-2Z: 0 slides of CCRCC, 0 slides of CHRCC, and 23 slides of PRCC.-5P: 0 slides of CCRCC, 0 slides of CHRCC, and 15 slides of PRCC.In this part, we will define the sufficiency and consistency in the context of the concept information bottleneck (CIB) module.First, we define the mutual information between two random variables X and Y as:EQUATIONSufficiency: To quantify the requirement that the calibrated features should be maximally informative about the label information y, we define the sufficiency as the mutual information between the label information and the calibrated features:Sufficiency. α is sufficient for y ⇐⇒ I(x; y| α) = 0 ⇐⇒ I(x; y) = I( α; y).The sufficiency definition requires that the calibrated features α encapsulate all information about the label information y that is accessible from the original features x. In essence, the calibrated feature α of original feature x is sufficient for determining label y if and only if the amount of information regarding the specific task is unchanged after the transformation.Consistency: Since the image feature x and the concept anchor c are consistent in the pathology VLM, we posit that any representation containing all information accessible from both the original feature and the concept anchor also encompasses the necessary discriminative label information. Thus, we define the consistency between the concept anchor and the original feature as:Consistency. c is consistent with x for y ⇐⇒ I(y; x|c) = 0.The consistency definition requires that the concept anchor c and the original feature x should be consistent with the label information y.In this part, we will prove that the predictive information for the concept anchor can be maximized by maximizing the InfoNCE mutual information lower bound, which is defined in Equation 4. As in [27] , the optimal value for f (c, αi ) should be proportional to the density ratio:EQUATIONAnd we can get:EQUATIONwhere C neg denotes the negative concepts, including both negative class concepts and other type concepts. We can see that the InfoNCE estimation is a lower bound of the mutual information between the concept anchor and the calibrated features:EQUATIONThus, the predictive information of calibrated features for the concept anchor can be maximized by maximizing the InfoNCE mutual information lower bound.As shown in Equation 3, the mutual information between the original features and the calibrated features I(x; α) can be decomposed into the mutual information I( α; c) between the calibrated features and the concept anchor and the mutual information I(x; α|c) between the concept anchor and the calibrated features:EQUATIONAs introduced in section 3.2, the first item is the predictive information for the concept anchor, which can be maximized by maximizing the InfoNCE mutual information lower bound. Thus, the second superfluous information item can be minimized by minimizing I(x; α). In practice, the superfluous information minimization is conducted on all patches in the subset x, since the taskirrelevant information is distributed across all patches. Thus, the superfluous information can be minimized by minimizing I(x; α). As in [1] , this mutual information can be represented as:EQUATIONHowever, the computing of the marginal distribution p (α) = p (α|x) p (x) dx is intractable. Thus, we can let the Gaussian distribution r (α) approximate the marginal distribution p (α). Since the Kullback-Leibler divergence between two distributions is non-negative, we can get:EQUATIONThus, the mutual information I(x; α) have the upper bound:EQUATIONIn practice, to compute this upper bound, we approximate distribution p (x) with the empirical distribution:EQUATIONThus, the upper bound of the mutual information I(x; α) can be computed as:EQUATIONwhere q θ (α|x) is a variational distribution with parameter θ to approximate p (α|x), and we implement this variational distribution with MLP:EQUATIONwhere M LP µ and M LP Σ are implemented with MLPs which output the mean µ and covariance matrix Σ of the Gaussian distribution. In practice, we utilize the reparameterization trick [18] to sample from the Gaussian distribution to get an unbiased estimate of the gradient to optimize the opjective. Thus, the upper bound can be optimized by minimizing the KL divergence between the variational distribution q θ (α|x) and the Gaussian distribution r (α):EQUATIONTo better understand the concept information bottleneck (CIB) method, we provide an explanation of CIB with the information plane [8, 6] . As shown in Figure 6 , the information plane is a two-dimensional space where the x-axis represents the mutual information between the original features the calibrated features of CIB module I(x; α), and the y-axis represents the mutual information between the label information and the calibrated features I(y; x). The line colored orange in the information plane defines sufficiency, indicating where the calibrated features are maximally informative about the label information y.Feasible RegionThe consistency definition falls to the left of the sufficiency line, denoting a state where the calibrated features, containing all information from both the original feature and the concept anchor, also include the label information. The InfoMAX principle [9] states that the optimal calibrated feature should be maximally informative about the original feature.Ideally, a good calibrated feature α should be maximally informative about the label information y (the sufficiency line in the information plane) while minimally informative about the original features x (i.e., the ideal situation in the information plane). This situation can be achieved by directly supervised learning. However, the weakly supervised nature of MIL methods complicates this goal.As demonstrated in the ablation analysis in Section 4.4, if we directly perform superfluous information minimization without the guidance of concept anchor (i.e., predictive information maximization), the calibrated features will be less informative about the label information y and the model performance will be degraded significantly. To address this issue, the proposed CIB module introduces the concept anchor c to guide the learning process of information bottleneck for each instance. As shown in Figure 6 , the CIB module could effectively reduce the mutual information between the original features and the calibrated features I(x; α) while preserving the mutual information between the label information and the calibrated features I(y; α).In this section, we provide the detailed prompt templates and concepts for the datasets, which are used to generate the concept anchor for the CIB module. The prompt templates are shown in Table 9 , and the concepts for TCGA-BRCA, TCGA-NSCLC, and TCGA-RCC are shown in Table 10 , Table 11 , and Table 12 , respectively. The templates and the class-agnostic prompts are referred from the original paper of CONCH [25] , and the class-specific prompts are generated by querying the LLM with the question 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 LLM generated prompts has been demonstrated in several recent studies [29, 22] . 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. To further illustrate the effectiveness of the proposed CATE-MIL, we provide more visualization results in this section. The visualization results for IDC and ILC in TCGA-BRCA are shown in Figure 7 and Figure 8 , respectively. Justification:The paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material. Guidelines:• The answer NA means that paper does not include experiments requiring code. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.https://docs.gdc.cancer.gov/Encyclopedia/pages/TCGA_Barcode/ 


Decision: The paper introduces Concept Anchor-guided Task-specific feature Enhancement (CATE), a novel approach that enhances the performance of pathology foundation models by tailoring them to specific downstream tasks. By leveraging task-specific concepts derived from vision-language models and incorporating a Concept-guided Information Bottleneck module, the method effectively maximizes the relevance of extracted features while minimizing irrelevant information. Adding a Concept-Feature Interference module further refined the discriminativeness of these features. The reviewers acknowledged the novelty and potential impact of this approach, especially its ability to incorporate domain-specific knowledge into whole slide image (WSI) classification. While some concerns were raised regarding the dependency on vision-language model quality and potential overfitting, the additional experiments and clarifications provided by the authors in their rebuttal were sufficient to address most of these issues. The consensus is to accept this paper.

Reviews: review: paper_summary: paper_summary: The authors propose an approach to enhance the performance of Multiple Instance Learning (MIL) models by incorporating concept prompts within the context of pathological Vision-Language Models (VLMs). Two modules, i.e., the Concept-guided Information Bottleneck (CIB) module and the Concept-Feature Interference (CFI) module, are designed to calibrate and inject similarity characteristics between features and concepts. Experiments are conducted on three Whole Slide Image (WSI) datasets, yielding interesting results. The paper includes qualitative and ablation analyses.
strength_weakness: strength_weakness: Dependence on Text Prompts: The performance heavily relies on the quality of text prompts/concepts. A more principled approach to generating expert-designed prompts would be beneficial. Only the hard-crafted ones make the  "free lunch" not that free. 

Experimental Settings: The choice of the number of IND and OOD sites for different datasets is unclear. How to choose the number of IND and OOD sites? It would be helpful to see performance under more traditional settings on NSCLC and RCC. What's more, it is noticed that the gain of CATE decreases with an increasing number of sites, raising questions about the practicality of the main experimental settings in real-world WSI analysis.


Attention Maps: The attention maps are not that convincing to me. Since CATE-MIL provides almost identical attention maps to the baseline ABMIL except for intensity. If ABMIL predicts incorrect attention, CATE-MIL could potentially worsen the situation.


Underperformance on NSCLC and RCC: The authors claim that the underperformance of CATE-MIL on NSCLC and RCC due to the elimination of site-specific patterns is concerning. Clarification on why site-specific information for the IID test is not considered task-relevant is needed, especially as this issue does not appear with BRCA.
questions: questions: The paper is technically sound, but I am not convinced by the provided experimental results.

1) Could the authors suggest more principled ways to obtain expert-designed prompts? Relying on hand-crafted prompts may not be scalable or generalizable.

2) Please refer to the weakness and clarify the experimental settings. More detailed qualitative and quantitative analyses are needed.

3) Lack of comparison with ref [1], which also uses the Information Bottleneck theory to improve the feature quality of MIL.


[1] Task-specific Fine-tuning via Variational Information Bottleneck for Weakly-supervised Pathology Whole Slide Image Classification. CVPR 2023.
limitations: limitations: N/A
paper_summary: paper_summary: The authors introduce a new tool called Concept Anchor-guided Task-specific Feature Enhancement (CATE) for analysis of whole slide images.  The authors are taking advantage of the open source weight now available for the CONCH pathology vision language model trained on paired captions of images from journals and other publicly available image-caption pairs.  The high level conceptulization is that language prompts related to a task at the whole slide image level can be enrich for the features most relevant to the downstream task.  The title cleverly asserts this is a free lunch when using a vision language encoder for feature extraction the expertly curated concepts can "bring foward" the most relevant features when aggregating at the whole slide level.  

Interestingly, they split up the CATE task into Predictive Information Maximization to enhance features associated with related concepts and Superfluous Information Minimization is used to supress the concepts likely to be associated with unimportant features. 

The outputs of PIM and SIM are used during aggregation task to determine weighting that MIL should use for a given feature vector derived from the tile. 

Finally a module called concept-feature interference is used calibrate the CIB scores. 

These module provide enhanced features to the aggregation function that provides relative weighting based on the relationship between the semantic concepts relationship to the downstream task. 

CATE and weights are not backpropogated through the feature embedding tasks and they are static in the MIL function (to the best of my understanding). 

The authors chose RCC, NSCLC, and BRCA tasks from the TCGA to demonstrate performance
strength_weakness: strength_weakness: I agree with the weaknesses the authors have put forth.  
The model scope is ~mostly~ limited to downstream task for the the concepts are part of the VLM training.  The tasks that are chosen for benchmarking are actually tasks that are likely well represented in VLM training sets with robust descriptions of SCC vs. LUAD, etc.  For tasks that source material does not contain common descriptions, say more challenging tasks like predicting mutations, it is unclear to what extent, this work would provide advantage over other approaches.  
My presumption is that VLM encoders are less robust for feature extraction than the now available large foundation models trained with 10^6 plus whole slide images. If one is not able to use these models and must use an inferior model, then one might argue it is not a "free lunch".
questions: questions: Out of distribution (OOD) is a term that means many different things to different people.  Will you please spare a paragraph to better explain what you mean by OOD. If you find that OOD is not the best term for what you mean, I think it would be better to use a less loaded term. 

Have you considered a hybrid approach where the vision feature vectors can be extracted with a different foundation model but the VLM is use for the concept guidance?

How the normal distribution of SIM in equation (8) align with Bernoulli distribution of ABMIL? Are there alternative distributions that could be used here? (may be Poisson distribution to define random occurance of a non-informative patch?)

Minor comments for improvement in readability: 
1) In figure 2a, the image encoder cartoon is larger on the side of image which implies the feature vectors are larger than the original patches. 
2) Equation (6) LHS should be I(x; alpha | c)
limitations: limitations: Limitations section is reasonable.  I have provided further feedback in the weaknesses section.
paper_summary: paper_summary: The authors introduce CATE, a novel approach designed to enhance the generalizability of histopathology models by leveraging task-specific concepts derived from the text endoder of pathology vision language models(VLM). CATE includes two modules: CIB and CIF, which work synergistically to improve model robustness. The CIB module identifies and retains features that contribute positively to predicting the task outcome, while it eliminates features that contain irrelevant or superfluous information.  Additionally, CIF specializes in creating task-specific features by leveraging similarities between the enhanced image features the and predefined concept anchors. The authors demonstrate the effectiveness of CATE across an extensive series of histopathology datasets, showcasing superior results in model performance and generalization.
strength_weakness: strength_weakness: - CATE's performance hinges largely on the quality of its concept anchors, which in turn depends on domain expertise and the quality of the pre-trained pathology VLM, a weakness that is also highlighted by the authors.
questions: questions: - Why did you use accuracy as an evaluation metric? Accuracy is straightforward to interpret, but it can be misleading in imbalanced datasets. Metrics like F1-score, weighted accuracy, precision, and recall offer more nuanced insights into model performance, particularly in contexts where class distribution is skewed such as the BRCA dataset.

- What strategy was used to split the dataset into in-domain and out-of-domain test sets? Understanding this process ensures transparency in evaluating the model's generalizability across different clinical settings or data sources. 

- How do you define the class-agnostic concepts?

- What is the intuition behind maziming the distance between the enhanced featrures and the class-agnostic concepts? When features are overly tailored to the training samples by maximizing their separation from class-agnostic concepts, the model may become too specialized and prone to overfitting.
limitations: limitations: - Exploring the potential of concept anchors for applications beyond subtyping, such as mutation predictio could be an intersting avenue to explore.
paper_summary: paper_summary: The paper proposes a way to enhance image features from VLM pathology foundation models for downstream tasks by aligning them with task-specific text prompts. The authors use two modules -- one which extracts an image representation aligned with task concepts and other which measures similarity of representation with tasks, to enrich image representations for downstream tasks.
The authors show results for the cancer subtyping task across 3 datasets indicating the usefulness of the proposed approach.
strength_weakness: strength_weakness: The main weaknesses are around some of the assumptions made in the paper listed below which limit the applicability of the approach.


**Aligning patches with tasks**

In WSIs the labels are only available at slide-level and not patch-level and for many problems only a fraction of the patches in the slide have task-relevant info and for many others the information may need to be aggregated from multiple patches.

Even though there is a filtering step to filter patches, the step of aligning patches to tasks is noisy when many patches dont have task info and it may not even be appropriate for problems which need aggregation of information across patches (as the approach homogenizes representations of all patches belonging to the same class).

The authors show results on only one type of task -- subtyping where most patches on the slide have the sub-type information and is also typically a simpler MIL task indicated by the strong results on the benchmarks for most methods. 

Instead of aligning patch representations, it would be better to align slide-level representations with task-concepts, which addresses both the drawbacks listed above. There is also evidence from prior works around usefulness of this like SC-MIL (https://arxiv.org/abs/2303.13405 which align representations of slides belonging to same-class) and they also show aligning patch-level representations (patch SC-MIL) does poorly compared to aligning at the slide-level


**Unclear Improvements**

It is unclear how useful the task-alignment objective with the information bottleneck layer is. 
The key part seems like the filtering step which extracts relevant patches for the task and I suspect if thats contributing to majority of the improvements.
Once filtering is done, is the alignment step really needed as it seems like a dimensionality reduction step to extract task specific features. 
A sufficiently strong model should be able to able to extract task relevant features from the embedding dimensions with cross-entropy loss.
The alignment step might also be unsuitable as it homogenizes patch representations which may not be good for many MIL tasks and directly optimizing with cross entropy on task-labels might be more flexible and general.

Its also unclear if the improvements are coming from a bigger model due to the CIB and CFI modules. Scaling ABMIL to similar number of params would provide more context on this.


**Task Concept Alignment**
The other key assumption is using text prompt embeddings as task anchors which is very dependent on factors like the quality of the pre-trained VLM, the textual prompts and descriptions used and the idea that it is possible to write morphological descriptors for tasks which is typically not possible for many MIL tasks like survival or molecular signature prediction.

The sufficiency and consistency constraints imposed could be limiting specially when any of the factors above break. The SIM objective specifically could be problematic when text doesnt fully capture the required features for the task and other relevant image features could be suppressed.
questions: questions: Its unclear how much improvements are coming from the patch filtering step and the increased number of parameters and if those alone are sufficient enough.

I think its worth trying the alignment on the slide-level representations instead of patch-level representations (similar to SC-MIL) which makes the approach more generic and limits the drawbacks with using slide-level labels for patches.
limitations: limitations: Highlighted the limitations in weaknesses.

Provided some suggestions and improvements in the questions. 

Given the weaknesses and the lack of evidence and limited applicability of the approach to more complex MIL problems like scoring (gleason grading, HER2 grading) which need aggregation of information across patches or other analytical tasks like survival prediction, molecular signature prediction where there might not be possible to write text descriptions, I am basing my decision

Comments: comment: I appreciate the authors running  additional experiments and ablations and providing clarifications to my concerns. I will keep my rating as it is already high, as I find this work to be both novel and promising in advancing the field of computational pathology.
We apologize for the confusion caused by the term "bottleneck". We will clarify this in the revised manuscript to avoid any misunderstanding. Thank you for pointing this out.
Dear Reviewer 9fs8,

Thank you very much for recognizing the novelty and potential of our work. We agree that tasks such as survival prediction and biomarker prediction are indeed more challenging due to the difficulty of defining reliable features, and we will explore these tasks in our future work.

Additionally, there might be a potential solution to address this challenge. For instance, we could leverage LLMs or retrieval-based LLMs to generate descriptive prompts about the general morphological appearance of WSIs for specific cancer types. By asking targeted questions, we can summarize reliable and general morphological descriptions associated with different survival outcomes or biomarker expressions (e.g., “Displays significant lymphocyte infiltration within both the neoplastic epithelium and the stroma” for EBV-positive tumors) and further verify these prompts with pathologists. These prompts can then be used to generate morphological concepts that serve as concept anchors, guiding the alignment of patch features and enhancing the task-relevant morphological features.

Thank you for your insightful comments, which have encouraged us to think more deeply about potential future improvements to our method. We will include the above clarification and discussion in the revised manuscript.
Thanks for clarifying the question about the feature dimension.
It makes sense that the dimension is kept same to match text embedding, however the usage of "bottleneck" implicated a smaller dimension, worth mentioning this distinction in the paper.
Thank you for pointing out my mistake on OOD being out-of-domain rather than the more common use of OOD.  I maintain my weak accept characterization.  I think that this work is very novel and interesting and will contribute to furthering the field of computational pathology. I also do not see the comments addressing my concerns to be comprehensive enough to warrant upgrading my initial impression.  The most useful tasks for computational pathology are tasks for which pathologists do not yet have reliable features for predictions, such as survival and biomarker prediction. Gleason grading is yet another task that pathologists are very good at and the features are widely known.
Dear Reviewer VcAb,

Thank you very much for your thoughtful and constructive feedback. We greatly appreciate your positive comments, and your insights have significantly contributed to strengthening our work. We would like to address your remaining minor concerns regarding the feature dimension of the output from the CIB module.

The encoder of CIB module does not change the feature dimension of the input patch features. Since the patch features and concept anchors are generated by the image encoder and text encoder of the pathology VLM respectively, they share the same feature dimension.
Therefore, the number of parameters in the classifiers and attention layers in CATE-MIL is identical to that in vanilla ABMIL.

We hope this clarifies your concerns. And we will add these clarifications and the results of the additional experiments to the revised manuscript or the appendix.
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.
Dear Reviewer ZzEi,

We appreciate your detailed reviews and valuable comments. We would like to address your additional concerns below:

> From my perspective, using LLMs just transforms the manually designed prompts into questions...

Thank you for your feedback. In response to your comments:
- The questions we designed are more general and do not require an extensive pathological background. They are intended to allow the LLM or retrieval-based LLM to utilize embedded domain knowledge in the model or the knowledge in the literature for high-quality prompts.
- We understand that designing such questions is not without effort and we will relax the "free-lunch" claim and discuss it in our revised manuscript.
- We want to emphasize that medical imaging analysis and computational pathology are complex fields requiring significant domain expertise. Incorporating such knowledge/background into the method design could enhance the performance and robustness of the framework. Our approach offers a principled approach to integrating domain knowledge/background into framework design.

> The site-specific signatures have been thoroughly discussed in ref[1]. I recommend using the site-preserved CV...

Thank you for your constructive suggestion. We conducted experiments using the suggested site-preserved CV, following the original code and settings from ref [1] to split the data into three folds based on the source sites. The results show that CATE-MIL consistently outperforms vanilla ABMIL across all datasets. We will add these results to the revised paper or appendix (if space is not allowed).

||Method|AUC|
|---|---|---|
|BRCA|ABMIL|0.912±0.012|
||CATE-MIL|0.935±0.014|
|NSCLC|ABMIL|0.953±0.018|
||CATE-MIL|0.965±0.011|
|RCC|ABMIL|0.980±0.001|
||CATE-MIL|0.983±0.000|

> For the performance, the authors attribute to the class-specific features and the site-specific...

Thanks for your comments. We try to provide more evidence from the following two aspects:
- To evaluate the influence of site-specific features, we selected two sites (site IDs 49 and 55, respectively) with the largest number of samples of **the same subtype (LUAD)** from the NSCLC dataset.
We used the corresponding site ID as labels for the samples and trained an ABMIL model to predict which site the samples originated from. The results show that the ABMIL model can easily predict the site with an AUC of **0.988**. This indicates that the model can easily use site-specific features to distinguish different sites, which may be used as shortcuts for subtyping.
- We also conducted a comparative experiment on the NSCLC dataset using patch features extracted after **stain normalization**, which helps reduce the impact of staining styles and imaging characteristics.
As shown in the following table, the IND performance is reduced. This demonstrates the effect of site-specific features on shortcutting subtyping and hindering the generalization of the model. Meanwhile, the OOD performance of ABMIL is enhanced after stain normalization, indicating that without site-specific features, the model can generalize better to OOD data, which is consistent with CATE-MIL.
Notably, CATE-MIL still outperforms ABMIL in OOD performance after stain normalization, validating that CATE can better eliminate site-specific features to enhance the generalization of the model compared to stain normalization.

||Method|OOD-AUC|IND-AUC|
|---|---|---|---|
|NSCLC($N_{\text{IND}}$=2)|ABMIL w/o stain norm|0.874±0.021|0.997±0.004|
||ABMIL w/ stain norm|0.889±0.019|0.995±0.006|
||CATE-MIL|0.945±0.016|0.985±0.011|
|NSCLC ($N_{\text{IND}}$=4)|ABMIL w/o stain norm|0.951±0.023|0.974±0.018|
||ABMIL w/ stain norm|0.959±0.021|0.966±0.026|
||CATE-MIL|0.969±0.003|0.967±0.019|

> Besides this cherry-picked sample, I am still not convinced. I think the the attention maps of...

We apologize for any remaining doubts regarding the attention maps. As discussed in the previous response and supported by the visual evidence in the attached PDF, we have shown that the attention maps of CATE-MIL and ABMIL are not identical. 
The attention map of CATE-MIL is more focused on cancerous patches and less on non-cancerous patches. To further clarify this, we calculated the average **Jaccard index** between the binarized attention maps of CATE-MIL and ABMIL for **all samples** to provide the **quantitative evidence**.

We first binarized the attention maps by setting the threshold to 0.5, and then calculated the Jaccard index by dividing the intersection of the two binarized attention maps by their union. The results show that the average Jaccard index on the BRCA dataset is **0.699**, indicating that the overlap between the attention maps of CATE-MIL and ABMIL is not very high and that the attention maps of CATE-MIL are indeed different from those of ABMIL.

[1] The impact of site-specific digital histology signatures on deep learning model accuracy and bias, nature communications
I appreciate the detailed responses provided. Based on the responses, I still have further questions:
1. From my perspective, using LLMs just transforms the manually designed prompts into questions.  I assume that both of them need a certain level of pathological background, which is totally different from natural images and makes the proposed method not that free.
2. The site-specific signatures have been thoroughly discussed in ref[1]. I recommend using the site-preserved CV for more convincing results instead of 10 runs of Monte-Carlo CV, where higher means and fewer stds are expected for better performance.
3. For the performance, the authors attribute to the class-specific features and the site-specific features. Any more direct evidence to support this argument?
4. Besides this cherry-picked sample, I am still not convinced. I think the the attention maps of CATE-MIL and ABMIL are still almost identical after binarized.

[1] The impact of site-specific digital histology  signatures on deep learning model accuracy and  bias, nature communications
Dear Reviewer VcAb, 

Thank you very much for your thorough review and for taking the time to read our rebuttal carefully. We appreciate your insightful comments and the opportunity to address them further. We would like to address your additional comments below:

> In CATE-MIL w/o concept alignment, do you filter patches using similarity and then directly use the patch features from the image encoder for ABMIL? In the ABMIL baseline I assume we do not have the patch-filtering step but everything should be similar to the CATE-MIL w/o concept alignment setup right? Any ideas why AB-MIL is doing better than CATE-MIL w/o concept alignment then?

- To ease your understanding, we will further clarify the key idea of our concept-guided Information Bottleneck (CIB) module. 
In CATE-MIL, the CIB module functions as an encoder (similar to the encoder of VAE) to acquire the enhanced features of **all** the original patch features, where the module parameters are optimized via the concept alignment objectives (i.e., PIM loss and SIM loss).
As not all patches contain task-relevant information, we filter out a subset of patches that contain task-specific information and only calculate the objective losses on these selected patches to avoid instability during the optimization.
However, in the prediction process, we take **all** enhanced patch features into ABMIL to get the final slide-level prediction. Therefore, there is no filtering step in the testing phase.
In summary, the "patch filtering-like" step only **occurs in the training phase, not the testing phase**. 

- In CATE-MIL w/o concept alignment, we did not remove the CIB module directly, but only removed the concept alignment objectives (i.e., PIM loss and SIM loss) and used the cross-entropy loss on task labels to optimize the encoder parameters. 
**All the original patch features** are still input to the CIB module and then to ABMIL, **without the filtering-like step**. Therefore, the process of CATE-MIL w/o concept alignment resembles ABMIL with an encoder of VAE appended.

- The results presented in the rebuttal indicate that such a setting does not enhance the performance of ABMIL; instead, it diminishes performance. This could be attributed to the increased difficulty in training ABMIL and the heightened risk of overfitting caused by the additional layers added before ABMIL. Consequently, the performance of CATE-MIL w/o concept alignment is lower than that of ABMIL, which also validates the effectiveness of the proposed concept alignment objectives.

> In Replace CIB with MLP, do you project the image features to a bottleneck dimension using an MLP and then use these features for ABMIL? Agree without concept alignment the features may not always capture relevent information but surprised to see it doing so poorly compared to ABMIL. Any hypothesis as to why this is the case.

- Yes, we project the image features to a bottleneck dimension that matches the enhanced feature dimension in the CIB module using an MLP and then use these features for ABMIL.
As we discussed above, the CIB module can be regarded as a encoder of VAE with concept alignment objectives regularizing the training of the encoder. Without concept alignment objectives, the CIB module functions merely as an encoder, similar to an MLP.
In this case, **directly adding such an encoder or MLP before ABMIL without proper "regularization" will increase the difficulty of the training of ABMIL and potentially lead to overfitting.** Therefore, the performance of CATE-MIL w/o concept alignment is similar to that of replacing the CIB module with an MLP, and both underperform compared to ABMIL.

> Summary

The "patch filtering-like" step is only used for calculating alignment losses during training and does not affect the prediction process during inference. The performance improvement of our CATE-MIL model comes from the task/concept alignment objectives, not from the patch filtering step. Additionally, the improvement is not due to an increased number of parameters, as shown by the reduced results when replacing the CIB module with a vanilla MLP having a similar number of parameters.

We hope this clarifies your concerns. If you have any further questions or need additional clarification, please feel free to let us know.
> Aligning patches with tasks

Thanks to the authors for exploring the suggestion of aligning with slides. It makes sense that without patch-level alignment, the slide-level alignment can be noisy and has a harder task. Good to see the slide-level alignment helping and complementing the patch-level alignment.


> Unclear improvements

Thanks for runnning the ablations. In CATE-MIL w/o concept alignment, do you filter patches using similarity and then directly use the patch features from the image encoder for ABMIL? 
In the ABMIL baseline I assume we do not have the patch-filtering step but everything should be similar to the CATE-MIL w/o concept alignment setup right? Any ideas why AB-MIL is doing better than CATE-MIL w/o concept alignment then?

In Replace CIB with MLP, do you project the image features to a bottleneck dimension using an MLP and then use these features for ABMIL? Agree without concept alignment the features may not always capture relevent information but surprised to see it doing so poorly compared to ABMIL. Any hypothesis as to why this is the case.

> Task concept alignment and applicability to more complex MIL problems

Thanks for running the experiments for gleason grading. This is helpful.

> Summary

Thanks to the authors for running these additional experiments and strengthening the paper. Could you confirm the setup for the ablations discussed above?