brca-wsi-classification-mil 🧬🔬

Dual-Stream Multiple Instance Learning with Pathology Foundation Models for BRCA1/BRCA2 Expression Stratification from Whole Slide Images

---

🔬 Overview & Objective

A PyTorch implementation of Gated Dual-Stream Multiple Instance Learning (Gated-DSMIL) for breast cancer gene expression classification on Whole Slide Images (WSIs)

brca-wsi-classification-mil evaluates whether Dual-Stream Multiple Instance Learning (DSMIL) on H&E Whole Slide Images (WSIs) can stratify BRCA1 and BRCA2 gene expression status (over-expressed vs. under-expressed based on median bulk RNA-seq FPKM).

We perform a controlled ablation study investigating the impact of adding Gated Attention (Ilse et al., 2018) to the query projection of DSMIL (Li et al., CVPR 2021) using two state-of-the-art computational pathology foundation models: UNI (1024-dim) and CONCH (768-dim).

Original DSMIL :  Q = tanh(W · x)
Modified DSMIL :  Q = tanh(W · x) ⊙ sigmoid(U · x)   [+ Gated Attention]

Note on Task Scope: This benchmark predicts transcriptomic expression levels (FPKM) from histology images, not somatic/germline genetic variant calling.

---

📊 Dataset & Validation Setup

• Cohort: TCGA-BRCA (2,787 Whole Slide Images across 1,059 unique patients).
• Patient-Grouped CV: Strict 4-fold cross-validation grouped by patient ID (GroupKFold) — zero patient overlap between train and validation sets.
• Encoders: Evaluated independently (no arbitrary feature concatenation):
  • UNI: ViT-Large/16, $D=1024$
  • CONCH: Vision-Language CoCa, $D=768$

flowchart LR
    WSI["WSI (H&E)"] --> CLAM["CLAM Tiling"]
    CLAM --> Enc["UNI / CONCH Encoders"]
    Enc --> Bag["Bag X: (N, D)"]
    Bag --> Ins["Instance Stream\n(Critical Patch q_crit)"]
    Bag --> BagStream["Bag Stream (Q & V)\nOriginal vs Gated Q"]
    Ins & BagStream --> Loss["Dual-Stream Loss: 0.5 L_bag + 0.5 L_inst"]

Loading

---

📈 Key Results

Validation Performance (Mean ± Std over 4 Folds)

Gene	Model Architecture	Feature Encoder	Accuracy (%)	Precision (%)	Recall (%)	F1-Score
BRCA1	Original DSMIL	CONCH	$58.81 \pm 2.82$	$55.87 \pm 3.17$	$\mathbf{83.86 \pm 3.77}$	$0.6704 \pm 0.0317$
BRCA1	Modified DSMIL	CONCH	$\mathbf{59.67 \pm 2.57}$	$\mathbf{56.65 \pm 3.13}$	$83.31 \pm 3.80$	$\mathbf{0.6737 \pm 0.0246}$
BRCA1	Original DSMIL	UNI	$57.27 \pm 3.94$	$55.46 \pm 4.41$	$\mathbf{79.30 \pm 5.84}$	$0.6501 \pm 0.0157$
BRCA1	Modified DSMIL	UNI	$\mathbf{59.92 \pm 2.24}$	$\mathbf{57.87 \pm 3.54}$	$74.85 \pm 7.82$	$\mathbf{0.6504 \pm 0.0279}$
BRCA2	Original DSMIL	CONCH	$59.13 \pm 4.24$	$56.13 \pm 4.67$	$\mathbf{84.24 \pm 1.88}$	$0.6731 \pm 0.0379$
BRCA2	Modified DSMIL	CONCH	$\mathbf{59.96 \pm 4.51}$	$\mathbf{57.08 \pm 5.01}$	$82.94 \pm 5.15$	$\mathbf{0.6743 \pm 0.0322}$
BRCA2	Original DSMIL	UNI	$\mathbf{59.31 \pm 3.42}$	$\mathbf{56.72 \pm 4.55}$	$79.22 \pm 1.74$	$0.6604 \pm 0.0327$
BRCA2	Modified DSMIL	UNI	$58.99 \pm 3.60$	$56.70 \pm 5.25$	$\mathbf{81.04 \pm 6.47}$	$\mathbf{0.6636 \pm 0.0171}$

Statistical & Overfitting Findings

1. Paired Fold Comparison: Paired $t$-tests ($\Delta \text{F1} = \text{Modified} - \text{Original}$) yield $p = 0.62 - 0.97$. Modified DSMIL leads on 7 folds, Original DSMIL on 7 folds, with 2 ties $\rightarrow$ Statistically indistinguishable.
   • Why? Gated attention only modifies bag query projections, while 50% of the dual-stream loss supervises the unmodified instance classifier.
2. Encoder Overfitting: Models trained on UNI features exhibit a significantly larger generalization gap ($\mathcal{L}{\text{val}} - \mathcal{L}{\text{train}} \approx 0.41 - 0.45$) than CONCH features ($\approx 0.030$), attributable to higher embedding dimensionality and feature space properties.

---

📁 Repository Structure

BRCA-project/
├── .gitignore                       # Git ignore configuration
├── Code/
│   ├── build_brca_training_csv.py   # Cohort dataset builder
│   ├── data_utils.py                # Dataset loader & GroupKFold splitting
│   ├── train_dsmil.py               # Main training script (Original vs Modified)
│   ├── results.py                   # Metrics aggregation & plot generator
│   ├── dsmil.sh / dsmil_mod.sh      # SLURM execution scripts
│   ├── models/                      # Saved fold weights (.pt)
│   └── results/                     # Metric logs (.csv, .jsonl, .png)
└── README.md                        # Project documentation

---

🚀 Quick Start

1. Setup

# Clone & install dependencies
git clone https://github.com/shazimjaved/brca-wsi-classification-mil.git
cd brca-wsi-classification-mil/Code
pip install torch torchvision pandas numpy scikit-learn matplotlib

2. Prepare Dataset

python build_brca_training_csv.py

3. Train Models

# Train Original DSMIL
python train_dsmil.py --variant original --feature_mode conch --label_column BRCA1_label

# Train Modified DSMIL (Gated Attention)
python train_dsmil.py --variant modified --feature_mode conch --label_column BRCA1_label

4. Evaluate & Plot Results

python results.py