[YHC66]

Description

ActivationCache.logit_attrs (and so direct_logit_attribution) dots each component with W_U directions. Some adapters apply a constant factor after the unembedding, in apply_output_logits_transform: Granite divides by logits_scaling, MiniCPM by d_model / dim_model_base, Falcon-H1 and Muse Glimmer multiply by lm_head_multiplier / output_multiplier. That factor is not in W_U, so on these models the attributions sum to the logit difference times the inverse of the factor. Released Granite checkpoints set logits_scaling well above 1, and MiniCPM4-0.5B has 1024 / 256 = 4, so DLA comes out several times too large.

The DLA readout check does not catch this: it applies the transform before comparing with the model's logits, so it passes. Cohere is not affected, because its adapter folds logit_scale into W_U in compatibility mode.

import torch
from transformers import GraniteConfig, GraniteForCausalLM
from transformer_lens.model_bridge.sources import build_bridge_from_module
from transformer_lens.tools.analysis import direct_logit_attribution

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, logits_scaling=8.0)
model = GraniteForCausalLM(config).eval()
bridge = build_bridge_from_module(model, "GraniteForCausalLM", hf_config=config)
bridge.enable_compatibility_mode()

logits, cache = bridge.run_with_cache(torch.tensor([[3, 4, 5]]))
dla = direct_logit_attribution(bridge, answer_tokens=7, incorrect_tokens=9, cache=cache)
print(dla.attribution.sum().item(), (logits[0, -1, 7] - logits[0, -1, 9]).item())
# dev:     0.15515334904193878 0.019394170492887497  (8x)
# this PR: 0.019394168630242348 0.019394170492887497

Changes:

• ArchitectureAdapter.output_logits_scale() returns the constant factor the output transform applies before any soft cap (1.0 by default). Granite, MiniCPM, Falcon-H1, Muse Glimmer and Cohere return the same factor their apply_output_logits_transform uses (for Cohere, 1.0 once the scale is folded). Granite MoE and Granite MoE Hybrid inherit it from Granite. The existing transforms are unchanged, except that Cohere's now calls the new method.
• logit_attrs multiplies the logit directions by this factor. tokens_to_residual_directions still returns the plain W_U columns. HookedTransformer has no adapter, so it is not affected.
• Tests: a Granite DLA test with logits_scaling=8 (component and head units) next to the existing one, which uses logits_scaling=1.0; and a contract test that output_logits_scale() matches apply_output_logits_transform for Granite, Falcon-H1 and MiniCPM.

With a soft cap (Gemma 2, Muse Glimmer) the logits are not linear in the residual stream, so DLA still attributes the pre-cap logits. This PR does not change that.

Type of change

• Bug fix (non-breaking change which fixes an issue)

Checklist:

• I have commented my code, particularly in hard-to-understand areas
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes
• I have not rewritten tests relating to key interfaces which would affect backward compatibility

Tested on macOS (Apple Silicon), Python 3.12, CPU, torch 2.11, transformers 5.13. Without the fix the 5 new tests fail; with it they pass. tests/unit -m "not slow": 8360 passed, 0 failed. test_granite_embedding_decomposition.py, test_minicpm_tiny.py and test_ssm_hybrid_resid_decomposition.py: 213 passed. I also checked DLA sums against the logit difference on tiny GraniteMoe, MiniCPM (4 and 1/2 code), Cohere and Llama models in compatibility mode. black, isort, pycln and mypy pass on the changed files.

[jlarson4]

@YHC66 Thanks for putting this together! Just one comment on a potential improvement. Let me know if you have any questions

[jlarson4]

Every adapter now states its factor twice, here and in its apply_output_logits_transform. If a future adapter only overrides the transform this bug will resurface. Can the base transform apply output_logits_scale() before the soft cap, in order to eliminate the per-adapter transform overrides?

[YHC66]

Good call, done in dac9ee3. The base apply_output_logits_transform now multiplies by output_logits_scale() before the soft cap, and the Granite, Falcon-H1, MiniCPM, Muse Glimmer and Cohere transform overrides are gone; each adapter only overrides output_logits_scale(). The existing per-adapter transform tests (including Muse Glimmer's multiplier-before-softcap and Cohere's folded-scale cases) still pass unchanged.

[jlarson4]

Looks good, thanks for making those updates @YHC66