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

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

Note the contrast between the two entries: the layer-17 transcoder loses ~11 EV points off-domain while the layer-29 transcoder loses ~4, 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.}
}