DeltaNet Attention

DeltaNet Attention.

Implements the DeltaNet attention mechanism from Yang et al. (2024), arXiv:2406.06484. DeltaNet applies a delta learning rule to the recurrent linear attention state, performing targeted error-correcting writes instead of purely additive memory updates. The state update follows:

S_t = S_{t-1} (I - β_t k_t k_t^T) + β_t v_t k_t^T

where β_t ∈ (0, 1) is a per-head write strength learned from the input, and k_t is L2-normalized. This formulation achieves perfect recall on the Multi-Query Associative Recall (MQAR) task by removing conflicting key-value associations before writing new ones.

Reference:

Yang, S., Kailash, B., Zhang, Y., & Kim, Y. (2024). Parallelizing Linear Transformers with the Delta Rule over Sequence Length. arXiv:2406.06484.

class src.model.attention.gated.deltanet_attn.DeltaNetAttention(*args: Any, **kwargs: Any)

Bases: Module

DeltaNet attention with delta-rule state updates.

Maintains a per-head matrix-valued recurrent state S_t ∈ R^{d×d}. At each timestep, a learned write strength β_t controls how much of the existing key-associated content is removed before the new key-value pair is written. Queries and keys are L2-normalized to ensure stable inner products. A silu-gated output projection provides channel-wise modulation after the recurrent readout.

Parameters:

config –

Model configuration object with the following relevant attributes: hidden_size (int): Dimensionality of input embeddings. num_heads (int): Number of attention heads. Must divide

hidden_size evenly.

dropout (float): Dropout probability applied after the

output gate.

use_bitnet (bool): If True, uses BitLinear for Q/K/V/G/O

projections instead of nn.Linear.

mode (str, optional): "encoder" or "decoder".

Defaults to "encoder".

hidden_size

Input embedding dimensionality.

Type:

int

num_heads

Number of attention heads.

Type:

int

head_dim

Dimensionality per head (hidden_size // num_heads).

Type:

int

total_dim

Total Q/K/V dimensionality (head_dim * num_heads).

Type:

int

q_proj

Query projection.

Type:

nn.Module

k_proj

Key projection.

Type:

nn.Module

v_proj

Value projection.

Type:

nn.Module

beta_proj

Per-head write strength projection.

Type:

nn.Linear

g_proj

Output gate projection (silu-gated).

Type:

nn.Module

out_proj

Output projection.

Type:

nn.Module

norm

Layer normalization applied to the recurrent readout before the output gate.

Type:

nn.LayerNorm

dropout

Dropout layer.

Type:

nn.Dropout

mode

"encoder" or "decoder".

Type:

str

Raises:

ValueError – If hidden_size is not divisible by num_heads.

Reference:

Yang, S., Kailash, B., Zhang, Y., & Kim, Y. (2024). Parallelizing Linear Transformers with the Delta Rule over Sequence Length. arXiv:2406.06484.

__init__(config)

Initialize DeltaNetAttention.

Parameters:

config – Model configuration object. See class docstring for required attributes.

Raises:

ValueError – If hidden_size is not divisible by num_heads.

forward(x: torch.Tensor, logical_layer_idx: int | None = None) → torch.Tensor

Compute DeltaNet attention over the input sequence.

Processes the sequence token-by-token with a delta-rule recurrent matrix state. At each step t, the state is updated by removing the projection onto k_t (weighted by β_t) and then adding the new association β_t · v_t k_t^T. Queries and keys are L2-normalized before use.

Parameters:

• x – Input tensor of shape (batch_size, seq_len, hidden_size).

• logical_layer_idx – Unused; accepted for interface compatibility with other attention mixers.

Returns:

Output tensor of shape (batch_size, seq_len, hidden_size).