saillab/medgemma-circuit-tools

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

import sys, torch
sys.path.insert(0, "circuit_tracing")
from train_layer17_transcoder import (
    MEDGEMMA_CONFIG, Layer17ActivationCollector, TranscoderLayer17, load_model,
)

model, processor = load_model(device="cuda")
collector = Layer17ActivationCollector(model, layer=29)  # or 17

ckpt = torch.load("layer29_transcoder_final.pt", map_location="cpu",
                  weights_only=False)
transcoder = TranscoderLayer17(d_model=2560, expansion_factor=8, top_k=64,
                               dtype=torch.float32, device="cuda")
transcoder.load_state_dict(ckpt["state_dict"])
transcoder.eval()

# mlp_input / mlp_output: [batch, seq, 2560] activations for your prompts
mlp_in, mlp_out = collector.get_activations(inputs)
reconstructed, 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

@misc{sadanandan2026medgemma_circuit_tools,
  title  = {MedGemma Circuit Tools: Domain-Adapted Layer-17 and Layer-29
            Transcoders (ReXGradient-trained)},
  author = {Sadanandan, Binesh and Behzadan, Vahid},
  year   = {2026},
  note   = {Domain-adapted transcoders for circuit tracing of paraphrase
            sensitivity in medical VLMs. SAIL Lab, University of New Haven.}
}