PaperlayerView exact revision on arXiv

Source-linked AI summary

Scaling Vision Transformers to Gigapixel Images via Hierarchical Self-Supervised Learning

Richard J. Chen, Chengkuan Chen, Yicong Li, Tiffany Y. Chen, Andrew D. Trister, Rahul G. Krishnan, Faisal Mahmood

arXiv:2206.02647v1cs.CV

TL;DR

Gigapixel whole-slide images pose challenges from enormous sequence lengths and hierarchical, fixed-scale visual tokens. HIPT addresses these challenges with hierarchical Transformer aggregation and self-supervised pretraining, achieving stronger slide-level performance than conventional weakly supervised approaches across evaluated tasks.

  • Problem

    Gigapixel whole-slide images contain hierarchical visual tokens across resolutions, while small-token processing makes self-attention intractable and large-token processing loses detail.

  • Method

    HIPT uses Transformer-based hierarchical aggregation and recursive self-supervised pretraining to learn slide-level representations for gigapixel pathology images.

  • Results

    HIPT consistently achieves the highest macro-averaged AUC across slide-level classification tasks and outperforms conventional MIL approaches, especially for context-aware survival prediction.

  • Takeaways & Limitations

    Hierarchical pretraining is important for HIPT, while self-supervised ViT representations support slide-level weakly supervised and K-nearest-neighbor evaluation.

  • Takeaways & Limitations

    HIPT is difficult to pretrain end-to-end, and the study is limited by TCGA-only data, excluded slides, and evaluation lacking independent test cohorts.

Abstract

from arXiv · show

Vision Transformers (ViTs) and their multi-scale and hierarchical variations have been successful at capturing image representations but their use has been generally studied for low-resolution images (e.g. - 256x256, 384384). For gigapixel whole-slide imaging (WSI) in computational pathology, WSIs can be as large as 150000x150000 pixels at 20X magnification and exhibit a hierarchical structure of visual tokens across varying resolutions: from 16x16 images capture spatial patterns among cells, to 4096x4096 images characterizing interactions within the tissue microenvironment. We introduce a new ViT architecture called the Hierarchical Image Pyramid Transformer (HIPT), which leverages the natural hierarchical structure inherent in WSIs using two levels of self-supervised learning to learn high-resolution image representations. HIPT is pretrained across 33 cancer types using 10,678 gigapixel WSIs, 408,218 4096x4096 images, and 104M 256x256 images. We benchmark HIPT representations on 9 slide-level tasks, and demonstrate that: 1) HIPT with hierarchical pretraining outperforms current state-of-the-art methods for cancer subtyping and survival prediction, 2) self-supervised ViTs are able to model important inductive biases about the hierarchical structure of phenotypes in the tumor microenvironment.

1. Introduction

Whole-slide images pose a scale and context challenge for representation learning: fixed-scale visual tokens form hierarchies, while conventional MIL relies on limited patches and global pooling. HIPT addresses this with hierarchical Transformer aggregation and self-supervised representations for slide-level pathology tasks.

  • WSIs can reach 150000 × 150000 pixels and require tissue phenotyping for diagnosis, prognosis, and treatment-response estimation.
  • Conventional MIL patches WSIs at one magnification, extracts patch embeddings, and globally pools them under slide-level supervision.
  • 256 × 256 windows capture fine morphology but provide limited context for invasion, tumor size, immune infiltrates, and tissue organization.
  • WSIs contain fixed-scale visual tokens arranged hierarchically from cell-level 16 × 16 images to broader 4096 × 4096 tissue regions.
  • HIPT performs bottom-up Transformer aggregation from 16 × 16 tokens through 256 × 256 and 4096 × 4096 windows to form slide-level representations.
  • HIPT outperforms conventional MIL, especially on context-aware survival prediction, while self-supervised attention models fine- and coarse-grained histopathology concepts.

2. Related Work

Prior work established set-based MIL for gigapixel histology and increasingly incorporated pretrained instance features and richer aggregation. Vision Transformer research separately advanced patch-sequence modeling and multiscale image pyramids, motivating their combination for WSIs.

  • Multiple Instance Learning in WSIs: Set-based deep learning and MIL architectures enabled weakly supervised learning from histology regions and gigapixel WSIs.
  • Multiple Instance Learning in WSIs: MIL variants incorporated VAE-GAN, SimCLR, and MOCO pretraining, while later methods expanded aggregation layers and scoring functions.
  • Vision Transformers and Image Pyramids: Vision Transformers formulate 256 × 256 images as sequences of 16 × 16 visual tokens, while related work integrates multiscale information through image-pyramid designs.

3. Method

