Feature Re-Embedding: Towards Foundation Model-Level Performance in Computational Pathology

TL;DR

A Re-embedded Regional Transformer (R2T) for re-embedding the instance features online, which captures fine-grained local features and establishes connections across different regions, validate that feature re-embedding improves the performance of MIL models based on ResNet-50 features to the level of foundation model features, and further enhances the performance of foundation model features.

Abstract

Multiple instance learning (MIL) is the most widely used framework in computational pathology, encompassing sub-typing, diagnosis, prognosis, and more. However, the existing MIL paradigm typically requires an offline instance feature extractor, such as a pre-trained ResNet or a foundation model. This approach lacks the capability for feature fine-tuning within the specific downstream tasks, limiting its adaptability and performance. To address this issue, we propose a Re-embedded Regional Transformer (R$^2$T) for re-embedding the instance features online, which captures fine-grained local features and establishes connections across different regions. Unlike existing works that focus on pre-training powerful feature extractor or designing sophisticated instance aggregator, R$^2$T is tailored to re-embed instance features online. It serves as a portable module that can seamlessly integrate into mainstream MIL models. Extensive experimental results on common computational pathology tasks validate that: 1) feature re-embedding improves the performance of MIL models based on ResNet-50 features to the level of foundation model features, and further enhances the performance of foundation model features; 2) the R$^2$T can introduce more significant performance improvements to various MIL models; 3) R$^2$T-MIL, as an R$^2$T-enhanced AB-MIL, outperforms other latest methods by a large margin.The code is available at: https://github.com/DearCaat/RRT-MIL.

Figures (11)

• Figure 1: Top: The conventional MIL paradigm lacks fine-tuning of the offline embedded instance features. Bottom: The proposed MIL paradigm that introduces instance feature re-embedding to provide more discriminative features for the MIL model.
• Figure 2: Overview of proposed R$^2$T-MIL. A set of patches is first cropped from the tissue regions of a slide and embedded in features by an offline extractor. Then, the sequence is processed with the R$^2$T module: (1) region partition, (2) feature re-embedding within each region, and (3) cross-region feature fusion. Finally, a MIL model predicts the bag labels using the re-embedded instance features.
• Figure 3: Illustration of Embedded Position Encoding Generator.
• Figure 4: Performance improvement by adding R$^2$T. Features re-embedded by R$^2$T online outperform PLIP offline features on most tasks.
• Figure 5: The tSNE van2008tsne visualization of instance features from the CAMELYON-16 dataset, comparing (a) features extracted by ResNet-50 pre-trained on ImageNet-1k, (b) features extracted by PLIP, (c) features after N-MSA re-embedding, and (d) features after R$^2$T re-embedding. In (a), we obtain instance-level labels from the tumor annotations and report the instance numbers of different labels.