Upload README.md with huggingface_hub

README.md

@@ -12,34 +12,61 @@ tags:
 
 # MedGemma Circuit Tools
 
-Circuit-tracing artifacts for MedGemma, from SAIL Lab. First entry:
+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).
 
-## Layer-17 Transcoder (ReXGradient-trained)
+Both are top-k transcoders trained on MLP activations of
+[`google/medgemma-4b-it`](https://huggingface.co/google/medgemma-4b-it) collected
+from radiology-report text
+([ReXGradient-160K](https://huggingface.co/datasets/rajpurkarlab/ReXGradient-160K)),
+for circuit tracing and mechanistic analysis of paraphrase sensitivity in medical
+vision-language models.
 
-A top-k transcoder for **layer 17** of [`google/medgemma-4b-it`](https://huggingface.co/google/medgemma-4b-it),
-trained on layer-17 MLP activations collected from radiology-report text
-([ReXGradient-160K](https://huggingface.co/datasets/rajpurkarlab/ReXGradient-160K)).
-Built for circuit tracing and mechanistic analysis of paraphrase sensitivity in
-medical vision-language models (PSF-Med, Sadanandan et al. 2026).
-
-## What it is
+## Shared configuration
 
 - **Architecture:** top-k transcoder, 2,560 → 20,480 features (8x expansion), `top_k = 64`
-- **Hookpoint:** layer 17 of MedGemma-4B's language backbone (Gemma 3 4B),
-  reconstructing MLP output from MLP input
+- **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
-- **Motivation:** layer 17 carries Feature 3818, the "clinical query register gate"
-  identified in PSF-Med circuit analysis using Gemma Scope 2. This transcoder is
-  the domain-adapted counterpart for MedGemma-specific mechanistic work.
 
-## Validation (post-training, 64 documents per slice)
+## 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 | domain-specific; not for general text |
+| 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
 
@@ -51,9 +78,10 @@ from train_layer17_transcoder import (
 )
 
 model, processor = load_model(device="cuda")
-collector = Layer17ActivationCollector(model, layer=17)
+collector = Layer17ActivationCollector(model, layer=29)  # or 17
 
-ckpt = torch.load("final.pt", map_location="cpu", weights_only=False)
+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"])
@@ -64,31 +92,35 @@ mlp_in, mlp_out = collector.get_activations(inputs)
 reconstructed, features = transcoder(mlp_in.view(-1, 2560).float())
 ```
 
-Training and validation code: `circuit_tracing/train_layer17_transcoder.py`,
+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](https://github.com/UNHSAILLab/medical-vlm-robustness) repository.
 
 ## Limitations
 
-- **Single layer.** This is a layer-17 transcoder only, not a full transcoder
-  suite. For other layers or general-domain text, use
+- **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](https://huggingface.co/google/gemma-scope-2-4b-it).
-- **Domain-specific.** Reconstruction degrades on non-medical text
-  (EV 0.884 on WikiText). That specialization is intentional.
+- **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; this artifact is derived from MedGemma activations.
+  Foundations terms; these artifacts are derived from MedGemma activations.
 
 ## Citation
 
 ```bibtex
-@misc{sadanandan2026medgemma_l17_transcoder,
-  title  = {MedGemma-4B Layer-17 Transcoder (ReXGradient-trained)},
+@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 transcoder for circuit tracing of paraphrase
+  note   = {Domain-adapted transcoders for circuit tracing of paraphrase
             sensitivity in medical VLMs. SAIL Lab, University of New Haven.}
 }
 ```