Titan Attention

Titans memory-augmented attention.

Implements the core attention component of the Titans architecture, which augments standard multi-head attention with a neural memory module that learns to memorize at test time. This module provides the attention pathway that interacts with the surprise-driven long-term memory. Uses HoPE or RoPE positional encoding on query and key projections.

Reference:

Behrouz et al. (2025), “Titans: Learning to Memorize at Test Time”, arXiv:2501.00663.

class src.model.attention.titan.TitanAttention(*args: Any, **kwargs: Any)

Bases: Module

Multi-head attention with positional encoding for Titans architecture.

Projects the input into query, key, and value tensors, applies HoPE or RoPE positional encoding to queries and keys, computes scaled dot-product attention with softmax, and aggregates values. Supports both encoder (bidirectional) and decoder (causal) modes. Designed to work alongside Titans’ neural memory module for handling contexts beyond 2M tokens.

Reference:

Behrouz et al. (2025), “Titans: Learning to Memorize at Test Time”, arXiv:2501.00663.

Parameters:

config – Model configuration object with attributes hidden_size, num_heads, dropout, use_bitnet, positional_encoding ("hope" or "rope"), hope_base, hope_damping, rope_base, rope_scaling, and optionally mode ("encoder" or "decoder").

hidden_size

Dimensionality of the input and output embeddings.

num_heads

Number of parallel attention heads.

head_dim

Dimensionality of each attention head (hidden_size // num_heads).

scale

Scaling factor 1 / sqrt(head_dim) applied to dot products.

q_proj

Linear (or BitLinear) projection for queries.

k_proj

Linear (or BitLinear) projection for keys.

v_proj

Linear (or BitLinear) projection for values.

out_proj

Linear (or BitLinear) output projection.

pos_encoder

Positional encoding module (HoPE or RoPE).

dropout

Dropout layer applied to attention weights.

mode

"encoder" for bidirectional attention, "decoder" for causal (upper-triangular) masking.

Raises:

ValueError – If hidden_size is not divisible by num_heads, or if positional_encoding is not one of {"hope", "rope"}.

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

Compute Titans multi-head attention with positional encoding.

Parameters:

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

• logical_layer_idx – Logical layer index passed to the positional encoder for layer-dependent scaling. Defaults to 0 if None.

Returns:

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