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

TL;DR

This work introduces two interconnected modules to dynamically calibrate the generic image features extracted by foundation models for certain tasks or cancer types, and designs 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.

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. The source code is available at https://github.com/HKU-MedAI/CATE.

Figures (8)

• Figure 1: (a) Illustration of the key idea of concept-guided information bottleneck to enhance the task-relevant information and discard the task-irrelevant information. (b) Task-specific model adaptation with CATE to enhance the generalization across different data sources.
• Figure 2: (a) Overview of CATE: the outputs of the CIB and CFI modules are concatenated to form the enhanced feature for downstream MIL models. (b) Task-relevant concept generation. (c) Concept-guided Information Bottleneck (CIB) module. (c) Concept-Feature Interference (CFI) module.
• Figure 3: (a) Attention heatmap of CATE-MIL. (b) Attention heatmap of the original ABMIL. (c) similarity between the calibrated features and the corresponding class concept feature. (d) similarity between the original features and the corresponding class concept feature. (e) Original WSI. (f) UMAP visualization of class concept features, original features, and enhanced features.
• Figure 4: Ablation study of the weight of PIM and SIM losses on the model performance.
• Figure 5: Ablation study of $k$.