fix: respect parallel MLP inputs in residual decomposition

tests/integration/model_bridge/test_parallel_mlp_input_decomposition.py

@@ -0,0 +1,121 @@
+"""MLP-input decomposition against real tiny NeoX normalization inputs, without Hub loads."""
+
+import pytest
+import torch
+from transformers import GPTNeoXConfig, GPTNeoXForCausalLM
+
+from transformer_lens.model_bridge.sources import build_bridge_from_module
+
+
+@pytest.fixture(scope="module", params=[False, True], ids=["sequential", "parallel"])
+def model_and_cache(request):
+    with torch.random.fork_rng():
+        torch.manual_seed(0)
+        config = GPTNeoXConfig(
+            vocab_size=32,
+            hidden_size=16,
+            intermediate_size=32,
+            num_hidden_layers=2,
+            num_attention_heads=2,
+            max_position_embeddings=16,
+            rotary_pct=0.5,
+            use_parallel_residual=request.param,
+        )
+        config._attn_implementation = "eager"
+        hf = GPTNeoXForCausalLM(config).eval()
+        with torch.no_grad():
+            for block in hf.gpt_neox.layers:
+                block.attention.dense.bias.copy_(torch.linspace(-0.4, 0.4, 16))
+                block.mlp.dense_4h_to_h.bias.copy_(torch.linspace(0.1, 0.2, 16))
+        bridge = build_bridge_from_module(
+            hf, "GPTNeoXForCausalLM", hf_config=config, dtype=torch.float32, device="cpu"
+        )
+        with torch.no_grad():
+            _, cache = bridge.run_with_cache(torch.tensor([[1, 2, 3], [3, 2, 1]]))
+    return bridge, cache, request.param
+
+
+@pytest.mark.parametrize("layer", [0, 1])
+@pytest.mark.parametrize("normalized", [False, True])
+@pytest.mark.parametrize("batchless", [False, True])
+def test_component_and_full_decomposition_match_actual_mlp_input(
+    model_and_cache, layer, normalized, batchless
+):
+    bridge, cache, parallel = model_and_cache
+    if batchless:
+        cache = cache.apply_slice_to_batch_dim(0)
+    target = cache[f"blocks.{layer}.ln2.hook_{'out' if normalized else 'in'}"]
+    components, labels = cache.decompose_resid(
+        layer=layer, mlp_input=True, apply_ln=normalized, return_labels=True
+    )
+    torch.testing.assert_close(components.sum(0), target)
+    assert (f"{layer}_attn_out" in labels) is not parallel
+    for expand in (False, True):
+        full = cache.get_full_resid_decomposition(
+            layer=layer, mlp_input=True, expand_neurons=expand, apply_ln=normalized
+        )
+        torch.testing.assert_close(full.sum(0), target, atol=2e-5, rtol=1e-5)
+        directions = torch.arange(32, dtype=torch.float32).reshape(16, 2) / 32
+        projected = cache.get_full_resid_decomposition(
+            layer=layer,
+            mlp_input=True,
+            expand_neurons=expand,
+            apply_ln=normalized,
+            project_output_onto=directions,
+        )
+        torch.testing.assert_close(projected.sum(0), target @ directions, atol=2e-5, rtol=1e-5)
+
+
+@pytest.mark.parametrize("layer", [0, 1])
+@pytest.mark.parametrize("include_mlp", [False, True])
+def test_accumulated_bias_stops_before_parallel_attention(model_and_cache, layer, include_mlp):
+    bridge, _, parallel = model_and_cache
+    expected = torch.zeros(16)
+    for block in bridge.blocks[:layer]:
+        expected += block.attn.b_O
+        if include_mlp:
+            expected += block.mlp.b_out
+    if not parallel:
+        expected += bridge.blocks[layer].attn.b_O
+    torch.testing.assert_close(
+        bridge.accumulated_bias(layer, mlp_input=True, include_mlp_biases=include_mlp), expected
+    )
+
+
+@pytest.mark.parametrize("layer", [0, 1])
+def test_accumulated_resid_uses_existing_streams(model_and_cache, layer):
+    _, cache, parallel = model_and_cache
+    parts, labels = cache.accumulated_resid(
+        layer=layer, mlp_input=True, incl_mid=True, return_labels=True
+    )
+    torch.testing.assert_close(parts[-1], cache[f"blocks.{layer}.ln2.hook_in"])
+    assert any(label.endswith("_mid") for label in labels) is not parallel
+
+
+@pytest.mark.parametrize("layer", [0, 1])
+@pytest.mark.parametrize("normalized", [False, True])
+@pytest.mark.parametrize("projected", ["none", "vector", "matrix"])
+@pytest.mark.parametrize("batchless", [False, True])
+def test_neuron_remainder_matches_mlp_input(
+    model_and_cache, layer, normalized, projected, batchless
+):
+    _, cache, _ = model_and_cache
+    if batchless:
+        cache = cache.apply_slice_to_batch_dim(0)
+    direction = None
+    if projected == "vector":
+        direction = torch.arange(16, dtype=torch.float32) / 16
+    elif projected == "matrix":
+        direction = torch.arange(32, dtype=torch.float32).reshape(16, 2) / 32
+    parts = cache.stack_neuron_results(
+        layer=layer,
+        mlp_input=True,
+        incl_remainder=True,
+        apply_ln=normalized,
+        pos_slice=-1,
+        project_output_onto=direction,
+    )
+    target = cache[f"blocks.{layer}.ln2.hook_{'out' if normalized else 'in'}"][..., -1, :]
+    if direction is not None:
+        target = target @ direction
+    torch.testing.assert_close(parts.sum(0), target, atol=2e-5, rtol=1e-5)

tests/integration/model_bridge/test_parallel_residual_identities.py

@@ -74,6 +74,18 @@
         "StableLmForCausalLM", StableLmConfig, dict(_GQA, use_parallel_residual=True), id="stablelm"
     ),
     pytest.param("FalconForCausalLM", FalconConfig, dict(_FALCON, parallel_attn=True), id="falcon"),
+    pytest.param(
+        "FalconForCausalLM",
+        FalconConfig,
+        dict(
+            _FALCON,
+            new_decoder_architecture=True,
+            num_ln_in_parallel_attn=2,
+            multi_query=False,
+            num_kv_heads=4,
+        ),
+        id="falcon-two-norms",
+    ),
 ]
 
 SEQUENTIAL_CASES = [
@@ -162,6 +174,59 @@ def test_sequential_variant_exposes_resid_mid(
         assert f"blocks.{layer}.hook_resid_mid" in cache
 
 
+@pytest.mark.parametrize(
+    "hf_architecture,config_cls,config_kwargs", PARALLEL_CASES + SEQUENTIAL_CASES
+)
+@pytest.mark.parametrize("normalized", [False, True])
+@pytest.mark.parametrize("batchless", [False, True])
+def test_mlp_input_decomposition_matches_norm_input(
+    hf_architecture: str,
+    config_cls: Any,
+    config_kwargs: dict[str, Any],
+    normalized: bool,
+    batchless: bool,
+) -> None:
+    bridge, cache = _run(hf_architecture, config_cls, config_kwargs)
+    if batchless:
+        cache = cache.apply_slice_to_batch_dim(0)
+    for layer in range(bridge.cfg.n_layers):
+        # Shared-norm parallel blocks feed both branches from ln1; NeoX uses a separate ln2.
+        norm = "ln2" if "ln2" in bridge.blocks[layer]._modules else "ln1"
+        target = cache[f"blocks.{layer}.{norm}.hook_{'out' if normalized else 'in'}"]
+        components = cache.decompose_resid(layer=layer, mlp_input=True, apply_ln=normalized)
+        torch.testing.assert_close(components.sum(0), target, atol=1e-5, rtol=1e-5)
+        for expand in (False, True):
+            full = cache.get_full_resid_decomposition(
+                layer=layer, mlp_input=True, apply_ln=normalized, expand_neurons=expand
+            )
+            torch.testing.assert_close(full.sum(0), target, atol=2e-5, rtol=1e-5)
+        for projected in (False, True):
+            direction = torch.arange(bridge.cfg.d_model, dtype=torch.float32) / bridge.cfg.d_model
+            neurons = cache.stack_neuron_results(
+                layer=layer,
+                mlp_input=True,
+                apply_ln=normalized,
+                incl_remainder=True,
+                pos_slice=-1,
+                project_output_onto=direction if projected else None,
+            )
+            expected = target[..., -1, :] @ direction if projected else target[..., -1, :]
+            torch.testing.assert_close(neurons.sum(0), expected, atol=2e-5, rtol=1e-5)
+
+
+def test_missing_independent_mlp_scale_does_not_use_attention_scale() -> None:
+    from transformer_lens.ActivationCache import ActivationCache
+
+    bridge, cache = _run(
+        "GPTNeoXForCausalLM", GPTNeoXConfig, dict(_MHA, use_parallel_residual=True)
+    )
+    incomplete = ActivationCache(
+        {key: value for key, value in cache.items() if key != "blocks.0.ln2.hook_scale"}, bridge
+    )
+    with pytest.raises(KeyError, match="blocks.0.ln2.hook_scale"):
+        incomplete.decompose_resid(layer=0, mlp_input=True, apply_ln=True)
+
+
 @pytest.mark.parametrize("use_parallel_residual", [True, False])
 def test_hf_use_parallel_residual_reaches_the_bridge_config(use_parallel_residual: bool) -> None:
     """The adapter reads HF's own flag, so the passthrough must carry it — the

transformer_lens/ActivationCache.py

@@ -373,6 +373,12 @@ def apply_slice_to_batch_dim(self, batch_slice: Union[Slice, SliceInput]) -> Act
         }
         return ActivationCache(new_cache_dict, self.model, has_batch_dim=still_has_batch_dim)
 
+    def _mlp_input_includes_attn(self, layer: int) -> bool:
+        """Whether the target block's MLP reads the current attention contribution."""
+        return layer < self.model.cfg.n_layers and bool(
+            getattr(self.model.blocks[layer], "mlp_input_includes_attn", True)
+        )
+
     def accumulated_resid(
         self,
         layer: Optional[int] = None,
@@ -484,7 +490,8 @@ def accumulated_resid(
                 immediately pre logits). The indices are taken such that this gives the accumulated
                 streams up to the input to layer l.
             incl_mid:
-                Whether to return `resid_mid` for all previous layers.
+                Whether to return `resid_mid` for previous sequential-residual layers.
+                Parallel-residual layers have no distinct mid stream.
             apply_ln:
                 Whether to apply the final layer norm to the stack. When True, applies
                 `model.ln_final`, which recomputes normalization statistics (mean and
@@ -493,8 +500,8 @@ def accumulated_resid(
             pos_slice:
                 A slice object to apply to the pos dimension. Defaults to None, do nothing.
             mlp_input:
-                Whether to include resid_mid for the current layer. This essentially gives the MLP
-                input rather than the attention input.
+                Target the MLP input: include the current resid_mid for sequential blocks;
+                parallel blocks already expose this input as resid_pre.
             return_labels:
                 Whether to return a list of labels for the residual stream components. Useful for
                 labelling graphs.
@@ -518,7 +525,9 @@ def accumulated_resid(
                 continue
             components_list.append(self[("resid_pre", l)])
             labels.append(f"{l}_pre")
-            if (incl_mid and l < layer) or (mlp_input and l == layer):
+            if ((incl_mid and l < layer) or (mlp_input and l == layer)) and (
+                self._mlp_input_includes_attn(l)
+            ):
                 components_list.append(self[("resid_mid", l)])
                 labels.append(f"{l}_mid")
         components_list = [pos_slice.apply(c, dim=-2) for c in components_list]
@@ -713,9 +722,8 @@ def decompose_resid(
                 means just embed and pos_embed. The indices are taken such that this gives the
                 accumulated streams up to the input to layer l
             mlp_input:
-                Whether to include attn_out for the current
-                layer - essentially decomposing the residual stream that's input to the MLP input
-                rather than the Attn input.
+                Decompose the MLP input rather than the attention input. The current layer's
+                attention contribution is included only for sequential-residual blocks.
             mode:
                 Values are "all", "mlp" or "attn". "all" returns all
                 components, "mlp" returns only the MLP components, and "attn" returns only the
@@ -763,7 +771,7 @@ def decompose_resid(
             if incl_mlp:
                 components_list.append(self[("mlp_out", l)])
                 labels.append(f"{l}_mlp_out")
-        if mlp_input and incl_attn:
+        if mlp_input and incl_attn and self._mlp_input_includes_attn(layer):
             components_list.append(self[("attn_out", layer)])
             labels.append(f"{layer}_attn_out")
         components_list = [pos_slice.apply(c, dim=-2) for c in components_list]
@@ -1199,16 +1207,30 @@ def _get_cached_ln_scale(
         when the expected hook isn't in the cache (some non-decoder-only architectures expose
         LN scale at a different path or not at all).
         """
+        norm = "ln2" if mlp_input else "ln1"
+        blocks = getattr(self.model, "blocks", None)
+        block = (
+            blocks[layer]
+            if blocks is not None and layer is not None and layer != self.model.cfg.n_layers
+            else None
+        )
+        if (
+            mlp_input
+            and block is not None
+            and layer is not None
+            and not self._mlp_input_includes_attn(layer)
+            and getattr(block, "ln2", None) is None
+            and getattr(block, "ln2_post", None) is None
+        ):
+            # Shared-norm parallel blocks feed the MLP from ln1, not a separate ln2.
+            norm = "ln1"
         if layer == self.model.cfg.n_layers or layer is None:
             key = "ln_final.hook_scale"
         else:
-            key = f"blocks.{layer}.ln{2 if mlp_input else 1}.hook_scale"
+            key = f"blocks.{layer}.{norm}.hook_scale"
         try:
             scale = self[key]
         except KeyError as e:
-            norm = f"ln{2 if mlp_input else 1}"
-            blocks = getattr(self.model, "blocks", None)
-            block = blocks[layer] if blocks is not None and layer is not None else None
             if (
                 block is not None
                 and getattr(block, norm, None) is None
@@ -1242,8 +1264,8 @@ def _stack_neuron_results_apply_ln_projected(
 
         Analytical formula (LN models, cached scale ``s``):
             ``LN_s(a_n * W_out_n) @ p = (a_n / s) * (W_out_n @ p - mean(W_out_n) * sum_p)``
-        RMS models drop the ``mean(W_out_n) * sum_p`` term (no centering). ``s`` is the ln2
-        scale of ``layer`` when ``mlp_input`` is set, else the ln1 scale.
+        RMS models drop the ``mean(W_out_n) * sum_p`` term (no centering). ``s`` is the
+        MLP-input norm scale when ``mlp_input`` is set, else the ln1 scale.
         """
         scale = self._get_cached_ln_scale(layer, mlp_input=mlp_input, pos_slice=pos_slice)
 
@@ -1334,8 +1356,10 @@ def stack_neuron_results(
             mlp_input:
                 Treat the stack as the input to ``layer``'s MLP rather than its attention, as in
                 ``decompose_resid``: with ``incl_remainder=True`` the remainder fills the stack up to
-                ``resid_mid`` instead of ``resid_pre``, and with ``apply_ln=True`` the LN scale of
-                ln2 is used instead of ln1. The neurons stacked are unchanged.
+                ``resid_mid`` for sequential blocks or ``resid_pre`` for parallel blocks.
+                With ``apply_ln=True`` the scale of the norm feeding the MLP is used
+                (ln2, or ln1 on shared-norm parallel blocks).
+                The neurons stacked are unchanged.
         """
         if layer is None or layer == -1:
             # Default to the residual stream immediately pre unembed
@@ -1366,9 +1390,11 @@ def stack_neuron_results(
         neuron_labels = neuron_slice.apply(torch.arange(d_mlp), dim=0)
 
         labels = [f"L{l}N{h}" for l in range(layer) for h in neuron_labels]
-        # The remainder fills the stack up to the residual stream that layer's LN normalizes:
-        # resid_mid (the MLP input) with mlp_input, else resid_pre (resid_post of layer - 1).
-        remainder_target = ("resid_mid", layer) if mlp_input else ("resid_post", layer - 1)
+        if mlp_input and layer < self.model.cfg.n_layers:
+            stream = "resid_mid" if self._mlp_input_includes_attn(layer) else "resid_pre"
+            remainder_target = (stream, layer)
+        else:
+            remainder_target = ("resid_post", layer - 1)
         components: Any
         ln_folded = apply_ln and project_2d is not None
         if ln_folded:
@@ -1479,7 +1505,8 @@ def apply_ln_to_stack(
                 The layer we're taking the input to. In [0, n_layers], n_layers means the unembed.
                 None maps to the n_layers case, ie the unembed.
             mlp_input:
-                Whether the input is to the MLP or attn (ie ln2 vs ln1). Defaults to False, ie ln1.
+                Whether the input is to the MLP or attention. The MLP uses ln2 when present,
+                or ln1 on shared-norm parallel blocks. Defaults to False, ie ln1.
                 If layer==n_layers, must be False, and we use ln_final
             pos_slice:
                 The slice to take of positions, if residual_stack is not over the full context, None
@@ -1614,7 +1641,9 @@ def _ln_then_project(stack: torch.Tensor) -> torch.Tensor:
             return stack @ project_2d if project_2d is not None else stack
 
         head_stack, head_labels = self.stack_head_results(
-            layer + (1 if mlp_input else 0), pos_slice=pos_slice, return_labels=True
+            layer + int(mlp_input and self._mlp_input_includes_attn(layer)),
+            pos_slice=pos_slice,
+            return_labels=True,
         )
         if ln_folded:
             head_stack = _ln_then_project(head_stack)

transformer_lens/model_bridge/generalized_components/block.py

@@ -38,6 +38,7 @@ class BlockBridge(GeneralizedComponent):
 
     is_list_item: bool = True
     hook_out_is_single_residual_stream: bool = True
+    mlp_input_includes_attn: bool = True
     # hook_mlp_in is a direct HookPoint on this class (not aliased) so it can
     # fire on the MLP-branch entry (pre-ln2, or the MLP input on post-norm
     # blocks); see __init__. The normalized mlp input stays at block.mlp.hook_in.
@@ -466,6 +467,8 @@ class ParallelBlockBridge(BlockBridge):
     of the adapter sees ``ParallelBlockBridge`` and knows the hook is absent.
     """
 
+    mlp_input_includes_attn: bool = False
+
     def __init__(
         self,
         name: str,