A graph-transformer for whole slide image classification

In plain words

A whole slide image is a microscope scan of a tissue sample, so large that it holds thousands of small tiles. The task here is to label the whole slide as normal lung tissue, lung adenocarcinoma, or lung squamous cell carcinoma. Most earlier methods looked at tiles one at a time, or gave every tile the slide's label, which adds noise.

The method keeps the layout. Each tile becomes a point described by a feature vector from a network that was first trained, without labels, to recognise that two differently altered versions of the same tile belong together. Tiles that touch on the slide are linked, so the slide becomes a graph. The graph is first smoothed by neighbours, then squeezed into a smaller set of groups, and a transformer reads those groups to produce one label. A companion method, GraphCAM, colours the slide to show which regions pushed the decision towards each label, so a pathologist can check the result.

The authors report strong accuracy on their main data and a drop on a second, external dataset. The numbers below are theirs.

TL;DR

GTP is a graph convolution, min-cut pooling and a small transformer stacked on features from a contrastive-learning ResNet18. It reaches 91.2% three-class accuracy on CPTAC under 5-fold cross-validation and 82.3% on TCGA. GraphCAM gives class-specific heatmaps. The official code differs from the paper in several training and architecture details, listed under Paper vs code.

Problem

Patch-level models give every tile the slide label, which injects label noise, and they miss how tiles relate across the slide §1. Transformers on all tiles would need quadratic memory in the number of tiles, and a slide can contain thousands of them §1 §2.2. ImageNet pretraining is common, but the paper's ablation finds ImageNet features weaker than its contrastive features §1 tab. 3. The paper uses slide-level labels from CPTAC, TCGA and NLST §2.1 tab. 1.

Method

The pipeline has three stages: patch feature extraction, graph construction with graph-transformer inference, and GraphCAM for explanation fig. 1.

Patch features. Each slide is cut into non-overlapping 512 × 512 patches at 20× magnification, and patches that are more than half background are discarded §3.1. A ResNet18 is trained with a contrastive loss (NT-Xent): two random augmentations of the same patch form a positive pair, and the other 2K − 1 views in the mini-batch are negatives §2.2 eq. 1 fig. 1c. A projection MLP sits on top during training; the 512-d backbone output is used as the node feature §2.2 fig. 1c. NLST data (about 1.8 million patches) is used only for this stage §3.

Graph construction. Each retained patch is a node. Two nodes are joined by an edge when their patches are neighbours in the 8-neighbourhood on the slide, giving a binary adjacency matrix A with at most 8 ones per row §2.2 fig. 1b.

Graph-transformer. Three pieces run in sequence §2.2:

1. One graph-convolution (GC) layer propagates features over the graph using the normalised adjacency with self-loops eq. 2a–2b.
2. A dense learned assignment S maps the N nodes onto a fixed number of clusters, following min-cut pooling. Pooled features are S-transposed times the node features; the pooling layer has its own min-cut and orthogonality losses that are added to the classification loss §2.2 fig. 1a.
3. A learnable class token is prepended to the pooled cluster tokens and a vision-transformer encoder processes them eq. 3a–3d eq. 4a–4c. No positional embedding is added; spatial information enters only through the adjacency used in the GC layers and pooling §2.2. The class token, after a layer norm, gives the prediction eq. 5.

# GTP forward pass for one WSI  [§2.2] [eq. 2a-2b] [eq. 3a-3d] [eq. 5]
# Input:  features (N, 512), adjacency (N, N) of 0/1 (8-neighbour graph)
# Output: logits (3,), training loss (scalar)

def gtp_forward(features, adjacency, label):
    # Add self-loops and symmetric-normalise the adjacency  [eq. 2b]
    loops = adjacency + identity(N)
    degree = diagonal(row sums of loops)
    adj_hat = inv_sqrt(degree) @ loops @ inv_sqrt(degree)
    # One graph-convolution layer with ReLU  [eq. 2a] [§3.1]
    # Code: no ReLU and no degree normalisation  [code: models/gcn.py#L364-L372]
    hidden = relu(adj_hat @ features @ weight_gc)       # (N, 64)
    # Dense soft assignment of the N nodes to 100 clusters  [§2.2]
    assign = softmax(hidden @ weight_s, axis=1)         # (N, 100)
    pooled = assign.T @ hidden                          # (100, 64)
    # Pooling losses; the paper gives no formula  [§2.2]
    # Code: PyG dense_mincut_pool  [code: models/GraphTransformer.py#L67]
    cut_loss, ortho_loss = mincut_terms(adjacency, assign)
    # Prepend the learnable class token  [§2.2] [eq. 4a]
    tokens = concatenate(class_token, pooled)           # (101, 64)
    # Three pre-norm blocks, 8 heads, MLP width 128  [§3.1] [eq. 4b-4c]
    for block in range(3):
        tokens = tokens + attention(layer_norm(tokens)) # head size 8
        tokens = tokens + mlp(layer_norm(tokens))       # GELU in code
    # Read out the class token  [eq. 5]
    # Code adds one linear head after the layer norm  [code: models/ViT.py#L338]
    logits = linear_head(layer_norm(tokens[0]))
    loss = cross_entropy(logits, label) + cut_loss + ortho_loss
    # Code: unit weights on both pooling terms  [code: models/GraphTransformer.py#L76]
    return logits, loss

