fix: carry the output logit scale into logit_attrs and DLA

tests/integration/model_bridge/test_granite_embedding_decomposition.py

@@ -85,6 +85,40 @@ def test_component_dla_reconstructs_logit_difference(granite):
     )
 
 
+@pytest.mark.parametrize("unit", ["component", "head"])
+def test_dla_reconstructs_logit_difference_under_logits_scaling(unit):
+    """Granite divides lm_head's output by logits_scaling; DLA must carry that factor."""
+    torch.manual_seed(0)
+    config = GraniteConfig(
+        vocab_size=32,
+        hidden_size=16,
+        intermediate_size=32,
+        num_hidden_layers=2,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        max_position_embeddings=32,
+        logits_scaling=8.0,
+        attn_implementation="eager",
+    )
+    model = GraniteForCausalLM(config).eval()
+    bridge = build_bridge_from_module(model, "GraniteForCausalLM", hf_config=config, device="cpu")
+    try:
+        bridge.enable_compatibility_mode()
+        with torch.no_grad():
+            logits, cache = bridge.run_with_cache(torch.tensor([[3, 4, 5]]))
+        result = direct_logit_attribution(
+            bridge, answer_tokens=7, incorrect_tokens=9, unit=unit, cache=cache
+        )
+        torch.testing.assert_close(
+            result.attribution.sum(0),
+            logits[:, -1, 7] - logits[:, -1, 9],
+            rtol=1e-5,
+            atol=1e-6,
+        )
+    finally:
+        bridge.close()
+
+
 def test_embedding_hook_exposes_residual_contribution_without_changing_forward(granite):
     bridge, reference, multiplier = granite
     tokens = torch.tensor([[3, 4, 5]])

tests/unit/model_bridge/test_output_logits_contract.py

@@ -12,6 +12,9 @@
 from transformer_lens.model_bridge.supported_architectures.granite import (
     GraniteArchitectureAdapter,
 )
+from transformer_lens.model_bridge.supported_architectures.minicpm import (
+    MiniCPMArchitectureAdapter,
+)
 from transformer_lens.model_bridge.supported_architectures.muse_glimmer import (
     MuseGlimmerArchitectureAdapter,
 )
@@ -135,3 +138,30 @@ def test_muse_glimmer_applies_multiplier_before_the_softcap() -> None:
         MuseGlimmerArchitectureAdapter(cfg).apply_output_logits_transform(logits),
         20.0 * torch.tanh(logits * 0.5 / 20.0),
     )
+
+
+@pytest.mark.parametrize(
+    ("adapter_cls", "architecture", "field", "value"),
+    [
+        (GraniteArchitectureAdapter, "GraniteForCausalLM", "logits_scaling", 4.0),
+        (FalconH1ArchitectureAdapter, "FalconH1ForCausalLM", "lm_head_multiplier", 0.25),
+        (MiniCPMArchitectureAdapter, "MiniCPMForCausalLM", "dim_model_base", 4),
+    ],
+)
+def test_output_logits_scale_matches_the_transform(
+    adapter_cls: type, architecture: str, field: str, value: float
+) -> None:
+    """Linear readouts (logit_attrs) use output_logits_scale in place of the transform."""
+    cfg = build_bridge_config_from_hf(
+        _text_config(**{field: value}),
+        architecture,
+        "tiny-output-scale-contract",
+        torch.float32,
+    )
+    adapter = adapter_cls(cfg)
+    logits = torch.tensor([[-8.0, 4.0]])
+
+    assert adapter.output_logits_scale() != 1.0
+    torch.testing.assert_close(
+        adapter.apply_output_logits_transform(logits), logits * adapter.output_logits_scale()
+    )

transformer_lens/ActivationCache.py

@@ -789,6 +789,14 @@ def logit_attrs(
                 incorrect_tokens
             )
 
