From 007c7ee5054e1150c29663abb42165a686748647 Mon Sep 17 00:00:00 2001 From: emerard <113128214+emerardd@users.noreply.github.com> Date: Sun, 25 Oct 2026 16:01:45 +0800 Subject: [PATCH 1/2] fix: respect parallel MLP inputs in residual decomposition --- .../test_parallel_mlp_input_decomposition.py | 121 ++++++++++++++++++ .../test_parallel_residual_identities.py | 16 +++ transformer_lens/ActivationCache.py | 41 ++++-- .../generalized_components/block.py | 3 + .../model_bridge/transformer_bridge.py | 11 +- 5 files changed, 173 insertions(+), 19 deletions(-) create mode 100644 tests/integration/model_bridge/test_parallel_mlp_input_decomposition.py diff --git a/tests/integration/model_bridge/test_parallel_mlp_input_decomposition.py b/tests/integration/model_bridge/test_parallel_mlp_input_decomposition.py new file mode 100644 index 0000000000..ecd80324d3 --- /dev/null +++ b/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) diff --git a/tests/integration/model_bridge/test_parallel_residual_identities.py b/tests/integration/model_bridge/test_parallel_residual_identities.py index f9dbcb68a8..015457c66a 100644 --- a/tests/integration/model_bridge/test_parallel_residual_identities.py +++ b/tests/integration/model_bridge/test_parallel_residual_identities.py @@ -162,6 +162,22 @@ 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 +) +def test_mlp_input_decomposition_matches_norm_input( + hf_architecture: str, config_cls: Any, config_kwargs: dict[str, Any] +) -> None: + bridge, cache = _run(hf_architecture, config_cls, config_kwargs) + 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" + components = cache.decompose_resid(layer=layer, mlp_input=True) + torch.testing.assert_close( + components.sum(0), cache[f"blocks.{layer}.{norm}.hook_in"], atol=1e-5, rtol=1e-5 + ) + + @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 diff --git a/transformer_lens/ActivationCache.py b/transformer_lens/ActivationCache.py index bc6d1c7a7b..edc9726b6d 100644 --- a/transformer_lens/ActivationCache.py +++ b/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] @@ -1334,8 +1342,9 @@ 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 LN scale of ln2 is used instead of ln1. + The neurons stacked are unchanged. """ if layer is None or layer == -1: # Default to the residual stream immediately pre unembed @@ -1366,9 +1375,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: @@ -1614,7 +1625,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) diff --git a/transformer_lens/model_bridge/generalized_components/block.py b/transformer_lens/model_bridge/generalized_components/block.py index 085958bb10..950b979880 100644 --- a/transformer_lens/model_bridge/generalized_components/block.py +++ b/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, diff --git a/transformer_lens/model_bridge/transformer_bridge.py b/transformer_lens/model_bridge/transformer_bridge.py index 081e700685..930bbd3784 100644 --- a/transformer_lens/model_bridge/transformer_bridge.py +++ b/transformer_lens/model_bridge/transformer_bridge.py @@ -1740,7 +1740,7 @@ def accumulated_bias( """Sum of variant + MLP output biases through the residual stream up to `layer`. Includes all layer types (attn, SSM, linear-attn). Set mlp_input=True - to include the variant bias of the target layer itself. + to target its MLP input: the current variant bias is included only for sequential blocks. """ accumulated = torch.zeros(self.cfg.d_model, device=self.cfg.device) for i in range(layer): @@ -1756,10 +1756,11 @@ def accumulated_bias( if mlp_input: assert layer < self.cfg.n_layers, "Cannot include attn_bias from beyond the final layer" block = self.blocks[layer] - scale = getattr(block, "residual_contribution_scale", 1.0) - b_O = self._get_block_variant_bias(block) - if b_O is not None: - accumulated = accumulated + b_O.to(accumulated.device) * scale + if getattr(block, "mlp_input_includes_attn", True): + scale = getattr(block, "residual_contribution_scale", 1.0) + b_O = self._get_block_variant_bias(block) + if b_O is not None: + accumulated = accumulated + b_O.to(accumulated.device) * scale return accumulated def all_composition_scores(self, mode: str) -> CompositionScores: From 558597d39e682b20cd56f6dec57e0be6aa3ae802 Mon Sep 17 00:00:00 2001 From: emerard <113128214+emerardd@users.noreply.github.com> Date: Tue, 27 Oct 2026 00:16:34 +0800 Subject: [PATCH 2/2] fix: use shared MLP normalization scales in cache decomposition --- .../test_parallel_residual_identities.py | 59 +++++++++++++++++-- transformer_lens/ActivationCache.py | 32 +++++++--- 2 files changed, 78 insertions(+), 13 deletions(-) diff --git a/tests/integration/model_bridge/test_parallel_residual_identities.py b/tests/integration/model_bridge/test_parallel_residual_identities.py index 015457c66a..b3a65008b9 100644 --- a/tests/integration/model_bridge/test_parallel_residual_identities.py +++ b/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 = [ @@ -165,17 +177,54 @@ def test_sequential_variant_exposes_resid_mid( @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] + 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" - components = cache.decompose_resid(layer=layer, mlp_input=True) - torch.testing.assert_close( - components.sum(0), cache[f"blocks.{layer}.{norm}.hook_in"], atol=1e-5, rtol=1e-5 - ) + 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]) diff --git a/transformer_lens/ActivationCache.py b/transformer_lens/ActivationCache.py index edc9726b6d..40e391b408 100644 --- a/transformer_lens/ActivationCache.py +++ b/transformer_lens/ActivationCache.py @@ -1207,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 @@ -1250,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) @@ -1343,7 +1357,8 @@ def stack_neuron_results( 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`` for sequential blocks or ``resid_pre`` for parallel blocks. - With ``apply_ln=True`` the LN scale of ln2 is used instead of ln1. + 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: @@ -1490,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