Tucker Attention: A generalization of approximate attention mechanisms

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

cs.LGMar 31, 2026

Tucker Attention: A generalization of approximate attention mechanisms

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

Organizations: Department of Mathematics; Otto von Guericke University Magdeburg; 39106 Magdeburg; Germany. · Scientific Computing, Norwegian University of Life Sciences, ˚As, Norway · Max Planck Institute for Dynamics of Complex Technical Systems; 39106 Magdeburg; Germany. · Computer Science and Mathematics Division, Oak Ridge National Laboratory, Oak Ridge, TN 37980, USA.

Abstract

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.

Figures & tables

Panel of Figure 1 — Figure 1: Illustration of the pre-softmax attention tensor ${\mathcal{W}}$ under existing factorizations; colored edges denote contraction dimensions. MHA: ${\mathcal{W}}$ is formed by contracting $d_{\rm{model}}\times d_{\rm{H}}$ query and key matrices along $d_{\rm{H}}$ (orange) for each of the ${n_{\rm{H}}}$ heads, yielding tensor rank $({n_{\rm{H}}},d_{\rm{model}},d_{\rm{model}})$ . GQA: Queries follow the same parametrization as MHA, but only $n_{\rm{KV}}$ key matrices are used. During contraction along $d_{\rm{H}}$ (orange), each key matrix is broadcasted to all queries in its head group, yielding rank $({n_{\rm{H}}},d_{\rm{model}},n_{\rm{KV}}d_{\rm{H}})$ . MLA: Key and query matrices are replaced by down- and up-projections, with up-projections reshaped to ${n_{\rm{H}}}\times d_{\rm{c}}\times d_{\rm{H}}$ . Contractions occur along $d_{\rm{c}}^{Q},d_{\rm{c}}^{K}$ (red/blue), and along $d_{\rm{H}}$ per head, giving rank $({n_{\rm{H}}},d_{\rm{c}}^{\rm{Q}},d_{\rm{c}}^{\rm{K}})$ . At inference, contractions along $d_{\rm{c}}^{\rm{Q}}$ and $d_{\rm{H}}$ are precomputed. The post-softmax tensor $\widetilde{\mathcal{W}}$ is parametrized analogously for MHA and GQA; for MLA, only the value matrix has low-rank factorization while the output matrix remains full-rank.

Figure 2: Normalized singular spectrum of pre-softmax head mode,

\Mat_{1}(\mathcal{W})

, for GPT2 pretraining. See Figure 5 for a detailed description. MLA shares a similar low-rank behavior across heads as Tucker Attention, but is not able to leverage this fact to reduce parameters.

Method	#Params (train)	#Params (inference)	KV Cache
Eq.( 6 )	$4{n_{\rm{H}}}d_{\rm{model}}^{2}$	$4d_{\rm{model}}^{2}$	$4N{n_{\rm{H}}}d_{\rm{model}}$
MHA	$4d_{\rm{model}}^{2}$	same as training	$2Nd_{\rm{model}}$
MQA	$2d_{\rm{model}}^{2}+2d_{\rm{H}}d_{\rm{model}}$	same as training	$2Nd_{\rm{H}}$
GQA- $n_{\rm{KV}}$	$2d_{\rm{model}}^{2}+2n_{\rm{KV}}d_{\rm{H}}d_{\rm{model}}$	same as training	$2Nn_{\rm{KV}}d_{\rm{H}}$
MLA (separated KV)	$d_{\rm{model}}^{2}+6d_{\rm{model}}d_{\rm{c}}$	$d_{\rm{model}}^{2}+4d_{\rm{model}}d_{\rm{c}}$	$2Nd_{\rm{c}}$
MLA (shared KV)	$d_{\rm{model}}^{2}+5d_{\rm{model}}d_{\rm{c}}$	$d_{\rm{model}}^{2}+2d_{\rm{model}}d_{\rm{c}}$	$Nd_{\rm{c}}$

Table 1: Comparison of the memory footprint of MHA, MQA, GQA with

n_{\rm{KV}}

KV heads, MLA, and the proposed Tucker Attention. We count the total parameters of the key, query, value, and output weight matrices. MLA fuses query down and up projections during inference, thus the parameter count changes. We assume

d={n_{\rm{H}}}d_{\rm{H}}

for MHA, MQA, GQA,

n_{\rm{KV}}

is the number of kv heads for GQA,

d_{\rm{c}}

is the latent dim/rank of MLA, and

(r_{\rm{1}},r,r)

are the Tucker ranks of

\mathcal{W}

and

\widetilde{\mathcal{W}}

in Tucker Attention. Sequence length is denoted by

N

. Tucker Attention enjoys a favorable

\mathcal{O}(rd_{\rm{model}})

scaling compared to the other methods.

