Proposal
Make TransformerBridge.boot_transformers apply the same logit/weight normalization preprocessing that HookedTransformer.from_pretrained applies by default -specifically fold_ln, center_writing_weights, and center_unembed -or expose them as constructor flags and document the difference prominently in the v3 migration guide.
Currently the bridge skips all three by default, so for the same model and same input, TransformerBridge and HookedTransformer return numerically different logits and cached activations, even though the README presents them as interchangeable.
Motivation
The README and v3 migration guide describe TransformerBridge as the recommended replacement for HookedTransformer.from_pretrained, with the legacy API "still available through a compatibility layer." A natural reading is that the two paths produce equivalent results for the same model.
In practice they do not. For gpt2 (verified across four different prompts):
TransformerBridge matches raw HuggingFace GPT2LMHeadModel to fp32 noise (max abs diff on logits in the 5e-5 to 9e-5 range).
HookedTransformer differs from both by max 116–240 / mean 72–97 on the same inputs -and the difference is a single constant per token position (std-of-diff across the vocab dimension is ~1e-5).
Because the difference is a per-row constant, softmax / argmax / cross-entropy are unaffected, so generation and loss behave identically and the discrepancy is invisible in casual use. But anything consuming raw logit or activation values -logit lens, calibration, KL between models, ensembling, cached unembed.hook_out, residual-stream–based analyses -will silently produce different numbers depending on which API loaded the model.
When run_with_cache results are compared on shared hook names, ~half of the legacy hooks (100/210 for gpt2) hold numerically different tensors on the bridge, with magnitudes growing with depth - consistent with the absent centering propagating through the residual stream. So this is not just a logit-layer concern; it affects the intermediate activations that mech-interp work is actually about.
The root cause is configuration, not implementation: HookedTransformer.from_pretrained defaults fold_ln=True, center_writing_weights=True, center_unembed=True (HookedTransformer.py L1153–1155). TransformerBridge.boot_transformers does not apply them. Both paths are individually correct; they disagree because their defaults differ.
Pitch
Two acceptable resolutions, in preference order:
Option A -change the bridge default to match legacy. Have TransformerBridge.boot_transformers apply fold_ln / center_writing_weights / center_unembed by default for architectures where they are valid (skipping shortformer, OLMo 2 post-norm, models with logit softcap, etc., consistent with the existing exclusion logic in HookedTransformer.py lines 1384–1416). This preserves the v3 migration framing of the bridge as a drop-in replacement, keeps the residual stream "clean" in the way mech-interp work has historically assumed, and means tutorials and reusable analyses written against HookedTransformer continue to produce identical numbers under the bridge.
Option B -expose the flags and document the change. Add fold_ln, center_writing_weights, center_unembed as keyword arguments to boot_transformers (defaulting to False, i.e. current behavior), and add a clearly-labeled section to the v3 migration guide titled e.g. "Bridge returns un-centered logits and activations by default" explaining the implication for raw activation values. Users can opt in to legacy-equivalent behavior with boot_transformers("gpt2", fold_ln=True, center_writing_weights=True, center_unembed=True).
Alternatives
Document only, change nothing. Add a clear note to the migration guide that the bridge returns raw HF-equivalent logits and activations, and that this differs from legacy. Cheapest, but leaves the "drop-in replacement" implication intact and silently invalidates any analysis ported across the two APIs.
Additional context
Reproduction (TransformerLens HEAD 58b007f, torch 2.5.1+cu121, A40, transformers ≥4.56):
import torch
from transformers import GPT2LMHeadModel
from transformer_lens import HookedTransformer
from transformer_lens.model_bridge import TransformerBridge
legacy = HookedTransformer.from_pretrained("gpt2", device="cuda").eval()
bridge = TransformerBridge.boot_transformers("gpt2", device="cuda"); bridge.eval()
hf = GPT2LMHeadModel.from_pretrained("gpt2").to("cuda").eval()
ids = legacy.to_tokens("The capital of France is")
with torch.no_grad():
hf_l, leg_l, br_l = hf(ids).logits, legacy(ids), bridge(ids)
print("HF vs bridge max:", (hf_l - br_l ).abs().max().item()) # ~5e-5
print("HF vs legacy max:", (hf_l - leg_l).abs().max().item()) # ~116
print("leg vs bridge max:", (leg_l - br_l ).abs().max().item()) # ~116
the diff is a per-row constant, not noise:
d = (br_l - leg_l)[0]
print("per-row std (vocab dim):", d.std(dim=-1).max().item()) # ~1e-5
Visible-by-eye exhibit — first 5 logits at the last position for prompt "The capital of France is":
HF : -111.72, -110.79, -114.99, -115.99, -114.74
bridge : -111.72, -110.79, -114.99, -115.99, -114.74 (matches HF)
legacy : 4.73, 5.66, 1.46, 0.46, 1.71 (= HF + 116.45)
Disabling the legacy preprocessing makes all three paths converge:
pythonlegacy_raw = HookedTransformer.from_pretrained(
"gpt2", device="cuda",
fold_ln=False, center_writing_weights=False, center_unembed=False,
).eval()
legacy_raw vs bridge max = 6.10e-05
legacy_raw vs HF max = 6.87e-05
i.e. when legacy is configured to not preprocess, it numerically matches both the bridge and raw HF, confirming the centering defaults are the entire source of the divergence.
A separate but related observation: when run_with_cache is run on both APIs and the resulting caches are compared on the 210 shared hook names, blocks.{i}.attn.hook_attn_scores differs by inf on every layer, which looks like one path caches scores pre–causal-mask and the other post-mask. I'll file that as a separate issue if it isn't a duplicate, since it's an independent concern from the centering question.
Checklist
- [Y ] I have checked that there is no similar issue in the repo (required)
Proposal
Make TransformerBridge.boot_transformers apply the same logit/weight normalization preprocessing that HookedTransformer.from_pretrained applies by default -specifically fold_ln, center_writing_weights, and center_unembed -or expose them as constructor flags and document the difference prominently in the v3 migration guide.
Currently the bridge skips all three by default, so for the same model and same input, TransformerBridge and HookedTransformer return numerically different logits and cached activations, even though the README presents them as interchangeable.
Motivation
The README and v3 migration guide describe TransformerBridge as the recommended replacement for HookedTransformer.from_pretrained, with the legacy API "still available through a compatibility layer." A natural reading is that the two paths produce equivalent results for the same model.
In practice they do not. For gpt2 (verified across four different prompts):
TransformerBridge matches raw HuggingFace GPT2LMHeadModel to fp32 noise (max abs diff on logits in the 5e-5 to 9e-5 range).
HookedTransformer differs from both by max 116–240 / mean 72–97 on the same inputs -and the difference is a single constant per token position (std-of-diff across the vocab dimension is ~1e-5).
Because the difference is a per-row constant, softmax / argmax / cross-entropy are unaffected, so generation and loss behave identically and the discrepancy is invisible in casual use. But anything consuming raw logit or activation values -logit lens, calibration, KL between models, ensembling, cached unembed.hook_out, residual-stream–based analyses -will silently produce different numbers depending on which API loaded the model.
When run_with_cache results are compared on shared hook names, ~half of the legacy hooks (100/210 for gpt2) hold numerically different tensors on the bridge, with magnitudes growing with depth - consistent with the absent centering propagating through the residual stream. So this is not just a logit-layer concern; it affects the intermediate activations that mech-interp work is actually about.
The root cause is configuration, not implementation: HookedTransformer.from_pretrained defaults fold_ln=True, center_writing_weights=True, center_unembed=True (HookedTransformer.py L1153–1155). TransformerBridge.boot_transformers does not apply them. Both paths are individually correct; they disagree because their defaults differ.
Pitch
Two acceptable resolutions, in preference order:
Option A -change the bridge default to match legacy. Have TransformerBridge.boot_transformers apply fold_ln / center_writing_weights / center_unembed by default for architectures where they are valid (skipping shortformer, OLMo 2 post-norm, models with logit softcap, etc., consistent with the existing exclusion logic in HookedTransformer.py lines 1384–1416). This preserves the v3 migration framing of the bridge as a drop-in replacement, keeps the residual stream "clean" in the way mech-interp work has historically assumed, and means tutorials and reusable analyses written against HookedTransformer continue to produce identical numbers under the bridge.
Option B -expose the flags and document the change. Add fold_ln, center_writing_weights, center_unembed as keyword arguments to boot_transformers (defaulting to False, i.e. current behavior), and add a clearly-labeled section to the v3 migration guide titled e.g. "Bridge returns un-centered logits and activations by default" explaining the implication for raw activation values. Users can opt in to legacy-equivalent behavior with boot_transformers("gpt2", fold_ln=True, center_writing_weights=True, center_unembed=True).
Alternatives
Document only, change nothing. Add a clear note to the migration guide that the bridge returns raw HF-equivalent logits and activations, and that this differs from legacy. Cheapest, but leaves the "drop-in replacement" implication intact and silently invalidates any analysis ported across the two APIs.
Additional context
Reproduction (TransformerLens HEAD 58b007f, torch 2.5.1+cu121, A40, transformers ≥4.56):
import torch
from transformers import GPT2LMHeadModel
from transformer_lens import HookedTransformer
from transformer_lens.model_bridge import TransformerBridge
legacy = HookedTransformer.from_pretrained("gpt2", device="cuda").eval()
bridge = TransformerBridge.boot_transformers("gpt2", device="cuda"); bridge.eval()
hf = GPT2LMHeadModel.from_pretrained("gpt2").to("cuda").eval()
ids = legacy.to_tokens("The capital of France is")
with torch.no_grad():
hf_l, leg_l, br_l = hf(ids).logits, legacy(ids), bridge(ids)
print("HF vs bridge max:", (hf_l - br_l ).abs().max().item()) # ~5e-5
print("HF vs legacy max:", (hf_l - leg_l).abs().max().item()) # ~116
print("leg vs bridge max:", (leg_l - br_l ).abs().max().item()) # ~116
the diff is a per-row constant, not noise:
d = (br_l - leg_l)[0]
print("per-row std (vocab dim):", d.std(dim=-1).max().item()) # ~1e-5
Visible-by-eye exhibit — first 5 logits at the last position for prompt "The capital of France is":
HF : -111.72, -110.79, -114.99, -115.99, -114.74
bridge : -111.72, -110.79, -114.99, -115.99, -114.74 (matches HF)
legacy : 4.73, 5.66, 1.46, 0.46, 1.71 (= HF + 116.45)
Disabling the legacy preprocessing makes all three paths converge:
pythonlegacy_raw = HookedTransformer.from_pretrained(
"gpt2", device="cuda",
fold_ln=False, center_writing_weights=False, center_unembed=False,
).eval()
legacy_raw vs bridge max = 6.10e-05
legacy_raw vs HF max = 6.87e-05
i.e. when legacy is configured to not preprocess, it numerically matches both the bridge and raw HF, confirming the centering defaults are the entire source of the divergence.
A separate but related observation: when run_with_cache is run on both APIs and the resulting caches are compared on the 210 shared hook names, blocks.{i}.attn.hook_attn_scores differs by inf on every layer, which looks like one path caches scores pre–causal-mask and the other post-mask. I'll file that as a separate issue if it isn't a duplicate, since it's an independent concern from the centering question.
Checklist