PaperlayerView exact revision on arXiv

Source-linked AI summary

Tucker Attention: A generalization of approximate attention mechanisms

Timon Klein, Jonas Kusch, Sebastian Sager, Stefan Schnake, Steffen Schotthöfer

arXiv:2603.30033v1cs.LGcs.AI

TL;DR

MHA’s attention-weight matrices impose substantial parameter and memory costs. Tucker Attention models pre- and post-softmax weights as tensor objects and factorizes them directly; it is often by an order of magnitude more parameter efficient than GQA, MLA, or related methods in the reported evaluations.

  • Problem

    MHA’s key, query, value, and output matrices create substantial parameter and memory costs.

    Source: §1 · p. 1

  • Method

    Tucker Attention treats pre- and post-softmax weights as tensor objects and applies a direct low-rank Tucker decomposition.

    Source: §1.1 · p. 2, §2.1 · p. 4

  • Results

    Often by an order of magnitude, Tucker Attention is more parameter efficient than GQA, MLA, or related methods in vision-transformer transfer learning and GPT-2 pretraining.

    Source: §1.1 · p. 2

  • Takeaways & Limitations

    The framework identifies uncompressed head-mode redundancy in existing methods and allows independent rank selection across attention modes.

    Source: §1.1 · p. 2

    The evaluation uses an existing FlashAttention-compatible strategy; exploring custom kernels that exploit the factorized core is left to future work.

    Source: §3.3 · p. 7

Abstract

from arXiv · show

The pursuit of reducing the memory footprint of the self-attention mechanism in multi-headed self attention (MHA) spawned a rich portfolio of methods, e.g., group-query attention (GQA) and multi-head latent attention (MLA). The methods leverage specialized low-rank factorizations across embedding dimensions or attention heads. From the point of view of classical low-rank approximation, these methods are unconventional and raise questions of which objects they really approximate and how to interpret the low-rank behavior of the resulting representations. To answer these questions, this work proposes a generalized view on the weight objects in the self-attention layer and a factorization strategy, which allows us to construct a parameter efficient scheme, called Tucker Attention. Tucker Attention requires an order of magnitude fewer parameters for comparable validation metrics, compared to GQA and MLA, as evaluated in LLM and ViT test cases. Additionally, Tucker Attention~encompasses GQA, MLA, MHA as special cases and is fully compatible with flash-attention and rotary position embeddings (RoPE). This generalization strategy yields insights of the actual ranks achieved by MHA, GQA, and MLA, and further enables simplifications for MLA.

1. Introduction

MHA’s attention weights create substantial memory costs, and designing compressed representations that address training, inference, and attention-computation bottlenecks remains an open challenge. Tucker Attention treats pre- and post-softmax weights as tensor objects and factorizes them directly, unifying several existing methods while offering parameter savings and compatibility with common attention implementations.

1. Introduction

  • MHA’s key, query, value, and output matrices create substantial parameter and memory costs, while addressing training, inference, and attention-computation bottlenecks together remains an open challenge.Existing approaches include GQA’s limited key/value heads and MLA’s low-rank query, key, and value weights.Source: §1 · pp. 1–2

1.1. Contribution

  • The paper represents pre-softmax query/key weights and post-softmax value/output weights as separate order-three tensors, rather than analyzing the matrices individually.The former yields a bilinear form over the token sequence; the latter combines post-softmax attention matrices into the output.Source: §1.1 · p. 2
  • Assuming low-rank structure in these tensors, Tucker Attention applies direct low-rank Tucker decompositions; MHA is most efficient if both tensors are full rank.The paper frames this as an alternative to approximating MHA indirectly.Source: §1.1 · p. 2
  • MHA, MQA, GQA, and MLA are special cases of Tucker Attention, and none compresses the attention tensor’s head mode.The methods correspond to particular Tucker ranks.Source: §1.1 · p. 2
  • Up to 9× parameter reduction at comparable accuracy follows from selecting ranks independently across head, query, key, value, and output modes.The framework is also compatible with KV caching, a RoPE variation, and FlashAttention.Source: §1.1 · p. 2
  • Tucker Attention is often by an order of magnitude more parameter efficient than GQA, MLA, or related methods in vision-transformer transfer learning and GPT-2 pretraining.The paper reports comparable validation metrics.Source: §1.1 · p. 2

2. Tensor attention reformulation

The section recasts multi-head attention as tensor-valued operations, then applies Tucker decomposition to compress attention tensors across modes. This framework relates MQA, GQA, and MLA to Tucker Attention and reveals where their factorizations do and do not capture low-rank structure.

