transformer_lens.model_bridge.generalized_components.gated_delta_net module

GatedDeltaNet bridge for Qwen3.5/Qwen3Next linear-attention layers.

Reimplements forward (prefill only) to expose mech-interp-relevant intermediate states. Falls back to HF native forward during autoregressive generation where cache state management is required.

class transformer_lens.model_bridge.generalized_components.gated_delta_net.GatedDeltaNetBridge(name: str, config: Any | None = None, submodules: Dict[str, GeneralizedComponent] | None = None, **kwargs)

Bases: SSMStateHookMixin, GeneralizedComponent

Bridge for GatedDeltaNet linear-attention with full hook decomposition.

Hooks (prefill, in execution order):

hook_in: input hidden_states [batch, seq, d_model] hook_q_pre_conv: Q after projection, before conv [batch, seq, n_k_heads, head_k_dim] hook_k_pre_conv: K after projection, before conv [batch, seq, n_k_heads, head_k_dim] hook_v_pre_conv: V after projection, before conv [batch, seq, n_v_heads, head_v_dim] hook_q: Q after conv, pre-GQA-expansion [batch, seq, n_k_heads, head_k_dim]

Note: on standard attn layers, hook_q is post-projection. Here it’s post-conv — use hook_q_pre_conv for the projection-only output.

hook_k: K after conv [batch, seq, n_k_heads, head_k_dim] hook_v: V after conv [batch, seq, n_v_heads, head_v_dim] hook_beta_logit: pre-sigmoid write gate logit, per v-head [batch, seq, n_v_heads] hook_beta: write strength sigmoid(b), per v-head [batch, seq, n_v_heads] hook_log_decay: log-space decay g (NEGATIVE; multiplicative decay = exp(g)),

per v-head [batch, seq, n_v_heads]

hook_recurrence_out: output of linear recurrence [batch, seq, n_v_heads, head_v_dim] hook_gate_input: z tensor (pre-silu) for GatedRMSNorm [batch, seq, n_v_heads, head_v_dim] hook_ssm_state: recurrent state trajectory S_t [batch, seq, n_v_heads, head_k_dim,

head_v_dim] — fires ONLY on the opt-in eager-scan path (eager_scan=True), which swaps the fused kernel for a Python delta-rule scan so S_t can be read/patched. hook_ssm_write (alias -> hook_beta) is the write strength and propagates through the scan.

hook_out: final output to residual stream [batch, seq, d_model]

During generation (cache_params present), only hook_in/hook_out fire.

Property aliases:

W_in_proj_qkvz, W_in_proj_ba, W_out_proj, A_log, dt_bias

compute_effective_attention(cache: ActivationCache, layer_idx: int) → Tensor

Materialize a heuristic effective-attention matrix from cached hooks.

Uses the gated-linear-attention form of the recurrence (the exact gated delta rule additionally removes the key being written):

S_t ≈ exp(g_t) * S_{t-1} + beta_t * v_t @ k_t^T
o_t = S_t^T @ q_t
M[i,j] = (q_i^T @ k_j) * beta_j * prod_{t=j+1}^{i} exp(g_t)

so M is an interpretability heuristic, not a faithful output decomposition.

Requires the interior hooks (hook_q/k/beta/log_decay), which fire only on the hooked prefill path: call run_with_cache(tokens, use_cache=False) so cache_params is None. The default cached path exposes only hook_in/hook_out and this method then raises.

Measured divergence (tiny random-init test fixture, seed-stable):

  • L2-norm gap. The fused kernel L2-normalizes Q/K internally (use_qk_l2norm_in_kernel=True) but the hooked Q/K are pre-norm, so M differs from the normalized form by ≈1.0 relative when Q/K norms are small/non-uniform (the random-init regime); the gap shrinks toward 0 as norms equalize after training.

  • Delta-rule omission. Even with normalized Q/K, M @ V reconstructs the fused-kernel hook_recurrence_out only to O(1) relative error because the key-removal term is dropped.

Parameters:

  • cache – ActivationCache from run_with_cache(..., use_cache=False).

  • layer_idx – Block index for this linear_attn layer.

Returns:

[batch, n_v_heads, seq, seq] causal matrix (upper triangle zero).

Cost is O(batch * n_heads * seq^2); use on short sequences.

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

Reconstruct the recurrent state S of the gated delta rule from cache.

Read-only: replays the eager delta-rule scan (_gated_delta_scan) on the cached hook_q/k/v/beta/log_decay — no forward() re-run. Faithful (the full delta rule, key-removal included), unlike compute_effective_attention which is a gated-linear-attention heuristic that drops key removal.

Requires the interior hooks, which fire only on the hooked prefill path: call run_with_cache(tokens, use_cache=False) so cache_params is None.

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(..., use_cache=False).

  • layer_idx – Block index for this linear_attn layer.

  • time_step – If given, return only S at that position ([batch, n_v_heads, head_k_dim, head_v_dim]). None returns every step.

Returns:

[batch, seq, n_v_heads, head_k_dim, head_v_dim] for all steps, or [batch, n_v_heads, head_k_dim, head_v_dim] for a single time_step.

Memory is O(batch · n_v_heads · seq · head_k_dim · head_v_dim); pass time_step (or short sequences) when that is too large.

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

Generic forward pass for bridge components with input/output hooks.

hook_aliases: Dict[str, str | List[str]] = {'hook_linear_attn_in': 'hook_in', 'hook_linear_attn_out': 'hook_out', 'hook_ssm_B': 'hook_k', 'hook_ssm_C': 'hook_q', 'hook_ssm_decay': 'hook_log_decay', 'hook_ssm_out': 'hook_out', 'hook_ssm_write': 'hook_beta'}
property_aliases: Dict[str, str] = {'A_log': 'A_log', 'W_in_proj_ba': 'in_proj_ba.weight', 'W_in_proj_qkvz': 'in_proj_qkvz.weight', 'W_out_proj': 'out_proj.weight', 'dt_bias': 'dt_bias'}