+        # W_U omits any constant factor the architecture applies to the unembedding
+        # output (e.g. Granite's logits_scaling); include it so attributions sum to logits.
+        output_logits_scale = getattr(
+            getattr(self.model, "adapter", None), "output_logits_scale", None
+        )
+        if callable(output_logits_scale):
+            logit_directions = logit_directions * output_logits_scale()
+
         scaled_residual_stack = self.apply_ln_to_stack(
             residual_stack,
             layer=-1,

transformer_lens/model_bridge/architecture_adapter.py

@@ -294,12 +294,26 @@ def _set_rms_rotary_defaults(self, *, final_rms: bool = True, gated: bool = True
     def apply_output_logits_transform(self, logits: torch.Tensor) -> torch.Tensor:
         """Apply the architecture's declared post-unembedding transform.
 
-        Most causal decoders only apply the shared optional Gemma-style soft cap.
-        Adapters with additional scaling override this method instead of relying
-        on similarly named configuration fields used by unrelated architectures.
+        Scales the logits by ``output_logits_scale()``, then applies the shared
+        optional Gemma-style soft cap. Adapters with a constant output scale
+        override ``output_logits_scale`` rather than this method, so linear
+        readouts that use the scale stay in sync with the forward path.
         """
+        scale = self.output_logits_scale()
+        if scale != 1.0:
+            logits = logits * scale
         return apply_softcap(logits, getattr(self.cfg, "output_logits_soft_cap", None))
 
+    def output_logits_scale(self) -> float:
+        """Return the constant factor applied to the logits before any soft cap.
+
+        ``apply_output_logits_transform`` multiplies by it, and linear readouts
+        such as ``ActivationCache.logit_attrs`` multiply ``W_U`` directions by it
+        so their attributions add up to the model's logits. Adapters whose
+        architecture scales the logits override this instead of the transform.
+        """
+        return 1.0
+
     def validate_output_logits_transform(self) -> None:
         """Validate that the adapter can reproduce its post-unembedding path."""
 

transformer_lens/model_bridge/supported_architectures/cohere.py

@@ -202,15 +202,14 @@ def postprocess_weights(self, bridge: Any) -> None:
         self._logit_scale_already_folded = True
         self._logit_scale_fold_pending = False
 
-    def apply_output_logits_transform(self, logits: torch.Tensor) -> torch.Tensor:
+    def output_logits_scale(self) -> float:
         """Match Cohere's ``lm_head -> logit_scale -> optional softcap`` path.
 
         Once preprocess_weights has folded the scale into the live unembed, the
         head output already carries it, so applying cfg.logit_scale again would
         double-scale every reader of this transform.
         """
-        scale: float = 1.0 if self._logit_scale_already_folded else getattr(self.cfg, "logit_scale")
-        return super().apply_output_logits_transform(logits * scale)
+        return 1.0 if self._logit_scale_already_folded else getattr(self.cfg, "logit_scale")
 
 
 class _Cohere2AttentionBridge(PositionEmbeddingsAttentionBridge):

transformer_lens/model_bridge/supported_architectures/falcon_h1.py

@@ -41,8 +41,6 @@
 
 from typing import Any
 
-import torch
-
 from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter
 from transformer_lens.model_bridge.generalized_components import (
     BlockBridge,
@@ -159,7 +157,6 @@ def __init__(self, cfg: Any) -> None:
             "unembed": UnembeddingBridge(name="lm_head", config=self.cfg),
         }
 
-    def apply_output_logits_transform(self, logits: torch.Tensor) -> torch.Tensor:
+    def output_logits_scale(self) -> float:
         """Match Falcon-H1's post-unembedding multiplier."""
-        multiplier = float(getattr(self.cfg, "lm_head_multiplier", 1.0))
-        return super().apply_output_logits_transform(logits * multiplier)
+        return float(getattr(self.cfg, "lm_head_multiplier", 1.0))

transformer_lens/model_bridge/supported_architectures/granite.py

@@ -6,8 +6,6 @@
 
 from typing import Any
 
-import torch
-
 from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter
 from transformer_lens.model_bridge.generalized_components import (
     AttentionBridge,
@@ -112,7 +110,6 @@ def _build_component_mapping(self) -> dict:
             "unembed": UnembeddingBridge(name="lm_head", config=self.cfg),
         }
 
-    def apply_output_logits_transform(self, logits: torch.Tensor) -> torch.Tensor:
+    def output_logits_scale(self) -> float:
         """Match Granite's ``lm_head / logits_scaling`` output path."""
-        scaling = float(getattr(self.cfg, "logits_scaling", 1.0))
-        return super().apply_output_logits_transform(logits / scaling)
+        return 1.0 / float(getattr(self.cfg, "logits_scaling", 1.0))

transformer_lens/model_bridge/supported_architectures/minicpm.py

@@ -70,7 +70,7 @@ class MiniCPMArchitectureAdapter(ArchitectureAdapter):
     - each sublayer output is multiplied by ``scale_depth / sqrt(num_hidden_layers)``
       before the residual add (hook_attn_out / hook_mlp_out expose the product);
     - the final hidden state is divided by ``hidden_size / dim_model_base`` before
-      lm_head (reproduced by apply_output_logits_transform).
+      lm_head (reproduced by output_logits_scale).
 
     Nothing is folded into weights: HF applies the scales on the raw and the
     processed path alike, and weight processing is linear, so it commutes.
@@ -115,13 +115,12 @@ def __init__(self, cfg: Any) -> None:
             "unembed": UnembeddingBridge(name="lm_head", config=self.cfg),
         }
 
-    def apply_output_logits_transform(self, logits: torch.Tensor) -> torch.Tensor:
+    def output_logits_scale(self) -> float:
         """Match MiniCPM's ``lm_head(h / (hidden_size / dim_model_base))`` output path.
 
         lm_head is bias-free, so dividing its output equals dividing its input.
         """
-        divisor = self.cfg.d_model / float(getattr(self.cfg, "dim_model_base", 1.0))
-        return super().apply_output_logits_transform(logits / divisor)
+        return float(getattr(self.cfg, "dim_model_base", 1.0)) / self.cfg.d_model
 
     def prepare_loading(self, model_name: str, model_kwargs: dict) -> None:
         """Patch the remote code's transformers-v5 incompatibilities before loading."""

transformer_lens/model_bridge/supported_architectures/muse_glimmer.py

@@ -2,8 +2,6 @@
 
 from typing import Any
 
-import torch
-
 from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter
 from transformer_lens.model_bridge.generalized_components import (
     AttentionBridge,
@@ -77,6 +75,5 @@ def __init__(self, cfg: Any) -> None:
             "unembed": UnembeddingBridge(name="lm_head", config=self.cfg),
         }
 
-    def apply_output_logits_transform(self, logits: torch.Tensor) -> torch.Tensor:
-        multiplier = float(getattr(self.cfg, "output_multiplier", 1.0))
-        return super().apply_output_logits_transform(logits * multiplier)
+    def output_logits_scale(self) -> float:
+        return float(getattr(self.cfg, "output_multiplier", 1.0))