Medical LLM Interpretability

Code for the paper "Medical Interpretability and Knowledge Maps of Large Language Models" (ICLR 2026).

This repository provides tools to study how LLMs represent and process medical knowledge through four interpretability techniques:

1. UMAP projections of intermediate-layer activations
2. Gradient-based weight saliency
3. Layer lesioning (replacing a layer with the identity function)
4. Activation patching (swapping layer outputs between clean and corrupted prompts)

These techniques are applied across multiple medical knowledge areas (patient age, symptoms, diseases, disease progression, drug treatments, drug dosages) to build LLM knowledge maps showing where medical knowledge is stored in the model.

Key Findings

• Age is encoded non-linearly at intermediate layers, with a discontinuity at age 18 (teenager vs. adult)
• Disease progression is represented circularly and non-monotonically — late-stage embeddings loop back toward early-stage embeddings
• Drugs cluster more by medical specialty than mechanism of action (especially in Llama3.3-70B)
• Gemma/MedGemma activations collapse at intermediate layers but recover by the final layers
• For Llama3.3-70B, most medical knowledge is processed in the first half of the model's layers

Supported Models

Installation

conda create -n medical-interp python=3.10 -y
conda activate medical-interp
pip install -r requirements.txt

Hardware Requirements

• Running the 70B models requires 4-8 GPUs (A40 or A100 recommended)
• Smaller models (1B, 8B, 13B) can run on a single GPU
• The code uses HuggingFace device_map for automatic multi-GPU sharding

Repository Structure

.
├── runAct.py              # Main entry point — loads model, dispatches analyses
├── runActSomeGPUs.py      # Variant for running on a subset of GPUs
├── act.py                 # Base Act class: hooks, UMAP, saliency, lesioning, patching
├── ageAct.py              # Age analysis (1–100 year sweep)
├── genderAct.py           # Gender representation analysis
├── diseaseAct.py          # Disease representation by specialty
├── symptomAct.py          # Symptom clustering analysis
├── drugAct.py             # Drug analysis (mechanism & specialty clustering)
├── top100DrugAct.py       # Top-100 US drugs analysis
├── dosageAct.py           # Drug dosage analysis
├── progressionAct.py      # Disease progression (9-stage trajectories)
├── maps.py                # LLM map generation and metrics aggregation
├── saliency.py            # Gradient-based weight saliency utilities
├── utils.py               # Silhouette score computation with bootstrap CIs
├── Makefile               # Convenience targets for running analyses
├── requirements.txt       # Python dependencies
├── runall.sh              # Bash script to run analyses across multiple models
├── gemma.py               # Standalone test script for Gemma/MedGemma
└── l4.py                  # Standalone test script for Llama-4

Usage

Download a Model

# Example: download Llama-3.3-70B-Instruct
huggingface-cli download meta-llama/Llama-3.3-70B-Instruct \
    --local-dir /data/hf/meta-llama/Llama-3.3-70B-Instruct

Run Individual Analyses

All analyses are run via runAct.py. You select the model with --model-dir and enable analyses with flags:

# UMAP analysis on age and disease
python runAct.py --model-dir /data/hf/meta-llama/Llama-3.3-70B-Instruct \
    --do_umap --age --disease

# Saliency analysis on symptoms and drugs
python runAct.py --model-dir /data/hf/meta-llama/Llama-3.3-70B-Instruct \
    --do_saliency --symptom --drug

# Layer lesioning on diseases
python runAct.py --model-dir /data/hf/meta-llama/Llama-3.3-70B-Instruct \
    --do_lesioning --disease

# Activation patching on drugs and dosages
python runAct.py --model-dir /data/hf/meta-llama/Llama-3.3-70B-Instruct \
    --do_activation_patching --drug --dosage

# Generate LLM knowledge maps (requires previous analyses to be cached)
python runAct.py --model-dir /data/hf/meta-llama/Llama-3.3-70B-Instruct \
    --do_maps --age --disease --symptom --drug --dosage

Run All Analyses at Once

python runAct.py --model-dir /data/hf/meta-llama/Llama-3.3-70B-Instruct \
    --do_umap --do_saliency --do_lesioning --do_activation_patching --do_maps \
    --age --disease --symptom --drug --dosage

Using the Makefile

The Makefile provides shortcuts. Edit the MODEL_DIR variable at the top, then enable the desired knowledge areas by uncommenting them:

# Uncomment the analyses you want to run:
age=--age
disease=--disease
symptom=--symptom
drug=--drug
dosage=--dosage

Then run:

make umap            # UMAP analysis
make saliency        # Weight saliency
make lesioning       # Layer lesioning
make activation_patching  # Activation patching
make maps            # Generate LLM maps (from cached results)
make all_analyses    # Run everything

Generate Metrics Tables

To produce LaTeX tables and per-layer metric plots from cached results (no GPU needed):

python runAct.py --model-dir /data/hf/meta-llama/Llama-3.2-1B-Instruct --metrics_table

Caching

By default, results are cached in the results/ directory and reused on subsequent runs. To force re-computation:

python runAct.py --model-dir /path/to/model --do_umap --age --no_cache

Analysis Flags

Knowledge Area Flags

Citation

@inproceedings{marinescu2026medical,
  title={Medical Interpretability and Knowledge Maps of Large Language Models},
  author={Marinescu, Razvan and Gruber, Victoria-Elisabeth and Fajardo, Diego},
  booktitle={International Conference on Learning Representations (ICLR)},
  year={2026}
}

License

This project is for research purposes. Please cite the paper if you use this code.