MAD-MIL: Multi-head Attention MIL

• The paper introduces a multi-head attention MIL framework that enhances feature aggregation and computational efficiency compared to single-head gated models.
• MAD-MIL partitions whole slide images into patches and processes them via a transformer-style, gated multi-head attention module for interpretable bag-level embeddings.
• The approach achieves 20-30% reductions in trainable parameters and FLOPs over ABMIL while delivering competitive AUC and F1 scores across multiple pathology datasets.

Multi-head Attention MIL (MAD-MIL) is a multiple instance learning (MIL) framework designed for weakly supervised classification tasks in digital pathology, particularly for whole slide images (WSIs). MAD-MIL generalizes the single-head gated attention mechanism of Attention-based Deep MIL (ABMIL) to a multi-head formulation inspired by Transformer architectures. The model emphasizes efficient, interpretable, and accurate aggregation of information from large sets of image patches, reducing computational footprint and increasing representational diversity relative to prior state-of-the-art MIL approaches (Keshvarikhojasteh et al., 2024).

1. Model Architecture and Computational Flow

MAD-MIL replaces the single gated-attention module of ABMIL with an M-headed attention block, introducing architectural parallels to Transformer-style multi-head attention. The model comprises four primary components:

1. Instance Feature Extraction: The WSI is partitioned into $N$ tiles $\{p_1, \ldots, p_N\}$. Each tile $p_i$ is processed by a pretrained CNN (e.g., ResNet-50) to yield high-dimensional features $h_i \in \mathbb{R}^{D_0}$, which are compressed via a learnable fully connected (FC) layer to $f_i \in \mathbb{R}^D$.
2. Multi-head Attention Module: The feature vector $f_i$ is split evenly into $M$ sub-vectors along the feature dimension: $f_i = [f_{i,1}; f_{i,2}; \ldots; f_{i,M}]$, with $f_{i,m} \in \mathbb{R}^{D/M}$. Each sub-vector is processed by a distinct gated attention head, yielding per-head attention weights $a_{i,m}$ and aggregated vectors $\{p_1, \ldots, p_N\}$0.
3. Aggregation Layer (Bag-level Embedding): For each head $\{p_1, \ldots, p_N\}$1, the representation is aggregated as $\{p_1, \ldots, p_N\}$2, where $\{p_1, \ldots, p_N\}$3 and $\{p_1, \ldots, p_N\}$4. The outputs from $\{p_1, \ldots, p_N\}$5 heads are concatenated to form the slide-level embedding $\{p_1, \ldots, p_N\}$6.
4. Classifier: A final FC layer $\{p_1, \ldots, p_N\}$7 computes predictions $\{p_1, \ldots, p_N\}$8, using sigmoid activation for binary and softmax for multiclass tasks.

2. Multi-head Attention Mechanisms

MAD-MIL supports two conceptualizations for multi-head aggregation: its practical implementation and a Transformer-style formulation.

A. Transformer-style Multi-head (for context):

• Each embedding $\{p_1, \ldots, p_N\}$9 is used to compute queries ($p_i$0), keys ($p_i$1), and values ($p_i$2) via linear projections.
• For each head $p_i$3:

  $p_i$4

• Dot-product attention weights and per-head outputs are computed, concatenated, and pooled for the final bag embedding.

B. Gated Multi-head Attention (MAD-MIL implementation):

• Each split feature $p_i$5 is processed by a gated attention module:

  $p_i$6

  where $p_i$7, $p_i$8, and $p_i$9 is the sigmoid function.

• Per-head bag embedding: $h_i \in \mathbb{R}^{D_0}$0.

3. Model Complexity and Efficiency

MAD-MIL is designed to reduce both trainable parameters and floating point operations per bag relative to existing deep MIL architectures such as ABMIL and DS-MIL, without loss of accuracy. Parameter and computational requirements across representative tasks are summarized below.

Dataset	Method	Params	FLOPs
MNIST-BAGS	ABMIL	167.1 K	19.9 M
	MAD-MIL/6	107.1 K	12.7 M
TUPAC16	ABMIL	788.7 K	94.4 M
	MAD-MIL/3	614.8 K	73.5 M
	DS-MIL	1.186 M	142.0 M
TCGA BRCA	ABMIL	788.7 K	94.4 M
	MAD-MIL/2	657.6 K	78.6 M
TCGA LUNG	MAD-MIL/8	559.3 K	66.8 M
TCGA KIDNEY	MAD-MIL/5	582.7 K	69.6 M

