PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Wenhao Tang, Fengtao Zhou, Sheng Huang, Xiang Zhu, Yi Zhang, Bo Liu

arXiv:2402.17228v4cs.CV

TL;DR

MIL relies on offline instance features that lack downstream task-specific fine-tuning, while full end-to-end learning is constrained by memory costs. The paper introduces R2T, a plug-in regional Transformer that re-embeds features online. Across computational pathology benchmarks, re-embedding improves MIL performance, and R2T-MIL outperforms compared methods.

  • Problem

    MIL commonly relies on offline instance features without downstream task-specific fine-tuning, while full end-to-end learning is infeasible because of prohibitive memory costs.

    Source: §1 · p. 1

  • Method

    R2T re-embeds offline instance features between feature input and aggregation using local self-attention, cross-region fusion, and positional encoding.

    Source: §1 · p. 2, §3.2 · p. 3

  • Results

    Across diagnosis, subtyping, and survival benchmarks, R2T-MIL outperforms compared methods, including a C-index of 67.19% on LUAD.

    Source: §1 · p. 2, §4.2.2 · p. 6, §4.2.1 · p. 6

  • Takeaways & Limitations

    Transformer-based re-embedding consistently boosts diverse MIL architectures, and the paper concludes that local Transformers remain versatile alongside foundation models.

    Source: §5 · p. 8

    Re-embedded features can lose original patch-label information, limiting accurate weakly supervised tumor localization or segmentation.

    Source: F. Limitation · p. 15

Abstract

from arXiv · show

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.

1. Introduction

Conventional MIL uses offline instance features that cannot be fine-tuned for downstream tasks, while full end-to-end learning is prohibitively memory-intensive. The paper introduces R2T for online feature re-embedding and reports state-of-the-art performance for R2T-MIL across computational pathology benchmarks.

  • Conventional MIL extracts instance features offline, aggregates them, and predicts a bag label, with most prior work emphasizing aggregation and prediction.Source: §1 · p. 1
  • Full end-to-end learning is infeasible at WSI scale because of memory costs, while offline features lack task-specific fine-tuning and foundation-model pretraining does not resolve that gap.Source: §1 · pp. 1–2
  • Global Transformer attention misses pathology-relevant local detail; linear MSA reduces memory demands but has high computational cost and suboptimal performance.Source: §1 · p. 2
  • R2T applies native MSA within local regions, fuses regions with Cross-region MSA, and encodes patch positions with EPEG while integrating into mainstream MIL models.Source: §1 · p. 2
  • Feature re-embedding improves MIL performance to a level competitive with the latest foundation model.Source: §1 · p. 2
  • R2T-MIL, an R2T-enhanced AB-MIL, achieves state-of-the-art performance on various computational pathology benchmarks.Source: §1 · p. 2

2. Related Work

Computational pathology applies computational methods to pathology images, while MIL provides a weakly supervised framework for whole-slide analysis. Prior MIL work has developed approaches to feature extraction and aggregation, but typically leaves offline instance features without downstream task-specific fine-tuning.

2.1. Computational Pathology

  • Computational pathology combines pathology, image analysis, and computer science to develop computational methods for analyzing and interpreting pathology images.Source: §2.1 · p. 2

2.2. Multiple Instance Learning

  • MIL commonly uses pre-trained feature extractors such as ImageNet-pretrained ResNet-50, but pathology-image differences can make the extracted features poorly discriminative.Source: §2.2 · p. 2
  • Prior methods use instance-label fusion or instance-feature fusion, and Transformer blocks have been used to aggregate instance features and model mutual information.Source: §2.2 · p. 3
  • Most existing methods emphasize aggregating discriminative information from pre-extracted features, which lack fine-tuning for specific downstream tasks.Source: §2.2 · p. 3

3. Methodology

The method inserts online re-embedding between offline instance feature extraction and MIL aggregation, using regional attention and cross-region fusion to create task-adapted instance representations. The resulting module can be paired with MIL aggregators and classifiers, with AB-MIL used by default in R2T-MIL.

3.1. Preliminary

  • MIL separates feature extraction, instance aggregation, and bag classification, but offline, non-fine-tuned features can yield sub-optimal performance.The proposed re-embedding step follows feature extraction and updates instance representations before aggregation.Source: §3.1 · p. 3

3.2. Re-embedded Regional Transformer

  • R2T re-embeds input instance features by applying self-attention within local regions, then fusing information across regions.Regional attention uses instance ordering to reduce computation and highlight salient local features; cross-region attention connects regional representations.Source: §3.2 · pp. 3–4

3.3. R2Transformer-based MIL

  • R2T can be inserted between feature input and instance aggregation, and its re-embedding is optimized jointly with aggregation and classification.R2T-MIL uses AB-MIL's instance aggregation method and bag classifier by default.Source: §3.2 · p. 3, §3.3 · p. 5

3.2. Re-embedded Regional Transformer

  • EPEG embeds a lightweight positional encoding operation within regional self-attention to encode patch positions.It incorporates a 1-D convolution into the attention module.Source: §3.2 · p. 4

