HGRN2 Attention

HGRN2 Attention.

Implements the HGRN2 (Hierarchically Gated Recurrent Network v2) attention mechanism from Qin et al. (2024), arXiv:2404.07904 (COLM 2024). HGRN2 uses outer-product state expansion with hierarchically lower-bounded forget gates to combine recurrent memory efficiency with richer matrix-valued state updates. The state update follows:

S_t = diag(g_t) · S_{t-1} + v_t k_t^T

where g_t ∈ [lower_bound, 1]^d is a per-channel forget gate with a configurable lower bound that prevents complete memory erasure.

Reference:

Qin, Z., Yang, S., Zhong, Y., Shen, Y., & Sun, M. (2024). HGRN2: Gated Linear RNNs with State Expansion. arXiv:2404.07904. COLM 2024.

class src.model.attention.gated.hgrn2_attn.HGRN2Attention(*args: Any, **kwargs: Any)

Bases: Module

HGRN2 attention with lower-bounded forget gates.

Maintains a per-head matrix-valued recurrent state S_t ∈ R^{d×d}. At each timestep, a per-channel forget gate g_t (sigmoid-activated and lower-bounded) scales the previous state element-wise before adding the outer product v_t k_t^T. The lower bound prevents complete memory erasure, enabling hierarchical memory retention. 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.

hgrn2_lower_bound (float, optional): Minimum value for the

forget gate. Defaults to 0.0.

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

lower_bound

Minimum forget gate value.

Type:

float

q_proj

Query projection.

Type:

nn.Module

k_proj

Key projection.

Type:

nn.Module

v_proj

Value projection.

Type:

nn.Module

forget_proj

Per-channel forget gate 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:

Qin, Z., Yang, S., Zhong, Y., Shen, Y., & Sun, M. (2024). HGRN2: Gated Linear RNNs with State Expansion. arXiv:2404.07904. COLM 2024.

__init__(config)

Initialize HGRN2Attention.

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 HGRN2 attention over the input sequence.

Processes the sequence token-by-token with a lower-bounded forget-gated recurrent matrix state. At each step t, the per-channel forget gate g_t (clamped to [lower_bound, 1]) scales the previous state before adding v_t k_t^T.

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