Across datasets, MAD-MIL achieves $h_i \in \mathbb{R}^{D_0}$1 reductions in trainable parameters and FLOPs relative to ABMIL, and over $h_i \in \mathbb{R}^{D_0}$2 reduction versus DS-MIL (Keshvarikhojasteh et al., 2024).

4. Experimental Protocol and Evaluation

Empirical validation covered both synthetic (MNIST-BAGS) and real-world WSI datasets:

• Datasets:
  • MNIST-BAGS: 20-instance bags, binary classification of digit ‘8’ under controlled positive/negative instance ratios.
  • TUPAC16: 821 WSIs (H&E), binary proliferation grading.
  • TCGA BRCA: 1,038 slides, subtype classification (IDC vs ILC).
  • TCGA LUNG: 1,046 slides, LUAD vs LUSC.
  • TCGA KIDNEY: 918 slides, three-class subtyping.
• Feature Extraction:
  • MNIST: Flatten $h_i \in \mathbb{R}^{D_0}$3 images, project to $h_i \in \mathbb{R}^{D_0}$4.
  • WSIs: Patch extraction ($h_i \in \mathbb{R}^{D_0}$5 at $h_i \in \mathbb{R}^{D_0}$6), ResNet-50 to 1,024-d, then FC to $h_i \in \mathbb{R}^{D_0}$7.
• Training:
  • Adam optimizer.
  • Task-specific epochs (MNIST: 20; TUPAC16/TCGA: 50).
  • Hyperparameters: Validation-based selection, 10-fold cross-validation (TCGA).
  • Head count $h_i \in \mathbb{R}^{D_0}$8 optimized via validation loss.
• Performance Metrics: AUC (ROC), F1-score (binary), macro-F1 (multi-class).

5. Comparative Performance Analysis

Experimental results demonstrate that MAD-MIL consistently surpasses ABMIL, and in most cases matches or narrowly trails the highest-performing, but more complex, methods such as DS-MIL and ACMIL.

Dataset/Task	Method	AUC	F1
MNIST-BAGS	ABMIL	$h_i \in \mathbb{R}^{D_0}$9	$f_i \in \mathbb{R}^D$0
	MAD-MIL/7	$f_i \in \mathbb{R}^D$1	$f_i \in \mathbb{R}^D$2
TUPAC16	ABMIL	$f_i \in \mathbb{R}^D$3	$f_i \in \mathbb{R}^D$4
	MAD-MIL/3	$f_i \in \mathbb{R}^D$5	$f_i \in \mathbb{R}^D$6
	CLAM-MB	$f_i \in \mathbb{R}^D$7	$f_i \in \mathbb{R}^D$8
TCGA BRCA	ABMIL	$f_i \in \mathbb{R}^D$9	$f_i$0
	MAD-MIL/2	$f_i$1	$f_i$2
	DS-MIL	$f_i$3	$f_i$4
TCGA LUNG	ABMIL	$f_i$5	$f_i$6
	MAD-MIL/8	$f_i$7	$f_i$8
TCGA KIDNEY	ABMIL	$f_i$9	$M$0
	MAD-MIL/5	$M$1	$M$2
	DS-MIL	$M$3	$M$4

A consistent AUC and F1-score improvement is observed over ABMIL, with competitive ranking alongside other transformer-inspired methods, but at a lower computational and parameter budget (Keshvarikhojasteh et al., 2024).

6. Interpretability Features

MAD-MIL generates per-head attention heatmaps, enhancing transparency of slide-level predictions:

• Each attention head $M$5 produces an attention score map $M$6, which can be spatially registered to patch locations.
• Heatmaps derived from these scores can be up-scaled and superimposed on original WSIs.
• Empirical visualization (e.g., on LUAD slides) shows that MAD-MIL’s eight attention heads yield complementary highlight regions: tumor, stroma, necrosis, and lymphocyte infiltration.
• A plausible implication is that the diversity of M-heads offers greater opportunity for fine-grained, multi-faceted clinical interpretability and pathologist trust, compared to single-head models or those producing only a single map.

7. Implementation and Prospective Extensions

The published implementation includes modular code (PyTorch) with data preprocessing, model modules, and visualization tools (heatmap overlay) [GitHub: https://github.com/tueimage/MAD-MIL]:

• Feature extraction and tiling can be decoupled (offline), enabling low-latency batch inference.
• Moderate memory footprint due to reduced multilayer perceptron (MLP) sizes.
• Multi-head outputs support integration into graphical user interfaces for interactive slide review.
• Potential extensions identified in the original source include replacement of gated attention with dot-product multi-head attention, self-supervised pretraining of the feature encoder, and algorithmic head pruning or regularization to maximize information diversity for a given model size (Keshvarikhojasteh et al., 2024).