fix: preserve final RMS predictions during weight processing

tests/unit/model_bridge/test_boot_native.py

@@ -573,6 +573,7 @@ def test_native_adapter_declares_fold_ln_support():
     "overrides",
     [
         pytest.param({"normalization_type": "LN"}, id="LN"),
+        pytest.param({"normalization_type": "LN", "final_rms": True}, id="LN-final-RMS"),
         pytest.param({"normalization_type": "RMS"}, id="RMS"),
         pytest.param({"normalization_type": "LNPre"}, id="LNPre"),
         pytest.param({"normalization_type": "RMSPre"}, id="RMSPre"),
@@ -646,6 +647,54 @@ def test_boot_native_weight_processing_preserves_predictions(overrides, recwarn)
     torch.testing.assert_close(after, before, atol=1e-4, rtol=1e-4)
 
 
+@pytest.mark.parametrize("nonidentity_scale", [False, True])
+def test_mixed_final_rms_fold_only_preserves_logits_and_unembed(nonidentity_scale):
+    bridge = TransformerBridge.boot_native(
+        _cfg(n_layers=2, normalization_type="LN", final_rms=True)
+    ).eval()
+    with torch.no_grad():
+        if nonidentity_scale:
+            bridge.ln_final.weight.copy_(torch.linspace(0.5, 1.5, bridge.cfg.d_model))
+        weight = bridge.W_U.detach().clone()
+        scale = bridge.ln_final.weight.detach().clone()
+        tokens = torch.tensor([[2, 3, 4]])
+        before = bridge(tokens)
+    bridge.enable_compatibility_mode(
+        fold_ln=True,
+        center_writing_weights=False,
+        center_unembed=False,
+        fold_value_biases=False,
+        refactor_factored_attn_matrices=False,
+    )
+    with torch.no_grad():
+        after = bridge(tokens)
+    torch.testing.assert_close(bridge.W_U, weight * scale[:, None])
+    torch.testing.assert_close(bridge.ln_final.weight, torch.ones_like(scale))
+    torch.testing.assert_close(after, before, atol=1e-5, rtol=1e-5)
+
+
+@pytest.mark.parametrize("fold_ln", [False, True])
+def test_mixed_final_rms_value_bias_folding_preserves_predictions(fold_ln):
+    bridge = TransformerBridge.boot_native(
+        _cfg(n_layers=2, normalization_type="LN", final_rms=True)
+    ).eval()
+    with torch.no_grad():
+        for block in bridge.blocks:
+            block.attn.v.bias.fill_(0.2)
+        tokens = torch.tensor([[2, 3, 4]])
+        before = bridge(tokens).log_softmax(-1)
+    bridge.enable_compatibility_mode(
+        fold_ln=fold_ln,
+        center_writing_weights=True,
+        center_unembed=False,
+        fold_value_biases=True,
+        refactor_factored_attn_matrices=False,
+    )
+    with torch.no_grad():
+        after = bridge(tokens).log_softmax(-1)
+    torch.testing.assert_close(after, before, atol=1e-5, rtol=1e-5)
+
+
 def test_boot_native_fold_ln_zeroes_norm_parameters():
     """fold_ln must actually absorb the norm into the downstream weights — an
     untouched ln1.weight means the fold silently skipped the layer."""

tests/unit/model_bridge/test_mixed_norm_decomposition.py

@@ -0,0 +1,110 @@
+"""Mixed block-LN/final-RMS predictions and cached decompositions must both reconcile."""
+
+import pytest
+import torch
+
+from transformer_lens.config import TransformerBridgeConfig
+from transformer_lens.model_bridge import TransformerBridge
+from transformer_lens.tools.analysis import direct_logit_attribution
+
+
+@pytest.fixture(params=["LN", "LNPre"])
+def mixed_bridge(request):
+    model = TransformerBridge.boot_native(
+        TransformerBridgeConfig(
+            d_model=16,
+            d_head=8,
+            n_heads=2,
+            n_layers=2,
+            n_ctx=8,
+            d_vocab=32,
+            d_mlp=24,
+            act_fn="gelu",
+            normalization_type=request.param,
+            final_rms=True,
+            seed=3,
+        )
+    ).eval()
+    return model
+
+
+@pytest.mark.parametrize("processing", ["none", "fold_only", "default"])
+@pytest.mark.parametrize("batchless", [False, True])
+def test_mixed_final_rms_decompositions_reconcile(mixed_bridge, processing, batchless):
+    model = mixed_bridge
+    tokens = torch.tensor([[2, 3, 4], [4, 5, 6]])
+    with torch.no_grad():
+        if processing != "none" and hasattr(model.ln_final, "weight"):
+            model.ln_final.weight.copy_(torch.linspace(0.5, 1.5, model.cfg.d_model))
+        before = model(tokens).log_softmax(-1)
+    if processing == "none":
+        model.enable_compatibility_mode(no_processing=True)
+    elif processing == "fold_only":
+        model.enable_compatibility_mode(
+            center_writing_weights=False, center_unembed=False, fold_value_biases=False
+        )
+    else:
+        model.enable_compatibility_mode()
+    if batchless:
+        tokens = tokens[:1]
+    with torch.no_grad():
+        logits, cache = model.run_with_cache(tokens, remove_batch_dim=batchless)
+    torch.testing.assert_close(
+        logits.log_softmax(-1), before[: tokens.shape[0]], atol=1e-5, rtol=1e-5
+    )
+    residual = cache["blocks.1.hook_resid_post"]
+    normalized = cache.apply_ln_to_stack(
+        residual.unsqueeze(0), layer=-1, has_batch_dim=not batchless
+    )[0]
+    torch.testing.assert_close(normalized, residual / cache["ln_final.hook_scale"])
+    direction = model.W_U[:, 7] - model.W_U[:, 8]
+    expected = logits[..., 7] - logits[..., 8] - (model.b_U[7] - model.b_U[8])
+    if batchless:
+        expected = expected[0]
+    for expand_neurons in (False, True):
+        dense = cache.get_full_resid_decomposition(expand_neurons=expand_neurons, apply_ln=True)
+        projected = cache.get_full_resid_decomposition(
+            expand_neurons=expand_neurons, apply_ln=True, project_output_onto=direction
+        )
+        torch.testing.assert_close(projected, dense @ direction, atol=1e-5, rtol=1e-5)
+        torch.testing.assert_close(projected.sum(0), expected, atol=1e-5, rtol=1e-5)
+    if not batchless:
+        result = direct_logit_attribution(model, answer_tokens=7, incorrect_tokens=8, cache=cache)
+        torch.testing.assert_close(result.attribution.sum(0), expected[:, -1], atol=1e-5, rtol=1e-5)
+
+
+@pytest.mark.parametrize("mlp_input", [False, True])
+def test_intermediate_layernorm_still_centers(mixed_bridge, mlp_input):
+    model = mixed_bridge
+    model.enable_compatibility_mode(no_processing=True)
+    with torch.no_grad():
+        _, cache = model.run_with_cache(torch.tensor([[2, 3, 4]]))
+    residual = cache["blocks.1.hook_resid_mid" if mlp_input else "blocks.1.hook_resid_pre"]
+    expected = (residual - residual.mean(-1, keepdim=True)) / cache[
+        f"blocks.1.ln{2 if mlp_input else 1}.hook_scale"
+    ]
+    actual = cache.apply_ln_to_stack(residual.unsqueeze(0), layer=1, mlp_input=mlp_input)[0]
+    torch.testing.assert_close(actual, expected)
+    direction = torch.linspace(0.1, 1.0, model.cfg.d_model)
+    dense = cache.stack_neuron_results(layer=1, apply_ln=True, mlp_input=mlp_input)
+    projected = cache.stack_neuron_results(
+        layer=1, apply_ln=True, mlp_input=mlp_input, project_output_onto=direction
+    )
+    torch.testing.assert_close(projected, dense @ direction, atol=1e-5, rtol=1e-5)
+
+
+@pytest.mark.parametrize("matrix_projection", [False, True])
+def test_final_rms_projected_neurons_match_uncentered_oracle(mixed_bridge, matrix_projection):
+    model = mixed_bridge
+    model.enable_compatibility_mode(no_processing=True)
+    with torch.no_grad():
+        _, cache = model.run_with_cache(torch.tensor([[2, 3, 4], [4, 5, 6]]))
+    direction = torch.linspace(0.1, 1.0, model.cfg.d_model)
+    if matrix_projection:
+        direction = torch.stack([direction, direction.square()], dim=-1)
+    raw = cache.stack_neuron_results(layer=model.cfg.n_layers, pos_slice=-1, apply_ln=False)
+    expected = (raw / cache["ln_final.hook_scale"][:, -1]) @ direction
+    actual = cache.stack_neuron_results(
+        layer=model.cfg.n_layers, pos_slice=-1, apply_ln=True, project_output_onto=direction
+    )
+    torch.testing.assert_close(actual, expected, atol=1e-5, rtol=1e-5)

transformer_lens/ActivationCache.py

@@ -1179,6 +1179,13 @@ def get_neuron_results(
             return neuron_acts * projected
         return neuron_acts[..., None] * projected
 
+    def _uses_ln_centering(self, layer: Optional[int]) -> bool:
+        """Use the target norm's semantics, not just the block normalization family."""
+        is_final = layer is None or layer in (-1, self.model.cfg.n_layers)
+        return self.model.cfg.normalization_type in ["LN", "LNPre"] and not (
+            is_final and getattr(self.model.cfg, "final_rms", False)
+        )
+
     def _get_cached_ln_scale(
         self,
         layer: Optional[int],
@@ -1238,7 +1245,7 @@ def _stack_neuron_results_apply_ln_projected(
         """
         scale = self._get_cached_ln_scale(layer, mlp_input=mlp_input, pos_slice=pos_slice)
 
-        apply_centering = self.model.cfg.normalization_type in ["LN", "LNPre"]
+        apply_centering = self._uses_ln_centering(layer)
         sum_p = project_2d.sum(dim=0) if apply_centering else None  # [n_outs]
 
         components: list = []
@@ -1508,8 +1515,8 @@ def apply_ln_to_stack(
                 results.append(out.reshape(original_shape))
             return torch.stack(results, dim=0)
 
-        # Center the stack onlny if the model uses LayerNorm
-        if self.model.cfg.normalization_type in ["LN", "LNPre"]:
+        # Final RMSNorm must not inherit the block LayerNorm centering convention.
+        if self._uses_ln_centering(layer):
             residual_stack = residual_stack - residual_stack.mean(dim=-1, keepdim=True)
 
         # Shape is [batch, position, 1] or [position, 1]; final dim is a dummy for broadcasting.

transformer_lens/weight_processing.py

@@ -1038,7 +1038,8 @@ def _fold_unembed_layer_norm(
         )
         if alternate_final_w_key != ln_final_w_key and alternate_final_w_key in state_dict:
             state_dict[alternate_final_w_key] = identity_val
-        if center_weights:
+        # RMSNorm outputs are not zero-mean, so centering would change predictions.
+        if center_weights and not getattr(cfg, "final_rms", False):
             unembed_weight_centered = ProcessWeights.convert_tensor_to_tl_format(
                 unembed_W_U_key, state_dict, state_dict.get(unembed_W_U_key), cfg, adapter, None
             )
@@ -1701,19 +1702,18 @@ def process_weights(
                 UserWarning,
             )
             center_writing_weights = False
-        if center_writing_weights:
-            if getattr(cfg, "normalization_type", "LN") in ["LN", "LNPre"] and (
-                not getattr(cfg, "final_rms", False)
-            ):
-                state_dict = ProcessWeights.center_writing_weights(state_dict, cfg, adapter=adapter)
+        should_center_writing_weights = (
+            center_writing_weights
+            and getattr(cfg, "normalization_type", "LN") in ["LN", "LNPre"]
+            and not getattr(cfg, "final_rms", False)
+        )
+        if should_center_writing_weights:
+            state_dict = ProcessWeights.center_writing_weights(state_dict, cfg, adapter=adapter)
         if center_unembed:
             state_dict = ProcessWeights.center_unembed(state_dict, cfg=cfg, adapter=adapter)
         if fold_value_biases:
             state_dict = ProcessWeights.fold_value_biases(state_dict, cfg, adapter=adapter)
-            if center_writing_weights and getattr(cfg, "normalization_type", "LN") in [
-                "LN",
-                "LNPre",
-            ]:
+            if should_center_writing_weights:
                 for layer_idx in range(cfg.n_layers):
                     b_O_key = ProcessWeights._get_param_key(f"blocks.{layer_idx}.attn.b_O", adapter)
                     if b_O_key in state_dict: