[Bug Report] Compatibility mode legacy hook aliases do not preserve HookedTransformer hook semantics

[SamuelePunzo]

TransformerBridge.enable_compatibility_mode() appears to expose some legacy HookedTransformer hook names without preserving the same computation points.

In particular, the legacy aliases for attention input hooks seem to be semantically different from the legacy HookedTransformer hooks:

• blocks.{layer}.hook_q_input
• blocks.{layer}.hook_k_input
• blocks.{layer}.hook_v_input

My understanding is that in legacy HookedTransformer, these correspond to pre-LN residual-stream forks, while in current TransformerBridge compatibility mode they fire on post-LN tensors instead.

This means compatibility mode may currently provide:

• matching or near-matching logits,
• matching hook names,
• matching hook shapes,

while still not preserving the same hooked activations and backward gradients.

Observed behavior

For GPT-2, the current behavior appears to be:

• logits are close,
• legacy hook aliases can be registered,
• hook shapes match,
• but Q/K/V input hook values do not match legacy HookedTransformer,
• and downstream attribution-style scores diverge substantially.

There may also be a similar issue for:

• blocks.{layer}.hook_mlp_in

Checklist

• I have checked that there is no similar [issue]

It seems to me like one of these should be made explicit:

1. Compatibility mode should preserve legacy hook semantics.
2. If exact semantic parity is not intended, the aliases/docs should say they are name/shape compatible rather than semantically equivalent.
3. Both should exist:
   • bridge-native canonical post-LN hooks,
   • legacy-compatible aliases for legacy pre-LN hook semantics.

My preference would be option 3, since it preserves bridge-native behavior while making compatibility mode meaningful for legacy tooling.

I’d be happy to work on this and open a PR if this direction sounds right to maintainers.

[jlarson4]

Hi @SamuelePunzo! Thanks for the detailed write up.

This is a known bug, it is documented within the code, and I've actually already started a fix for this specific feature. I haven't touched it in about a week, but I will pick it up today and see if I can't get it wrapped up and included in the next release. It is closest to your option 1 (compat mode preserves legacy semantics), but effectively gets you option 3's benefit too (both surfaces remain available, just with different names). I'll send another ping your way when the PR is up

[SamuelePunzo]

Thanks, that’s very helpful to know.
Glad the write-up matched the intended direction!

I’ll keep an eye out for the PR.

[jlarson4]

Pr is now up at #1319, take a look!

[SamuelePunzo]

I tested PR #1319 locally against my EAP-IG parity checks. The hook-level parity looks much better now, but my tiny EAP / EAP-IG score parity tests are still not passing, so I’m still investigating the remaining end-to-end gap on my side.

[jlarson4]

Glad to hear it is working better, sorry it doesn't cover everything. If you can provide a little more detail on what you're seeing, I can also look into it from my end. I have a couple hypotheses about what it could be.

[jlarson4]

@SamuelePunzo #1319 is merged now, and I will be doing the 3.3.0 release later today. If you end up finding a solution for your end-to-end gap in party tests, feel free to put up a PR and once its merged I'll make that its own patch release in case any other users are hitting this same bug. Thanks for your work on this, keep me updated if you run into any issues or have any questions.

[SamuelePunzo]

Thanks, I reran everything against the merged #1319 code locally. For my GPT-2 parity case, the original TransformerBridge compatibility problem now looks fixed: hook value parity passes, hook gradient parity passes, and plain EAP parity now passes between HookedTransformer and TransformerBridge.

The only remaining mismatch is in my EAP-IG-inputs end-to-end path. At this point the leftover gap looks like an EAP-IG-side issue rather than the original TransformerLens hook-semantics issue.