transformer_lens.model_bridge.generalized_components.ssm_mixer module

Wrap-don’t-reimplement bridge for HF’s MambaMixer (Mamba-1), plus S6 effective attention.

class transformer_lens.model_bridge.generalized_components.ssm_mixer.SSMMixerBridge(*args: Any, **kwargs: Any)

Bases: SSMStateHookMixin, GeneralizedComponent

Opaque wrapper around Mamba-1’s MambaMixer.

Submodules (in_proj, conv1d, x_proj, dt_proj, out_proj) are swapped into the HF mixer by replace_remote_component, so their hooks fire when slow_forward accesses them. A_log and D reach the user via GeneralizedComponent.__getattr__ delegation.

Decode-step caveat: conv1d.hook_out fires only on prefill during stateful generation; see DepthwiseConv1DBridge for the reason.

compute_effective_attention(cache: ActivationCache, layer_idx: int, include_dt_scaling: bool = False, per_state_coord: bool = False) → Tensor

Materialize Mamba-1’s per-channel effective attention from cached hooks.

Mamba-1’s S6 selective scan is equivalent to causal attention with a per-channel, per-state-coordinate learned decay — see “The Hidden Attention of Mamba” (Ali et al., ACL 2025). Unlike Mamba-2 there is no head grouping: each of the intermediate_size channels has its own A row, so the “head” axis here is the channel axis:

M[c, i, j] = sum_n C[i, n] · prod_{k=j+1..i} exp(A[c,n]·dt[c,k]) · B[j, n]

Reads B/C from x_proj.hook_out (or post-norm b_layernorm / c_layernorm hooks when present, as on Jamba) and dt from dt_proj.hook_out (softplus of that output); A via __getattr__. Read-only: no forward() re-run.

Parameters:

  • cache – ActivationCache from run_with_cache with this layer’s x_proj and dt_proj hooks.

  • layer_idx – Block index for this mixer.

  • include_dt_scaling – False (default) returns the attention-like form; True multiplies column j by dt[c, j], giving the reconstruction form satisfying y[c,i] = sum_j M[c,i,j]·x[c,j] + D[c]·x[c,i] (x is the post-conv SiLU input; y the pre-gate scan output).

  • per_state_coord – False (default) sums over the state coordinate and returns [batch, channels, seq, seq]. True returns the unsummed [batch, channels, state, seq, seq] tensor (the paper’s D·N matrices) — OFF by default.

Returns:

[batch, intermediate_size, seq, seq], or [batch, intermediate_size, state_size, seq, seq] when per_state_coord is True.

Peak memory is O(batch · intermediate_size · state_size · seq²) — the per-(channel, state) decay tensor — even for the summed default; use on short sequences.

compute_ssm_state(cache: ActivationCache, layer_idx: int, time_step: int | None = None) → Tensor

Reconstruct Mamba-1’s recurrent state S from cached hook values.

S6 recurrence h_t[c,s] = exp(A[c,s]·dt_t[c])·h_{t-1}[c,s] + dt_t[c]·x_t[c]·B_t[s] unrolls to:

S_t[c, s] = sum_{j<=t} decay[c, s, t, j] · dt_j[c] · x_j[c] · B_j[s]

with the same per-(channel, state) decay used by compute_effective_attention. x is the post-conv SiLU input (SiLU(conv1d.hook_out)). Read-only: no forward() re-run. Verify with y_t[c] = sum_s C_t[s]·S_t[c,s] + D[c]·x_t[c].

On padded batches the cached hooks are unmasked, so S is exact only at non-pad positions; pad-position state is out of contract.

Parameters:

  • cache – ActivationCache from run_with_cache with this layer’s x_proj, dt_proj, and conv1d hooks.

  • layer_idx – Block index for this mixer.

  • time_step – If given, return only S at that position ([batch, channels, state]); None returns every step.

Returns:

[batch, channels, seq, state] for all steps, or [batch, channels, state] for a single time_step.

Peak memory is O(batch · channels · state · seq²) — the per-(channel, state) decay tensor, built in full by _s6_terms regardless of time_step (which bounds only the returned tensor, not this peak); use on short sequences.

forward(*args: Any, **kwargs: Any) → Any

Hook the input, run HF slow_forward (or the eager scan), hook the output.

hook_aliases: Dict[str, str | List[str]] = {'hook_conv': 'conv1d.hook_out', 'hook_dt_proj': 'dt_proj.hook_out', 'hook_in_proj': 'in_proj.hook_out', 'hook_ssm_dt': 'dt_proj.hook_out', 'hook_ssm_out': 'hook_out', 'hook_x_proj': 'x_proj.hook_out'}
real_components: Dict[str, tuple]
training: bool