TransMIL: Transformer based Correlated Multiple Instance Learning for Whole Slide Image Classification

In plain words

A whole-slide image is a scan of a tissue sample so large that it holds thousands of small patches. Usually only the whole slide has a diagnosis label, not each patch, which makes the task "weakly supervised". Common methods look at each patch on its own, or pool patch scores with a simple rule, and so ignore how patches relate to one another.

TransMIL lets the patches of one slide attend to each other with a Transformer, so the decision can use the context around a region and the relations between regions. That is the idea the authors want to capture from how a pathologist reads a slide. Plain attention would be too expensive for this many patches, so the model uses an approximation whose cost grows with the number of patches. A convolution step gives the model a sense of where each patch sits in the sequence.

The authors report higher accuracy and AUC than the compared methods on three public datasets, two cancer-subtype tasks and one tumour-detection task.

TL;DR

• Standard multiple instance learning treats the patches of a slide as independent. TransMIL replaces that with a two-layer Transformer over all patch features, so each patch's representation depends on the others §3.1 §3.3.
• Self-attention uses the Nystrom approximation, and a convolutional position encoder (PPEG) adds spatial information §3.3 eq. 9 alg. 3.
• Author-reported AUC: 0.9309 on CAMELYON16, 0.9603 on TCGA-NSCLC and 0.9882 on TCGA-RCC, each the highest of the compared methods tab. 1.
• The official code is for non-commercial academic use. It differs from the paper in batch size, classifier head, and class-token order. Read "Paper vs code" before you reimplement.

Problem

• Many MIL methods assume the instances in a bag are independent and identically distributed. The authors argue this is often not valid, because pathologists use both the context around an area and the relation between areas §1.
• Plain Transformers handle sequences of roughly a thousand tokens. A WSI bag has far more: CAMELYON16 has about 8,800 patches per bag on average at x20 §1 §4.
• Assumptions: each bag has one label. Binary bags are positive when at least one instance is positive. Instance-level labels are unknown §3.1 eq. 1.

Method

Formulation. A binary bag label is 0 only if every instance label is 0, and 1 otherwise eq. 1. The prediction is a scoring function S of the bag eq. 2.

Correlated MIL. The authors approximate a continuous set function S by g(P{f(x)+h(x)}), where f extracts morphological information, h extracts spatial information, and P is a pooling matrix that can mix instances eq. 7 alg. 1. They prove an approximation result (Theorem 1 eq. 3, with the proof in app. A) and an entropy result (Theorem 2 eq. 8). Theorem 2 says the correlated assumption has no more entropy than the independent one §3.1.

Generic three-step method. alg. 1

# Generic three-step approach under correlated MIL  [alg. 1]
# Input:  bag X of n instances
# Output: bag-level prediction

def correlated_mil(bag):
    # 1. Morphological and spatial features of every instance  [alg. 1 step 1]
    morph = f(bag)                  # (n, d)
    spatial = h(bag)                # (n, d)
    combined = morph + spatial      # (n, d)
    # 2. Aggregate across instances with pooling matrix P  [alg. 1 step 2]
    pooled = P @ combined           # P: (n, n)
    # 3. Map the pooled features to the bag label  [alg. 1 step 3]
    return g(pooled)

Transformer view. Bags go from the bag space to a Transformer space and then to labels. The mapping into the Transformer space is the hard part app. B. The Transformer layers use pre-norm: LayerNorm, then multi-head self-attention, then a residual add app. B TransMIL.py.

TPT module. alg. 2 TransMIL has two Transformer layers with a PPEG step in between §3.3 fig. 3.

# TPT module processing flow  [alg. 2]
# Input:  bag of projected features, shape (n, 512)
# Output: bag-level logits

def tpt(features):
    # 1. Squaring: pad to side x side and prepend a class token  [alg. 2 step 1]
    side = ceil(sqrt(n))
    padding = features[: side * side - n]   # copies of the first tokens
    tokens = concat(class_token, features, padding)   # (side*side+1, 512)
    # 2. Correlation modelling with the first attention layer  [alg. 2 step 2]
    tokens = tokens + nystrom_attention(layer_norm(tokens))
    # 3. Conditional position encoding and local fusion  [alg. 2 step 3]
    tokens = ppeg(tokens, side)
    # 4. Second attention layer for deep aggregation  [alg. 2 step 4]
    tokens = tokens + nystrom_attention(layer_norm(tokens))
    # 5. Read the class token and predict  [alg. 2 step 5]
    class_out = layer_norm(tokens)[0]       # (512,)
    return classifier(class_out)

