[Bug Report] BLOOM `TransformerBridge` maps `hook_attn_out` / `hook_mlp_out` to residual-added states, breaking HookedTransformer residual identities

[hijohnnylin]

Versions: transformer-lens 3.7.0, also reproduced on 3.6.0; transformers 5.14.1; torch 2.13.0; CPU; float32

Summary

On bigscience/bloom-560m, the TransformerBridge compatibility hooks:

blocks.{layer}.hook_attn_out
blocks.{layer}.hook_mlp_out

contain residual-added states rather than additive attention and MLP contributions.

In HookedTransformer, these hook names represent the tensors added to the residual stream, so the following identities hold:

resid_pre + attn_out = resid_mid
resid_mid + mlp_out  = resid_post

On the BLOOM bridge, the aliases instead expose the Hugging Face attention and MLP module outputs. BLOOM performs its residual additions inside those modules, so for this checkpoint:

hook_attn_out = resid_mid
hook_mlp_out  = resid_post

Using those values as additive contributions double-counts an entire residual stream.

The converted HookedTransformer implementation is unaffected.

Scope of the issue

This issue is specifically about the legacy TransformerLens-compatible aliases:

blocks.{layer}.hook_attn_out
blocks.{layer}.hook_mlp_out

A canonical architecture-shaped hook such as blocks.{layer}.attn.hook_out may intentionally expose the literal output of the underlying Hugging Face submodule.

However, aliases named hook_attn_out and hook_mlp_out should preserve the established HookedTransformer semantics associated with those names: the additive components of the residual-stream decomposition.

Root cause

BLOOM performs the residual addition inside each sublayer:

class BloomAttention(nn.Module):
    def forward(self, hidden_states, residual, ...):
        ...
        output_tensor = self.dense(context_layer)
        output_tensor = dropout_add(
            output_tensor,
            residual,
            self.hidden_dropout,
            self.training,
        )
        return output_tensor, ...

class BloomMLP(nn.Module):
    def forward(self, hidden_states, residual):
        ...
        intermediate_output = self.dense_4h_to_h(hidden_states)
        output = dropout_add(
            intermediate_output,
            residual,
            self.hidden_dropout,
            self.training,
        )
        return output

The general relationship is:

module output = residual argument + dropout(projected sublayer output)

For bigscience/bloom-560m in this reproduction:

• the attention residual argument is resid_pre;
• the MLP residual argument is resid_mid;
• hidden dropout is zero.

Therefore:

BloomAttention module output = resid_mid
BloomMLP module output       = resid_post

Mapping hook_attn_out and hook_mlp_out to those module outputs changes their meaning from additive contributions to accumulated residual states.

A residual parameter in a forward signature is a useful warning sign, but the generic condition is whether the returned tensor already contains a residual addition.

Expected behavior

For every sequential-residual block, the TransformerLens decomposition should satisfy:

torch.testing.assert_close(
    cache[f"blocks.{layer}.hook_resid_pre"]
    + cache[f"blocks.{layer}.hook_attn_out"],
    cache[f"blocks.{layer}.hook_resid_mid"],
)

torch.testing.assert_close(
    cache[f"blocks.{layer}.hook_resid_mid"]
    + cache[f"blocks.{layer}.hook_mlp_out"],
    cache[f"blocks.{layer}.hook_resid_post"],
)

This holds exactly for the converted BLOOM HookedTransformer.

Actual behavior

For the BLOOM TransformerBridge, the decomposition fails:

bigscience/bloom-560m / HookedTransformer:

  L0   resid_pre + attn_out == resid_mid
       cos=+1.000000  rel=0.0000  max_abs=0

  L0   resid_mid + mlp_out == resid_post
       cos=+1.000000  rel=0.0000  max_abs=0

  L12  resid_pre + attn_out == resid_mid
       cos=+1.000000  rel=0.0000  max_abs=0

  L12  resid_mid + mlp_out == resid_post
       cos=+1.000000  rel=0.0000  max_abs=0

  L23  resid_pre + attn_out == resid_mid
       cos=+1.000000  rel=0.0000  max_abs=0

  L23  resid_mid + mlp_out == resid_post
       cos=+1.000000  rel=0.0000  max_abs=0


bigscience/bloom-560m / TransformerBridge:

  L0   resid_pre + attn_out == resid_mid
       cos=+0.977455  rel=0.9037  max_abs=8.486

  L0   resid_mid + mlp_out == resid_post
       cos=+0.966881  rel=0.7104  max_abs=8.604

  L12  resid_pre + attn_out == resid_mid
       cos=+0.999993  rel=0.9996  max_abs=516.9

  L12  resid_mid + mlp_out == resid_post
       cos=+0.999952  rel=0.9979  max_abs=517.0

  L23  resid_pre + attn_out == resid_mid
       cos=+0.999157  rel=0.9727  max_abs=50.21

  L23  resid_mid + mlp_out == resid_post
       cos=+0.999485  rel=0.9930  max_abs=52.75