HIPT represents gigapixel WSIs through nested visual-token aggregation from cell- to patch- to region- and slide-level contexts. Its hierarchical Transformer design supports self-supervised pretraining across resolutions, addressing the fixed-scale and long-sequence constraints of conventional MIL.

  • Motivation: HIPT addresses conventional MIL limitations by modeling coarse tissue organization alongside fine morphology and by replacing purely global pooling with hierarchical Transformer aggregation.Fixed 256 × 256 windows have limited context for invasion, tumor size, lymphocytic infiltrates, and broader phenotype organization, while long WSI sequences constrain standard self-attention.
  • Hierarchical architecture: HIPT recursively aggregates [16 × 16] cell tokens within [256 × 256] patches, [256 × 256] patches within [4096 × 4096] regions, and regions into a slide representation.ViT256-16 produces [CLS]256 tokens, ViT4096-256 aggregates them into [CLS]4096 tokens, and ViTWSI-4096 aggregates region tokens at slide level.
  • Hierarchical architecture: The architecture uses Transformer self-attention as a permutation-equivariant aggregation layer to model dependencies between visual concepts at each resolution.The design treats WSIs as nested sequences whose tokens recursively break down to the cell level.
  • Hierarchical architecture: ViT256-16 processes 256 tokens per 256 × 256 window, while ViT4096-256 also processes 256 patch tokens per 4096 × 4096 region; slide-level sequence lengths are usually below 256.This preserves a manageable forward-pass sequence length while increasing spatial context.
  • Hierarchical pretraining: Hierarchical pretraining first applies DINO to ViT256-16, then freezes its weights and reuses it to embed inputs for a second DINO stage training ViT4096-256.The procedure transfers self-supervised ViT pretraining across two nested resolution levels; it does not reach the slide level.
  • Hierarchical pretraining: Pretrained x4096 representations are competitive with supervised methods for slide-level subtyping, while two-stage hierarchical pretraining achieves state-of-the-art performance.The paper reports these outcomes for slide-level representation learning in gigapixel histopathological images.

4. Experiments

HIPT is evaluated through slide-level classification, survival prediction, embedding comparisons, and attention-map analyses. Across these experiments, hierarchical pretraining and multi-scale representations support strong performance and reveal tissue phenotypes at different spatial resolutions.

  • Experimental setup: HIPT is compared with weakly supervised architectures, pretrained and unpretrained variants, and self-supervised embeddings across slide-level classification and survival tasks.The experiments include attention-map visualization of learned morphological features.
  • Slide-Level Classification: HIPT achieves the highest macro-averaged AUC across classification tasks and training-data fractions, with improvements over CLAM-SB on BRCA, NSCLC, and RCC subtyping.Using 100% of training data, the increases are 1.86%, 2.59%, and 0.72%; using 25%, they are 3.14%, 8.33%, and 1.78%, respectively.
  • Embedding Comparisons: ViT4096-256 region-level embeddings outperform ViT256-16 patch embeddings across all evaluated tasks and surpass weakly supervised approaches on BRCA and RCC subtyping with full training data.The broader regions provide contexts closer to the slide-level view used in evaluation.
  • Survival Prediction: HIPT achieves the best c-Index on IDC, COADREAD, CCRCC, and STAD survival tasks, with the largest reported values of 0.634 and 0.608 on IDC and COADREAD.The paper emphasizes these tasks because context-aware relationships are important for survival prediction.
  • Self-Supervised ViTs Find Unique Morphological Phenotypes: ViT256-16 attention heads isolate stroma, red blood cells, cells, lymphocytes, and white spaces, while hierarchical maps connect tumor cells with stroma and poorly differentiated glands.The maps combine fine-grained and coarse-grained attention distributions in colorectal cancer.
  • Further Ablation Experiments: Hierarchical pretraining is important because HIPT variants without pretraining overfit in MIL tasks, while pan-cancer pretraining improves cell localization.The ablations also find strong representation quality for ViT256-16 image patches.

5. Conclusion

HIPT advances self-supervised slide-level representation learning, but its training procedure and study design remain constrained by computational and dataset limitations.

  • HIPT features achieve superior performance in weakly-supervised and KNN evaluations.
  • Pretraining the last aggregation layer is difficult because only a small number of WSI data points are available.
  • End-to-end hierarchical HIPT pretraining is computationally intractable on commercial workstations and must be performed in stages.
  • The study excludes some slides, lacks independent test cohorts, and is limited to TCGA, which overrepresents patients with European ancestry.

A. Multi-Head Self-Attention