Nystrom attention. The softmax attention is approximated with m landmark tokens selected from Q and K eq. 9 §3.3. The paper states that the approximation reduces the cost from O(n^2) to O(n) §3.3.

\hat{S}=\operatorname{softmax}\left(\frac{Q\tilde{K}^{T}}{\sqrt{d_q}}\right)\left(\operatorname{softmax}\left(\frac{\tilde{Q}\tilde{K}^{T}}{\sqrt{d_q}}\right)\right)^{+}\operatorname{softmax}\left(\frac{\tilde{Q}K^{T}}{\sqrt{d_q}}\right) \qquad [\text{eq. 9}]

PPEG (positional encoding). alg. 3 fig. 4 PPEG is applied to the patch tokens after the first attention layer.

# Pyramid Position Encoding Generator  [alg. 3]
# Input:  tokens (side*side + 1, 512), class token first
# Output: tokens with the same shape

def ppeg(tokens, side):
    # 1. Split off the class token  [alg. 3 step 1]
    class_token = tokens[0]
    patches = tokens[1:]
    # 2. Restore patches to a side x side grid, in order  [alg. 3 step 2]
    grid = reshape(patches, (side, side, 512))
    # 3. Convolutions with kernels 7, 5 and 3, zero-padded  [alg. 3 step 3]
    # [code: models/TransMIL.py#L32-L34]
    # 4. Sum the grid with the three convolution outputs  [alg. 3 step 4]
    fused = grid + conv7(grid) + conv5(grid) + conv3(grid)
    # 5. Flatten to a sequence and put the class token first  [alg. 3 steps 5-6]
    return concat(class_token, flatten(fused))

Training and inference. Cross-entropy loss, Lookahead optimizer, learning rate 2e-4, weight decay 1e-5, batch size 1 §4. Inference applies softmax to the scores §4.

Hyperparameters

Values are as the paper states them. "not stated" means the paper does not give a value. Code values are from the official repository (see Paper vs code).

Results

Author-reported. Setup §4: CAMELYON16 is binary tumour detection, scored by AUC on the official test set. TCGA-NSCLC is binary LUAD vs LUSC, scored by mean AUC over 4-fold CV. TCGA-RCC is 3-class (KICH, KIRC, KIRP), scored by macro one-vs-rest AUC over 4-fold CV. Accuracy uses a threshold of 0.5. Hardware: one RTX 3090 §4. The paper reports no standard deviations.

Table 1 reproduces the paper's main results. tab. 1

Notes from the text §4.1: MIL-RNN is not applied to multi-class RCC. On NSCLC the paper reports TransMIL is 2.16 points higher in accuracy and 1.40 points higher in AUC than the second-best method for each metric tab. 1.

Ablation on the position encoding tab. 2 §4.2.1. The paper says the ablation is run on CAMELYON16 and NSCLC, scored by AUC.

The sinusoidal encoding is scaled by 0.001 §4.2.1. The paper does not define the "both" row tab. 2 (see Pitfalls).

Other results, described in words (figures are not in this card):

• Interpretability: attention maps line up with an expert annotation on RCC §4.3 fig. 6.
• Convergence: the paper reports about two to three times fewer training epochs than ABMIL, DSMIL and CLAM. No epoch counts are given §4.4 fig. 7.

Paper vs code

Official repository: github.com/szc19990412/TransMIL, at commit 9d6aee57c7c72375fb9132dc58cd8c9b0f0a949c. The repository has no LICENSE file. The README says it is "GPLv3 ... available for non-commercial academic purposes" README.md. No code is quoted in this card.