4. Experiments and Results

Experiments evaluate diagnosis, sub-typing, and survival across CAMELYON-16 and TCGA datasets, using task-appropriate metrics and cross-validation. R2T-MIL leads the reported comparisons, while experiments examine re-embedding, local attention, positional encoding, cross-region fusion, and feature distributions.

4.1. Datasets and Evaluation Metrics

  • The evaluation covers C16, TCGA-BRCA, and TCGA-NSCLC for diagnosis and sub-typing, plus TCGA-LUAD, TCGA-LUSC, and TCGA-BLCA for survival prediction.Diagnosis and sub-typing use Accuracy, AUC, and F1-score; survival uses C-index, with cross-validation protocols varying by dataset.Source: §4.1 · p. 5

4.2.1 Cancer Diagnosis, and Sub-typing

  • R2T-MIL leads all reported metrics on C16 and BRCA, with C16 gains over second-best methods of 0.59% Accuracy, 1.18% AUC, and 0.69% F1-score.On BRCA, its AUC improvement is 0.69%; it also improves on AB-MIL and foundation-model features in the reported comparisons.Source: §4.2.1 · pp. 5–6

4.2.2 Survival Prediction

  • R2T-MIL reaches C-index values of 61.13% on BLCA, 67.19% on LUAD, and 60.95% on LUSC, beating second-best methods by 2.23%, 3.08%, and 1.77%, respectively.The reported re-embedding gains over AB-MIL are 4.8%, 3.85%, and 3.3% on the three datasets, respectively.Source: Table 3 · p. 6, §4.2.2 · p. 6

4.3.1 Re-Embedding Matters in Pathology

  • Online R2T features outperform PLIP offline features on most tasks, and R2T boosts all tested MIL frameworks across tasks.Under AB-MIL, R2T gains 2.80% on C16 and 8.41% on LUAD, exceeding the second-best re-embedding approaches by 0.85% and 2.63% while using one-fifth their inference time.Source: §4.3.1 · pp. 6–7, Figure 4 · p. 7, Table 4 · p. 7

4.3.2 Local is more Appropriate than Global

  • Local self-attention outperforms global N-MSA for re-embedding, but requires around four times more training time; R2T reduces the computational cost.The authors also report that local attention increases re-embedded feature diversity.Source: §4.3.2 · p. 7

4.3.3 Effects of EPEG

  • EPEG improves results by 0.50% on C16, 0.39% on NSCLC, and 1.74% on LUAD.The authors contrast EPEG’s lightweight 1-D convolution of the attention matrix with conventional positional encodings that do not effectively improve performance.Source: Table 5 · p. 7, §4.3.3 · pp. 7–8

4.4. Visualization

  • Visualizations show R2T features becoming more discriminative and less homogeneous than offline or globally re-embedded features.The authors note that offline features remain insufficiently discriminative under extreme instance imbalance, including a 1:224 positive-to-negative ratio.Source: Figure 5 · p. 8, §4.4 · p. 8

4.3.4 Impact of Cross-Region MSA

  • Cross-region MSA improves feature distributions and quantitative results; multi-region semantics are reported as especially important for survival prediction.The authors connect this to survival being case-based, with bags containing multiple slides.Source: Table 6 · p. 8, §4.3.4 · p. 8

5. Conclusion

The conclusion emphasizes instance-feature re-embedding as a response to under-learning in conventional MIL and presents R2T as a versatile local Transformer module.

  • Instance-feature re-embedding alleviates under-learning of instance features in conventional MIL computational pathology.Source: §5 · p. 8
  • Transformer-based re-embedding modules consistently boost MIL methods across different architectures.Source: §5 · p. 8
  • R2T’s main contribution is a local Transformer with CR-MSA and EPEG, presented as a versatile re-embedding module in the foundation-model era.Source: §5 · p. 8

A. Additional Method Detail

The additional method detail describes the R-MSA representation with EPEG, its shared projection parameters, and the CR-MSA pseudocode.

A.1. Attention Matrix

  • The appendix gives an R-MSA representation that includes EPEG.Source: §A.1 · p. 11
  • R-MSA projection matrices are shared across regions, and Nhead denotes the number of heads.Source: §A.1 · p. 11

A.2. Pseudocodes of CR-MSA

B. Dataset Description

The appendix describes the datasets and slide-patch preparation used in the experiments, including C16’s cross-validation and weak-label setup.

  • C16 contains 400 WSIs, officially split into 270 training and 130 testing slides, with two test slides excluded as mislabeled.The study uses three-times three-fold cross-validation and slide-level labels for weak supervision.Source: B. Dataset Description · p. 11
  • TCGA-NSCLC includes LUAD and LUSC, with 541 and 512 diagnostic slides, respectively, while TCGA-BRCA includes 779 IDC and 198 ILC slides.Source: B. Dataset Description · p. 11
  • TCGA-BLCA contains 376 cases of Bladder Urothelial Carcinoma.Source: B. Dataset Description · p. 11
  • Each WSI is cropped into non-overlapping 256 × 256 patches at 20X magnification, with background regions discarded.Source: B. Dataset Description · p. 11

