From eb5a53cbdb905ce9100a08f0f0734b5788204319 Mon Sep 17 00:00:00 2001 From: Zhe Wang <53119724+flyingdoubleG@users.noreply.github.com> Date: Sun, 8 Nov 2026 21:36:14 +0800 Subject: [PATCH 1/2] fix: residual decompositions on Gemma 4 (layer_scalar, per-layer inputs, norms) Gemma 4 layers multiply the whole residual stream by a learned layer_scalar and E2B/E4B add a per-layer input branch after the MLP; the decompositions now carry each contribution through later layer scalars and add the per-layer input rows. The Gemma 4 norms are bridged as native RMS norm bridges so heads and neurons map through the post norms and apply_ln has a cached scale. Neuron expansion on MoE layers raises instead of returning a wrong stack. --- .../test_gemma4_resid_decomposition.py | 226 ++++++++++++++++++ .../test_gemma4_adapter.py | 29 +++ transformer_lens/ActivationCache.py | 181 ++++++++++++-- .../supported_architectures/gemma4.py | 69 +++++- 4 files changed, 480 insertions(+), 25 deletions(-) create mode 100644 tests/integration/model_bridge/test_gemma4_resid_decomposition.py diff --git a/tests/integration/model_bridge/test_gemma4_resid_decomposition.py b/tests/integration/model_bridge/test_gemma4_resid_decomposition.py new file mode 100644 index 0000000000..0cf814de70 --- /dev/null +++ b/tests/integration/model_bridge/test_gemma4_resid_decomposition.py @@ -0,0 +1,226 @@ +"""Residual-decomposition APIs on Gemma 4. + +Each Gemma 4 decoder layer normalizes its attention and MLP outputs (sandwich norms), +E2B/E4B add a per-layer input branch after the MLP, and every layer ends with +``hidden_states *= layer_scalar``, a learned per-layer scale on the whole residual stream. +The decompositions must still sum to the residual stream they decompose. +""" +from unittest.mock import MagicMock + +import pytest +import torch + +from transformer_lens.model_bridge.bridge import TransformerBridge +from transformer_lens.model_bridge.sources._bridge_builder import ( + build_bridge_config_from_hf, +) + +transformers = pytest.importorskip("transformers") +try: + from transformers.models.gemma4 import Gemma4TextConfig +except ImportError: # pragma: no cover - older transformers + pytest.skip("transformers has no Gemma 4", allow_module_level=True) + +ATOL = 1e-5 +TOKENS = torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8], [8, 7, 6, 5, 4, 3, 2, 1]]) +_BASE = dict( + vocab_size=128, + hidden_size=64, + intermediate_size=128, + num_hidden_layers=4, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=16, + global_head_dim=16, + sliding_window=8, + max_position_embeddings=64, + layer_types=["sliding_attention", "full_attention", "sliding_attention", "full_attention"], +) +_VARIANTS = { + # Per-layer inputs and KV sharing, as on E2B / E4B. + "e2b_like": dict( + hidden_size_per_layer_input=16, vocab_size_per_layer_input=128, num_kv_shared_layers=2 + ), + # K == V global attention, as on 31B. + "31b_like": dict( + hidden_size_per_layer_input=0, attention_k_eq_v=True, num_global_key_value_heads=1 + ), + # Dense MLP plus experts, as on 26B-A4B. + "moe_like": dict( + hidden_size_per_layer_input=0, + enable_moe_block=True, + num_experts=4, + top_k_experts=2, + moe_intermediate_size=32, + attention_k_eq_v=True, + num_global_key_value_heads=1, + ), +} +# Distinct per-layer scales, so a contribution carried through the wrong layers is caught. +_LAYER_SCALARS = [0.3, 0.7, 1.1, 1.5] + + +def _gemma4(variant: str) -> TransformerBridge: + from transformers import AutoModelForCausalLM + + from transformer_lens.factories.architecture_adapter_factory import ( + ArchitectureAdapterFactory, + ) + + torch.manual_seed(0) + c = Gemma4TextConfig(**_BASE, **_VARIANTS[variant]) + c.architectures = ["Gemma4ForCausalLM"] + hf = AutoModelForCausalLM.from_config(c, attn_implementation="eager").eval() + with torch.no_grad(): + for layer, scalar in zip(hf.model.layers, _LAYER_SCALARS): + layer.layer_scalar.fill_(scalar) + for name, param in layer.named_parameters(): + if "norm" in name and param.ndim == 1: + param.normal_(1.0, 0.3) + cfg = build_bridge_config_from_hf(hf.config, "Gemma4ForCausalLM", "tiny", torch.float32) + adapter = ArchitectureAdapterFactory.select_architecture_adapter(cfg) + return TransformerBridge(model=hf, adapter=adapter, tokenizer=MagicMock()) + + +@pytest.fixture(scope="module", params=sorted(_VARIANTS)) +def bridge_and_cache(request): + bridge = _gemma4(request.param) + with torch.no_grad(): + _, cache = bridge.run_with_cache(TOKENS) + return request.param, bridge, cache + + +def _final(bridge, cache): + return cache[f"blocks.{bridge.cfg.n_layers - 1}.hook_resid_post"] + + +def _target(bridge, cache, layer, mlp_input=False): + if mlp_input: + return cache[f"blocks.{layer}.hook_resid_mid"] + if layer == bridge.cfg.n_layers: + return _final(bridge, cache) + return cache[f"blocks.{layer}.hook_resid_pre"] + + +def test_layer_scalar_rescales_the_residual_stream(bridge_and_cache): + """The setup this file relies on: each block's output is scaled after its residual adds.""" + _, bridge, cache = bridge_and_cache + write = cache["blocks.0.hook_resid_mid"] - cache["blocks.0.hook_resid_pre"] + torch.testing.assert_close(write, cache["blocks.0.hook_attn_out"], atol=ATOL, rtol=0) + unscaled = cache["blocks.0.hook_resid_post"] / _LAYER_SCALARS[0] + assert not torch.allclose(unscaled, cache["blocks.0.hook_resid_post"]) + + +class TestDecomposeResid: + @pytest.mark.parametrize("layer", [0, 1, 2, 3, 4]) + def test_sums_to_resid_pre(self, bridge_and_cache, layer): + _, bridge, cache = bridge_and_cache + comps = cache.decompose_resid(layer=layer) + torch.testing.assert_close(comps.sum(0), _target(bridge, cache, layer), atol=ATOL, rtol=0) + + @pytest.mark.parametrize("layer", [0, 1, 2, 3]) + def test_mlp_input_sums_to_resid_mid(self, bridge_and_cache, layer): + _, bridge, cache = bridge_and_cache + comps = cache.decompose_resid(layer=layer, mlp_input=True) + torch.testing.assert_close( + comps.sum(0), _target(bridge, cache, layer, mlp_input=True), atol=ATOL, rtol=0 + ) + + def test_labels(self, bridge_and_cache): + name, bridge, cache = bridge_and_cache + _, labels = cache.decompose_resid(return_labels=True) + per_layer = ["attn_out", "mlp_out"] + if name == "e2b_like": + per_layer.append("per_layer_input_out") + assert labels == ["embed"] + [ + f"{l}_{kind}" for l in range(bridge.cfg.n_layers) for kind in per_layer + ] + + def test_attn_and_mlp_modes_partition_the_stream(self, bridge_and_cache): + _, bridge, cache = bridge_and_cache + attn = cache.decompose_resid(mode="attn", incl_embeds=False).sum(0) + mlp = cache.decompose_resid(mode="mlp", incl_embeds=False).sum(0) + embed = cache.decompose_resid(mode="attn", return_labels=False)[0] + torch.testing.assert_close(attn + mlp + embed, _final(bridge, cache), atol=ATOL, rtol=0) + + +class TestHeadsAndNeurons: + @pytest.mark.parametrize("layer", [1, 2, 3, 4]) + def test_heads_plus_remainder(self, bridge_and_cache, layer): + _, bridge, cache = bridge_and_cache + stack = cache.stack_head_results(layer=layer, incl_remainder=True) + torch.testing.assert_close( + stack.sum(0), cache[f"blocks.{layer - 1}.hook_resid_post"], atol=ATOL, rtol=0 + ) + + def test_heads_sum_to_scaled_attn_out(self, bridge_and_cache): + """Through ln1_post and the layer's own layer_scalar.""" + _, bridge, cache = bridge_and_cache + n_heads = bridge.cfg.n_heads + for layer in range(bridge.cfg.n_layers): + heads = cache.stack_head_results(layer=layer + 1)[layer * n_heads :] + torch.testing.assert_close( + heads.sum(0), + cache[f"blocks.{layer}.hook_attn_out"] * _LAYER_SCALARS[layer], + atol=ATOL, + rtol=0, + ) + + def test_neurons_plus_remainder(self, bridge_and_cache): + name, bridge, cache = bridge_and_cache + if name == "moe_like": + with pytest.raises(NotImplementedError, match="post_feedforward_layernorm_1"): + cache.stack_neuron_results(layer=bridge.cfg.n_layers, incl_remainder=True) + return + stack = cache.stack_neuron_results(layer=bridge.cfg.n_layers, incl_remainder=True) + torch.testing.assert_close(stack.sum(0), _final(bridge, cache), atol=ATOL, rtol=0) + + +class TestFullResidDecomposition: + @pytest.mark.parametrize("expand_neurons", [False, True]) + def test_sums_to_final_residual(self, bridge_and_cache, expand_neurons): + name, bridge, cache = bridge_and_cache + if name == "moe_like" and expand_neurons: + with pytest.raises(NotImplementedError, match="experts"): + cache.get_full_resid_decomposition(expand_neurons=True) + return + stack, labels = cache.get_full_resid_decomposition( + expand_neurons=expand_neurons, return_labels=True + ) + assert stack.shape[0] == len(labels) + if name == "e2b_like": + assert "0_per_layer_input_out" in labels + torch.testing.assert_close(stack.sum(0), _final(bridge, cache), atol=ATOL, rtol=0) + + @pytest.mark.parametrize("expand_neurons", [False, True]) + @pytest.mark.parametrize("layer", [1, 3]) + def test_mlp_input_sums_to_resid_mid(self, bridge_and_cache, expand_neurons, layer): + name, bridge, cache = bridge_and_cache + if name == "moe_like" and expand_neurons: + pytest.skip("neuron expansion is refused on MoE layers") + stack = cache.get_full_resid_decomposition( + layer=layer, mlp_input=True, expand_neurons=expand_neurons + ) + torch.testing.assert_close( + stack.sum(0), _target(bridge, cache, layer, mlp_input=True), atol=ATOL, rtol=0 + ) + + @pytest.mark.parametrize("expand_neurons", [False, True]) + def test_ln_and_projection(self, bridge_and_cache, expand_neurons): + name, bridge, cache = bridge_and_cache + if name == "moe_like" and expand_neurons: + pytest.skip("neuron expansion is refused on MoE layers") + torch.manual_seed(1) + directions = torch.randn(bridge.cfg.d_model, 3) + stack = cache.get_full_resid_decomposition( + expand_neurons=expand_neurons, apply_ln=True, project_output_onto=directions + ) + n = bridge.cfg.n_layers + expected = cache.apply_ln_to_stack(_final(bridge, cache)[None], n)[0] @ directions + torch.testing.assert_close(stack.sum(0), expected, atol=1e-4, rtol=0) + + def test_logit_attrs_runs(self, bridge_and_cache): + _, bridge, cache = bridge_and_cache + resid = cache.decompose_resid() + attrs = cache.logit_attrs(resid, tokens=torch.tensor([3, 3])) + assert attrs.shape[0] == resid.shape[0] diff --git a/tests/unit/model_bridge/supported_architectures/test_gemma4_adapter.py b/tests/unit/model_bridge/supported_architectures/test_gemma4_adapter.py index 839b5ce663..0cd13c15d4 100644 --- a/tests/unit/model_bridge/supported_architectures/test_gemma4_adapter.py +++ b/tests/unit/model_bridge/supported_architectures/test_gemma4_adapter.py @@ -15,6 +15,7 @@ EmbeddingBridge, GatedMLPBridge, LinearBridge, + RMSNormalizationBridge, RotaryEmbeddingBridge, UnembeddingBridge, ) @@ -23,6 +24,7 @@ ) from transformer_lens.model_bridge.supported_architectures.gemma4 import ( Gemma4ArchitectureAdapter, + Gemma4BlockBridge, ) ARCH = "Gemma4ForConditionalGeneration" @@ -118,6 +120,33 @@ def test_block_decomposition(): assert blocks.submodules[norm].optional is False +def test_norms_are_native_norm_bridges(): + """Norm bridges on HF's own forward: hook_scale for the residual decompositions.""" + m = _adapter().component_mapping + norms = [m["blocks"].submodules[n] for n in ("ln1", "ln1_post", "ln2", "ln2_post")] + for norm in norms + [m["ln_final"]]: + assert isinstance(norm, RMSNormalizationBridge) + assert norm.use_native_layernorm_autograd is True + + +def test_block_exposes_layer_scalar_and_per_layer_input_write(): + blocks = _adapter().component_mapping["blocks"] + assert isinstance(blocks, Gemma4BlockBridge) + assert blocks.residual_side_outputs == {"post_per_layer_input_norm": "per_layer_input"} + + layer = torch.nn.Module() + layer.layer_scalar = torch.nn.Buffer(torch.tensor([0.25])) + blocks.set_original_component(layer) + assert torch.equal(blocks.residual_output_scale, torch.tensor([0.25])) + assert blocks.neuron_expansion_unsupported_reason is None + + +def test_moe_blocks_refuse_neuron_expansion(): + blocks = _adapter().component_mapping["blocks"] + blocks.add_module("experts", torch.nn.Identity()) + assert "post_feedforward_layernorm_1" in blocks.neuron_expansion_unsupported_reason + + def test_split_qkv_fork_aliases_absent(): """Attention is delegated wholesale to HF; per-layer structure is heterogeneous (KV-shared layers have no k/v projections), so the split-qkv fork aliases diff --git a/transformer_lens/ActivationCache.py b/transformer_lens/ActivationCache.py index 2fa617dc67..b23fb3e4ef 100644 --- a/transformer_lens/ActivationCache.py +++ b/transformer_lens/ActivationCache.py @@ -494,6 +494,93 @@ def _mixer_stack( labels.append(self._mixer_label(layer, name)) return components, labels + def _residual_output_scales(self) -> List[Optional[torch.Tensor]]: + """Per block, the scale it applies to the whole residual stream after its residual adds. + + Gemma 4 multiplies each layer's output by a learned ``layer_scalar``, so everything + written before it is rescaled too. None for blocks without such a scale (or a scale of + 1). Read once per cache. + """ + memo = self.__dict__.get("_residual_output_scales_memo") + if memo is not None: + return memo + scales: List[Optional[torch.Tensor]] = [] + for block in list(getattr(self.model, "blocks", None) or []): + scale = getattr(block, "residual_output_scale", None) + if scale is not None and not isinstance(scale, torch.Tensor): + scale = torch.tensor(float(scale)) + if scale is not None and bool(torch.all(scale == 1)): + scale = None + scales.append(None if scale is None else scale.detach()) + self.__dict__["_residual_output_scales_memo"] = scales + return scales + + def _downstream_scale(self, layer: int, target: int) -> Optional[torch.Tensor]: + """Factor by which a contribution written in block ``layer`` reaches ``target``. + + The product of the residual output scales of blocks ``layer`` to ``target - 1``; for the + embeddings use ``layer=0``. None when it is 1. + """ + total: Optional[torch.Tensor] = None + for scale in self._residual_output_scales()[max(layer, 0) : max(target, 0)]: + if scale is not None: + total = scale if total is None else total * scale + return total + + @staticmethod + def _rescale(x: torch.Tensor, scale: Optional[torch.Tensor]) -> torch.Tensor: + return x if scale is None else x * scale.to(device=x.device, dtype=x.dtype) + + def _side_outputs(self, layer: int, pos_slice: Slice) -> Tuple[List[torch.Tensor], List[str]]: + """Residual writes of block ``layer`` after its MLP, e.g. Gemma 4's per-layer input. + + A block lists them in ``residual_side_outputs`` as ``{submodule: label}``; each writes + its submodule's ``hook_out`` and is labelled ``{layer}_{label}_out``. + """ + block = self.model.blocks[layer] + outputs: List[torch.Tensor] = [] + labels: List[str] = [] + for name, label in (getattr(block, "residual_side_outputs", None) or {}).items(): + if name in block._modules: + outputs.append(pos_slice.apply(self[f"blocks.{layer}.{name}.hook_out"], dim=-2)) + labels.append(f"{layer}_{label}_out") + return outputs, labels + + def _check_neurons_expandable(self, layer: int) -> None: + """Raise when block ``layer``'s ``mlp`` neurons do not map linearly onto its MLP output. + + E.g. Gemma 4's MoE layers normalize the dense MLP output and add the experts' output + before the post-feedforward norm. + """ + reason = getattr(self.model.blocks[layer], "neuron_expansion_unsupported_reason", None) + if reason: + raise NotImplementedError( + f"blocks.{layer}: {reason} Use expand_neurons=False or decompose_resid instead." + ) + + def _scaled_accumulated_bias( + self, layer: int, mlp_input: bool, include_mlp_biases: bool + ) -> torch.Tensor: + """``accumulated_bias`` with each layer's biases carried through later residual scales.""" + if self._downstream_scale(0, layer) is None: + return self.model.accumulated_bias( + layer, mlp_input, include_mlp_biases=include_mlp_biases + ) + total = self.model.accumulated_bias(0, mlp_input=False) + previous = total + for l in range(layer): + current = self.model.accumulated_bias( + l + 1, mlp_input=False, include_mlp_biases=include_mlp_biases + ) + total = total + self._rescale(current - previous, self._downstream_scale(l, layer)) + previous = current + if mlp_input: + with_attn = self.model.accumulated_bias( + layer, mlp_input=True, include_mlp_biases=include_mlp_biases + ) + total = total + (with_attn - previous) + return total + def accumulated_resid( self, layer: Optional[int] = None, @@ -845,7 +932,9 @@ def decompose_resid( attention components. Defaults to "all". On SSM / hybrid models, blocks without attention contribute their sequence-mixer output (e.g. ``0_linear_attn_out``, ``0_mixer_out``) in the attention slot, and blocks without an MLP contribute no - ``mlp_out``. + ``mlp_out``. Residual writes after the MLP (Gemma 4's ``0_per_layer_input_out``) + are in the MLP slot. When later blocks rescale the whole residual stream + (Gemma 4's ``layer_scalar``), each component is carried through those scales. apply_ln: Whether to apply LayerNorm to the stack. pos_slice: @@ -874,23 +963,29 @@ def decompose_resid( incl_mlp = mode != "attn" and not self.model.cfg.attn_only components_list = [] labels = [] + # Block each component is written in, for residual scales applied after it (embeds: 0). + sources: List[int] = [] if incl_embeds: if self.has_embed: components_list = [self["hook_embed"]] labels.append("embed") + sources.append(0) if self.has_pos_embed: components_list.append(self["hook_pos_embed"]) labels.append("pos_embed") + sources.append(0) def add_mixer(l: int) -> None: mixers, mixer_labels = self._mixer_stack(l, Slice(None)) components_list.extend(mixers) labels.extend(mixer_labels) + sources.extend([l] * len(mixers)) def add_attn_slot(l: int) -> None: if self._has_attn(l): components_list.append(self[("attn_out", l)]) labels.append(f"{l}_attn_out") + sources.append(l) elif self._mixer_role(l) == "attn": add_mixer(l) @@ -901,11 +996,19 @@ def add_attn_slot(l: int) -> None: if self._has_mlp(l): components_list.append(self[("mlp_out", l)]) labels.append(f"{l}_mlp_out") + sources.append(l) elif not self._has_attn(l) and self._mixer_role(l) == "mlp": add_mixer(l) + side_outputs, side_labels = self._side_outputs(l, Slice(None)) + components_list.extend(side_outputs) + labels.extend(side_labels) + sources.extend([l] * len(side_outputs)) if mlp_input and incl_attn and self._mlp_input_includes_attn(layer): add_attn_slot(layer) - components_list = [pos_slice.apply(c, dim=-2) for c in components_list] + components_list = [ + self._rescale(pos_slice.apply(c, dim=-2), self._downstream_scale(src, layer)) + for c, src in zip(components_list, sources) + ] if components_list: components = torch.stack(components_list, dim=0) else: @@ -1231,6 +1334,7 @@ def _post_norm_bias_correction( if norm_bias is not None: mapped = mapped + norm_bias.to(device=scale.device) delta = (mapped - raw) * getattr(block, "residual_contribution_scale", 1.0) + delta = self._rescale(delta, self._downstream_scale(l, layer)) correction = delta if correction is None else correction + delta return correction @@ -1281,7 +1385,24 @@ def stack_head_results( raise ValueError( f"layer must be >= 0 with incl_remainder=True (-1/None mean n_layers); got {layer}." ) + return self._stack_head_results( + layer, return_labels, incl_remainder, pos_slice, apply_ln, scale_target=layer + ) + def _stack_head_results( + self, + layer: int, + return_labels: bool, + incl_remainder: bool, + pos_slice: Slice, + apply_ln: bool, + scale_target: int, + ) -> Any: + """``stack_head_results`` with heads carried to the residual stream at ``scale_target``. + + ``get_full_resid_decomposition(mlp_input=True)`` stacks the heads of ``layer`` itself, + which reach ``resid_mid`` before that block's residual output scale. + """ # Idempotent; cleans up stale Bridge entries without the 'already cached' warning, # which is only meaningful for direct compute_head_results callers. self._ensure_head_results() @@ -1298,7 +1419,7 @@ def stack_head_results( if post is not None: scale, weight, centers, _ = post result = self._apply_post_norm_map(result, scale.unsqueeze(-2), weight, centers) - components.append(result) + components.append(self._rescale(result, self._downstream_scale(l, scale_target))) labels.extend([f"L{l}H{h}" for h in range(self.model.cfg.n_heads)]) if components: components = torch.cat(components, dim=-2) @@ -1405,6 +1526,7 @@ def get_neuron_results( if not isinstance(pos_slice, Slice): pos_slice = Slice(pos_slice) + self._check_neurons_expandable(layer) neuron_acts = self[("post", layer, "mlp")] block = self.model.blocks[layer] W_out = block.mlp.W_out @@ -1544,7 +1666,11 @@ def _stack_neuron_results_apply_ln_projected( if post is not None: a_l = a_l / post_scale # (a_l / s)[..., None] is [..., d_mlp, 1]; broadcast with lin_form_l [d_mlp, n_outs] - components.append((a_l / scale)[..., None] * lin_form_l) + components.append( + self._rescale( + (a_l / scale)[..., None] * lin_form_l, self._downstream_scale(l, layer) + ) + ) if not components: empty_src = pos_slice.apply(self["hook_embed"], dim=-2) return torch.zeros( @@ -1653,6 +1779,7 @@ def stack_neuron_results( f"({names}), so its neurons cannot be expanded. Use " "expand_neurons=False or decompose_resid instead." ) + self._check_neurons_expandable(l) labels = [f"L{l}N{h}" for l in mlp_layers for h in neuron_labels] if mlp_input and layer < self.model.cfg.n_layers: stream = "resid_mid" if self._mlp_input_includes_attn(layer) else "resid_pre" @@ -1681,14 +1808,13 @@ def stack_neuron_results( else: per_layer: list = [] for l in mlp_layers: - per_layer.append( - self.get_neuron_results( - l, - pos_slice=pos_slice, - neuron_slice=neuron_slice, - project_output_onto=project_2d, - ) + neuron_results = self.get_neuron_results( + l, + pos_slice=pos_slice, + neuron_slice=neuron_slice, + project_output_onto=project_2d, ) + per_layer.append(self._rescale(neuron_results, self._downstream_scale(l, layer))) if per_layer: components = torch.cat(per_layer, dim=-2) components = einops.rearrange( @@ -1910,10 +2036,13 @@ def _ln_then_project(stack: torch.Tensor) -> torch.Tensor: stack = self.apply_ln_to_stack(stack, layer, pos_slice=pos_slice, mlp_input=mlp_input) return stack @ project_2d if project_2d is not None else stack - head_stack, head_labels = self.stack_head_results( + head_stack, head_labels = self._stack_head_results( layer + int(mlp_input and self._mlp_input_includes_attn(layer)), - pos_slice=pos_slice, return_labels=True, + incl_remainder=False, + pos_slice=pos_slice, + apply_ln=False, + scale_target=layer, ) if ln_folded: head_stack = _ln_then_project(head_stack) @@ -1930,7 +2059,7 @@ def _ln_then_project(stack: torch.Tensor) -> torch.Tensor: mixer_list: List[torch.Tensor] = [] for l in mixer_layers: mixers, mixer_labels = self._mixer_stack(l, pos_slice, subtract_bias=True) - mixer_list.extend(mixers) + mixer_list.extend(self._rescale(m, self._downstream_scale(l, layer)) for m in mixers) labels.extend(mixer_labels) if mixer_list: mixer_stack = torch.stack(mixer_list, dim=0) @@ -1953,6 +2082,21 @@ def _ln_then_project(stack: torch.Tensor) -> torch.Tensor: ) labels.extend(neuron_labels) components.append(neuron_stack) + # Residual writes besides the MLP; decompose_resid(mode="mlp") has them otherwise. + side_list: List[torch.Tensor] = [] + for l in range(layer): + side_outputs, side_labels = self._side_outputs(l, pos_slice) + side_list.extend( + self._rescale(o, self._downstream_scale(l, layer)) for o in side_outputs + ) + labels.extend(side_labels) + if side_list: + side_stack = torch.stack(side_list, dim=0) + if ln_folded: + side_stack = _ln_then_project(side_stack) + elif project_2d is not None: + side_stack = side_stack @ project_2d + components.append(side_stack) else: # Get the stack of just the MLP outputs # mlp_input included for completeness, but it doesn't actually matter, since it's @@ -1972,8 +2116,9 @@ def _ln_then_project(stack: torch.Tensor) -> torch.Tensor: labels.extend(mlp_labels) components.append(mlp_stack) + embed_scale = self._downstream_scale(0, layer) if self.has_embed: - embed = pos_slice.apply(self["hook_embed"], -2)[None] + embed = self._rescale(pos_slice.apply(self["hook_embed"], -2)[None], embed_scale) if ln_folded: embed = _ln_then_project(embed) elif project_2d is not None: @@ -1981,7 +2126,7 @@ def _ln_then_project(stack: torch.Tensor) -> torch.Tensor: labels.append("embed") components.append(embed) if self.has_pos_embed: - pos_embed = pos_slice.apply(self["pos_embed"], -2)[None] + pos_embed = self._rescale(pos_slice.apply(self["pos_embed"], -2)[None], embed_scale) if ln_folded: pos_embed = _ln_then_project(pos_embed) elif project_2d is not None: @@ -1989,7 +2134,9 @@ def _ln_then_project(stack: torch.Tensor) -> torch.Tensor: labels.append("pos_embed") components.append(pos_embed) # If we didn't expand the neurons, the MLP biases are already included in the MLP outputs. - bias_full = self.model.accumulated_bias(layer, mlp_input, include_mlp_biases=expand_neurons) + bias_full = self._scaled_accumulated_bias( + layer, mlp_input, include_mlp_biases=expand_neurons + ) bias_correction = self._post_norm_bias_correction( layer, mlp_input, include_mlp_biases=expand_neurons, pos_slice=pos_slice ) diff --git a/transformer_lens/model_bridge/supported_architectures/gemma4.py b/transformer_lens/model_bridge/supported_architectures/gemma4.py index 555e4fe36d..fba460106d 100644 --- a/transformer_lens/model_bridge/supported_architectures/gemma4.py +++ b/transformer_lens/model_bridge/supported_architectures/gemma4.py @@ -27,7 +27,9 @@ ``(1.0 + weight)`` offset. """ -from typing import Any +from typing import Any, Optional + +import torch from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter from transformer_lens.model_bridge.generalized_components import ( @@ -35,6 +37,7 @@ DelegatedAttentionBlockBridge, EmbeddingBridge, LinearBridge, + RMSNormalizationBridge, RotaryEmbeddingBridge, UnembeddingBridge, ) @@ -43,6 +46,34 @@ ) +class Gemma4BlockBridge(DelegatedAttentionBlockBridge): + """Gemma 4 decoder layer. + + After its residual adds the layer writes the per-layer input branch + (``post_per_layer_input_norm``, E2B/E4B) and then multiplies the whole residual + stream by its learned ``layer_scalar``. The residual decompositions read both. + """ + + residual_side_outputs = {"post_per_layer_input_norm": "per_layer_input"} + + @property + def residual_output_scale(self) -> Optional[torch.Tensor]: + """HF's ``layer_scalar`` buffer, applied to the residual stream at the end of the layer.""" + original = self._modules.get("_original_component") + return getattr(original, "layer_scalar", None) + + @property + def neuron_expansion_unsupported_reason(self) -> Optional[str]: + """MoE layers normalize the dense MLP output and add the experts' before ln2_post.""" + if "experts" not in self._modules: + return None + return ( + "the dense MLP output goes through post_feedforward_layernorm_1 and is summed with " + "the MoE experts' output before post_feedforward_layernorm, so its neurons cannot " + "be mapped onto the residual stream." + ) + + class Gemma4ArchitectureAdapter(ArchitectureAdapter): """Adapter for Gemma 4 (`Gemma4ForConditionalGeneration` — multimodal, or `Gemma4UnifiedForConditionalGeneration` — text-only 12B).""" @@ -100,15 +131,33 @@ def __init__(self, cfg: Any) -> None: # Single rotary module serving both layer types (full / sliding) via a # per-layer-type forward kwarg, with separate rope parameters per type. "rotary_emb": RotaryEmbeddingBridge(name="model.language_model.rotary_emb"), - "blocks": DelegatedAttentionBlockBridge( + "blocks": Gemma4BlockBridge( name="model.language_model.layers", submodules={ # Sandwich norms: ln1/ln1_post around attention, ln2/ln2_post - # around the MLP (same shape as Gemma 2/3). - "ln1": GeneralizedComponent(name="input_layernorm"), - "ln1_post": GeneralizedComponent(name="post_attention_layernorm"), - "ln2": GeneralizedComponent(name="pre_feedforward_layernorm"), - "ln2_post": GeneralizedComponent(name="post_feedforward_layernorm"), + # around the MLP (same shape as Gemma 2/3). Norm bridges on HF's own + # forward: their hook_scale lets the residual decompositions map heads + # and neurons through the post norms and apply the cached-scale LN. + "ln1": RMSNormalizationBridge( + name="input_layernorm", + config=self.cfg, + use_native_layernorm_autograd=True, + ), + "ln1_post": RMSNormalizationBridge( + name="post_attention_layernorm", + config=self.cfg, + use_native_layernorm_autograd=True, + ), + "ln2": RMSNormalizationBridge( + name="pre_feedforward_layernorm", + config=self.cfg, + use_native_layernorm_autograd=True, + ), + "ln2_post": RMSNormalizationBridge( + name="post_feedforward_layernorm", + config=self.cfg, + use_native_layernorm_autograd=True, + ), # PLE residual mix — present only when hidden_size_per_layer_input > 0 # (E2B/E4B; absent on 31B and 26B-A4B). "per_layer_input_gate": GeneralizedComponent( @@ -155,6 +204,10 @@ def __init__(self, cfg: Any) -> None: "mlp": self._gated_mlp(), }, ), - "ln_final": GeneralizedComponent(name="model.language_model.norm"), + "ln_final": RMSNormalizationBridge( + name="model.language_model.norm", + config=self.cfg, + use_native_layernorm_autograd=True, + ), "unembed": UnembeddingBridge(name="lm_head"), } From 2a30ba197779a943c8f12536bb43f1a6e3a201e0 Mon Sep 17 00:00:00 2001 From: Zhe Wang <53119724+flyingdoubleG@users.noreply.github.com> Date: Mon, 9 Nov 2026 01:22:07 +0800 Subject: [PATCH 2/2] test: Gemma 4 neuron results now map through ln2_post --- .../model_bridge/test_gemma4_text_adapter.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/tests/integration/model_bridge/test_gemma4_text_adapter.py b/tests/integration/model_bridge/test_gemma4_text_adapter.py index d81bc6488c..7937c85fc9 100644 --- a/tests/integration/model_bridge/test_gemma4_text_adapter.py +++ b/tests/integration/model_bridge/test_gemma4_text_adapter.py @@ -77,11 +77,13 @@ def grab(tensor, hook): for name, shape in expected.items(): assert captured.get(name) == shape, f"{name}: {captured.get(name)}" - def test_neuron_results_reject_unbridged_post_norm(self) -> None: - """ln2_post is a plain component with no cached scale; pre-norm values would be wrong.""" + def test_neuron_results_map_through_post_norm(self) -> None: + """ln2_post caches hook_scale, so the neurons sum to the post-norm MLP output.""" bridge, _ = _tiny_gemma4_text_pair() ids = torch.randint(3, VOCAB, (1, 8), generator=torch.Generator().manual_seed(0)) with torch.no_grad(): _, cache = bridge.run_with_cache(ids) - with pytest.raises(KeyError, match=r"blocks\.0\.ln2_post .*no hook_scale"): - cache.get_neuron_results(0) + neurons = cache.get_neuron_results(0) + torch.testing.assert_close( + neurons.sum(-2), cache["blocks.0.hook_mlp_out"], atol=1e-5, rtol=0 + )