Here, relative error is:

||lhs - rhs||₂ / ||rhs||₂

A relative error near 1 is the expected signature when the sum is wrong by approximately one entire residual stream.

Cosine similarity remains above 0.96 in every failing case and above 0.999 in several cases. A direction-only comparison can therefore make these tensors appear nearly identical even though the residual decomposition is incorrect.

As a control, both implementations satisfy all six identities exactly for GPT-2:

gpt2 / HookedTransformer: all six exact, max_abs=0
gpt2 / TransformerBridge: all six exact, max_abs=0

Minimal reproduction

import torch
from transformer_lens import HookedTransformer
from transformer_lens.model_bridge import TransformerBridge

MODEL = "bigscience/bloom-560m"
PROMPT = "The capital of France is Paris."
LAYERS = (0, 12, 23)


def max_abs_error(left, right):
    return (left - right).abs().max().item()


models = (
    (
        "HookedTransformer",
        lambda: HookedTransformer.from_pretrained_no_processing(
            MODEL,
            device="cpu",
            dtype=torch.float32,
        ),
    ),
    (
        "TransformerBridge",
        lambda: TransformerBridge.boot_transformers(
            MODEL,
            device="cpu",
            dtype=torch.float32,
        ),
    ),
)

for implementation, load in models:
    model = load()

    with torch.inference_mode():
        _, cache = model.run_with_cache(PROMPT)

    print(f"\n{implementation}")

    for layer in LAYERS:
        resid_pre = cache[f"blocks.{layer}.hook_resid_pre"]
        attn_out = cache[f"blocks.{layer}.hook_attn_out"]
        resid_mid = cache[f"blocks.{layer}.hook_resid_mid"]
        mlp_out = cache[f"blocks.{layer}.hook_mlp_out"]
        resid_post = cache[f"blocks.{layer}.hook_resid_post"]

        print(
            layer,
            "resid_pre + attn_out -> resid_mid:",
            max_abs_error(resid_pre + attn_out, resid_mid),
        )
        print(
            layer,
            "resid_mid + mlp_out -> resid_post:",
            max_abs_error(resid_mid + mlp_out, resid_post),
        )

        # These additionally show the semantic collision on the bridge.
        print(
            layer,
            "attn_out -> resid_mid:",
            max_abs_error(attn_out, resid_mid),
        )
        print(
            layer,
            "mlp_out -> resid_post:",
            max_abs_error(mlp_out, resid_post),
        )

The key decomposition errors are:

HookedTransformer:
  layer 0:  0.0
  layer 12: 0.0
  layer 23: 0.0

TransformerBridge, attention identity:
  layer 0:  approximately 8.49
  layer 12: approximately 516.9
  layer 23: approximately 50.2

Hooking the bridge's blocks[layer].attn or blocks[layer].mlp submodules directly produces the same residual-added tensors. The problem is therefore the semantic location of the hook boundary, not the mechanism used to register the hook.

Why this matters

Many TransformerLens analyses rely on the residual decomposition:

resid_mid  = resid_pre + attn_out
resid_post = resid_mid + mlp_out

Examples include:

• residual-stream attribution;
• direct-logit attribution;
• activation patching;
• component ablations;
• SAE training on attention or MLP outputs;
• comparisons between converted models and bridged Hugging Face models.

On BLOOM, the bridge returns tensors with the expected names, shapes, and dtypes, but their semantics are different. Downstream code can therefore produce incorrect results without raising an exception.

Suggested fix

The bridge should distinguish between:

1. the literal tensor returned by the Hugging Face attention or MLP module; and
2. the additive contribution inserted into the residual stream.

For BLOOM, the contribution hook should run after output projection and hidden dropout but before the residual addition.

The legacy aliases:

blocks.{layer}.hook_attn_out
blocks.{layer}.hook_mlp_out

should target these contribution hooks so that they preserve HookedTransformer semantics.

Canonical architecture-shaped hooks may continue to expose the literal residual-added Hugging Face module outputs under separate names.

For read-only cache construction, the contributions can alternatively be derived by subtraction:

attn_out = resid_mid - resid_pre
mlp_out = resid_post - resid_mid

That restores the decomposition by construction. For interventions, however, subtraction alone is insufficient: a modified contribution must be reinserted into the computation before reconstructing the residual-added module output.

A regression test should assert the following identities for each sequential-residual architecture supported by the bridge:

torch.testing.assert_close(resid_pre + attn_out, resid_mid)
torch.testing.assert_close(resid_mid + mlp_out, resid_post)

Checklist

• I have checked that there is no similar issue in the repo (required)

[jlarson4]

Resolved and released as part of #1643, now available as v3.7.1