Contrastive pretraining. The loss for one positive pair (i, j) among the 2K augmented views is the negative log-probability of the positive, with cosine similarity scaled by a temperature τ §2.2 eq. 1:

# NT-Xent loss for one positive pair  [§2.2] [eq. 1]
# z: L2-normalised projections of the 2K augmented views of a mini-batch
# Output: scalar loss for the pair (i, j)

def pair_loss(z, i, j, tau):
    # Scaled cosine similarity of view i to every other view  [eq. 1]
    scores = [dot(z[i], z[k]) / tau for k in range(2 * K) if k != i]
    # Negative log-probability of the positive view j  [eq. 1]
    return logsumexp(scores) - dot(z[i], z[j]) / tau

The paper computes the loss over both (i, j) and (j, i) in a mini-batch §2.2.

GraphCAM. GraphCAM adapts the transformer relevance method of Chefer et al. §2.3. For each attention block l, it takes the gradient of the attention map and the layer relevance, multiplies them element-wise, averages over heads, adds the identity, and multiplies the block matrices together eq. 6a–6b. The resulting transformer relevance is pushed back to the graph nodes through the pooling assignment S and drawn onto the slide using the adjacency and patch coordinates §2.3.

# GraphCAM for one class  [§2.3] [eq. 6a-6b]
# Input:  per block l: attention A_l (heads, T, T), gradient dA_l,
#         relevance R_l; T = 101 tokens; assign S (N, 100)
# Output: relevance per patch, shape (N,)

def graphcam_patch_map(attn, grad_attn, relevance, assign):
    block_maps = []
    for layer in range(L):
        # Head-averaged Hadamard product plus identity  [eq. 6b]
        # Code clamps the product at zero first  [code: models/ViT.py#L382]
        mixed = mean_over_heads(grad_attn[layer] * relevance[layer])
        block_maps.append(mixed + identity(T))
    # Product of block maps; code multiplies later blocks on the left  [eq. 6a]
    rollout = product(block_maps)
    # Class-token row gives cluster relevances  [code: models/ViT.py#L385]
    cluster_relevance = rollout[0, 1:]
    # Spread cluster relevance back to patches with the assignment  [§2.3]
    return assign @ cluster_relevance

Pitfall: the paper describes the relevance rules only by reference to Chefer et al., but the official code implements its own relevance propagation for each module (linear, attention, layer norm, add, clone), which the paper does not specify ViT.py.

Hyperparameters

Values are as stated in the paper unless marked "code". Where the paper is silent, the table says so. Citations point to the paper or the official code at the commit in the front matter.

Table 4 hyperparameter ablation (all rows trained and tested on a 70/30 split of CPTAC; accuracy on the held-out CPTAC cases, from tab. 4):

Results

Author-reported. Setup: three-class task (normal, LUAD, LSCC). Training on CPTAC (2,071 WSIs) with 5-fold cross-validation; TCGA (2,082 WSIs) is the external test set, with the same hyperparameters §3 tab. 1.

Headline numbers §1 tab. 2:

Full accuracy and AUC table, including the graph-only variant GTP* (no transformer) tab. 2:

Per-class precision, recall and specificity for GTP (mean, SD in parentheses; tab. 2 panel a; the baseline rows are in the same table):

Baselines in the same panel, in order of the table's rows: TransMIL CPTAC precision 89.7 / 81.0 / 87.1, recall 90.4 / 81.2 / 85.9; TransMIL TCGA precision 76.6 / 64.3 / 80.4, recall 87.3 / 75.8 / 55.6; AttPool CPTAC precision 88.4 / 77.1 / 80.6, recall 85.9 / 78.0 / 81.6; AttPool TCGA precision 89.1 / 69.9 / 81.4, recall 87.8 / 79.4 / 71.3; GTP* CPTAC precision 82.9 / 81.6 / 86.5, recall 93.6 / 74.4 / 80.3; GTP* TCGA precision 72.5 / 69.1 / 82.0, recall 90.7 / 57.7 / 69.9 tab. 2.