• Batch size. The paper uses a mini-batch of 1 §4. The config uses batch size 1 and gradient accumulation of 2, so 2 bags per update TransMIL.yaml TransMIL.yaml train.py.
• Classifier head. The paper uses an MLP on the class token alg. 2 step 5 app. B. The code applies LayerNorm and one linear layer TransMIL.py TransMIL.py.
• Input projection. The paper says a fully connected layer §4. The code adds a ReLU TransMIL.py.
• Class token position in PPEG. Algorithm 3 step 6 concatenates the class token last alg. 3 step 6. The code puts it first TransMIL.py, and so does Algorithm 2 alg. 2 step 1.
• Nystrom value residual. The code sets residual on TransMIL.py. Eq. 9 has no such term eq. 9.
• Inner optimizer. The paper says Lookahead §4 and does not name the inner optimizer. The code wraps RAdam optim_factory.py optim_factory.py.
• Training length and stopping. Not stated in the paper. The code trains up to 200 epochs with early stopping, patience 10 TransMIL.yaml TransMIL.yaml.
• Mixed precision. Not stated in the paper. The code uses fp16 with AMP level O2 TransMIL.yaml.
• CAMELYON16 split. The paper says about 90/10 train/val of 270 slides §4. The code's fold 0 has 216 train and 54 val (80/20) fold0.csv. Its test column lists 129 slides (test_049 is absent), but the paper says the test set has 130 slides §4 fold0.csv.
• Learning-rate schedule. Not stated in the paper. The code uses none model_interface.py.
• Device. The code calls .cuda() on the class token TransMIL.py.
• Ablation code and feature extraction. The ablation scripts and the feature extractor are not in the repo; the README points to CLAM README.md.

Pitfalls

• "Spatial" means sequence order. PPEG reshapes the 1-D patch sequence into a square grid in input order TransMIL.py. The code uses no (x, y) coordinates. Whether the feature files are spatially sorted is not in this repo, so the model's spatial signal depends on how those files were written README.md.
• Padding. Algorithm 2 step 1 writes N = ceil(sqrt(n)) and M = N - n, which reads as a negative count. The code pads to side^2 tokens by repeating the first tokens TransMIL.py. Implement the padding to side^2. The repeated tokens are visible to the PPEG convolutions.
• Padding size. Derived from ceil(sqrt(n)): padding adds at most about 2 sqrt(n) tokens, so a bag of 8,800 patches gets about 94 x 94 = 8,836 slots.
• No feed-forward block. The Transformer layers have attention and a residual, but no MLP (feed-forward) sublayer app. B TransMIL.py. This differs from a standard ViT block.
• Appendix embedding. Appendix B adds a positional embedding E_pos to the input app. B. TransMIL has no learned or additive positional embedding; PPEG is the position encoding TransMIL.py.
• Nystrom details. The paper does not state the landmark method, the number of landmarks, the pseudo-inverse iterations or the head count eq. 9 app. B. The code uses an external package pinned to nystrom-attention 0.0.9 requirements.txt. The pseudo-inverse is an iterative approximation, and fp16 with AMP level O2 TransMIL.yaml may be sensitive. Check it on your data.
• Undefined ablation row. The "both" row in tab. 2 is not defined in the text.
• Sinusoidal scaling. The sinusoidal encoding is multiplied by 0.001 §4.2.1. The ablation code is not in the repo, so the scaling cannot be checked.
• Variance. The TCGA numbers are 4-fold averages, and CAMELYON16 is a single official split. No standard deviations are given §4.
• Epoch claim. The two-to-three-times faster convergence claim has no epoch counts §4.4 fig. 7.
• License of the code. The official code is GPLv3 with non-commercial academic use README.md. This card does not quote the code. A commercial implementation should follow the paper and this card, not the code.
• Paper's RCC note. MIL-RNN is excluded from RCC because it does not handle multi-class tasks §4.1 tab. 1.
• Batch and memory. The paper says memory and computation grow with higher magnification, and it leaves that to future work §5.

Lineage

• Builds on: the Transformer for NLP §2.2; the Nystrom approximation of self-attention (Nystromformer, cited as xiong2021nystromformer) §3.3; the conditional position encoding of Chu et al., cited as chu2021conditional §3.3; ABMIL §1; DSMIL (non-local attention in MIL) §1; CLAM §4; Lookahead (cited as Zhang2019LookaheadOK) §4.
• Compared against: Mean-pooling, Max-pooling, ABMIL, PT-MTA, MIL-RNN, DSMIL, CLAM-SB, CLAM-MB §4 tab. 1.
• Superseded by: no information in this wiki.