Panel of Figure 3 — Figure 3: ViT32l (left) and ViT14g (right) top-5 validation performance on ImageNet1k for MHA and approximate attention mechanisms over total attention parameter count (top-left is best). GQA is presented with $n_{\rm{KV}}=1$ to $4$ KV heads, where GQA-1 corresponds to MQA. MLA is presented with latent dimension 16-64 with individual kv-weights. Tucker Attention is presented with ranks $[r_{\rm{1}},r_{\rm{2}},r_{\rm{3}}]$ . Is it apparent that Tucker Attention requires an order of magnitude fewer parameters for comparable accuracy levels and thus shifts the parameter-accuracy pareto-frontier to the top left.

	Method	Attn params	KV Cache	Val loss	WikiText2	LAMBADA	Mean	LAMBADA	CBT-CN	CBT-NE	HellaSwag	WinoGrande	PIQA	ARC-E	ARC-C	Mean
	Metric	[MB]	[MB]	(lower is better)	Perplexity (lower is better)			Accuracy [%] (higher is better)
w/o RoPE	MHA	56.62	37.74	2.861	37.15	16.95	27.05	27.17	82.44	62.84	31.28	51.30	64.25	41.96	24.40	48.20
	GQA [ $n_{\rm{KV}}=4$ ]	37.74	12.85	2.896	42.97	22.30	32.63	25.13	82.24	61.52	30.77	51.78	63.17	42.17	25.94	47.84
	GQA [ $n_{\rm{KV}}=2$ ]	33.02	6.29	2.869	40.63	21.67	31.15	24.37	81.44	59.80	30.42	50.99	63.38	41.25	25.94	47.20
	MLA [GQA, $n_{\rm{KV}}=2$ equiv.]	37.74	6.29	2.882	40.32	23.68	32.00	22.38	82.04	61.60	30.66	51.93	62.79	40.91	25.17	47.19
	MLA shared K,V [ $d_{\rm{c}}^{\rm{K}},d_{\rm{c}}^{\rm{Q}}=128$ ]	26.00	3.14	2.937	43.30	24.01	33.66	22.13	80.92	59.72	29.76	50.59	60.93	41.30	25.09	46.31

Table 2: Zero-shot results for GPT-2 [

d_{\rm{model}}=768,{n_{\rm{H}}}=12,d_{\rm{H}}=64

] variants with

N=1024

trained on OpenWebText with identical training hyperparameters, without fine-tuning. For each RoPE setting, we mark the and lowest KV cache and parameter storage (measured in MB and assuming BF16 encoding) of GQA, MLA, and Tucker Attention in bold. Tucker Attention achieves competitive validation accuracy and perplexity scores with a fewer trainable parameters and KV cache than the baseline methods.

		Eval		Training Efficiency				Decode Efficiency
Method	Attn Params [MB]	PPL $\downarrow$	Val loss $\downarrow$	Fwd [ms] $\downarrow$	Bwd [ms] $\downarrow$	Opt [ms] $\downarrow$	Train Iter [ms] $\downarrow$	KV Cache [MB]	Peak Mem [GB]	Latency [ms] $\downarrow$
MHA	268	12.87	2.55	1.58	3.51	0.101	5.19	268	6.5	1.379
GQA [ $n_{\rm{KV}}=8$ ]	168	13.62	2.61	1.47	3.30	0.087	4.87	67	5.7	1.310
GQA [ $n_{\rm{KV}}=2$ ]	142	13.87	2.63	1.44	3.20	0.086	4.73	16.7	5.7	1.305
MQA [ $n_{\rm{KV}}=1$ ]	138	14.27	2.64	1.43	3.13	0.083	4.67	8.38	5.6	1.296
MLA [ $d_{\rm{c}}^{\rm{K}},d_{\rm{c}}^{\rm{Q}}=128$ ]	93.3	14.54	2.68	1.37	2.96	0.079	4.42	16.7	5.9	1.341

Table 3: Validation performance and per-iteration training timing breakdown for LLaMA3-1B variants on OpenWebText2 with

N=4096

. Note LLaMa3 1B default setup is

n_{\rm{KV}}=8

. Training: Timing is reported per training iteration on 2 cluster nodes with 2 H100 GPUs each. Decode: We use query sequence length 1, with cached KV sequence length

N=4096

. Timing is reported on 1 H100 GPU. Peak decode GPU memory load is measured in GB with batch size 1. All timings denote the mean over 10 iterations in [ms]. We remark that this implementation is non-optimized research code. Tucker Attention achieves competitive validation loss and perplexity scores at significant parameter and latency reductions in training and decoding mode.

Panel of Figure 4 — Figure 4: Llama3-1B training cross-entropy loss (left) and final time validation cross-entropy loss (right) on OpenWebtext2 with RoPE in Bf16 floating point accuracy. Approximate attention methods are MLA GQA, and Tucker Attention. Tucker Attention converges well, despite having only a fraction of the trainable parameters of the other methods, and does not exhibits instabilities.

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

