Tucker Attention

Tucker Attention: generalised low-rank attention.

Implements Tucker Attention, arXiv:2603.30033 (Klein et al. 2026). Tucker Attention provides a unified low-rank view of Multi-Head Attention (MHA), Grouped-Query Attention (GQA) and Multi-Head Latent Attention (MLA): each is recovered as a special case of a Tucker-style factorisation of the query, key and value weight tensors along the num_heads and hidden_size axes. Concretely, the canonical W_q, W_k, W_v of shape hidden_size -> num_heads * head_dim are replaced by a small core tensor contracted against per-mode factor matrices, exposing the actual ranks achieved by MHA, GQA and MLA.

Special cases:

  • MHA : query_rank = key_rank = value_rank = hidden_size.

  • GQA : key_rank = value_rank < hidden_size (shared KV).

  • MLA : key_rank = value_rank = latent_rank (joint KV latent).

The paper reports an order of magnitude fewer parameters than GQA and MLA for comparable validation metrics on LLM and ViT test cases. Tucker Attention is fully compatible with FlashAttention and RoPE.

Reference:

Klein, T., Kusch, J., Sager, S., Schnake, S., & Schotthöfer, S. (2026). “Tucker Attention: A generalization of approximate attention mechanisms”. arXiv:2603.30033.

class src.model.attention.latent.tucker_attn.TuckerAttention(*args: Any, **kwargs: Any)[source]

Bases: Module

Tucker-factorised low-rank attention (Klein et al. 2026).

Parameters:

config – Configuration object. Relevant attributes: hidden_size, num_heads, dropout, use_bitnet, mode, and optional tucker_query_rank, tucker_key_rank, tucker_value_rank (default hidden_size // 2 for K and V, hidden_size for Q).

hidden_size

Input dimensionality.

num_heads

Number of attention heads.

head_dim

Per-head dimensionality.

query_rank, key_rank, value_rank

Tucker ranks along the hidden axis for Q, K, V respectively.

q_factor, k_factor, v_factor

Down-projections hidden_size -> rank.

q_core, k_core, v_core

Per-head core projections rank -> num_heads * head_dim.

out_proj

Output projection.

dropout, mode

As in the rest of the family.

Raises:

ValueError – If hidden_size not divisible by num_heads.

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

Compute Tucker-factorised attention.

Parameters:

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

  • logical_layer_idx – Unused; accepted for interface compatibility.

Returns:

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