[Paper Review] MHAttnSurv: Multi-Head Attention for Survival Prediction Using Whole-Slide Pathology Images

Shuai Jiang, Arief A. Suriawinata+1 more|arXiv (Cornell University)|Oct 22, 2021

AI in cancer detection44 references4 citations

TL;DR

MHAttnSurv proposes a multi-head attention mechanism to improve survival prediction from whole-slide pathology images (WSIs), focusing on diverse morphological patterns across tumor slides. Evaluated on four TCGA cancer types, it achieves a c-index of 0.640, outperforming prior state-of-the-art methods with c-indices of 0.603 and 0.619.

ABSTRACT

In pathology, whole-slide images (WSI) based survival prediction has attracted increasing interest. However, given the large size of WSIs and the lack of pathologist annotations, extracting the prognostic information from WSIs remains a challenging task. Previous studies have used multiple instance learning approaches to combine the information from multiple randomly sampled patches, but different visual patterns may contribute differently to prognosis prediction. In this study, we developed a multi-head attention approach to focus on various parts of a tumor slide, for more comprehensive information extraction from WSIs. We evaluated our approach on four cancer types from The Cancer Genome Atlas database. Our model achieved an average c-index of 0.640, outperforming two existing state-of-the-art approaches for WSI-based survival prediction, which have an average c-index of 0.603 and 0.619 on these datasets. Visualization of our attention maps reveals each attention head focuses synergistically on different morphological patterns.

Motivation & Objective

• To address the challenge of extracting prognostic information from large, unannotated whole-slide images (WSIs) in pathology.
• To overcome limitations of random patch sampling in multiple instance learning by focusing on biologically relevant morphological patterns.
• To develop a deep learning model that adaptively attends to diverse regions of WSIs for more comprehensive feature representation.
• To improve survival prediction performance over existing WSI-based approaches using attention mechanisms.
• To provide interpretable attention maps that reveal which tissue patterns contribute most to survival predictions.

Proposed method

• A multi-head attention mechanism is applied to patch-level features extracted from whole-slide images (WSIs) to learn dynamic, attention-weighted representations.
• Patch features are extracted using a pre-trained convolutional neural network (e.g., ResNet) applied to random crops from WSIs.
• The attention mechanism computes query, key, and value matrices across patches, allowing the model to focus on relevant spatial patterns.
• Multiple attention heads allow the model to attend to different morphological patterns (e.g., tumor cells, stroma, necrosis) in parallel.
• A global pooling layer aggregates the attended features, which are then passed through fully connected layers to predict survival outcomes.
• The model is trained end-to-end using a Cox proportional hazards loss to optimize survival prediction.

Experimental results

Research questions

• RQ1 Can multi-head attention mechanisms effectively identify and prioritize diverse morphological patterns in whole-slide images for survival prediction?
• RQ2 How does attention-based feature aggregation compare to random patch sampling or global average pooling in WSI-based survival modeling?
• RQ3 To what extent do individual attention heads focus on biologically meaningful tissue structures in WSIs?
• RQ4 Does the proposed attention mechanism improve c-index performance compared to existing state-of-the-art WSI survival models?
• RQ5 Can attention maps provide interpretable insights into the model’s decision-making process for survival prediction?

Key findings

• MHAttnSurv achieved a mean c-index of 0.640 across four cancer types from The Cancer Genome Atlas (TCGA), outperforming two state-of-the-art methods with c-indices of 0.603 and 0.619.
• Visualization of attention maps revealed that different attention heads focus on distinct morphological patterns, such as tumor cell clusters, stromal regions, and necrotic areas.
• The model demonstrated improved robustness and generalization by attending to multiple relevant tissue regions rather than relying on a single patch or global average.
• The attention mechanism enabled more interpretable predictions, as attention maps highlighted pathologically relevant structures in the slides.
• The performance gain was consistent across all four evaluated cancer types, indicating broad applicability of the method.
• The results suggest that multi-head attention effectively captures heterogeneous prognostic features in WSIs, leading to better survival prediction.