[flyingdoubleG]

Description

Falcon-H1 runs attention and a Mamba-2 mixer in parallel in every block, and HF scales several outputs inline before their residual adds:

# FalconH1DecoderLayer.forward
mamba_hidden_states = self.mamba(...) * self.ssm_out_multiplier
attention_hidden_states = self.self_attn(...)[0] * self.attn_out_multiplier
hidden_states = residual + mamba_hidden_states + attention_hidden_states
# FalconH1MLP.forward: down_proj(...) * down_multiplier (mlp_multipliers[1])
# FalconH1Model.forward: embed_tokens(input_ids) * embedding_multiplier

The bridge hooks see the outputs before these multipliers, and #1987 skips mixers in blocks that have attention. So on Falcon-H1, decompose_resid, get_full_resid_decomposition (with or without expand_neurons) and stack_neuron_results silently miss part of the residual stream:

• the Mamba output is not in any decomposition;
• hook_attn_out, the head results and the attention bias are not multiplied by attention_out_multiplier;
• hook_embed is not multiplied by embedding_multiplier;
• the neuron expansion (post @ W_out) and b_out are not multiplied by mlp_multipliers[1] (hook_mlp_out already is, since it is the module output).

The released checkpoints set all of these to values other than 1. On the seeded tiny random Falcon-H1 of the new test (3 layers, d_model 64, multipliers set to values other than 1, all biases on), comparing the sum with resid_post (max abs error relative to max abs resid_post):

	dev	this PR
decompose_resid()	0.83	2e-7
get_full_resid_decomposition(expand_neurons=False)	0.83	2e-7
get_full_resid_decomposition(expand_neurons=True)	0.86	3e-7

With all multipliers at 1 the Mamba output is still missing (3e-3 on random weights, where Mamba writes little; more on trained weights).

Fix.

