diff --git a/tests/integration/model_bridge/test_granite_embedding_decomposition.py b/tests/integration/model_bridge/test_granite_embedding_decomposition.py index 0a516572b..459f8319a 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 436252318..f858c1428 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 08c743317..b592d964c 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 e7fd9033c..fc7e21d79 100644 --- a/transformer_lens/model_bridge/architecture_adapter.py +++ b/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.""" diff --git a/transformer_lens/model_bridge/supported_architectures/cohere.py b/transformer_lens/model_bridge/supported_architectures/cohere.py index 9080bb16e..7d00b65a5 100644 --- a/transformer_lens/model_bridge/supported_architectures/cohere.py +++ b/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): diff --git a/transformer_lens/model_bridge/supported_architectures/falcon_h1.py b/transformer_lens/model_bridge/supported_architectures/falcon_h1.py index b633cd42f..eea3759d2 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,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)) diff --git a/transformer_lens/model_bridge/supported_architectures/granite.py b/transformer_lens/model_bridge/supported_architectures/granite.py index 1aada0b28..c8b9558cb 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,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)) diff --git a/transformer_lens/model_bridge/supported_architectures/minicpm.py b/transformer_lens/model_bridge/supported_architectures/minicpm.py index cbf72ec03..360700fac 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,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.""" diff --git a/transformer_lens/model_bridge/supported_architectures/muse_glimmer.py b/transformer_lens/model_bridge/supported_architectures/muse_glimmer.py index 313ec6896..18ac9e1ab 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,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))