C. Implementation Details

The implementation uses projected patch features and a fixed training setup to train R2T with ResNet-50 or PLIP features.

C. Implementation Details

  • ResNet-50 patch features are reduced from 1024 to 512 dimensions, while PLIP’s 512-dimensional features are mapped to 512 dimensions.ResNet-50 uses ImageNet pretraining; its last convolutional module is removed and global average pooling produces the initial feature vector.Source: C. Implementation Details · p. 11
  • R-MSA uses 8 heads, and training uses Adam with a 2 × 10−4 learning rate and 1 × 10−5 weight decay.The learning rate follows a cosine schedule; models train for 200 epochs with early stopping and batch size 1.Source: C. Implementation Details · p. 11
  • Experiments use NVIDIA GPUs without gradient cropping or accumulation; early-stopping patience is 30 for CAMELYON-16 and 20 for TCGA.The paper points to Section G for code and pretrained PLIP weights.Source: C. Implementation Details · p. 11

D. Additional Quantitative Experiments

These experiments examine R2T with foundation-model features, EPEG variants, and regional partition choices. Results favor local convolutional positional encoding and moderate two-dimensional regions.

D.1. More on Foundation Model Features

  • The CR-MSA pseudocode forms representative features per region, applies native MSA, then redistributes and combines the outputs using dispatch weights.The pseudocode also computes min-max-normalized dispatch weights.Source: §D.1 · p. 12

ABMIL DSMIL CLAM DTFD IBMIL (b) Sub‐typing (NSCLC)

D.2. More on EPEG

  • Most 1-D convolution kernels improve local Transformer re-embedding, while larger kernels typically perform worse, attributed to overfitting from excess parameters on limited slides.Figure 10 discusses kernel sizes across datasets.Source: §D.2 · p. 12
  • The EPEG “value” variant is feasible but performs significantly below the default, especially on C16; the default places positional information in the core attention matrix.Table 7 compares EPEG variants; Figure 8 illustrates their structures, with the left variant as default.Source: §D.2 · p. 12, Table 7 · p. 13, Figure 8 · p. 13

D.3. Region Partition Strategy

  • Two-dimensional region partitioning outperforms 1-D partitioning, while regions that are too small or too large degrade performance; moderate regions are optimal for R-MSA.Small regions reduce receptive field, whereas large regions damage re-embedded feature diversity.Source: §D.3 · p. 13, Figure 7 · p. 13
  • R2T maintains high performance with many regions, including 512 or 32×32 partitions, supporting a performance-efficiency trade-off on longer input sequences.More regions lower spatial and temporal costs; Figure 7 reports performance under different partition settings on two datasets.Source: §D.3 · p. 13, Figure 7 · p. 13

D.4. More Parameters are not Always Better

The additional analyses find that increasing model size need not improve performance, local MSA suits WSI characteristics, and regional re-embedding improves feature discriminativeness in visualizations.

D.4. More Parameters are not Always Better

  • R2T-MIL has equal or fewer parameters than TransMIL with better performance, and re-embedding can yield higher performance at lower cost than the matched “+N-MSA” structure.The paper also reports that R2T-MIL can reach foundation-model performance with 2× DTFD’s parameters.Source: §D.4 · p. 13, Table 8 · p. 14
  • Adding FFN and layers increases R2T-MIL parameters by 7.35M but drops LUAD performance by 4.12%.The efficiency analysis also contrasts R2T-MIL’s parameter size and performance with Transformer-based methods and DTFD.Source: §D.4 · pp. 13–14, Table 8 · p. 14

D.5. More on Local Transformer

  • More use of global MSA worsens performance on both datasets; local Transformers better adapt to WSI scale and small tumor areas.The global-MSA strategy adds hyperparameters that reduce generalization ability; Table 9 compares global and local MSA.Source: §D.5 · p. 14, Table 9 · p. 14

D.6. Discussion of Hyper-parameter in CR-MSA

  • R2T performs at a high level across different CR-MSA parameter values, with similar consistency across offline features.Figure 10’s bottom results support generality of the preset optimal parameter across scenarios.Source: §D.6 · p. 14

E. Additional Visualization

  • When tumor prediction confidence is too low, visualizations use original instance labels rather than attention scores to color instances.Attention scores after softmax otherwise label instances in the visualizations.Source: E. Additional Visualization · pp. 14–15
  • Regional MSA produces more discriminative re-embedded features than global MSA when tumor-instance proportions are low, supporting slide classification.The visualizations describe high cohesion and low coupling for regional features; global MSA can lack feature diversity when instance counts are high and tumor-instance ratios are extremely low.Source: E. Additional Visualization · p. 15

F. Limitation

Re-embedding can erase original instance-label information, constraining the use of the final aggregation module for weakly supervised tumor localization or segmentation.

F. Limitation

  • Re-embedded features can lose original label information, limiting accurate weakly supervised tumor localization or segmentation through the final aggregation module.The authors identify improving detection or segmentation with re-embedding as future work.Source: F. Limitation · p. 15

Related papers

Loading 2402.17228v4…