• ScaledResidualBlockBridge takes an optional mlp_contribution_scale (default: residual_contribution_scale, so Granite is unchanged) and mixer_contribution_scales for a mixer that runs beside attention. hook_mlp_out uses the MLP scale.
• ActivationCache lists a mixer that sits beside attention in the same block (Falcon-H1's mamba) as a mixer row, {l}_mamba_out, in decompose_resid and get_full_resid_decomposition. It is scaled by its mixer_contribution_scales entry. Its bias is not subtracted, since accumulated_bias only counts the attention bias for such a block.
• MLP paths (neurons, b_out in accumulated_bias, the post-norm bias correction) use the MLP scale times an optional _neuron_contribution_scale on the MLP bridge. It is the multiplier between post @ W_out + b_out and hook_mlp_out, in the same way as _result_contribution_scale on attention.
• The Falcon-H1 adapter uses ScaledResidualBlockBridge with attention_out_multiplier, MLP scale 1 and {"mamba": ssm_out_multiplier}, sets the MLP's _neuron_contribution_scale to mlp_multipliers[1], and gives the embedding output_contribution_scale=embedding_multiplier, as Granite does. attention_out_multiplier, ssm_out_multiplier and mlp_multipliers are added to the config attributes copied from HF.

User-visible change: on Falcon-H1, hook_attn_out and hook_embed now carry the scaled contributions that are added to the residual stream (as on Granite), and the decompositions gain {l}_mamba_out rows. The attribution-patching check that refuses blocks with a parallel mixer is unchanged.

Type of change

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

Checklist:

• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• 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

tests/integration/model_bridge/test_ssm_hybrid_resid_decomposition.py gets a falcon_h1 factory: a seeded tiny FalconH1Config with every multiplier set to a value other than 1 and nonzero biases, built locally without hub access. It runs through the existing checks (decompositions sum to resid_pre / resid_mid / the final residual, the attn/mlp modes partition the stream, full decomposition with and without neurons, LN + projection), plus:

• labels include {l}_attn_out, {l}_mamba_out and {l}_mlp_out;
• hook_attn_out + mamba.hook_out * ssm_out_multiplier == resid_mid - resid_pre in every block;
• head results plus the scaled b_O sum to hook_attn_out (the existing test now applies the block's scale to b_O).

11 of these fail on dev and pass here.

Local runs (Windows, CPU, torch 2.9.1, --jaxtyping-packages=transformer_lens,beartype.beartype, PYTHONUTF8=1):

• tests/integration/model_bridge/test_ssm_hybrid_resid_decomposition.py: 82 passed, 20 skipped.
• tests/unit -m "not slow": 8173 passed, 14 failed. The same 14 fail on dev in my environment, with the same errors. Ten Gemma 3n adapter tests need the optional timm. Three test_sparse_probing.py tests check optimizer stop reasons. test_faithfulness_recovery_is_graded_not_binary is [Bug Report] test_faithfulness_recovery_is_graded_not_binary passes or fails on the sign of float32 rounding #1993.
• tests/integration/model_bridge/test_{direct_logit_attribution,granite_embedding_decomposition,granite_hook_semantics,granite_moe_hybrid_adapter,parallel_mlp_input_decomposition,parallel_residual_identities,post_norm_head_decomposition,residual_decomposition_identities,scaled_hook_autograd}.py and tests/unit/test_activation_cache.py: 650 passed, 8 failed. The 8 failures are in test_granite_moe_hybrid_adapter.py (SSD reconstruction finds no mixer.in_proj.hook_out in the cache) and fail the same way on dev here, which I take to be my local transformers version.

I also ran a sweep of tiny random models, comparing each decomposition's sum with the residual stream. For Jamba, Bamba, OLMo-Hybrid, Qwen3-Next, Mamba and FalconMamba nothing changes (all at 1e-7, except Qwen3-Next with expand_neurons=True, which raises on its MoE as before). For 38 dense architectures, including Granite and GraniteMoE, the results match dev. The sweep also showed that Zamba2, RecurrentGemma and RWKV raise KeyError in these APIs (LFM2 is handled in #2001); that is not changed here.

black (the locked 23.12.1) and isort --check-only pass. mypy reports no issues in the changed files.

This was prepared with AI assistance (Claude); I reviewed and tested the change.

[jlarson4]

This PR is approved! Once v4.4.0 goes live, this will be merged into dev for inclusion in the next release

[SBrouck]

Thanks for this! I wrote #1987, and you're right: it assumed a block writes through either attention or a sequence mixer, so Falcon-H1's parallel Mamba output was dropped. Listing the mixers that sit beside attention, each with its own scale, is the right fix.

Checked locally at 0f586cd (CPU, fp32) with an independent tiny Falcon-H1: 4 layers, batch 2, all biases on, and every multiplier different from 1 and from the ones in the test.

• Logits match HF exactly.
• decompose_resid() (and layer=l against each resid_pre) and get_full_resid_decomposition() with and without expand_neurons: 2e-7 to 4e-7, against 0.27 / 0.53 on dev.
• stack_head_results() plus b_O * attention_out_multiplier equals the sum of hook_attn_out (3e-7). stack_neuron_results() plus b_out * down_multiplier equals the sum of hook_mlp_out (5e-7).
• Doubling blocks.0.hook_attn_out moves hook_resid_mid by exactly hook_attn_out (2e-7), so the scaled hook writes in residual units.
• Emptying mixer_contribution_scales brings the error back to 0.37, so the test would catch a regression.

Since _mixer_out_names now lists other variant slots on blocks that have attention, I also checked the other adapters that declare attn plus another variant slot. On Bamba, Jamba and OLMo-Hybrid the layers alternate, so no block has both, and on tiny random models the decompositions and labels are identical to dev. GraniteMoeHybrid is in the test file. The existing decomposition, Granite, parallel-residual and ActivationCache tests pass unchanged.

One small, non-blocking thought: _neuron_contribution_scale is set from the adapter as a private attribute on the MLP bridge. A constructor argument on the gated MLP bridge, like output_contribution_scale on EmbeddingBridge, would make it easier to find for the next architecture with a multiplier after down_proj. Happy either way.