Circuit-Tracing MedGemma

1What we built#

Much of a transformer’s computation happens in its MLP layers, but their activations are dense vectors that do not line up with human concepts. A transcoder is trained to predict each MLP’s output from its input through a large, sparse dictionary of features. Of the 2,048 features per layer in ours, only about ten fire on a given token, and each tends to track one thing: the lung base, a chest tube, the “AP PORTABLE” stamp. With the transcoder standing in for the MLPs, the model’s computation can be read feature by feature.

Ours is cross-layer: a feature read at one layer writes to every layer after it, so a single feature can account for computation the model spreads across depth. This is the substrate the attribution-graph method was built around. The CLT uses TopK sparsity (k = 32) with 2,048 features on each of MedGemma 1.5’s 34 layers, 6.24B parameters in all, trained from scratch on one H100 on activations streamed live from the model over 50,000 NIH chest X-rays and CheXpert Plus reports. As a single tensor the decoder would exceed 2³¹ elements and overflow the int32 indexing in bitsandbytes’ 8-bit Adam, so it is stored as one tensor per write offset and decoded as a running sum, which also keeps activation memory flat.

2Images need their own dictionary#

MedGemma sees each X-ray as 256 image tokens placed ahead of the text. Their activations look nothing like text’s, and a transcoder trained the usual way, on text alone, does not transfer. Its fraction of variance unexplained (FVU, the share of the MLP output it fails to reconstruct) on image tokens is 1.13, worse than simply predicting the mean, and at layer 0 it is 354× worse than on text. Balancing image and text tokens in every training batch brings image FVU to 0.19 at no cost to text.

3How faithful it is#

The standard test is replacement: substitute the transcoder for all 34 MLPs, with no correction term, and check whether the model still behaves the same. On 60 held-out image/report pairs the replacement model predicts the original top-1 next token every time, with a mean KL divergence of 0.044 between the two output distributions.

This holds on the prompt the CLT was trained with. On a constrained yes/no question the text side degrades (KL 13.5), while image-token reconstruction is identical across prompts: under causal masking the image tokens come before the question and cannot see it. Retraining on mixed prompts cuts the probe KL 13.5×, while doubling the dictionary does not help, so the bottleneck is prompt coverage, not capacity.

4Exact attribution graphs#

To explain one answer, we build a local replacement model for that X-ray: freeze the attention patterns and normalization at the values they take on this input, and add back the small per-position gap between the CLT and the true MLPs. What is left is linear in the features, so every edge in the resulting attribution graph is exact, and the edges into the answer sum to its logit (the model’s raw score for that token). We checked each property: the logit matches the original exactly; scaling features moves it along a straight line (deviation 0.094, versus 11.2 with attention left live, about 120× closer); and the edges sum to the logit within bf16 precision.

5What the features see#

Labeling features by their top-activating image patches returns concrete radiology: lung bases, the mediastinum, tubes and catheters, shoulders and spine, and acquisition cues such as the “AP PORTABLE” stamp on bedside films. In layer 30, 25 of the top 30 features receive a specific anatomical or acquisition label. The features also know where the evidence is. Occluding small patches of an X-ray reveals which regions actually drive the model’s answer, and a linear probe on the features picks out those regions with AUC 0.88, against 0.62 for raw pixels and 0.48 for a shuffled-label control.

Costophrenic angle / lung baseL30 · f1045

Mediastinal contour & paraspinal regionL30 · f1434

Tubes & catheters crossing the chestL30 · f1439

Clavicle & shoulder regionL30 · f1344

Spine & paraspinal regionL30 · f1827

“AP / AP PORTABLE” technique textL24 · f565

6A readout artifact#

The usual way to test a feature is to intervene on it and watch the output. On films where MedGemma falsely calls cardiomegaly (an enlarged heart), we negated the 80 image features most associated with the call. P(yes), read as the softmax probability of the “yes” token, fell from 0.92 to 0.63: an apparent handle on the finding.

It is not one. P(no) did not rise; the lost probability went to other tokens. The intervention made the model less likely to answer in yes/no form, not less convinced. Renormalized over yes and no, P(yes) stays at 1.00, a change of 0.002. The raw readout inflates the effect about 150×.

Intervening on the input, by contrast, does move the decision. Blacking out the single most causal patch drops the normalized P(yes) by 0.72, while blacking out a patch at the border drops it by 0.001. The call is driven by a specific image region, the features encode where that region is, and yet editing the features does not change the decision.

The same kind of control caught one of our own earlier claims. We had localized a false cardiomegaly call on portable (AP) films to the attention pattern, because patching in a standard-film pattern removed it. Patching in the pattern from another portable film removes 87% as much, so the effect was mostly generic disruption rather than anything specific to film type. We have withdrawn it.

7Limits#

One model and one training run. The released CLT is faithful on its training prompt, not yet on arbitrary questions (§3). A per-layer transcoder trained on the same data reconstructs about as well at a seventeenth of the size, so the case for the cross-layer design rests on circuit tracing, not reconstruction. Feature labels are automatic and have not been validated by a radiologist.