SAGE — Semantic Attention Global Explanations

Code for SAGE: Semantic Explainability of Attention-Based Survival Models in Computational Pathology.

Attention maps tell you where an ABMIL model looks, not what it looks at. SAGE is a post-hoc, model-agnostic read-out that turns the attention of an already-trained ABMIL model into a global, language-grounded explanation: it scores every patch against a dictionary of 25 histology concepts with a pathology vision–language model, aggregates those scores with the model's own attention weights, and relates the resulting per-patient concept scores to the model's predictions across a cohort.

Nothing is retrained and no concept annotation is needed: the ABMIL predictor stays frozen, and only its attention weights and output scores are used.

         frozen ABMIL                     pathology VLM
   patches ──► attention a_i          patches ──► s_ij = cos(patch_i, concept_j)
                    │                                   │
                    └──────────►  S_j = Σ_i a_i · s_ij  ◄┘        (per patient)
                                          │
                    Spearman r( S_j , predicted prognosis )       (per cohort)
                                          │
                            global explanation = the SAGE profile

What is here

Path	Contents
sage/	The concept dictionary + prompt bank (concepts.py) and the path configuration (paths.py)
scripts/	The pipeline: text embeddings → ABMIL training → SAGE concept scores → statistics → ablations
figures/	One script per paper figure
configs/	Final hyperparameters, environment template, batch runners
results/	The paper's own result tables (see results/README.md)
docs/	REPRODUCE.md — step by step; FIGURES.md — figure/table ↔ code ↔ data

Install

git clone https://github.com/lotterlab/sage.git && cd sage
python -m venv .venv && source .venv/bin/activate
pip install -r requirements.txt
cp configs/env.example.sh configs/env.sh   # edit the paths, then:
source configs/env.sh

Every path is read from environment variables through sage/paths.py; no file in this repository needs to be edited to run it elsewhere.

Quickstart — no data required

The paper's result tables ship with the repository, so its figures and statistics can be regenerated immediately:

python figures/fig2_term_heatmap.py            # Fig. 2, the SAGE profile heatmap
python figures/fig4a_attention_dumbbell.py --cohorts KIRC BRCA
python figures/fig4b_vlm_agreement.py
python scripts/attention_ablation_merge.py --stats_dir results/stats --out_dir work/ablation

The last command reproduces the attention ablation of Section 4.5 — 175 concept–cohort pairs, attention sharpens |r| for 74.9 % of them, mean Δ|r| = +0.070, paired Wilcoxon p = 3.6e-13.

results/stats/ holds the per-concept statistics behind those numbers (cross-fold Spearman, bootstrap CI, permutation p), for both the attention-weighted and the uniform aggregation.

Full pipeline — from whole-slide images

See docs/REPRODUCE.md. In short:

# 1. concept text embeddings (needs CONCH weights)
python scripts/build_text_embeddings_conch.py --out_dir "$SAGE_TEXT_EMB_ROOT/conch_ensembled"   --mode mean_embedding
python scripts/build_text_embeddings_conch.py --out_dir "$SAGE_TEXT_EMB_ROOT/conch_template_bank" --mode template_bank

# 2. train the ABMIL survival models SAGE will explain
bash configs/run_train_all.sh

# 3. SAGE: per-patient concept scores (attention-weighted and uniform) + statistics
bash configs/run_sage_all.sh

# 4. figures
python figures/fig2_term_heatmap.py  --term_scores_dir "$SAGE_STATS_ROOT"
python figures/fig4a_attention_dumbbell.py --per_term_csv "$SAGE_FIG_ROOT/ablation/attention_ablation_per_term.csv"

Tiling and patch-embedding extraction are not part of this repository: the paper uses PathFMTools with foreground segmentation from the HEST pipeline, producing one HDF5 per slide. The expected HDF5 layout is documented in docs/REPRODUCE.md.

Applying SAGE to your own model

SAGE only needs (i) per-fold ABMIL checkpoints with a softmax attention branch, (ii) patch embeddings in a vision–language space, and (iii) text embeddings for your concepts. To query a different vocabulary, edit MEDICAL_TERMS in sage/concepts.py, regenerate the text embeddings, and rerun scripts/sage_concept_scores.py — the ABMIL model is untouched, so the same frozen predictor can be interrogated with as many dictionaries as you like.

Data

TCGA whole-slide images and clinical data are available from the GDC; the disease-specific survival endpoint follows the cBioPortal TCGA PanCancer Atlas tables. The seven cohorts used are BRCA, BLCA, CESC, COAD, KIRC, LGG and LUAD, one slide per patient.

Model weights are gated by their own licences: CONCH and UNI2 (MahmoodLab), MUSK (lilab-stanford).

Citation

@inproceedings{sage2026,
  title     = {SAGE: Semantic Explainability of Attention-Based Survival Models
               in Computational Pathology},
  author    = {Lamane, Abdallah and Diab, Abdul Rahman and Wu, Ren-Chin and
               Lotter, William},
  booktitle = {MICCAI Workshop on Interpretability of Machine Intelligence in
               Medical Image Computing (iMIMIC)},
  year      = {2026}
}

Released under the MIT licence (see LICENSE).