General conventions
Symbol	Description
$x,y$	Scalars and vectors (lowercase italic)
$W,X$	Matrices (uppercase italic)
$\mathcal{C},\mathcal{W}$	Tensors (uppercase calligraphic)
Dimensions
$N$	Sequence length (number of tokens)

Appendix

Table 4: Summary of notation.

Hyperparameter	ViT-B.16	ViT-L.32	ViT-G.14	GPT2	LLaMA3-1b
$d_{\rm{model}}$	768	1024	1664	768	2048
${n_{\rm{H}}}$	12	16	16	12	32
$d_{\rm{H}}$	64	64	104	64	64
$N$	196	49	256	1024 (256 for Shakespeare)	4096
Batch Size (Cifar10)	128	n.a.	n.a.	n.a.	n.a.
Batch Size (Cifar100)	128	n.a.	n.a.	n.a.	n.a.

Appendix

Table 5: Model and training hyperparameters

Panel of Figure 5 — Figure 5: Normalized singular spectrum of the transformer layers in GPT2 after training. The singular spectrum was computed via a matricization of the tensor along a single mode; specific information about which mode is given in each subplot. The ribbon plots are calculated as a sample distribution over all twelve transformer layers in GPT2.

Panel of Figure 6 — Figure 6: Vit-L 32 on Imagenet1k. Results denoting mean over 5 runs. The attention layer of ViT-L.32 has dimensions $d_{\rm{model}}=1024,{n_{\rm{H}}}=16,d_{\rm{H}}=64$ .

Panel of Figure 7 — Figure 7: Vit-B 16 on Cifar100. Results denoting mean over 5 runs. The attention layer of Vit-B 16 has dimensions $d_{\rm{model}}=768,{n_{\rm{H}}}=12,d_{\rm{H}}=64$ .

Panel of Figure 8 — Figure 8: Vit-B 16 on Cifar10. Results denoting mean over 5 runs. The attention layer of Vit-B 16 has dimensions $d_{\rm{model}}=768,{n_{\rm{H}}}=12,d_{\rm{H}}=64$ .

Figure 9: GPT2 validation cross-entropy loss on OpenWebtext without RoPE. Approximate attention methods are MLA with

d_{\rm{c}}^{\rm{Q}}=d_{\rm{c}}^{\rm{K}}=128

, GQA with

n_{\rm{KV}}=2

, Tucker Attention with

r=[8,64,128]

. Parameter counts are (in millions) MHA: 28.31, GQA: 16.51, MLA: 13.00, Tucker: 5.32. Tucker Attention converges well, despite having only a fraction of the trainable parameters of the other methods.

Tucker Decomposition

See all 14 papers

Sequential Functional Structured Tucker Compression for Large Language Model Attentions
Sep 30, 2026Jiangfeng Chen, Xinyu Wang, Tianshuo Yan +4Tucker Decomposition LLM Compression
University of Manitoba · McGill University · Simpleway +2
Tucker Bottleneck Attention for Multi-Dimensional Sequence Modeling
Oct 6, 2026Ryan Solgi, Parsa Madinei, Zheng ZhangAttention Mechanisms Low-Rank Attention
University of California-Santa Barbara Santa Barbara, CA, USA
CARE: Covariance-Aware and Rank-Enhanced Decomposition for Enabling Multi-Head Latent Attention
Mar 18, 2026Zhongzhu Zhou, Fengxiang Bie, Ziyan Chen +6Multi-Head Latent Attention Low-Rank Approximation
University of Sydney · Together AI · University of Texas at Austin +1

Sep 30, 2026J/K move · Enter open · S save

Sequential Functional Structured Tucker Compression for Large Language Model Attentions

Jiangfeng Chen, Xinyu Wang, Tianshuo Yan, Hanwei Wu+3

University of Manitoba · McGill University · Simpleway+2

Tucker Decomposition LLM Compression Decoder-Only Language Models Tensor Decomposition Low-Rank Compression

Open paper AskPDF arXiv

Post-training compression of LLM attention is often formulated as independent matrix approximation, ignoring both the shared structure among attention projections and the representation shift introduced by earlier compression. We propose FTC, a sequential structured compression framework that adapts the approximation to the current compressed model while jointly exploiting the native Q/K/V head structure under a fixed storage budget. The output projection is handled separately to account for the changed post-attention representation. FTC requires neither fine-tuning nor gradient-based recovery. Across seven decoder-only LLMs from 6B to 32B parameters, FTC achieves the lowest WikiText-2 perplexity among the compared methods at every tested keep ratio on five modern GQA models, with the largest gains under aggressive compression. The improvements transfer to downstream tasks and remain substantial at the 32B scale.

Abstract

Tucker Decomposition

Explore similar work

Sequential Functional Structured Tucker Compression for Large Language Model Attentions

Tucker Bottleneck Attention for Multi-Dimensional Sequence Modeling

CARE: Covariance-Aware and Rank-Enhanced Decomposition for Enabling Multi-Head Latent Attention