DeLong test, log10 p-value of AUC differences (threshold log10(0.05) = −1.301) tab. 2 panel c:

Feature-extractor ablation, accuracy and AUC tab. 3 panel b:

Per-class precision, recall and specificity for the feature-extractor ablation are in tab. 3 panel a; they are not copied here.

GraphCAM §4 fig. 4 fig. 5: maximum IoU against the pathologist's annotation was 0.857 at a threshold of 0.6 for one LUAD slide and 0.733 at 0.2 for one LSCC slide (CPTAC). On 20 TCGA slides (10 LUAD, 10 LSCC), the mean of the per-slide maximum IoU was 0.817 ± 0.165. GraphCAMs were consistent across the five cross-validation models fig. 5; in failure cases they were not interpretable fig. 6.

Compute §3.2 §3.3: PyTorch 1.9.0 on one NVIDIA 2080 Ti (11 GB). Training ran about 2.4 iterations per second and converged in under a day. Inference took under one second per slide at batch size 2. Training the transformer on raw patches, without the graph component, ran out of memory above 4,760 patches at batch size 1 on that card §3.2.

Paper vs code

Official code: github.com/vkola-lab/tmi2022 at commit 0d395a160a969578fa5ce66a8a7f0fd0817712b4, MIT licence, from the Kolachalama laboratory (the paper's corresponding group) LICENSE. The paper links it §2.4. The repository has no k-fold script and does not ship the split files train_set.txt, val_set.txt or test_set.txt, so the cross-validation protocol cannot be reproduced from the repository alone.

Differences found:

• GC layer. Paper uses ReLU and the symmetric normalised adjacency with self-loops eq. 2a–2b. Code adds self features to the raw adjacency sum, applies a linear map, L2-normalises each node, then batch-normalises, with no ReLU gcn.py GraphTransformer.py.
• GC layers. Paper text says one GC layer §3.1; Table 4 top row lists 3 tab. 4; code uses 1 GraphTransformer.py.
• Min-cut clusters. Table 4 top row (best) uses 120 tab. 4; code default is 100 GraphTransformer.py.
• Pooling losses. Paper describes a custom min-cut loss without a formula §2.2; code sums cut and orthogonality losses with unit weight into the classification loss GraphTransformer.py.
• Training length and schedule. Paper: 150 "iterations", LR 1e-3 with decay at 30 and 100 §3.1. Code: 120 epochs, MultiStepLR milestones at 20 and 100, gamma 0.1 option.py main.py. A comment in main.py reads "20,90,130 -> 150", which suggests earlier schedules.
• Optimiser for GTP. Paper names Adam only in the feature-extractor paragraph §3.1; code uses Adam with weight decay 5e-4 main.py.
• Checkpoint selection. Code saves the epoch with the best validation accuracy main.py. The paper does not describe a validation split.
• Feature-extractor optimiser. Paper: Adam, LR 1e-4 §3.1. Code: Adam at 1e-5, hard-coded simclr.py.
• Feature-extractor batch size. Paper: 512 §3.1. Code config: 256 config.yaml.
• Augmentations. Paper mentions random cropping with resizing §2.2. Code's random crop is commented out and replaced by a fixed resize; horizontal flip is added dataset_wrapper.py.
• Temperature, epochs, weight decay. Not in the paper; code uses τ = 0.5, 20 epochs, weight decay 1e-5 config.yaml config.yaml config.yaml.
• Feature-extractor schedule. Code: constant LR for 10 epochs, then cosine with T_max = 20 simclr.py. The paper says cosine annealing only.
• Input resolution. Code resizes 512 px patches to 224 px for both contrastive training and graph building build_graphs.py. The paper does not mention this step.
• Normalisation. Code initialises ResNet18 without pretrained weights and uses InstanceNorm2d resnet_simclr.py. Paper is silent on initialisation and normalisation.
• Transformer details. Code uses GELU in the MLP, LayerNorm eps 1e-6 in blocks, no query/key/value bias, and a linear head after the final layer norm ViT.py ViT.py ViT.py. Paper is silent on the activation and gives eq. 5 without a head.
• Feature weight loading. build_graphs.py loads the contrastive weights by pairing state-dict entries by position with strict=False build_graphs.py. The pairing depends on key order and may map projection-head weights onto the classifier. I did not run it.
• Cohort counts. Paper Table 1: CPTAC 2,071 WSIs (719 / 667 / 685 normal / LUAD / LSCC) and TCGA 2,082 WSIs (534 / 808 / 740) tab. 1. The shipped data/cptac.txt lists 2,082 slides (734 / 669 / 679) and data/tcga.txt lists 1,346 (420 / 476 / 450) cptac.txt tcga.txt. Which subset was used is not established.
• README cohort size. The README mentions 3,024 WSIs. The paper gives 4,818 across CPTAC, TCGA and NLST, with NLST (665) used only for contrastive pretraining tab. 1 §2.1.
• Feature extractor weights. A pretrained extractor is shipped as feature_extractor/model.pth model.pth. I did not check it against the paper's training or its reported numbers.
• GraphCAM. Paper gives no clamp, no class-probability scaling, and no row choice. Code clamps at zero before averaging heads ViT.py, reads the class-token row ViT.py, multiplies each class map by the predicted probability clipped to [0.4, 1] vis_graphcam.py vis_graphcam.py, and skips patches in the first 5 rows and the last 5 columns and rows vis_graphcam.py.
• Annotation and IoU. The repository has the expert annotations as image files but no IoU or binarisation code. The annotation readme says only that expert annotations were added for CPTAC readme. The reported IoU thresholds cannot be checked from the repository.

Pitfalls

• The "best" hyperparameter row is not the main result. Table 4 reports 92.5% accuracy on a 70/30 split of CPTAC with 3 GC layers and 120 clusters. The main 5-fold result is 91.2% tab. 2 tab. 4. Do not mix them.
• Which GC depth? The paper says one GC layer in the text and three in the best row of Table 4. The code uses one. Start from one and treat three as an ablation unless the authors confirm otherwise.
• Min-cut loss weights. The paper names a pooling loss but gives no formula or weight. The code's unit weights are the only specification available.
• Normalisation choice matters for reproduction. The paper's ReLU and symmetric normalisation are absent from the code. Pick one deliberately and report which.
• Validation leakage risk. Selecting the checkpoint by validation accuracy is standard, but the paper does not say which split was used. If the validation split is the CPTAC test fold, reported test numbers are optimistic.
• Splits are not in the repository. Fold membership, the train/val/test lists and the patch-level holdout for contrastive training are absent. Build your own, and state it.
• STL-10 versus NLST on CPTAC. Table 3 panel b shows STL-10 slightly above NLST on CPTAC accuracy and AUC (91.4 vs 91.2, 98.0 vs 97.7). NLST is clearly better on TCGA (71.9 vs 82.3 accuracy). The text's "significantly better" claim §4 should be read for the external test set.
• Patch graph size and memory. Slides can exceed 4,760 patches. A transformer over all patches ran out of memory on an 11 GB card §3.2, which is why pooling is needed. Pooling to 100 clusters keeps the transformer cheap, but the GC layer and the dense assignment still scale with N.
• Overlap versus non-overlap. A 10% overlap lowered CPTAC accuracy to 0.881 in Table 4 tab. 4. Non-overlapping tiles are the default in the paper and the code.
• GraphCAM is per slide. The README notes it handles one slide at a time [README]. Each patch value is a soft-assignment mix of the 100 cluster relevances, so the map's resolution is set by the cluster count.
• GraphCAM has no fixed binarisation. The paper reports IoU at several thresholds and gives the best value per slide §4. Choose a threshold on held-out slides, not on the annotated ones.
• Relevance rules are code-only. Reproducing GraphCAM means implementing the relevance propagation in models/ViT.py and models/layers.py, which the paper does not describe.
• Generalisation is weaker externally. Accuracy falls from 91.2% (CPTAC CV) to 82.3% (TCGA) tab. 2. The paper's own TCGA normal recall is 92.6%, but LSCC recall is 75.2%, and the paper says the model confused LSCC and LUAD on TCGA §4.

Lineage

• Builds on: contrastive learning with NT-Xent and a projection head §2.2; graph convolution §2.2 eq. 2a–2b; min-cut pooling §2.2; the Vision Transformer encoder §2.2 eq. 3a–3d eq. 4a–4c; Chefer et al.'s transformer relevance propagation §2.3.
• Compared against: TransMIL §3.2 tab. 2; AttPool §3.2 tab. 2; ResNet18 with and without fine-tuning and a convolutional autoencoder as feature extractors tab. 3.
• Superseded by: not stated in the paper; no later work is referenced here.