2. Tensor attention reformulation

  • MHA can be rewritten using stacked pre- and post-softmax attention tensors, with the head dimension as one mode and query/key or output/value computations as the others.The output is recovered by contracting the attention and projected-value tensors over head and sequence dimensions.Source: §2 · pp. 3–4

2.2. MQA, GQA and MLA are special cases of Tucker Attention

  • MQA, GQA, and MLA are special cases of Tucker Attention at specific ranks, inferred from low-rank factorizations of their query, key, and value weights.The discussion derives pre-softmax ranks; the post-softmax ranks and fuller theory are deferred to Appendix B.1.Source: §2.1 · p. 4, §2.2 · p. 5

2.1. Tucker Attention

  • Tucker Attention directly decomposes the pre- and post-softmax tensors, enabling compression across head, query, key, value, and output modes.Its basis matrices capture shared subspaces across heads, while separate head-mode bases enable compression across heads; its parameter count scales as O(rdmodel) under the stated equal-rank setting.Source: §2.1 · pp. 4–5
  • GPT-2 spectra show similar decay across MHA modes, lost head-mode low rank in GQA after query and output projections, and head-mode redundancy in MLA that its parametrization cannot exploit.The head-mode spectrum of MLA resembles Tucker Attention with r1 = 8.Source: §2.1 · pp. 4–5, Figure 2 · p. 5

2.2.1. MULTI AND GROUP QUERY ATTENTION

  • GQA shares key and value heads across query groups, bounding the key and value mode ranks by nKVdH; MQA is GQA with nKV = 1.The key matrices are broadcast across queries in each group.Source: Figure 1 · p. 3, §2.2.1 · p. 5

2.2.2. MULTIHEAD LATENT ATTENTION

  • MLA factorizes query and key heads through latent projections, and its up-projection tensors can be absorbed into the Tucker core for a more parameter-efficient representation.The key latent dimension is dc; the query heads are factorized similarly to the key heads.Source: §2.2.2 · p. 5

3. Compatibility of Tucker Attention

Tucker Attention is compatible with KV caching, a modified RoPE placement, and existing flash-attention implementations. Its latent-space rotation preserves relative-position behavior and can also simplify MLA’s RoPE pathway.

3.1. KV Caching

  • Tucker Attention is compatible with KV caching by computing and caching latent key and value representations, with storage cost 2Nr3.With a shared key-value down-projection, the cache can be reduced to Nr3; when r3 equals MLA’s latent dimension, the cache sizes are equal.Source: §3.1 · p. 6

3.2. Rotary Position Embedding (RoPE)

  • Moving RoPE from the head dimension to the latent key dimension yields latent RoPE, which preserves the relative-position property while changing rotation frequency.The rotation is computed once in the latent key space and shared across heads, reducing overhead compared with per-head RoPE in MHA and decoupled RoPE in MLA.Source: §3.2 · p. 6

3.2.1. COROLLARY: SIMPLIFIED ROPE FOR MLA

  • Applying latent RoPE directly to MLA offers an alternative to decoupled RoPE and removes the need for separate semantic and rotational pathways.The passage reports that this reduces MLA’s implementation complexity while preserving KV-cache efficiency.Source: §3.2.1 · p. 6

3.3. Flash Attention

  • Tucker Attention works with existing flash-attention implementations without custom kernels.Its shared key-value pair reduces chunk loads across query heads, while encoding chunks in the r3 latent dimension reduces SRAM pressure relative to MHA or GQA chunks encoded in dH.Source: §3.3 · p. 7
  • The flash-attention evaluation uses one possible strategy; custom kernels that exploit the factorized core C are left for future work.The passages identify such kernels as a possible route to more efficient use of the Tucker structure.Source: §3.3 · p. 7

4. Numerical Results

This section compares MHA, GQA, MLA, and Tucker Attention on vision transformers and GPT-2, with additional GPT-2 experiments on Shakespeare and LLaMA3-1B. Across these evaluations, Tucker Attention delivers competitive validation performance with fewer parameters and, in some settings, a smaller KV cache.

4.1. ViT16b on Cifar10/Cifar100 & ViT32l, ViT16g on ImageNet1k

  • On ViT tasks, Tucker Attention needs almost an order of magnitude fewer trainable parameters than GQA and MLA for similar performance.Results are means of five runs with standard deviation below 0.1% per setup; the low-rank methods were initialized by SVD-based approximation of MHA weights.Source: §4.1 · p. 7

4.2. GPT2 on OpenWebText

  • In GPT-2 zero-shot tests without RoPE, Tucker Attention with ranks (8, 128, 64) uses about 18% of MHA’s parameters and 39% of MLA’s while achieving similar mean perplexity and accuracy.Table 2 reports validation perplexity, loss, accuracy, attention-layer parameters, and inference KV cache.Source: Table 2 · p. 8, §4.2 · p. 8