Multi-head self-attention computes token relationships through learned query, key, and value projections, then aggregates information across multiple attention heads.

  • The input is a sequence of visual tokens extracted from an L × L patch, with sequence length M and embedding dimension d_l.
  • Self-attention projects each token into query, key, and value vectors using learned matrices W_q, W_k, and W_v.
  • Each token compares its query with other tokens' keys, normalizes the scores, and uses them to form a weighted sum of value vectors.
  • Using h attention heads increases feature expressiveness and supports permutation-equivariant aggregation of visual tokens.
  • Multi-head self-attention also learns representations of morphological phenotypes.

B. Additional Implementation Details

HIPT inference processes nested WSI regions bottom-up, while staged DINO pretraining and attention visualization support hierarchical representation learning and interpretation.

  • Inference uses a slide-level batch size of one and organizes inputs by 4096 regions, 256 patches, and 16-cell tokens.
  • DINO trains a student ViT256-16 to match a Siamese teacher's probability distribution using cross-entropy with local and global crops.
  • A second data loader batches 256-sized views, enabling bottom-up aggregation beginning with the ViT256-16 forward pass.
  • The pretraining corpus contains 10,687 FFPE H&E diagnostic slides from 33 cancer types, 408,218 4096 × 4096 regions, and 104M 256 × 256 patches.
  • Attention heatmaps interpolate token scores to 16 × 16 or 256 × 256 resolution, with overlapping stride 64 used for more granular ViT4096-256 maps.

C.1. Variations in HIPT Architecture

HIPT variants test attention, hierarchical pretraining, and freezing at different aggregation levels; the ablation evidence indicates that pretrained and frozen aggregation improves robustness on small WSI datasets.

  • Variations in HIPT Architecture: The ablation study evaluates Transformer attention, pretraining, and parameter freezing in slide-level classification and KNN evaluations.
  • Variations in HIPT Architecture: The attention-pooling variant removes ViT aggregation layers while retaining pre-extracted ViT256-16 features.
  • Variations in HIPT Architecture: One variant trains ViT4096-256 and ViTWSI-4096 aggregation layers from scratch.
  • Variations in HIPT Architecture: Another variant adds Stage 2 hierarchical pretraining to ViT4096-256, while a further variant also freezes that pretrained layer and finetunes only ViTWSI-4096.
  • Variations in HIPT Architecture: 0.923 to 0.652 and 0.952 to 0.820: without freezing ViT4096-256, NSCLC subtyping performance drops with 25% and 100% training data, respectively.

C.2. Assessing Quality of x256 Representations

ViT256-16 representations are robust across patch-level histology datasets and preserve meaningful global morphology, while slide-level advantages vary by cancer subtype and task difficulty.

  • Comparison with ImageNet Features: Global structures for stroma, tumor, normal, and mucous tissue are preserved across ViT models, unlike the poorly preserved structures in ResNet-50 features.The ViT representations retain class-level organization in UMAP visualizations with and without stain normalization.
  • Comparison with ImageNet Features: ResNet-50 features outperform ViT256-16 features on NSCLC and RCC subtyping, whereas ViT256-16 performs better on BRCA subtyping.The authors associate this pattern with easier subtype separation in NSCLC and RCC and greater stromal-context demands in BRCA.
  • Organ-Specific versus Pan-Cancer Image Pretraining: Pan-cancer and BRCA-pretrained ViT256-16 models perform similarly on most patch-level tasks, with pan-cancer pretraining slightly better on BCSS and BreastPathQ.BRCA-pretrained features preserve stroma and mucous structure better on CRC, while pan-cancer features localize tumor cells and lymphocytes more effectively.
  • Concluding Remarks on x256 Representations: The 384-dimensional ViT256-16 embedding keeps Stage 2 hierarchical pretraining and ViT4096-256 finetuning tractable on commercial workstations.Feature concatenation from the last four hidden stages instead produced a 1536-dimensional embedding without improving performance.

D. Additional Visualizations

Hierarchical attention visualizations show that HIPT combines cell-level and tissue-level attention to localize prognostically relevant tumor patterns in breast and colorectal cancer.

  • Additional Visualizations: Additional visualizations provide native-resolution attention maps and repository-based views of ViT256-16, ViT4096-256, and hierarchical attention.The repository supplements the space-limited figures with further model visualizations.
  • Hierarchical Attention Maps for Invasive Breast Carcinoma: Combined ViT256-16 and ViT4096-256 attention localizes invasive tumor cells in desmoplastic stroma and larger tumor nests in breast carcinoma.These patterns are relevant to distinguishing invasive ductal from lobular carcinoma and to survival outcomes.
  • Hierarchical Attention Maps for Colorectal Cancer: Combined hierarchical attention localizes invasive colorectal tumor cells in muscle and stromal regions and identifies poorly differentiated glands.Both patterns are described as prognostic histopathologic biomarkers for cancer staging and survival outcomes.
Loading 2206.02647v1…