GLA Attention

Gated Linear Attention (GLA).

Implements the Gated Linear Attention mechanism from Yang et al. (2023), arXiv:2312.06635. GLA introduces data-dependent diagonal gating on linear attention, replacing purely additive accumulation with a controlled memory retention scheme. The matrix-valued recurrent state evolves as:

S_t = G_t ⊙ S_{t-1} + v_t k_t^T

where G_t is a per-head, per-channel gate derived from the input via a low-rank projection followed by logsigmoid activation. This enables sub-quadratic training complexity and O(d²) inference memory, making it suitable for long-sequence modeling.

Reference:

Yang, S., Wang, B., Shen, Y., Panda, R., & Kim, Y. (2023). Gated Linear Attention Transformers with Hardware-Efficient Training. arXiv:2312.06635.

class src.model.attention.gated.gla_attn.GatedLinearAttention(*args: Any, **kwargs: Any)

Bases: Module

Gated Linear Attention with data-dependent diagonal gating.

Maintains a per-head matrix-valued recurrent state S_t ∈ R^{d×d} that is updated at each timestep with a learned forget gate. The gate is produced by a low-rank projection (gk_proj) followed by logsigmoid, scaled by 1/16 for numerical stability. A silu-gated output projection (g_proj) provides additional 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.

gla_gate_low_rank (int, optional): Rank of the low-rank

bottleneck in the gate projection. Defaults to 16.

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

g_proj

Output gate projection (silu-gated).

Type:

nn.Module

gk_proj

Low-rank gate projection producing per-head, per-channel forget logits.

Type:

nn.Sequential

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., Wang, B., Shen, Y., Panda, R., & Kim, Y. (2023). Gated Linear Attention Transformers with Hardware-Efficient Training. arXiv:2312.06635.

__init__(config)

Initialize GatedLinearAttention.

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 Gated Linear Attention over the input sequence.

Processes the sequence token-by-token with a recurrent matrix state. At each step t, the state is gated element-wise by exp(gk_t) before adding the outer product v_t k_t^T. The query reads from the state via inner product, and the result is modulated by a silu-gated output projection.

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).