MedGemma Circuit Tools
Circuit-tracing artifacts for MedGemma, from SAIL Lab. Two transcoders so far,
at the two ends of the candidate two-stage paraphrase-flip circuit identified in
PSF-Med (Sadanandan et al. 2026): the layer-17 register gate (Feature 3818) and
the layer-29 decision feature (Feature 12139).
Both are top-k transcoders trained on MLP activations of
google/medgemma-4b-it collected
from radiology-report text
(
ReXGradient-160K),
for circuit tracing and mechanistic analysis of paraphrase sensitivity in medical
vision-language models.
Shared configuration
- Architecture: top-k transcoder, 2,560 → 20,480 features (8x expansion),
top_k = 64
- Hookpoint: MLP output reconstructed from MLP input at the target layer of
MedGemma-4B's language backbone (Gemma 3 4B)
- Training: 50,000 steps, batch 32 x 512 tokens, ~1.6M documents (~11 epochs
over ReXGradient-160K Findings + Impression), AdamW lr 1e-4 with cosine decay
Entry 1: Layer-17 transcoder (layer17_transcoder_final.pt)
- Motivation: layer 17 carries Feature 3818, the "clinical query register
gate" identified in PSF-Med circuit analysis using Gemma Scope 2. This is the
domain-adapted counterpart for MedGemma-specific mechanistic work.
- Final training metrics: loss 5.8e-05, EV 0.996, L0 29.5
Validation (post-training, 64 documents per slice):
| Slice | Explained variance | L0 | Notes |
|---|
| ReXGradient reports (in-distribution) | 0.9964 | 28.8 | matches training |
| PSF-Med clinical questions (unseen) | 0.9648 | 55.0 | target distribution for circuit tracing |
| WikiText-103 (out-of-domain) | 0.8840 | 36.1 | strongly domain-specialized |
At n=256 documents per slice with per-document scoring, the out-of-domain EV is
0.9391 (see the paired domain-contrast table below).
Entry 2: Layer-29 transcoder (layer29_transcoder_final.pt)
- Motivation: layer 29 carries Feature 12139, the downstream yes/no decision
feature in the candidate two-stage 3818 → 12139 circuit. Training a matched
transcoder here allows the circuit account to be re-tested with domain-adapted
features at both stages instead of borrowed Gemma Scope 2 features.
- Final training metrics: loss 1.0e-04, EV 0.998, L0 27.3
Validation (post-training, 64 documents per slice):
| Slice | Explained variance | L0 | Notes |
|---|
| ReXGradient reports (in-distribution) | 0.9985 | 27.1 | matches training |
| PSF-Med clinical questions (unseen) | 0.9627 | 60.1 | target distribution for circuit tracing |
| WikiText-103 (out-of-domain) | 0.9611 | 34.7 | largely domain-general |
Domain specialization (n=256 per slice, paired by document)
Same 256 documents scored by both transcoders per slice; gaps are paired means
with bootstrap 95% CIs.
| Slice | L17 EV | L29 EV | Paired gap (L29 − L17) [95% CI] |
|---|
| ReXGradient (in-distribution) | 0.9977 | 0.9941 | −0.0036 [−0.0050, −0.0024] |
| PSF-Med questions (unseen) | 0.9735 | 0.9645 | −0.0091 [−0.0115, −0.0066] |
| WikiText-103 (out-of-domain) | 0.9391 | 0.9642 | +0.0251 [+0.0213, +0.0291] |
The EV ordering flips exactly at the domain boundary: the layer-17 (register)
transcoder reconstructs better on medical text, the layer-29 (decision)
transcoder generalizes better off-domain. The register computation is medically
specialized; the late decision computation is more generic.
Instrument-change test (2026-08-18)
These transcoders were substituted for Gemma Scope 2 in the PSF-Med canonical
two-stage circuit analysis (Feature 3818 at L17 → Feature 12139 at L29), with
identical pairs, arms, and statistics, and hypothesis features re-discovered on
the original discovery split:
- Single-feature restoration collapses from 17–58% (Gemma Scope 3818/12139) to
0–3% (local analogs) on identical pairs, while non-flip disruption stays ~0
for both instruments.
- Flip-separating information is present observationally (|delta| AUROC up to
0.885 frozen / 0.948 split-half, comparable to GS 3818's 0.72–0.79) but does
not restore causally at single-feature level.
- A weak mediation trace survives on MIMIC only (33/50 direction-coherent,
p=6e-04; PadChest at chance).
- Caveat: instruments differ in hookpoint (MLP block vs resid_post) and decoder
norm convention (unconstrained vs unit-norm), so the comparison confounds
feature causality with intervention geometry. Full analysis:
results/sae_analysis/LOCALSCOPE_INSTRUMENT_CHANGE_SUMMARY.md in the GitHub
repo.
The two-stage single-feature circuit as previously characterized is therefore
an artifact of the Gemma Scope instrumentation; the correlational signal
transfers across instruments, the causal single-feature account does not.
Note the contrast between the two entries: the layer-17 transcoder loses ~6 EV
points off-domain while the layer-29 transcoder loses ~4, and on the paired
n=256 contrast the ordering inverts exactly at the domain boundary (table
above), evidence that the
register computation at layer 17 is medically specialized while the late decision
computation at layer 29 is more generic.
Usage
1import sys, torch
2sys.path.insert(0, "circuit_tracing")
3from train_layer17_transcoder import (
4 MEDGEMMA_CONFIG, Layer17ActivationCollector, TranscoderLayer17, load_model,
5)
6
7model, processor = load_model(device="cuda")
8collector = Layer17ActivationCollector(model, layer=29) # or 17
9
10ckpt = torch.load("layer29_transcoder_final.pt", map_location="cpu",
11 weights_only=False)
12transcoder = TranscoderLayer17(d_model=2560, expansion_factor=8, top_k=64,
13 dtype=torch.float32, device="cuda")
14transcoder.load_state_dict(ckpt["state_dict"])
15transcoder.eval()
16
17# mlp_input / mlp_output: [batch, seq, 2560] activations for your prompts
18mlp_in, mlp_out = collector.get_activations(inputs)
19reconstructed, features = transcoder(mlp_in.view(-1, 2560).float())
Training and validation code (layer-selectable via
--layer):
circuit_tracing/train_layer17_transcoder.py,
circuit_tracing/validate_layer17_transcoder.py in the
medical-vlm-robustness repository.
Limitations
- Two layers only. These are layer-17 and layer-29 transcoders, not a full
transcoder suite. For other layers or general-domain text, use
Gemma Scope 2.
- Domain profile differs by layer. The layer-17 transcoder degrades on
non-medical text (EV 0.884 on WikiText) and should be treated as
medical-only; the layer-29 transcoder is more forgiving (EV 0.961) but was
still trained on medical text only.
- Text-only training stream. No image tokens appeared in training data;
reconstruction quality on image-token positions is untested.
- Research artifact. Not validated for, and not to be used in, any clinical
workflow. MedGemma itself is governed by Google's Health AI Developer
Foundations terms; these artifacts are derived from MedGemma activations.
Citation
1@misc{sadanandan2026medgemma_circuit_tools,
2 title = {MedGemma Circuit Tools: Domain-Adapted Layer-17 and Layer-29
3 Transcoders (ReXGradient-trained)},
4 author = {Sadanandan, Binesh and Behzadan, Vahid},
5 year = {2026},
6 note = {Domain-adapted transcoders for circuit tracing of paraphrase
7 sensitivity in medical VLMs. SAIL Lab, University of New Haven.}
8}