4.4. LLaMA3-1B on OpenWebText2

  • Tucker Attention uses 10–20% of the LLaMA3 GQA baseline’s parameters, with a 1–3% increase in validation cross-entropy loss, and reduces KV cache by 4–8×.It also decreases iteration wall time by up to 20% versus MHA and 15% versus GQA (nKV = 8), and outperforms MLA in validation loss and parameter count.Source: Table 3 · p. 8, §4.4 · p. 9

4.2. GPT2 on OpenWebText

  • With shared KV cache, Tucker Attention converges at the same rate as canonical methods and reaches lower loss than MLA; MLA with shared KV shows small instabilities and can diverge.The comparison uses Tucker Attention ranks (8, 128, 64) and equal KV-cache size for MLA and Tucker Attention.Source: §4.2 · p. 8
  • Tucker Attention is compatible with latent RoPE, and GPT-2 perplexity and test accuracy are slightly improved with RoPE, mirroring modest gains for MLA, GQA, and MHA.The latent RoPE implementation uses Equation (12).Source: §4.2 · pp. 8–9

4.3. GPT2 on Shakespeare

  • On Shakespeare, Tucker Attention achieves comparable validation loss with over an order of magnitude fewer attention parameters and KV cache than the baselines, including MLA.Shared KV down-projection is also compatible with Tucker Attention, with a slight increase in validation accuracy; MLA is compatible with latent RoPE.Source: §4.3 · p. 9

4. Numerical Results

  • Treating pre- and post-softmax weights as tensor-valued objects yields independent rank selection across head, query, key, value, and output modes, with up to 9× parameter reduction at comparable accuracy.The conclusion also reports compatibility with KV caching, latent RoPE, and FlashAttention.Source: §4.4 · p. 9

B.1. Tucker Representations of MQA, GQA, and MLA

This appendix formalizes Tucker representations of attention tensors and establishes maximal Tucker-rank results for MHA, GQA, and MLA. It also states that the expression in Equation (12) fulfills the condition in Equation (10).

B.1. Tucker Representations of MQA, GQA, and MLA

  • The pre-softmax attention tensor admits a Tucker decomposition with maximal rank (nH, rank(W Q), rank(W K)).The same general result is stated for a head-indexed tensor formed from the query and key matrices.Source: §B.1 · p. 13
  • GQA with nKV KV heads has maximal Tucker rank (nH, dmodel, nKVdH) for both pre-softmax and post-softmax attention tensors.MQA is identified as the special case of GQA with nKV = 1.Source: §B.1 · p. 13
  • MLA’s pre-softmax tensor has maximal Tucker rank (nH, dQ_c, dK_c), while its post-softmax tensor has maximal rank (nH, dmodel, dK_c).The theorem describes the pre-softmax ranks using latent query and key dimensions, and the post-softmax ranks using the latent value dimension.Source: §B.1 · p. 14

B.2. Rotary Position Embedding

  • Lemma B.5 states that the expression in Equation (12) fulfills the condition in Equation (10).The proof says this follows directly for token positions m and n.Source: §B.2 · p. 14

C.1.1. CIFAR-10/100 DATA

The CIFAR-10/100 data preparation uses training augmentation and normalization, while ViT inputs are resized to 224 × 224 pixels.

C.1.1. CIFAR-10/100 DATA

  • CIFAR training images receive random horizontal flipping and dataset-specific normalization; test images are only normalized, and ViT inputs are resized to 224 × 224.CIFAR-10 comprises 60,000 RGB images of size 32 × 32 across 10 classes; CIFAR-100 uses the stated corresponding augmentation strategy.Source: §C.1.1 · p. 14

C.1.2. IMAGENET-1K DATA

This setup describes ImageNet-1K preprocessing and ViT model initialization, then specifies how pretrained MHA parameters are converted to GQA, MLA, and Tucker variants.

C.1.2. IMAGENET-1K DATA

  • ImageNet-1K has 1,000 classes and over 1.2 million RGB training images; training uses random resized crops to 224 × 224 and normalization, while testing uses resize and center-crop.The test set is normalized after resizing and center-cropping.Source: §C.1.2 · p. 14

C.1.3. VIT MODELS

  • The listed ViT models are ViT-B.16, ViT-L.32, and ViT-G.14, with ImageNet21k weights used for initialization; model and training hyperparameters are in Table 5.The passage names Huggingface initialization endpoints for the models.Source: §C.1.3 · pp. 14–15

