FoX Attention

Forgetting Transformer (FoX) Attention.

Implements the Forgetting Transformer (FoX) attention mechanism from Lin et al. (2025), arXiv:2503.02130 (ICLR 2025). FoX injects a learned forget gate into the softmax logit space via cumulative log-bias terms, controlling recency while preserving full softmax attention expressiveness. The attention computation is:

O = softmax(Q K^T + D) V

where D_ij = Σ_{l=j}^{i} log f_l is a cumulative log-forget bias matrix constructed from per-head forget gates f_t ∈ (0, 1). This formulation is FlashAttention-compatible, enabling efficient hardware-aware implementations.

Reference:

Lin, Z., Gou, M., Gong, Y., Liu, X., Shen, Y., Xu, R., Lin, C., Yang, Y., Jiao, J., Duan, N., & Chen, W. (2025). Forgetting Transformer: Softmax Attention with a Forget Gate. arXiv:2503.02130. ICLR 2025.

class src.model.attention.gated.fox_attn.ForgettingAttention(*args: Any, **kwargs: Any)

Bases: Module

Forgetting Transformer attention with log-space forget gates.

Computes standard scaled dot-product attention with an additive forget bias D in the logit space. The bias is constructed from per-head forget gates f_t via cumulative log-sums, creating a lower-triangular matrix that progressively discounts past tokens. In decoder mode, an additional causal mask is applied.

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 to attention

weights.

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

projections instead of nn.Linear.

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

Defaults to "encoder". In decoder mode, a causal mask is applied in addition to the forget bias.

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

scale

Scaling factor for dot-product attention (head_dim ** -0.5).

Type:

float

q_proj

Query projection.

Type:

nn.Module

k_proj

Key projection.

Type:

nn.Module

v_proj

Value projection.

Type:

nn.Module

f_proj

Per-head forget gate projection.

Type:

nn.Linear

out_proj

Output projection.

Type:

nn.Module

dropout

Dropout layer applied to attention weights.

Type:

nn.Dropout

mode

"encoder" or "decoder".

Type:

str

Raises:

ValueError – If hidden_size is not divisible by num_heads.

Reference:

Lin, Z., Gou, M., Gong, Y., Liu, X., Shen, Y., Xu, R., Lin, C., Yang, Y., Jiao, J., Duan, N., & Chen, W. (2025). Forgetting Transformer: Softmax Attention with a Forget Gate. arXiv:2503.02130. ICLR 2025.

__init__(config)

Initialize ForgettingAttention.

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

Constructs a cumulative log-forget bias matrix D from per-head forget gates f_t, adds it to the scaled dot-product attention logits, and applies softmax. In decoder mode, a causal mask is applied on top of the forget bias.

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