From 70e74b603987bb0209141aae083449c8ee17202b Mon Sep 17 00:00:00 2001 From: Yihang Chen Date: Tue, 10 Nov 2026 10:23:01 -0800 Subject: [PATCH 1/2] fix: carry the output logit scale into logit_attrs and DLA Granite, MiniCPM, Falcon-H1 and Muse Glimmer scale lm_head's output by a constant that W_U does not contain, so logit_attrs (and therefore direct_logit_attribution) summed to the logit difference divided by that factor. Expose the factor as ArchitectureAdapter.output_logits_scale() and apply it to the logit directions in logit_attrs. --- .../test_granite_embedding_decomposition.py | 34 +++++++++++++++++++ .../test_output_logits_contract.py | 30 ++++++++++++++++ transformer_lens/ActivationCache.py | 8 +++++ .../model_bridge/architecture_adapter.py | 10 ++++++ .../supported_architectures/cohere.py | 6 ++-- .../supported_architectures/falcon_h1.py | 3 ++ .../supported_architectures/granite.py | 3 ++ .../supported_architectures/minicpm.py | 3 ++ .../supported_architectures/muse_glimmer.py | 3 ++ 9 files changed, 98 insertions(+), 2 deletions(-) diff --git a/tests/integration/model_bridge/test_granite_embedding_decomposition.py b/tests/integration/model_bridge/test_granite_embedding_decomposition.py index 0a516572b6..459f8319a4 100644 --- a/tests/integration/model_bridge/test_granite_embedding_decomposition.py +++ b/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]]) diff --git a/tests/unit/model_bridge/test_output_logits_contract.py b/tests/unit/model_bridge/test_output_logits_contract.py index 4362523188..f858c14283 100644 --- a/tests/unit/model_bridge/test_output_logits_contract.py +++ b/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() + ) diff --git a/transformer_lens/ActivationCache.py b/transformer_lens/ActivationCache.py index 08c743317d..b592d964c7 100644 --- a/transformer_lens/ActivationCache.py +++ b/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, diff --git a/transformer_lens/model_bridge/architecture_adapter.py b/transformer_lens/model_bridge/architecture_adapter.py index e7fd9033c0..d13a6fc86e 100644 --- a/transformer_lens/model_bridge/architecture_adapter.py +++ b/transformer_lens/model_bridge/architecture_adapter.py @@ -300,6 +300,16 @@ def apply_output_logits_transform(self, logits: torch.Tensor) -> torch.Tensor: """ return apply_softcap(logits, getattr(self.cfg, "output_logits_soft_cap", None)) + def output_logits_scale(self) -> float: + """Return the constant factor the output transform applies before any soft cap. + + 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 ``apply_output_logits_transform`` scales the logits + override this to return the same factor. + """ + return 1.0 + def validate_output_logits_transform(self) -> None: """Validate that the adapter can reproduce its post-unembedding path.""" diff --git a/transformer_lens/model_bridge/supported_architectures/cohere.py b/transformer_lens/model_bridge/supported_architectures/cohere.py index 9080bb16e7..2121e0ae53 100644 --- a/transformer_lens/model_bridge/supported_architectures/cohere.py +++ b/transformer_lens/model_bridge/supported_architectures/cohere.py @@ -209,8 +209,10 @@ def apply_output_logits_transform(self, logits: torch.Tensor) -> torch.Tensor: 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 super().apply_output_logits_transform(logits * self.output_logits_scale()) + + def output_logits_scale(self) -> float: + return 1.0 if self._logit_scale_already_folded else getattr(self.cfg, "logit_scale") class _Cohere2AttentionBridge(PositionEmbeddingsAttentionBridge): diff --git a/transformer_lens/model_bridge/supported_architectures/falcon_h1.py b/transformer_lens/model_bridge/supported_architectures/falcon_h1.py index b633cd42f9..d71e55d4a4 100644 --- a/transformer_lens/model_bridge/supported_architectures/falcon_h1.py +++ b/transformer_lens/model_bridge/supported_architectures/falcon_h1.py @@ -163,3 +163,6 @@ def apply_output_logits_transform(self, logits: torch.Tensor) -> torch.Tensor: """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) + + def output_logits_scale(self) -> float: + return float(getattr(self.cfg, "lm_head_multiplier", 1.0)) diff --git a/transformer_lens/model_bridge/supported_architectures/granite.py b/transformer_lens/model_bridge/supported_architectures/granite.py index 1aada0b289..b526bd932d 100644 --- a/transformer_lens/model_bridge/supported_architectures/granite.py +++ b/transformer_lens/model_bridge/supported_architectures/granite.py @@ -116,3 +116,6 @@ def apply_output_logits_transform(self, logits: torch.Tensor) -> torch.Tensor: """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) + + def output_logits_scale(self) -> float: + return 1.0 / float(getattr(self.cfg, "logits_scaling", 1.0)) diff --git a/transformer_lens/model_bridge/supported_architectures/minicpm.py b/transformer_lens/model_bridge/supported_architectures/minicpm.py index cbf72ec039..04d69887de 100644 --- a/transformer_lens/model_bridge/supported_architectures/minicpm.py +++ b/transformer_lens/model_bridge/supported_architectures/minicpm.py @@ -123,6 +123,9 @@ def apply_output_logits_transform(self, logits: torch.Tensor) -> torch.Tensor: divisor = self.cfg.d_model / float(getattr(self.cfg, "dim_model_base", 1.0)) return super().apply_output_logits_transform(logits / divisor) + def output_logits_scale(self) -> float: + 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.""" super().prepare_loading(model_name, model_kwargs) diff --git a/transformer_lens/model_bridge/supported_architectures/muse_glimmer.py b/transformer_lens/model_bridge/supported_architectures/muse_glimmer.py index 313ec6896a..a22a71d05e 100644 --- a/transformer_lens/model_bridge/supported_architectures/muse_glimmer.py +++ b/transformer_lens/model_bridge/supported_architectures/muse_glimmer.py @@ -80,3 +80,6 @@ def __init__(self, cfg: Any) -> None: 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)) From dac9ee327b0e18d3ddb860cf4c1ab06a04b1e403 Mon Sep 17 00:00:00 2001 From: Yihang Chen Date: Wed, 11 Nov 2026 08:41:30 -0800 Subject: [PATCH 2/2] refactor: apply output_logits_scale in the base output transform Adapters now only declare their constant factor via output_logits_scale(); the base apply_output_logits_transform multiplies by it before the soft cap, so the forward path and linear readouts like logit_attrs cannot drift apart. --- .../model_bridge/architecture_adapter.py | 20 +++++++++++-------- .../supported_architectures/cohere.py | 5 +---- .../supported_architectures/falcon_h1.py | 8 +------- .../supported_architectures/granite.py | 8 +------- .../supported_architectures/minicpm.py | 8 ++------ .../supported_architectures/muse_glimmer.py | 6 ------ 6 files changed, 17 insertions(+), 38 deletions(-) diff --git a/transformer_lens/model_bridge/architecture_adapter.py b/transformer_lens/model_bridge/architecture_adapter.py index d13a6fc86e..fc7e21d79d 100644 --- a/transformer_lens/model_bridge/architecture_adapter.py +++ b/transformer_lens/model_bridge/architecture_adapter.py @@ -294,19 +294,23 @@ 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 the output transform applies before any soft cap. + """Return the constant factor applied to the logits before any soft cap. - 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 ``apply_output_logits_transform`` scales the logits - override this to return the same factor. + ``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 diff --git a/transformer_lens/model_bridge/supported_architectures/cohere.py b/transformer_lens/model_bridge/supported_architectures/cohere.py index 2121e0ae53..7d00b65a5d 100644 --- a/transformer_lens/model_bridge/supported_architectures/cohere.py +++ b/transformer_lens/model_bridge/supported_architectures/cohere.py @@ -202,16 +202,13 @@ 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. """ - return super().apply_output_logits_transform(logits * self.output_logits_scale()) - - def output_logits_scale(self) -> float: return 1.0 if self._logit_scale_already_folded else getattr(self.cfg, "logit_scale") diff --git a/transformer_lens/model_bridge/supported_architectures/falcon_h1.py b/transformer_lens/model_bridge/supported_architectures/falcon_h1.py index d71e55d4a4..eea3759d29 100644 --- a/transformer_lens/model_bridge/supported_architectures/falcon_h1.py +++ b/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,10 +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: - """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) - def output_logits_scale(self) -> float: + """Match Falcon-H1's post-unembedding multiplier.""" return float(getattr(self.cfg, "lm_head_multiplier", 1.0)) diff --git a/transformer_lens/model_bridge/supported_architectures/granite.py b/transformer_lens/model_bridge/supported_architectures/granite.py index b526bd932d..c8b9558cb0 100644 --- a/transformer_lens/model_bridge/supported_architectures/granite.py +++ b/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,10 +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: - """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) - def output_logits_scale(self) -> float: + """Match Granite's ``lm_head / logits_scaling`` output path.""" return 1.0 / float(getattr(self.cfg, "logits_scaling", 1.0)) diff --git a/transformer_lens/model_bridge/supported_architectures/minicpm.py b/transformer_lens/model_bridge/supported_architectures/minicpm.py index 04d69887de..360700fac7 100644 --- a/transformer_lens/model_bridge/supported_architectures/minicpm.py +++ b/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,15 +115,11 @@ 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) - - def output_logits_scale(self) -> float: 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: diff --git a/transformer_lens/model_bridge/supported_architectures/muse_glimmer.py b/transformer_lens/model_bridge/supported_architectures/muse_glimmer.py index a22a71d05e..18ac9e1ab2 100644 --- a/transformer_lens/model_bridge/supported_architectures/muse_glimmer.py +++ b/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,9 +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))