C.1.4. CONVERTING PRETRAINED MHA PARAMETERS TO APPROXIMATE ATTENTION PARAMETERS

  • For GQA conversion, query and output matrices remain unchanged, while reshaped key and value matrices are reduced to nKV heads; MQA uses nKV = 1.The passage notes that selecting the retained heads can be refined using DEIM.Source: §C.1.4 · p. 15
  • For MLA conversion, truncated SVD of rank dc supplies down- and up-projections for query, key, and value; the key down-projection is shared, and WO is not compressed.The stated choice follows the DeepSeek implementation and DeepSeekV3 architecture (DeepSeek-AI et al., 2024).Source: §C.1.4 · p. 15

C.2. Training GPT2 from scratch

The GPT2-from-scratch setup specifies its corpora, implementation, and context length, then describes the RoPE variants and how rotational and semantic dimensions are handled.

C.2.1. OPENWEBTEXT

  • OpenWebText is used to pretrain GT2 variants, and OpenWebText and OpenWebText2 are tokenized with tiktoken v0.12.0 using GPT-2 encoding.Source: §C.2.1 · p. 15

C.2.2. TINYSHAKESPEARE

  • TinyShakespeare supplies character-level language-modeling data, with roughly 1 million characters tokenized using tiktoken v0.12.0 and GPT-2 encoding.The corpus is approximately 1MB of concatenated Shakespeare works.Source: §C.2.2 · p. 15

C.2.3. GPT2

  • The GPT2 implementation is based on NanoGPT, uses context length N = 1024, and takes its model and training hyperparameters from Table 5.The passage identifies the implementation as Andrey Karpathy’s NanoGPT.Source: §C.2.3 · p. 15, Table 5 · p. 16

C.2.4. USAGE OF ROPE

  • RoPE is applied to only a subset of query/key channels, leaving the rest semantic; Tucker evaluates its query and key RoPE matrices once across all heads.For GQA, the key RoPE matrix is evaluated once per head group and broadcast; the baseline GPT2 uses standard cosine position embeddings.Source: §C.2.4 · pp. 15–17

C.3. LLaMA-1B on OpenWebText2 training from scratch

The section describes a from-scratch LLaMA 3-style decoder-only training setup on OpenWebText2, including model configuration, optimization, and decoding timing conditions.

C.3. LLaMA-1B on OpenWebText2 training from scratch

  • The model has 16 layers, hidden size 2048, 32 query heads, and 8 GQA key/value heads, with RoPE and latent RoPE used as specified.It uses a 4096-token maximum context, full rotary coverage, and latent RoPE for Tucker Attention.Source: §C.3 · p. 17
  • Training runs for 100,000 iterations with shared optimization settings, BF16 precision, ZeRO stage 1, activation checkpointing, and FlashAttention2.The setup uses Adam, cosine learning-rate decay, weight decay 0.1, gradient clipping at 1.0, and disabled dropout.Source: §C.3 · p. 17
  • Decoding timings use a 16k KV cache and one query token to simulate autoregressive rollout with an existing cache.Timings are measured with the GPT-NeoX Megatron timer.Source: §C.3 · p. 17

C.4. Ablation study: GPT2 on Shakespeare

The Shakespeare ablation reports superior parameter-efficiency ratios for Tucker Attention. Additional figures show GPT-2 spectra and validation loss, alongside ViT results on ImageNet1k and CIFAR datasets.

D. Additional Figures

  • Figure 5 plots normalized singular spectra for GPT-2 transformer layers, computed by single-mode matricization and summarized across all 12 layers.Each subplot specifies the matricization mode.Source: Figure 5 · p. 18

C.4. Ablation study: GPT2 on Shakespeare

  • Tucker Attention achieves superior parameter-efficiency ratios among GQA, MLA, and Tucker variants in the GPT-2 Shakespeare ablation.The variants use identical training hyperparameters; MLA with Tucker-style RoPE refers to Equation (14).Source: §C.4 · p. 17, Table 6 · p. 19

D. Additional Figures

  • ViT figures report five-run means for ViT-L/32 on ImageNet1k and ViT-B/16 on CIFAR-100 and CIFAR-10, with model attention dimensions specified.The attention dimensions are 1024/16/64 for ViT-L/32 and 768/12/64 for both ViT-B/16 figures, respectively dmodel/nH/dH.Source: Figure 6 · p. 20, Figure 7 · p. 20, Figure 8 · p. 21
  • Without RoPE, GPT-2 validation-loss results report parameter counts of 28.31M for MHA, 16.51M for GQA, 13.00M for MLA, and 5.32M for Tucker Attention; Tucker converges well.The comparison uses MLA with dQ c = 128, GQA with nKV = 2, and Tucker ranks [8, 64, 128].Source: Figure 9 · p. 21
Loading 2603.30033v1…