Tensor Product Attention Is All You Need
Authors: Yifan Zhang
Yifeng Liu
Huizhuo Yuan
Zhen Qin
Yang Yuan
Quanquan Gu
Andrew C Yao
Venue: NeurIPS 2025 spotlight ()
Keywords: Attention mechanisms
Transformer
Primary area: deep_learning

Abstract: Scaling language models to handle longer input sequences typically necessitates large key-value (KV) caches, resulting in substantial memory overhead during inference. In this paper, we propose Tensor Product Attention (TPA), a novel attention mechanism that uses tensor decompositions to represent queries, keys, and values compactly, substantially shrinking the KV cache size at inference time. By factorizing these representations into contextual low-rank components and seamlessly integrating with Rotary Position Embedding (RoPE), TPA achieves improved model quality alongside memory efficiency. Based on TPA, we introduce the Tensor ProducT ATTenTion Transformer (T6), a new model architecture for sequence modeling. Through extensive empirical evaluation on language modeling tasks, we demonstrate that T6 surpasses or matches the performance of standard Transformer baselines including Multi-Head Attention (MHA), Multi-Query Attention (MQA), Grouped-Query Attention (GQA), and Multi-Head Latent Attention (MLA) across various metrics, including perplexity and a range of established evaluation benchmarks. Notably, TPA's memory efficiency and computational efficiency at decoding stage enables processing longer sequences under fixed resource constraints, addressing a critical scalability challenge in modern language models. Project Page: https://github.com/tensorgi/TPA.

Tensor Product Attention Is All You Need
Yifan Zhang 11,4 Yifeng Liu
Huizhuo Yuan Zhen Qin Andrew Chi-Chih Yao 1.2f
Yang Yuan 1,2 Quanquan Gu
Shangaai Qi Zhi Institute University of California, Los Angeles *Princeton University yi jifzhang@princeton edu, liuyifengecs ucla. edu qguêcs ucla. edu, andrewcyao@tssingua. edu. cn
'Ms, Tsinghua University
Abstract
Scaling language models to handle longer input sequences typically necessitates large key-value (KV) caches, resulting in substantial memory overhead during inference. In this paper, we propose Tensor Product Attention (TPA), a novel attention mechanism that uses tensor decompositions to represent queries, keys, and values compactly, substantially shrinking the KV cache size at inference time. By factorizing these representations into contextual low-rank components and seamlessly integrating with RoPE and any possible position encoding mechanisms, TPA achieves improved model quality alongside memory efficiency. Based on TPA, we introduce the Tensor ProducT ATTen Transformer (T6), new model architecture for sequence modeling. Through extensive empirical evaluation on language modeling tasks, we demonstrate that T6 surpasses matches the perfor- mance standard Transformer baselines, including Multi-Head Attention (MHA), Multi-Query Attention (MQA), Grouped-Query Attention (GQA), and Multi-Head Latent Attention (MLA) across various metrics, including perplexity and a range established evaluation benchmarks. Notably, TPA's memory efficiency and compu- tational efficiency at the decoding stage enable processing longer sequences under fixed resource constraints, addressing critical scalability challenge in modern language models. Project Page: https: ///gttttbbb//// com homm/tensorgi////
Introduction
Large language models (LLMs) have revolutionized natural language processing, demonstrating exceptional performance across tasks [5, 12, 58, 6]. As these models evolve, their ability to process longer contexts becomes increasingly important for sophisticated applications such as document analysis, complex reasoning, and code completion. However, managing longer sequences during inference poses significant computational and memory challenges, particularly due to the storage of key-value (KV) caches [70, 34]. Because memory consumption grows linearly with sequence length, the maximum context window limited by practical hardware constraints.
variety of solutions have been explored to address this memory bottleneck. Some approaches compress or selectively prune cached states through sparse attention patterns [10] or token eviction strategies [70, 62, 42], though such methods risk discarding tokens that may later prove important. Other work proposes off-chip storage of key-value states [17]. at the expense of increased I/O latency. Attention variants like Multi Multi-Query Attention (MQA) [46] and Grouped- -uuury-uurrruuuu Attention (GQA) [2] reduce per-token cache requirements by sharing keys and values across heads, but often compromise ficxibility or require significant architectural modifications. Meanwhile, low-rank weight factorization methods such as LoRA [20] effectively reduce fine-tuning memory, yet do not address the KV cache overhead that dominates inference at runtime. The recently introduced Multi- Head Latent Attention
Equual contribution; Project lead; Corresponding author.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Linear
Concat
Scaled Dot-Product Attention
Scale
Scale
Scale
Rx
R,
RoPE
RoPE
Aa
0
Ak
Bk
Av
By
Linear
Linear
Linear
Linear
Linear
Linear
Tensor Product Attention (TPA) within the Tensor ProducT ATTenTion Transformer (T6). In each TPA layer, the input hidden state x; is processed by linear layers to produce latent factor matrices for query (e.g., Ax(x)) (x,)), key (e.g.. Ax(x,) B (x:)), and value (e.g.. (x,), (x.)),,. Rotary Position Embedding (RoPE) is applied to the (x and (x, factors. The query, key, and value tensors for each attention head then formed by the tensor product of
Figure 1:
these factor matrices (e.g., Q = 1 (x (x,)). Finally, the TPA output is computed using scaled dot-produce attention, followed by a linear projection of the concatenated results from all heads.
Ro
(MLA) in Deepseek-V? V2 [32] caches compressed key-value representations but encounters difficulties with efficient Rotary Position Embedding (RoPH) [52] integration, necessitating additional position- encoded parameters per head.
To overcome the limitations of existing approaches, introduce Tensor Product Attention (TPA), illustrated in Figure 1. TPA is novel attention mechanism that employs tensor factorizations for queries (Q), keys (K), and values (V). By dynamically factorizing activations rather than static weights (as in LoRA), TPA constructs low-rank, contextual representations. This approach substantially reduces KV cache memory usage while offering improved representational capacity. In practice, TPA can decrease memory overhead by an order of magnitude compared to standard Multi-Head Attention (MHA), alongside achieving lower pretraining validation loss (perplexity) and better downstream performance. A key advantage of TPA is its native compatibility with rotary positional embeddings (RoPE) [52] and any possible position encodings, enabling straightforward drop-in replacement for multi-head attention (MHA) layers in modern L.M architectures such as L.aMA [58), Qwen [3], and Gemma [56].
Our main contributions are summarized as follows:
1. We propose Tensor Product Attention (TPA), mechanism that factorizes Q, K, and V activa- tions using contextual tensor decompositions. This achieves a substantial reduction in inference- time KV cache size relative to standard attention mechanisms [60], MHA. MQA, GQA, and MLA, while also improving performance. In addition, we analyze existing attention mechanisms and reveal that MHA, MQA, and GQA can be expressed as non-contextual variants of TPA.
We introduce the Tensor ProducT A ATTenTion Transformer (T6), a new TPA-based model archi- tecture for sequence modeling. In language modeling experiments, T6 consistently improves or matches validation perplexity and downstream evaluation performance, all while maintaining a reduced KV cache size.
3. We demonstrate that TPA integrates seamlessly with RoPE [52] and any possible position encod- ings as well as output gate and KV shifting, facilitating its easy adoption in popular foundation model architectures like LLaMA, Gemma, and Qwen.

---

4. We develop FPashTPA Decoding, an efficient autoregressive inference algorithm for TPA. Our empirical results show that FlashTPA Decoding can be faster than optimized MHA, MQA, GQA, and MIA decoding methods, particularly for long sequences.
2 Background
In this section, we briefly review Scaled Dot-Product Attention, Multi-Head Attention [60], and introduce key notations. Other attention mechanisms like Multi-Query Attention (MQA) [46], Grouped Query Attention (GQA) [2], Multi-head Latent Attention (MLA) [32, 33], and Rotary Position Embedding (RoPE) [52] are further discussed in the Appendix F.
Notations. We use bold uppercase letters (e.g., X. Q) for matrices, bold lowercase (e.g., a, b) for vectors, and italic uppercase (e.g.. W for learnable parameter matrices. We denote by [n] the set {l.....n) n for some positive integer n. We use T to denote the transpose of vector matrix. Let dmanel be the embedding dimension, h the number of attention heads, d,, the dimension per head, X2 Rd the input for the 1-th token at a given attention layer, X R R'' denotes the input embeddings for T tokens, and Q K, V E IR xxxdx denote the queries, keys, and values of h heads for T tokens. With a little abuse of notation, Qu K, V, 7 denote the i-th head of queries, keys, and values, and Q. K., V, E gaxda denote the heads of the query, key, and value for i-th token. Throughout the paper, Wº, W w denote projection matrices for queries, keys, and values, respectively. In multi-head attention, each head is associated with its own set of W.O. W.K. W.V. and each has dimension W.O, W.K. W. E R da Similarly, we have an output projection matrix W E (h--ds) ((h-ds))dduate
We define the tensor product of two vectors as follows: for vectors a E RT b E RR" the tensor product of and b is: a @ b C E Rmxx with C2 = a,by, where a; is the i-th element of a, b, is the j-th element of b, and Ci is the (i.j)-th entry of C. The vectorization of matrix C E Rm*" denoted vec(C) E Rmm, stacks the columns of C into single column vector. For example, [c1, C2 : Cn. where oj are columns, then vec(C) [c] [,] ....,,
C
2.1 Scaled Dot-Product Attention
Scaled dot-product attention [60] determines how focus on different parts of an input sequence by comparing queries (Q) and keys (K). It produces a weighted combination of the values (V). Formally, the attention output is:
Attention Q K, V) Softmax
OK
V.
dd,
where Q e Raxd, K E Rad,, and V Rnxd, for tokens. The softmax is applied row-wise over the n keys for each query.
2.2 Multi-Head Attention (MHA)
Multi-Head Attention (MHA) [60] extends scaled dot-product attention by dividing the model's s internal representation into several heads. Each head learns different projections for queries, keys, and values, allowing the model to attend to different types of information from different representational subspaces. For each token embedding X₃ E [ddinan MHA computes each head as follows:
Q = ( X, dh. K = (W:)² x, E Rd
Vt,: = (W.V) T x, E ER...
head; = Attention
where w W.K. W.V. E dl are learnable projection matrices for the i-th head, and Q K,, V, E Txxx are the query, key, and value matrices for the i-th head over T tokens. After computing each head's attention output, the results are concatenated and mapped back the model's original dimension via another learnable linear projection matrix D E g phd, Xdaaii
MHA(X) = Concat (head,,.... head,) W
MHA enables the model to capture rich set of dependencies by allowing each head to focus on different aspects of the input sequence. We also discuss how MHA, MQA, and GQA relate to TPA in the Section 4.
Often, h de dmooee, each head has query/key/value dimension da.
3

---

3 Tensor Product Attention
In this section, we provide a detailed description of our proposed Tensor Product Attention (TPA), which enables contextual low-rank factorization for queries, keys, and values. First, explain how TPA factorizes these components, specifying tensor shapes. Next, we describe TPA's integration into the multi-head attention framework and its benefits for reducing KV cache memory consumption during inference. Finally, we demonstrate RoPE's seamless integration with TPA, including a pre-rotated variant for efficiency.
3.1 Tensor Factorization of Queries, Keys, and Values
Let datm == dn denote the total attention projection dimension. Typically one sets dam = dmedal but this is not required: when datn # dmodel, the projection matrices W W W map from Rd into Rd and W maps Rd back to IR dndee, Standard attention projects the entire sequence into three tensors, Q. K. V E RTxxx,, where Qr, K,, Ve E R denote the slices for the t-th token.
Contextual Factorization. Instead of forming each head's query, key, or value via a single linear map, TPA factorizes each Qr, K,, V, into sum of (contextual) tensor products whose ranks are Rg. Rx, and Ry, respectively, and may differ. Specifically, for each token I, with a small abuse of notation, we define:
n2
Rg
Q
a
(xx)
(x1)
(x.);
K,
1
Σ
a
(xx)
& b (x:),
RQ
RK
By
V Rv
a (x1)
x by (xt);
(3.1)
where a (x,), K (x,). a. (x,) E R b& (x,) bY (x,) E Rd, Hence, for queries, each
tensor product a, (x(xx) & b, (x:): Re Rd - IR hxdi contributes to the query slice Q₁ E Axd, Analogous definitions apply the key slice K, and value slice V..
Latent Factor Maps. Each factor in the tensor product depends the token's 's hidden state xt. For example, for queries, we can write:
(x.) w X; E RRA
(x(()) W X, Rh,
where We Rixdd and W RRd xduata are learnable weight matrices. Similar linear maps produce the factors for keys and values.
One often merges the rank index into a single output dimension. For instance, for queries:
a (x,) w x E
R Rh b²(x)) W X4 R.
xh and Bx(x;) R Rgxds (where each row of (x:)) which are then reshaped into Aq (xt) Roxh corresponds to an a& (x,) and each row of B2(x;) to (x.) The query tensor for token can then be expressed as:
1
Q = AQ(x))
T Bo(x:) E gixds
RQ
This operation equivalent to Q - - Ro a& (x, where a& is the r-th column of (x(x)) T and (bę)T is the r-th row of (xx).. Repeating for all tokens reconstitutes Q Txxxd, Similar procedures are applied obtain K and V with ranks Rg and Rv, respectively. Scaled Dot-Product Attention. Once Q K, V are factorized, multi-head attention proceeds in standard Transformers. For each head E (1, (1,....hh)
The
head; = Softmax
Vaa Q (K.)) V.,
(3.2)
where Qu K, V, RTxd, are the slices along the head dimension. Concatenating these h heads along the last dimension yields an ((t-da))) tensor, which is projected back to dx by an output weight matrix wo E ((fh-d))xdaad
TPA(Q, K, V) Concat (head, head,,,,,,,,,, head,) W
(3.3)

---

Parameter Initialization. We use Xavier initialization [15] for the factor weight matrices; details are in the Appendix G.
3.2 RoPE Compatibility and Acceleration
In typical workflow of adding RoPE to standard multi-head attention, one first computes Q₁ K, E R of the t-th token and s-th token and then applies:
Q₂ Q2 RoPE:(QQ)) (Qr), K2 K2 Ro.E,(K.))
(3.4)
Direct Integration. A useful optimization is to integrate RoPE directly into the TPA factorization. For example, one can pre-rotate the ooken-dimension factors:
(x.) == RoPE, (Bx(xx)) = Bx(xx)TT,,
(3.5)
yielding pre-rotated key representation:
Rx
K. Σ
1 RK
K 7
(x+)
ROPF,,(; (xi)
1 Ax(x;) Bx(())-
Here, RoPE, is applied to each of (x,) (i.e., each by (x₂) vector). Thus, each cached key factor corresponds to RoPE-rotated key slice. This removes the need to rotate cached keys at decoding time; the current-step query (which is not cached) can still be rotated on the fly at negligible cost. Depending on hardware and performance requirements, different RoPE integration strategies can be adopted for training and inference.
Theorem 3.1 (RoPE's Compatibility with TPA). Let Qe be factorized by TPA as
1
Q
(x(x)) BQ(x,) E Rixds
RQ
where Aq(xi) IR Rg and BQ (x:) E IR xd,, Then we have:
1
RoPE, Qu) Q. = Re (x))) Bo(x)),
(3.6)
where (x(x)) Bx T, RoPE, (x+(x)) (RoPE applied row-wise (x()))). Further-
more, let Q = RoPF, (Qr) Q.T, and K, = RoPF,, = K,T, be the RoPF-transformed query/key slices. Then RoPE's standard relative-position identity is preserved:
Q K
Q.T.KKK,
equivalently
RoPE,-.(Q:) K. -
where T Te = T.T. In particular, for any head i (the i-th row), if qr,i ko, E RIdd and
q.. qr, Te, k,, = k,, T., then ir, K,, = q..T...k],
Theorem 3.1 indicates that TPA does not break RoPE's relative translational property. We prove in the Appendix D.I.
3.3 KV Caching and Memory Reduction
In autoregressive decoding, standard attention caches K,, V. Roxd, for each past token This accumulates to TTxxxd, for keys and 7xxxxd, for values, i.e., 2T h da total.
TPA Factorized KV Caching, Instead of storing the full K, and V,, TPA stores only their factor components. Specifically, for each past token we cache:
A Bx (xx) and Av(x), By;x,,,
where Ax(x)) E RR xh (x(x)) E IR Rx xda (pre-rotated), Av(xi) E Rv By e Ry xdh,
Hence, the memory cost per token is RK(h+ dh) + Ry + dh) ( (RK + Ry (h + dn..
For
For
Compared to the standard caching cost of 2 dr, the ratio is RRK+ (k+d..) For large h and dj
24 2h
(typically d, = or 128), setting Rx; Rv & h (e.g., rank 1 or 2) often yields substantial reduction of KV cache size. Table 1 provides a comparative overview of different attention mechanisms, including TPA and its variants, focusing on KV cache size per token and the number of parameters in an attention layer.
5

---

Table 1: Comparison of different attention mechanisms. Here, Ro. Rx. and Ry denote the ranks for queries, keys, and values in TPA, respectively. Variants of TPA, such as TPA (K Vonly)), TPA (Non- contextual A), and TPA (Non-contextual B), are detailed in the Appendix G. For MIA, di and dy the dimensions for RoPE and non-RoPE parts; d' and de are the dimensions of compressed vectors for query and key-value, respectively. The MLA parameter count includes the output projection.
METHOD
KV CACHE
PARAMETERS
QUERY HEADS
A KV HEADS
MHA
2hdr.
Adhest &.
MQA
2dase da (h
GQA
2Gd.
2dhente du (h + G)
d. (drese +
Rol))
+ 2hd.)
MLA
de
+- (hdn di")
(Rx Rv Y(h + de)
dusas Rx + Ry da) disael Add,
TPA (KVonly)
(Rx + Ry [hh + da)
devael Rx + Rv-(h du). + 2dhers hds
TPA (Non-cn Soon-conexxtual/ A)
(Rac Ry ids
(Rq + Ric Ry [dlooadss 1) drenn hdds
TPA (Non comtextull B)
(Ric + Ry
(Rl + Rlx + Ry Idawech da) + davan hdd.
Expressing MHA, MQA, GQA as Non-contextual TPA
We demonstrate that standard Multi-Hear Attention (MHA), Multi-Quee Query Attention (MQA), and Grouped- Query Attention (GQA) can be expressed as special, non-contextual variants of Tensor Product Attention (TPA). This is achieved by imposing specific constraints on the TPA factors, particularly by making head-dimension factors (a) independent of the input token (x:).
4.1 МНА as Non-contextual TPA
Standard Multi-Head Attention (MHA) can be precisely formulated TPA where the rank is equal to the number of heads (Rq = Rg = Rv = h), and the head-dimension factors are fixed, non-contextual basis vectors. To recover MHA, set the rank Rę h and define the factors for each head E [h] as follows:
Contextual token factor: This is the standard linear projection for the i-th head's query:
by (x:) (W:9) x; Rd
Non-contextual head factor: This factor a scaled standard basis vector, independent of x,:
e, R
where e; is the i-th standard basis vector (a vector of zeros with at the i-th position).
Substituting these into the TPA equation, the 1/RQ 1/h scaling factor cancels with the scaling of the a&? factor:
1 Q - Σ (h ex) ((W,)
Σ e (IW
The resulting tensor product, @ b Q (x:). produces an h x da matrix where only the i-th row is non-zero and contains the vector (b, (xx)) Summing these matrices for i ......u h assembles the complete query tensor Qr, where the i-th row is precisely the query vector for the i-th head in standard MHA. An analogous construction applies the key (K,) and value (V,) tensors.
Thus, MHA is equivalent non-contextual TPA where the head- head-dimension factors fixed and orthogonal, effectively assigning a dedicated rank component to each attention head.
4.2 MQA and GQA as Non-contextual TPA
Similarly, Multi-Query Attention (MQA) and Grouped-Query Attention (GQA) can be seen as non-contextual TPAs where the key and value tensors are formed with a rank lower than the number of heads.
MQA as Rank-1 TPA (for and V). In MQA, h query heads share a single key and value. This corresponds to a TPA with ranks Rx 1 and Ry 1. The key tensor K; is formed using

---

single, contexttll head- dimension factor a = 1n (a vector of all ones) and single contextual token-dimension factor b (x,) (WK) x,:
K, (1A &bb (1,@b*(x.)) (x+))
This creates an h x dn matrix where every row the same shared key vector (b* (x))). The same logic applies to the value tensor V,. The queries remain full-rank (Rq h) in MHA.
GQA as Rank-G TPA (for K and V). GQA is intermediate approach where h heads are divided into groups, with heads in the same group sharing key and value. This is equivalent to a TPA with ranks Rx- G and Ry G. The key tensor is formed by summing G components:
1 G
K:
K
by (xx)
Here, (xx) is the shared key vector for group j. The non-contextual factor a is scaled mask vector, defined as a G maskj, where the mask, vector has ones for heads belonging to group j and zeros elsewhere. This scaling cancels the 1/G pre-factor:
G D 1=1
K, = G
(G-mak)) ,) @ (x1) = Σ mask, @ by (x,)
1=1
For example, with h 4 8 heads and G = groups (2 KV heads), the factor for the first group of 4 heads would be a = 2- [1, 1, 1, 0, 0, 0, 0] This construction correctly assembles the final key tensor by broadcasting each group's ' shared key to its designated heads without any unintended extra scaling.
This perspective highlights that MHA, MQA, and GQA are specific instances of a more general TPA framework, where expressiveness and parameter sharing are controlled by the rank and the nature (contextual vs. non-contextual) of the tensor factors.
4.3 Model Architectures
We propose new architecture called Tensor ProducT ATTenTTion Transformer (T6), which uses our Tensor Product Attention (TPA) in place of standard MHA (multi-head attention) or GQA (grouped-query attention). Building upon the query, key, and value tensors Q K, V RTxxxxd defined in Section 3.1, T6 utilizes the overall architecture of LLaMA [58] while changing the self- attention block our TPA-based as in (47,5))
based version. The fred-forward network (FFN) adopts a SwiGLU layer,
Rotary Positional Embedding (RoPE). As discussed in Section 3.2, RoPE [52] is applied to the Q and K. Within TPA, we pre-rotate the factor by (xx) and bK (x.) directly, so that each K, is already rotated prior caching, see Equation (3.5) and Theorem 3.1.
SwIGLU Ferd-Forward Network. Following [47, 58]. our T6 uses SwiGLU-based Feed- Forward Network (FFN): (FN(x) = [a W1) @ (x W2) Ws. where a is the SiLU (a.k.a., swish) nonlinearity, is element-wist product, and W1, W2, Wa learnable parameters. Note that other activation functions can also be used.
Overall Block Structure. Putting everything together, one T6 block consists of:
f + TRM(NNM)))(()))
We place norm layers (e.g., RMSNorm) before each sub-layer. Stacking L such blocks yields T6 model architecture with L layers.
TPA Decoding Algorithm
For efficient autoregressive inference with Tensor Product Attention (TPA), we introduce FlashTPA Decoding, This algorithm is optimized for generating one token at time by leveraging the factorized
7

---

representation of queries, keys, and values. The core idea, illustrated in Figure 2, is to perform attention computations using a sequence of Einstein summations ("einsum") that operate directly on these factorized components. This avoids materializing the full query, key, and value tensors, which particularly beneficial as the Key-Value (KV) cache grows with sequence length. The detailed definitions of the input factorized components and the step-by-step pseudo-code for FlashTPA Decoding are provided in Algorithm An optimized Triton kernel implementation is outlined in Algorithm 3 (see Appendix B.
AG
Bo (R
(R0.00
RM Ro!
(RGMM)
(e) M)
(6.50
O
Kin
Figure 2: Data flow diagram for FlashTPA Decoding. Rectangles represent tensors (blue for inputs, yellow for intermediates, red for final output), circles with Σ or @ denote Einstein summation contractions or element-wirg products respectively, and the green rounded rectangle is the softmax operation. Shapes shown for single query (N 1) interacting with M cached items in the common rank-l setting Rx- Rv = 1. We use a head-first layout (H, MM for logits and attention weights; the cached head factors aace and a are shown transposed relative to their natural token-major layout for readability. H/ is the number of heads, Ri is the query rank, and D, E are respective feature dimensions for the Sg/bacche and bY cache factors. Scaling factors are omitted for visual clarity.
This sequence of factorized operations allows FlashTPA Decoding to compute the attention output efficiently. Consequently, TPA is not only memory- due to its smaller KV cache footprint but can also be computationally efficient during inference. The experimental results for Flash decoding time are presented in Section 6.2.
Experiments
6.1 Language Modeling Tasks
All experiments reported in this paper are implemented based on the nanoGPT codebase [24]. and we pretrain our models using the FineWeb-Edhe 100B dataset [37]. The dataset contains 100 billion tokens for training and 0.1 billion tokens for validation. We compare T6 against the baseline Llama architecture [58] with SwiGLU activation [47] and RoPE embeddings [52], as well as Llama variants Attention (MQA; [46]), Grouped Query Attention (GQA: [2I), or Multi-head Latent Attention (MLA; [32]). In our experiments, the number of heads h is adjusted for each attention mechanism ensure that all attention mechanisms have the same number of parameters as the standard Multi-Head Attention (MHA), which has 4 parameters per attention layer. We train models at four scales: small (1 24M parameters), medium (353M), large (773)), and XI. (1.5B). We pretrain all models for 50B tokens (roughly half an epoch over FineWeb-Edu--00 Details on architecture hyperparameters and training hardware are shown in Appendix H.l.
that replace Multi-Head Attention (MHA; [60]) with
Training & Validation Curves. Figure 4 compares validation loss curves for the medium (353M), large (773)), and XL (1.5B) models on FineWeb-Edd Breb-Edu-1009) 100B. Training loss curves are provided in Appendix Figure 3. Overall, TPA (red curves) and its simpler variant TPA-KVonle (pink curves) (see Appendix G) converge as fast faster than the baselines (MHA, MQA, GQA, MLA) while also achieving visibly lower final validation losses. For instance, in Figure 4(b), TPA and TPA-K Vonly remain below the MHA baseline in terms of validation loss nearly all training stages. Meanwhile, Multi-Head Latent Attention (MLA) [32] (blue curves) generally trains more slowly and yields higher validation losses.
Validation Perplexity. Figure (in the Appendix) shows the validation perplexities of the medium- and large-scale models. Mirroring the loss curves, TPA and TPA-K TPA-KVonly steadily outperform MHA, MQA, GQA, and MILA over the course of training. By the end of pretraining (around 49B tokens), baseb approaches achieve the lowest perplexities in most configurations.
Downstream Evaluation. We evaluate zero-shot and two-shot performance on standard bench- marks, including ARC [63]. BoolQ [13], HellaSwag [64]. OBQA [39], PIQA [4], WinoGrande [43],
8

---

and MMLU [18], using the in-evallattonnharee harness codebase [14]. For ARC-E, ARC-C, Hel- laSwag, OBQA, PIQA, and SciQ. we report accuracy norm; for other tasks, we report standard accuracy. Due to the page limitation, we only display the zero-shot evaluation results of medium and large models here Tables 2 and 3. Zero-shot evaluation of small and XI. models are dis- played in Tables 11 and 12 in the appendix. Moreover, we also present 2-shot evaluation results in Tables 13, 14, 15 and 16 in the appendix.
For the medium-size (353M) models (Table for O-shot and Table 14 in appendix for 2-shot), TPA generally ties or outperforms all competing methods, achieving. for example, an average of 51.41% zero-shot mode versus MHA's 50.11 %,, MQA'S 50.44%, and MLA's 50. 13.%. When given two-shot prompts, TPA again leads with 13% average accuracy. A similar trend appears for the large-size (773M) models (Table 3). where TPA-K Vonly attains the highest average (53.52% zero-shot). For the XL size models (1.5B) (Table 12 in the appendix), TPA-KV only achieves the highest average (55.03% zero-shot). Our experiments confirm that TPA consistently matches or exceeds the performance of established attention mechanisms (MHA, MQA, GQA, MLA) across medium and large model scales.
Medium Nodel. Iduloo
Large Model Re9Wbbedd1000
XL Vodal Poooie--
HCA
NOA
13-
TEA Conyy
THA-AAY
2.6
TRL ONLY
20
45
30
24-
(a) Medium models (353M)
(b) Large models (773M)
(c) XL. models (1.5B)
Figure 3: The training loss of medium-size (353M), large-size (773M) as well as XL-size (1.5B) models, with different attention mechanisms on the Fine Web-FuuE 100B dataset.
Meskum Marel RnelWe edu1008
Lange Hodel, Free//bbed:1000
Xi Mosel Rrekleb 100
2.
HOA
MLA
2.8
TTA
2.4
25
Toinnng Johura Et
Telm toheritti
(a) Medium models (353M)
(b) Large models (773M)
(c) XL models (1.5B)
Figure 4: The validation loss of medium-size (353M), large-size (773M) as well as XL-size (1.5B) models, with different attention mechanisms on the Fine Web-Edu 100B dataset.
Table 2: The evaluation results of medium models with different attention mechanisms pre-trained using Feb-WWebFad Edu 100B dataset (0-shot with Im-evaluation-harnes. The best scores in each column are bolded. Abbreviations: HellaSw. HellaSwag, W.G. = WinoGrande.
Method
ARC-E
ARC-C
BoolQ
HellaSw.
OBQA
PIQA
W.G.
MMLU
SeiQ
Avg.
MIIA
59,51
29.52
45.68
34.20
23.33
76.90
50.11
MQA
57.62
31.91
59.45
45.69
35.40
69.31
53.51
26.47
74.60
50.44
GQA
58.67
31.48
58.29
45.45
35.20
68.50
54.46
24.58
76.50
50.35
MLA
56.65
29.52
57.83
46.05
34.60
69.42
52.80
24.62
79.70
50.13
TPA-KVonly
58.01
30.12
58.01
45.95
35.60
69.10
53.12
25.39
75.10
50.04
TPA
58,38
31.57
59.39
46.83
37.00
70.02
54.06
25.52
79.90
51.41
6.2 Experimental Results on Flash TPA Decoding
This section presents an evaluation of FlashTPA's decoding time in comparison to several other optimized attention mechanisms. We benchmark FlashTPA against FlashMHA [45], FlashGQA, FlashMQA, and FlashMLA [23]. is important to note that our current FlashTPA inplementation utilizes Triton [57]. While the compared methods are typically available as highly optimized CUDA kernels, these experiments provide initial insights into FlashTPA's potential. Development of CUDA- based TPA kernel is ongoing and is expected to yield further performance improvements.

---

Table 3: The evaluation results of large models with different attention mechanisms pre- trainrrindd using the FineWeb-Edu 100B dataset (0-shot with Im-evaluaation-haresss)) The best scores in each column bolded. Abbreviations: HellaSw. HellaSwag, W.G. WinoGrande.
Mcthod
ARC-B
ARC-C
BoolQ
HollaSw.
OBQA
PIQA
W.G.
MMLU
SciQ
Avg.
MHA
59.93
33.62
61.93
50.63
36.00
71.06
55.41
22.87
81.20
52.52
MQA
60.73
33.62
57.34
50.09
37,00
69.97
55.49
25.30
79.60
52.13
GQA
61.66
34.30
58.72
49.85
38.40
71.16
53.75
25.23
77.60
52.30
MLA
63,55
32.85
60.95
51.72
38,80
70,51
55.01
24.55
81,90
53.32
TPn-KVonly
63,26
34.13
61.96
50,66
37:20
72.09
55.25
26.06
81.
53,52
TPA
63,22
35.58
60.03
51.26
36.80
71.44
55.56
24.77
79.60
53.10
Soa
E.bbediingg 2310
Daez Grossdiing 20:6
(a) Batch Size=1
(b) Batch Size=2
(c) Batch Size=4
Dascc Stn Croaddrrg Line 2040
Cercs Soe 25 Orbudding Onr, 2010
15
(d) Batch Sizc=8
(c) Batch Size=16
Figure 5: Decoding time comparison of different attention mechanisms with embedding dimension of 2048 and d, 64. The y-axis represents log, (time) in seconds, and the x-axis represents (0g2(sequenc)) length). Each subfigure corresponds different batch size.
The evaluations were performed with batch sizes selected from 21, 4,8, 16), model embedding dimensions (dmodal) chosen from (1024, 2048, 3072], and sequence lengths ranging from 2 (4,096) to 29 (524,288). For all experiments, the dimension per head (d,)) was fixed at 64. The ranks for TPA's factorized components (Rq, RK. Rv) were set to (16, 1, 1). and for GQA configurations the number of key-value head groups was 4. The decoding time per token, measured as logz (time) in seconds, is plotted against log, (sequence length). Lower values on the y- y-axis indicate faster decoding times. Results are presented in Figure 5 for embedding dimension of 2048 (corresponding to 32 attention heads). Additional results for embedding dimensions of 1024 (16 heads, Figure 8) and 3072 (48 heads, Figure 7) provided in Appendix B. Figure 5 depicts these speed comparisons for an embedding dimension of 2048. The results indicate that FlashTPA (blue line) is highly competitive and often outperforms other attention mechanisms, especially as the sequence length increases.
7 Conclusion
We introduced Tensor Product Attention (TPA), which factorizes query, key, and value matrices into rank-R R tensor products dependent on the token's 's hidden state. Storing only the factorized key/value components during autoregressive decoding substantially decreases the KV memory size with improved performance compared with MHA, MQA, GQA, and MLA. The approach is fully compatible with RoPE (and can store pre-rotated keys). Variants of TPA include factorizing only the key/value or sharing basis vectors across tokens. Overall, TPA offers powerful mechanism for compressing KV storage while improving the model performance, thereby enabling longer sequence contexts under constrained memory.
10

---

Acknowledgements
We thank the anonymous reviewers and area chairs for their helpful comments. We acknowledge the compute credits provided by Fetch.ai.
References
[1] Muhammad Adnan, Akhil Arunkumar, Gaurav Jain, Prashant Nair, llya Soloveychik, and Purushotham Kamath. Keyformer: Ky cache reduction through key tokens selection for efficient generative inference. Proceedings of Machine Learning and Systems, 6:1 6:114-127, 2024.
[2] Joshua Ainslie, James Lee-Thorp. Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai. GQA: training generalized multi-query transformer models from multi-head checkpoints. In Houda Bouamor, Juan Pino, and Kalika Bali, editors, Proceedings of the 2023 Conference Empirical Methods in Natural Language Processing, EMNLP 2023, Singapore, December 6-10, 2023, pages 4895-4901. Association for Computational Linguistics, 2023.
[3] Jinze Bai, Shuai Bai, Yunfei Chu, Zeyu Cui, Kai Dang, Xiaodong Deng, Yang Fan, Wenhin Ge, Yu Han, Fei Huang, et al. Qwen technical report. arXiv preprint arXiv:2309.16600) 2023.
[4] Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, and Yejin Choi. PIQA: reasoning about physical commonsense in natural language. In The Thirty-Fourth Fourth AAAI Conference on Artificial Intelligence, AAAI 2020, The Thirty-Secone Innovative Applications of Artificial Intelligence Conference, IAAI 2020, The Tenth AAAI Symposium on Educational Advances in Artificial Intelligence, EAAI 2020, New York, NY, USA, February 7-12, 2020, pages 7432-7439. AAAI Press, 2020.
[5] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D) Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33: 33:1877-1901, 2020.
[6] Sébastien Bubeck, Varun Chandrasekaran, Ronen Eldan, Johannes Gehrke, Eric Horvitz, Ece Kamar, Peter Lee, Yin Tat Lee, Yuanzhi Li, Scott Lundberg, et al. Sparks of artificial general intelligence: Early experiments with gpt-4, arXiv preprint arXiv:2303. 12712, 2023.
[7] Kerim Büyükakyüz.. Olora: Orthonormal low-rank adaptation of large language models. arXiv preprint arXiv:2406.01775, 2024.
[8] Zefan Cai, Yichi Zhang, Bofei Gao, Yuliang Liu, Tianyu Liu, Keming Lu, Wayne Xiong, Yue Dong, Baobao Chang. Junjie Hu, et al. Pyramidkv: Dynamic kv cache compression based on pyramidal information funneling. arXiv preprint arXiv:2406. 020669, 2024.
[9] Yukang Chen, Shengju Qian, Haotian Tang, Xin Lai, Zhijian Liu, Song Han, and Jiaya Jia. Lon- glora: Efficient fine-tuning of long-context large language models. In The Twelfth International Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024, 2024.
[10] Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv: 10509, 2019.
[11] Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamás Sarlós, Peter Hawkins, Jared Quincy Davis, Afroz Mohiuddin, Lukasz Kaiser, David Benjamin Belanger, Lucy J. Colwell, and Adrian Weller. Rethinking attention with performers. In 9th International Conference Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021.
[12] Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, Parker Schuh, Kensen Shi, Sasha Tsvyashchenko, Joshua Maynez, Abhishek Rao, Parker Barnes, Yi Tay, Noam Shazeer, Vinodkumar Prabhakaran, Emily Reif, Nan Du, Ben Hutchinson, Reiner Pope, James Bradbury, Jacob Austin, Michael Isard, Guy Ari, Pengcheng Yin,

---

Toju Duke, Anselm Levskaya, Sanjay Ghemawat, Sunipa Dev. Henryk Michalewski, Xavier Garcia, Vedant Misra, Kevin Robinson, Liam Fedus, Denny Zhou, Daphne Ippolito, David Luan, Hyeontaek Lim, Barret Zoph, Alexander Spiridonov, Ryan Sepassi, David Dohan, Shivani Agrawal, Mark Omernick, Andrew M. Dai, Thanumalayan Sankaranarayana Pillai, Marie Pellat, Aitor Lewkowyzz, Erica Moreira, Rewon Child, Oleksandr Polozov, Katherine Lee, Zongwei Zhou, Xuezhi Wang, Brennan Sacta, Mark Diaz, Orhan Firat, Michele Catasta, Jason Wei, Kathy Meier-Hellstern, Douglas Eck, Jeff Dean, Slav Petrov, and Noah Fiedel. Palm: Scaling language modeling with pathways. J. Mach. Learn. Res., 24:244: 1-240: 2023.
[13] Christopher Clark, Kenton Lee, Ming- Wing-Wei Chang. Tom Kwiatkowski, Michael Collins, and Kristina Toutanova. Boolq: Exploring the surprising difficulty of natural yes/no questions. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume (Long and Short Pupers),
pages 2924-29366 2936, 2019.
[14] Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le Noac"h, Haonan Li, Kyle McDonell, Niklas Muennighoff, Chris Ociepa, Jason Phang, Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang Sutawika, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. A framework for few-shot language model evaluation, 07 2024.
[15] Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedfor- ward neural networks. In Proceedings of the thirteenth international conference on artificial intelligence and statistics, pages 249 1699256..11 JMLR Workshop and Conference Proceedings, 2010.
[16] Insu Han, R Jayaram, A Karbasi, V Mirrokno, D Woodruff, and A Zandieh. Hyperatten- tion: Long-context attention in near-linear time. In International Conference on Learning Representations. International Conference on Learning Representations, 2024.
[17] Jiaao He and Jidong Zhai. Fastdecode: High-throughput gpu-efficient Ilm serving using heterogeneous pipelines. arXiv preprint arXiv:2403. 11421, 2024.
[18] Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding. In 9th International Conference on Learning Representations. ICLR 2021, Virtual Event, Austria, May 3-7, 2021, 2021.
[19] Coleman Hooper, Sehoon Kim, Hiva Mohammadzadeh, Michael W Mahoney, Yakun Sophia Shao, Kurt Keutzer, and Amir Gholami. Kvquant: Towards 10 million context length Ilm inference with kv cache quantization. arXiv preprint arXiv:2401.18073 2024.
[20] Edward J. Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models. In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022, 2022.
[21] Jingcheng Hu, Houyi Li. Yinmin Zhang, Zili Wang, Shuigeng Zhou Xiangyu Zhang, and Heung-Yeus Shum. Multi-matrix factorization attention. arXiv preprint 22.:2242..1955, 2024.
[22] Ting Jiang, Shaohan Huang, Shengyue Luo, Zihan Zhang, Haizhen Huang, Furu Wei. Weiwei Deng, Feng Sun, Qi Zhang, Deqing Wang, et al. Mora: High-rank updating for parameter- efficient fine-tuning, arXiv preprint arXiv:2405. 12130, 2024.
[23] Shengyu Liu Jiashi Li. Flashmla: Efficient mla decoding kernels. https://github. com/ aaapseekkaiiFFashhM... 2025.
[24] Andrej Karpathy. NanoGPT.
wwtphhab.ccaaaaappattaaaaaooo,. 2022.
[25] Angelos Katharopoulos, Apoorv Vyas. Nikolaos Pappas, and François Fleuret. Transformers are rnns: Fast autoregressive transformers with linear attention. In International conference on machine learning, pages 5156-5165, PMLR, 2020.
12

---

[26] Wonbeom Lee, Jungi Lee, Junghwan Seo, and Jaewoong Sim. {InfiniGen]: Efficient generative inference of large language models with dynamic (KV) cache management. 18th USENIX Symposium on Operating Systems Design and Implementation (OSDI 24), pages 155-172, 2024.
[27] Xiaoyu Li, Yingyu Liang, Zhenmei Shi, and Zhao Song. A tighter complexity analysis of sparsegpt. arXiv preprint arXiv:2408. 12151, 2024.
[28] Vladislav Lialin, Sherin Muckatira, Namrata Shivagunde, and Anna Rumshisky. Relora: High- rank training through low-rank updates. In The Twelfth International Conference on Learning Representations, 2023.
[29] Yan-Shuo Liang and Wu-Jun Li. Inflora: Interference-free low-rank adaptation for continual learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Puttern Recognition, pages
2024.
[30] Yingyu Liang, Heshan Liu, Zhenmei Shi, Zhao Song, Zhuoyan Xu, and Junze Yin. Conv-basis: new paradigm for efficient attention inference and gradient computation in transformers. arXiv preprini arXiv:2405.055219 20224.
[31] Yingyu Liang, Jiangxuan Long, Zhenmei Shi, Zhao Song, and Yufa Zhou. Beyond linear ap- proximations: A novel pruning approach for attention matrix, arXiv preprint 1r11
2024.
[32] Aixin Liu, Bei Feng, Bin Wang, Bingxuan Wang, Bo Liu, Chenggang Zhao, Chengqi Dengr, Chong Ruan, Damai Dai, Daya Guo, al. Deepseek-v2: A strong, economical, and efficient mixture-of-experts language model. arXiv preprint arXiv:2405.0444 04434, 2024.
[33] Aixin Liu, Bei Feng, Bing Xue, Bingxuan Wang, Bochao Wu, Chengda Lu, Chenggang Zhao, Chengqi Deng, Chenyu Zhang, Chong Ruan, et al. Deepseek-v3 technical report. arXiv preprint arXiv:2412.19433 2024.
[34] Zirui Liu, Jiayi Yuan, Hongye Jin, Shaochen Zhong, Zhaozhuo Xu, Vladimir Braverman, Beidi Chen, and Xia Hu. KIVI: A tuning-free free asymmetric 2bit quantization for KV cache. In Forty- first International Conference on Machine Learning, ICML 2024, Vienna, Austria, July 21-27, 2024, 2024.
[35] I Loshchilov. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05101, 2017.
[36] Ilya Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts. arXiv preprint arXiv: 1608v 03983, 2016.
[37] Anton Lozhkov, Loubna Ben Allal, Leandro von Werra, and Thomas Wolf, Fineweb-edu: the finest collection of educational content, 2024.
[38] Sadhika Malladi, Alexander Wettig, Dingli Yu, Danqi Chen, and Sanjeev Arora. kernel-based view of language model fine-tuning. In International Conference on Machine Learning, pages 23610-23641 PMLR, 2023.
[39] Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. Can suit of armor conduct electricity? a new dataset for open book question answering. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, pages 2381-2391, 2018.
[40] Myle Ott, Sergey Edunov, Alexei Baevski, Angela Fan, Sam Gross, Nathan Ng. David Grangier, and Michael Auli, fairseq: A fast, extensible toolkit for sequence modeling. In Waleed Ammar, Annie Louis, and Nasrin Mostafazadeh, editors, Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, NAACL-ILLT 2019, Minneupolis, MN, USA, June 2-7, 2019, Demonstrations, pages 48-53. Association for Computational Linguistics, 2019.
[41] Weijieying Ren, Xinlong Li, Lei Wang, Tianxiang Zhao, and Wei Qin. Analyzing and reducing catastrophic forgetting in parameter efficient tuning. arXiv preprint arXiv:2402.18866 2024.
[42] Luka Ribar, Ivan Chelombiev, Luke Hudlass-Galley, Charlie Blake, Carlo Luschi, and Douglas Orr. Sparq attention: Bandwidth- seffcidt LLM inference. Forty-first International Conference on Machine Learning, ICML 2024, Vienna, Austria, July 21-27, 2024, 2024.
13

---

[43] Keisuke Sakaguchi, Ronan Le Bras. Chandra Bhagavatula, and Yejin Choi. Winogrande: An adversarial winograd schema challenge at scale. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pages 4732-87400 Association for the Advancement of Artificial Intelligence (AAAI), 2020.
[44] Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. Linear transformers are secretly fast weight programmers. In International Conference on Machine Learning, pages 9355-9366 PMLR, 2021.
[45] Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, and Tri Dao. Flashattention-3: Fast and accurate attention with asynchrony and low-precision. Advances in Neural Information Processing Systems, 37:68658 68685, 2024.
[46] Noam Shazeer. Fast transformer decoding: One write-head all you need. arXiv preprint arXiv:1911.02150, 2019.
[47] Noam Shazeer. Glu variants improve transformer. arXiv preprint arXiv:2002.05002, 2020.
[48] Yiming Shi, Jiwei Wei, Yujia Wu, Ran Ran, Chengwei Sun, Shiyuan He, and Yang Yang. Loldu: Low-rank adaptation via lower-diag-upper decomposition for parameter-efficient fine-tuning. arXiv preprint arXiv:2410. 13618, 2024.
[49] Zhenmei Shi, Jiefeng Chen, Kunyang Li, Jayaram Raghuram, Xi Wu, Yingyu Liang, and Somesh Jha. The trade-off between universality and label efficiency of representations from contrastive learning. In The Eleventh International Conference on Learning Representations, 2023.
[50] Prajwal Singhania, Siddharth Singh, Shwai He, Soheil Feizi, and Abhinav Bhatele. Loki: Low-rank keys for efficient sparse attention. arXiv preprint 24X1v:2406.02542, 2024.
[51] Jianlin Su. The extreme pull between cache and effect: From MHA, MQA, GQA to MLA. https://spacees ac. cn/ chieeer May 2024.
[52] Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. Neurocomputingg 568:127063, 2024.
[53] Hanshi Sun, Len Chang, Wenlei Bao, Size Zheng, Ningxin Zheng, Xin Liu, Harry Dong, Yuejie Chi, and Beidi Chen. Shadowky: Kv cache in shadows for high-throughput long-context Ilm inference, arXiv preprint arXiv:2410.21465, 2024.
[54] Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma, Yuqing Xia, Jilong Xue, Jianyong Wang, and Furu Wei. Retentive network: A successor to transformer for large language models. arXiv preprint arXiv:2307.08621, 2023.
[55] Jiaming Tang, Yilong Zhao, Kan Zhu, Guangxuan Xiao, Baris Kasikei, and Song Han. QUEST: query-aware sparsity for efficient long-context LLM inference. In Forty-firs( International Conference on Machine Learning, ICML 2024, Vienna, Austria, July 21-27, 2024, 2024.
[56] Gemma Team, Thomas Mesnard, Cassidy Hardin, Robert Dadashi, Surya Bhupatiraju, Shreya Pathak, Laurent Sifre, Morgane Rivière, Mihir Sanjay Kale, Juliette Love, al. Gemma: Open models based on gemini research and technology. arXiv preprint arXiv:2403.08295 2024.
[57] Philippe Tillet, HT Kung, and David Cox. Triton: An intermediate language and compiler for tiled neural network computations. In ACM SIGPLAN International Workshop Machine Learning and Progrumming Languages co-located with PLDI. Association for Computing Machinery, 2019,
[58] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timo- thée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302. 2023.
14

---

[59] Yao-Hung Hubert Tsai, Shaojie Bai, Makoto Yamada, Louis-Philipee Morency, and Ruslan Salakhutdinov. Transformer dissection: An unified understanding for transformer's 's attention via the lens of kernel. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNIPP) pages 4344 4353, 2019.
[60] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[61] Guangxuan Xiao, Ji Lin, Mickael Seznec, Hao Wu, Julien Demouth, and Song Han. Smoothquant: Accurate and efficient post-training quantization for large language models. In International Conference on Machine Learning, pages 38087-38099, PMLR, 2023.
[62] Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. In The Twelfth International Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024, 2024.
[63] Vikas Yadav, Steven Bethard, and Mihai Surdeanu. Quick and (not so) dirty: Unsupervised selection of justification sentences for multi-hop question answering. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP), pages
2019.
[64] Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can a machine really finish your sentence? In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics. Association for Computational Linguistics, 2019.
[65] Yuchen Zeng and Kangwook Lee. The expressive power of low-rank adaptation. In The Twelfth International Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024, 2024.
[66] Hengyu Zhang. Sinklora: Enhanced efficiency and chat capabilities for long-context large language models. arXiv preprint arXiv:2406.0m 056788 2024.
[67] Michael Zhang, Kush Bhatia, Hermann Kumbong, and Christopher Re The hedgehog & the porcupine: Expressive linear attentions with softmax mimicry. In The Twelfth International Conference Learning Representations, 2024.
[68] Qingru Zhang, Minshuo Chen, Alexander Bukharin, Pengcheng He, Yu Cheng, Weizhu Chen, and Tuo Zhao. Adaptive budget allocation for parameter-efficient fine-tuning, In The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.net, 2023.
[69] Ruiqi Zhang, Spencer Frei, and Peter L Bartlett. Trained transformers learn linear models in-context. arXiv preprint urXiv:2306.099: 09927, 2023.
[70] Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Ré, Clark Barrett, et al. H2o: Heavy-hitter oracle for efficient generative inference of large language models. Advances in Neural Information Processing Systems, 114666611171... 10. 2023.
[71] Hongbo Zhao, Bolin Ni, Junsong Fan, Yuxi Wang, Yuntao Chen, Gaofeng Meng, and Zhaoxiang Zhang. Continual forgetting for pre-trained vision models. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pages 28631-28642, 2024.
15

---

Appendix
A Toward Faster Computation Without Materializing Q, K and V
17
A.1
Direct computation factor space
18
A.2
Complexity: materialized vs. specialized computation
18
A.3 Complexity of the specialized path
19
A.4
Inference-time decoding cost across mechanisms
19
B More FlashTPA Decoding Algorithm
24
B.I Triton FlashTPA Decoding Kernel
25
B.2 Additional Experimental Results
25
C Higher-Order Tensor Product Attention
28
C.1 RoPE Compatibility in Higher-Order TPA
28
Proofs of Theorems
29
D.I Proof of Theorem 3.1
D.2 Proof of Theorem C.I
30
More Related Work
32
More on Attention Mechanisms
32
FI1
Multi-Query Attention (MQA)
32
F.2
Grouped Query Attention (GQA)
33
F3
Multi-head Latent Attention (MLA)
33
F44
Multi-matrx Factorization Attention (MFA)
34
F.5
Rotary Position Embedding (RoPE)
34
More on TPA
35
H More on Experiments
36
H.l Experimental Settings
36
H.2 Additional Experimental Results
36
H.3 Ablation Studies on Learning Rates
39
Broader Impacts and I Limitations
39
16

---

A Toward Faster Computation Without Materializing Q, K and V
Our objective in this section is to compute attention without explicitly forming Q. K, V, by con- tracting their factorized representations in cache- and throughput-frient frienuly order. Recall from Equation (3.1) that each per-token slice Q₁, K,, V. E Rhxd2 is sum of rank-1 outer products. Unless otherwise stated
the per-factor normalizations vo=1/Bq- 8 =11/ 1v
We make the batch//umeerrea//rak/aallle dimensions explicit and introduce the shorthands D == de and E == dy (typically D-D):
AQ RR Bq E BxT,xxxX
Ak BR By E B
Av ERRXTTNNU
By B
Indices b, q, ki r,,,, 8, u, dl, denote batch, query position, key position, head, query-rank, key-rank, value-rank, feature (D), and value feature (E). We write T := Ty = & for full-sequence attention; in decoding, --=1 and we denote the cache length by M =Tk
Convention. For single token, the main text defines (x,)) E RR. and (x, R.R.D., with Q: (x,) (x,). Accordingly, throughout this appendix we index Aq as Aq q, g, (rank-major). Some implementations may store A, transposed as (H x R.) for memory layout, this equivalent, since all uses contract over the rank index.
High-level idea. We first compute head-shared feature-space dot products between BQ and BR, then mix them with head-specific Aq, AK obtain logits, apply the masked softmax, and finally aggregate values via Av, By. This ordering avoids materializing any T, x h x D queries/keys/valuss.
Phase Attention Score Computation
Bg
Ed
Bx
F2
Softmax
1/vRC nver key tokenx is
incl. *g.*x
Aq
Ax
Phuse 2: Value Aggregation
WAY
Li
sv
Figure 6: Specialized TPA computation without materializing Q K, V. Phase 1 (top): compute head-shared feature-space dot products Píb, q, k. r, =B= (Bqlb. g. r. Bx (b, k, s.:) and mix them with head-specific factors Aq, Ag to obtain logits Lbb, h, q, kj. Phase (bottom): apply the causal/padding mask and softmax to get adl, h q, k). then aggregate values via Bv. Scalings $Q, SK, sv and 1/ / are folded into the corresponding phases. Dropout is omitted for clarity. Batch B, heads H, ranks Ro, Rx, Ry, and feature dims D, E are indicated in the nodes.
17

---

A.1
Direct computation in factor space
Single head. For fixed head e [H] and token indices (q, using s0=1/Rq and sx= 1011/R Rx have
Rx
Q Q (K(N)
2.(r)
(Xx) akh (xx) (xx) (x((tr)(x)))
k.(o)
(xil)) (A.1)
q.k
Ro RK
q.h
and for values (with sv sy=1/Ry), V
Σ Dx= akk
By Vi(w) (xx) b bV((H) (xx).. The per- per-head attention
Rv
output query position y is then Σx softmax J Q (1) (K( )a:)) V
Multi-head with head-shared feature dot-products. Define head-shared feature-space dot products 208K,
Plb, q, k s] = q, r, Bxlb, (b, k s.:).
With S VD
we compute
Rg Rk
hl,,,, = 2 Aol r, k, kl,,q.
(A.2)
alb, 11 = Softmaxx (Clb, h q,
mask (b, k)),
Th
Ry
0 1 = 8V
alb h g. Av k, u,h By Bv k, e]
(A.3)
k=1==1
Here mask(b, q: k] E -x0 is an additive mask in logit space that enforces causality and padding. Eqs. (A.2)-(A.3) make explicit that (i) feature-space dot products P are head-shared, and (ii) the rank normalizations 8Q, 8K; 8y can be absorbed into the corresponding factor tensors (or into the scalar prefactors) without changing the computed attention output.
A.2 Complexity: materialized vs. specialized computation
We compare two execution strategies. (i) Natve/materialized:: form Q K, V explicitly and call standard kernels. (ii) Specialized: compute via Eq. (A.3) using head-shared feature-dot products and per-head rank contractions.
Standard MHA (baseline). Ignoring projections, full-sequence attention uses O ((HHTDD) FLOPS for scores and O(BHT=D) for value aggregation, i.e., Jmha 20BHTTD)) D).
TPA (materialized). Forming Q, K, V from factors costs O(BTHD RQ Rk Rk)) after the linear projections; subsequent attention uses the same 2 (0(HTTD MHA.
TPA (specialized).
Using Eqs, (A.2)-(A.3) and writing T--T the dominant FLOPs are
O RoRk D) + O(BHT RO + By
dots
per-head combine
value aggregation
Compared to FMHA, the specialized path reduces FLOPs whenever
Ro D + H RQRK + H Ry E < 2HD..
(A.4)
Dividing by HD yields (Ro (H) + (RoRx/Di + Ry RyyyyyE/en (E/D) 2. For E=D and small ranks (e.g., R the inequality holds for typical H, D 2 and the benefit grows with larger H or D.
Memory traffic and peak working set. For full-sequence attention the naive path streams Q, K, V of size O(BTHD) each. The specialized path streams factors only and needs the head-shared P tiles of size @ Tule qulle RoRx) plus per-head tiles for rank combine/value aggregation. In decoding with cache length M, the factorized KV cache uses (Rx+ + Rvy (h++(D))) numbers per token (cf. Section 3.3), 2hD for MHA; this reduction directly lowers memory bandwidth pressure.
18

---

Algorithm Specialized TPA (no explicit Q K, V: causal)
Require: AQ E B x BQ E x
Require: Aς E RBxx x Bk E R XRxD
Require: Ay E B
By E BBXTxXRXXE
Require: scales sq=1/Ro Sk=1/ Rg. sv= Ry; mask mask E {0, (0,-000BxTxxx0
Ensure:
x
2: [ (sQSk/ (ogsk/vD)= - BQ, Bk)
I: P f ""ogmmmm"budd.bbd->bbpop
DE R R
Aq. Ak, P)
3: C + + broadcasttmaskk
> causal/padding mask
4: a Softmax+
E RBXHXI, online/LSE in practice
5: O t sy. cinsumb
a, Av. Bv)
6: return transposc(O, "bhqe" "bqhe")
A.3 Complexity of the specialized path
Combining the terms above gives complexity FTPA-spo: = + R + O(BHT*Ry E), with the speed condition Eq. (A.4).
For single query (I,=1) against cache of length M, the specialized FLOPs are
{ M RO RK D) + M(BH RO + @(BH Ry E),
while MHA uses 2 (0HHMDD) This matches the asymptotics embodied in FLASHTPA (Section 5) and explains the regimes where & and =x=Ry [1,2] yield the largest gains.
We apply the causal mask before softmax and use an online log-sum-exp update for numerical stability (as in FlashAttention). The intermediate P E BBBT, xT, x R4 xRk is evaluated blockwise in T3 to keep peak memory linear in the block size; the same blocking naturally fuses with the masked softmax and the value aggregation step.
The constants SQ, SK, sy can be absorbed into either A()) or B()) at training time. We expose them explicitly only to make Eq. (A.4) transparent; The choice has no effect on softmax invariance or gradients.
The Triton kernel in Section implements the blocked computation of P. the masked online softmax over k and the fused value aggregation, mirroring Algorithm 1. This avoids creating any Q, K, or full T2xTT temporaries beyond working tiles.
Compared with 2 ((BHT2)) D) for MHA, the specialized path improves with small (Ro, Rx. Rv) and benefits further from pre-rotating Bx for RoPE (cf. Section 3.2), which removes per-step rotations decoding. Practical speed also depends on tiling, memory bandwidth, and kernel fusion; our measured gains in Section 6.2 align with the regime predicted by Eq. (A.4).
A.4 Inference-tie time decoding cost across mechanisms
In autoregressive decoding, we generate the output for the current token Xy given cached keys and values from T-1 previous tokens. We analyze the FLOPs for computing the attention output for this single query token and use M for the current cache length. For all mechanisms, we analyze the total Floating Point Operations (FLOPs) and the number of parameters in the attention layer, including the cost of projecting the current token's 's hidden state into its respective Query, Key, and Value representations. The parameter count formulas are taken from Table I.
For Multi-Head Attention (MHA), with H/ query heads and H distinet Key/Value heads, the complexity is determined by the dot-product attention and value aggregation steps.
Projection: Projecting to get query, key, and value vector for each of the H heads costs (dmmmte
Attention: Dot products and value aggregation over a cache of length cost ((2M (ignoring softmax constants).
Total MHA: The complexity is O(dmods + 2dH)
Multi-Query Attention (MQA) uses H query heads but shares single Key/Value head (Hgv 1). The arithmetic complexity remains the same MHA for the same number of query heads.
19

---

Projection: Projecting for H query heads and I shared K/V head costs (ddmode + 2dz)).
Attention: The interaction with the cache costs ( Hd,)..
Total MQA: The complexity is (ddmmdddda ( H + 2) + 2M AMH))..
Grouped- Querr Attention (GQA) uses H query heads and G Key/Value head groups (Hkv G). The arithmetic complexity is also identical MIA..
Projection:
Projecting for H query heads and G K/V head groups costs (cdmmddel (ddh + 2Gd,)).
Attention: The interaction with the cache costs (
. Total GQA: The complexity is (dmmaaddd, (H + 2G' + 2MHd,))
MQA and GQA significantly reduce the KV cache size and memory bandwidth compared to MHA. While the arithmetic FI.OP count for the core attention computation (dot products and weighted sums) is 2M Hd, for all three (for fixed H, d practical speedups for MQA/GQA arise from improved memory locality due to smaller K/V caches.
Multi-Head Latent Attention (MLA), as described Appendix F.3, uses H heads. Each head's (up-projected) query/key vectors have dimension d,, + dz. During decoding, however, the score computation against the cache can be decomposed into (i) a dot product in the cached latent space R&d, for the content part and (ii) additional RoPE dot product in Rd for the positional part. Concretely, MLA caches c E idd per past token 5, aggregates values in Rd, and then up-projects once per step.
Cached state: MLA caches the compressed KV latent kV Rd and the shared RoPE key component k.¹ e R per past token s.
Projection (current token): Computing the query latents and the new cache entry (up to constant factors) costs
(dmoderd + ddH((d + d + (mnoaa + +d d.h)).
corresponding to forming c QC, Q and computing/storing cV and kt for the current token. Attention (cache interaction): Using the identity q.TkC, (WUK q(,) Tcky the score against
each cached token can be computed via dot product in de plus the RoPE dot product in æ The latent value can be aggregated in Idd and then up-projected once. The dominant cache-dependent cost is
(MH(2dd + dfR))
up lower-order per-step terms such as (Hd,,dn)
Total MLA: (damate!d" + d' (dd, + ) + (dm + df. + +
TPA. We use the FlashTPA Decoding algorithm (Algorithm 2) for FLOPs analysis, with N = 1 query
token, M cached items, D as feature dimension for Bo/b# b (typically dz), and E for by (typically d,). For ranks (Rq. Rk. Rv):
Projection: Projecting the current token
XT
all
Q/K/V
factors
costs
O (dmmded +xx (H)))Hdd)))
Attention (cache interaction): Using Algorithm 2 with cache length M, the dominant cache- dependent FLOPs are
(M (RoRxD + HRQRR + HRy
up lower-order terms (masking/clementtwol) wise products and online-softmax bookkeeping).
Total for TPA decoding: @ (Rq Rx +dx) + M(RRRKD -RRRRK HRv
Example Comparison
We compare the total Floating Point Operations (FLOPs) required process single token during autoregressive inference. This analysis separates the initial, constant projection cost from the attention cost, which scales linearly with the cache length M.
The following parameters used for the comparison:
Model Dimension: dmodel 2048
Heads: H = 32
Head Dimension: d, 64 (so D = -== dr.)
20

---

GQA Groups: G = 4
MLA Dimensions: d, = 256, dff = 32, and d' = 768
MHA (16.8M parameters):
Parameters = model Hd,, 2048 (32- 64) 16.8 10
Projection = 3 dmodel H dh. = 1--0048--0 32 64 ~ 12.6 x 10
Attention 2- M H dg. 4096M
GQA (G = 4, 9.4M parameters):
Parameters = dmodal d (2H + 2G) 2048 64
2- ~ 9.4 x 10
Projection = dmmoeel H + 2G),,, = 2048 (82+8) -64 64
5.2 x 10
Attention = M H d, 4096.M
MLA (9.8M parameters):
Parameters 768(2048 + 2048 + 1024) + 2048(32 + 2048) + 256(2048 + 4096) P 9.8 x 10
Projection ~ dmodald" + d. (d,, + d(,)) + dmodel (d. + d + H dodh
2048 768 + 768 .32 (64 32) 2048 (256 32) 32 256 64
5.0 106
Attention M H (2dg + d = M 32-32-12 (512 + 32) = 17408M
TPA (Rq- 16, Rx Ry 7.7M parameters):
Parameters d 116+1+ + (H + dr)) + denoteel Hd,, (048(18)()6)) 20482 7.7 x 10 Projection = dmood [16 +1+ L) + dn) = 2048 (18) (96) = 3.5 x 10c
Attention M [1(1536) + 1(2048)] 3584M
TPA (Rq = 16, Rx = 2, Rly = 2, 8.1M parameters):
Parameters = dmaadee (16 2 (H + d,)) + dmnotel Hd,, = (048(20)(96) + 2048 ~ 8.1 x 106
Projection dmodel + (H d,)) 2048 (20) (96) == 3.9 x 1
Attention = (2(1536) + 2(2048) 7168M
TPA (Rq- 8, Rk Rv 1, 6.2M parameters):
Parameters dmncdel (8 + H) dx) + dmwste 2048(10)(96) + 2048 ~ 6.2 x 106
Projection dmooal (8 + () dr) = 2048 (10) (96) ~ 2.0 x 108
Attention M [1(768) + 1(2048)] 2816M
TPA (Ro = 8, Rx = 2, Rv = 2, 6.6M parameters):
Parameters dmustal (8 + 2) + d₂) + dmwsee Hd, = 2048(12)(96) + 2048 ~ 6.6 x 106
Projection dmooaa (8 + (H + dr.) 2048 (12) (96) ~ 2.4 x 108
21

---

Attention M (2(668)) + 2(2048) 5632M
The analysis shows that TPA with low ranks offers a favorable trade-off. Reducing the query rank (Rq) from 16 to 8 further decreases both the projection and attention costs, making the TPA (RQ=8, = Ry= configuration the most computationally efficient in this comparison. Increasing key/value ranks (e.g., to Rx=2, Ry=2) raises the attention cost linearly, remaining competitive with MHA for sufficiently long contexts where kernel fusion and blocking amortize memory traffic.
Example Comparison II.
We now repeat the analysis for larger model configuration to observe how these trade-offs scale.
The following parameters for larger model are used for this comparison:
Model Dimension: dmnddll = 4096
Heads: 32
Head Dimension: d₁ 128 (so D E d4,)
GQA Groups: G =
MLA Dimensions: d.- 512, dh = 64, and d' = 1536
MHA (67.1M parameters):
Parameters 4dmode Hd, = 4 40962 67.1 x 10
Projection 4096 32 128 = 50.3 x 10
Attention 2 . 128 8192M
GQA (G 4, 37.7M parameters):
Parameters = dm (2H + 2G) = 4096 128 (2-32 32 + 2-4 37.7 x 10
Projection 4096 (32 + 8) 128 R 21.0 x 105
Attention 2. .M 128 8192M
MLA (39.1M parameters):
Parameters = 1536(4096 + 4096 + 2048) + 4096 4096) 512(4096 + 8192) x 39.1 x 10
Projection z d'ooerd + H. (d,, + dR) + dmoddl (de d,) + Hd.dA
4096 1536 + 1536
(128 64) + 4096 (512 64) 32 512 128
20.2 x 10
Attention = M 32 (1024 + 64) 34816M
TPA (Rq = 16, Rx = 1, Ry = 1, 28.6M parameters):
Parameters 4096(18)(16)))) (160) + 40962 x= 28.6 x 106
Projection 4096 (16 + 1 1) (32 + 128) ~ 11.8 x 106
Attention (1(660)) + (44096) 6656M
TPA (Rq 16, Rx Ry 2, 29.9M parameters):
Parameters 4096(20) (160) + 4096 x 29.9 x 10
Projection 4096 (16 + 2 + - 002 128) & 13.1 x 10
Attention M (2(2560) + 2(4096) 13312M
TPA (Rq = 8, Rk 1, Ry 1, 23.3M parameters):
Parameters = 4096( 10)() + 40962 ~ 23.3 x 1
Projection 4096 (8 + +++1)) (32 + 128) ~ 6.6 1
22

---

Attention = . (1((280) + 144096) = 5376M
TPA (Rq 8, Rx 2, Rv 2, 24.6M parameters):
Parameters = (1)6( + 40962 = 24.6 x 10
Projection = 4096 (8 + + (32 + 128) R: 7.9 x 10
Attention .M [2(1280) 2(4096) = 10752M
For this larger configuration, TPA (Ro=8, Rx=1, 1y= Ry=1, efficiency, with the lowest projection and attention costs. This highlights the value of tuning TPA ranks balance expressiveness against compute.
remains the clear leader in computational
Example Comparison III.
Then we analyze a very large model configuration (e.g. MoE model with 1~2T parameters) to examine the scaling properties of each architecture, where dmosel * H dg to align MLA with other attention mechanisms. We also denote the number of parameters in the attention part for each layer.
The following parameters are used for this comparison:
Model Dimension: dmodel = 7168
Heads: H = 64
Head Dimension: d₃ = 128 (so D = E dl,))
GQA Groups: 8
MLA Dimensions: de 512, d 64, and d' 1536
MHA (235M parameters):
Parameters 4dmotel H d, = 4 7168 8192 235 10
Projection 3 7168 64 128 176.2 x 10
Attention 2 .M 64 128 16384M
GQA (G 8, 132M parameters):
Parameters dm (2Hl 2G) 7168 128 (2 64 2- (8) 132 10
Projection 7168 (64 + 16) 128 x 73.4 x 10
Attention = 2 M 64 128 16384M
MLA (101M parameters):
Parameters 1536(7168 + 8192 4096) 7168 + 512(7168 + 16384) 101 1
Projection dmoeadd Hd + d + dde di) Hd.d
7168 1536 + 1536 64 (128 + 64) + 7168 (512 + 64) - 64 512 128
= 38.2 x 10t
Attention M 64 (2 512 + 64) 69632M
TPA (Rq = 16, Rx = 1, Ry = 1, 83M parameters):
Parameters 7168(18) + 7168 8192 83.5 x 10
Projection 7168 (16 + 1) (64 128) 24.8 x 106
Attention M (1((2))) + 1(8192) 11264M
TPA (Rq = 16, Rx = 2, Ry = 2, 86.2M parameters):
Parameters 7168(20) (192) + 7168 8192 86.2 8
Projection 7168 (16+22 + 2) (64 128) R 27.5 x 106
Attention M (2(3072) + 2(8192) = 22528M
23

---

TPA (Rq = Rx = 1, Rv = 1, 72.3M parameters);
Parameters = (16)())))192)) + 7168 8192 x 72.5 x 10
Projection 7168 (8 + 1 (64 128) = 13.8 x 10
Attention M (1(1536) + 1(8192)) 9728M
TPA (Rq- 8, Rx 2, Ry 2, 75.2M parameters):
Parameters (168(12)(192)) + 7168 8192 ~ 75.2 x 10
Projection 7168 (8 + 2 (64 128) 16.5 10
Attention M [2(1536) + 2(8192)) = 19456M
At this very large scale, the cost of MHA projections becomes prohibitive. While MLA's projection cost can be competitive, its attention cost scales with (2d.+d and exceeds MHA for long sequences. TPA with low ranks (Rq=8, =1 =y=1) yields the lowest attention cost and a substantially smaller projection cost, strengthening its advantage as model size increases.
B More on Flash Decoding Algorithm
In this section, we present FlashTPA for decoding in hardware-friendly, numerically stable form and extend to general key/value ranks RK Rl 2 1. The algorithm computes attention without materializing Q, K, V or the full NxM attention matrix, by (i) forming head-shared feature-space dot products, (ii) mixing them with head-specifie factors to obtain logits as Eq. (A.2), and (iii) aggregating values in Eq. (A.3) in single online softmax pass.
Notation and shapes. We allow N query positions but decoding uses = Let B be batch, M the cache length, H heads, RQ, Rk, Ry ranks, and D, E feature sizes (typically D=E=d,). Inputs: Beache Rx Rx
Aq E R Bx
A cache ER Bx
g ER Bx Baache B
We use scalings sq=1/Rq. 9x=1/RK: Rk. sy=l/Ry, and 1wm=1/VD. Let mask E {0, [0.-o)]BxNxX> encode causality/padding. If RoPE pre-rotation is used (Section 3.2), Be BxNNM already includes positional phases; otherwise apply RoPE Bバ on load.
Algorithm 2 FlashTPA Decoding (general RK, Ry, masked, online-
Require: Aq. Bq, A gache Brache A cache Beache, mask; 8Q, 8K, 8V, Stota
Ensure: O Rx
1: Initialize y 0BxHxXXE S 0xHHN, m (-o0)BxHXN
accumulates
Σ cxxp(-); log-sum-exp is logs + m
2: for each cache block m:m+ Ams M do
3: Load B E B A.BxxmxXR., xD
Load A Ay.ik By AKM BBAAAMXRxX
4:
B E RB maskilic BRBxxxAm
5:
(1) Head-shared feature dots: P
einsum ("bard. BQ: Bx,b))
R BB Amx Ra x Rg
6:
(2) Per-head rank mixing to logits:
7:
Lank (Mum) einsum ("wnrr.bbshh....te.e.mmmm bww....hhbmssssbmntt Aq AK,NK, P)
Bxxxxxxmm
8:
Lask < Lalk + broadcast(masse
9:
(3) Online softmax update (no a materialization):
10:
mylk + maxx,(Cnn..)) pblk + exp( mbb); Sbik +- Σm Pblk
(4) Block value aggregation (fused over m, u):
12:
Yblk & einsum ("bhnm,bmon amuhhbmue-bbhet
Pblk;
Av.bk, BUNK))
D BBBxxXXX
(5) Fuse blocks with log-sum-exp:
14:
mmew &- max m, mok);
+ exp(m [ None Y + exp(maik
new I..., None ① Ybkk
15:
exp(m- Maew) O
mmew
O Sbiki m +- Macw
end for
17: return O <- 8V Ho..,onne permuted to (B, N, H, F)
24

---

Step (1)-(2) implements Eq. (A.2); step (4)-() implements Eq. (A.3) while fusing the masked softmax with value aggregation via online log-sum-exp (as in FlashAtrention), thereby avoiding any materialization. When =x=Ry== = the contractions reduce to the simpler einsums in Figure 2.
Complexity and working set. Per block of Am cache items, the dominant FLOPs are
( N Am D) + N Am RQ + ( H Am Ry F),
matching the specialized analysis in Appendix A.2 and the decoding bounds in Appendix A.4. Peak memory scales with tiles of Bx: Ak, Av, By and the small temporaries and ViRki neither Q,K, V nor the full N X. attention matrix is formed.
RoPE and masking. If keys are pre-rotated (Eq, (3.5)), Bx needs no decoding-tine rotation. Otherwise apply RoPE to B biikk row-wise before step (1). The mask mask (zeros or -oo) is added to logits in step (2) and supports both causal and padding masks.
B.1 Triton FLashTPA Decoding Kernel
We implement the experiments using Triton [57]; Algorithm 3 sketches the kernel corresponding Algorithm 2. The provided kernel outline specializes the frequently used case Rx=Ryy= general ranks follow by tiling over Rx, Rl and replacing the rank-1 vector-matrix products with the corresponding small GEMMs in steps S1/S2 and the value mixing path.
B.2
Additional Experimental Results
The following figures present additional speed comparisons for different embedding dimensions, with d = 64 maintained. The y-axis represents log, (time) in seconds (lower is faster), and the x-axis represents log, (sequence length).
Detailed Analysis of Figure (Embedding Dimension 2048): Figure in the main paper depicts speed comparisons for an embedding dimension of 2048. The results indicate that FlashTPA (blue line) is highly competitive. Across all tested batch sizes (1 16) for dymodel 2048:
MHA (orange line) is consistently the slowest mechanism, with its decoding time increasing most rapidly with sequence length.
MQA (purple line) and GQA (green line) offer significant speedups over MHA and perform very similarly to each other, often overlapping in the plots.
MLA (blue line) demonstrates strong performance, generally being faster than GQA, particularly
longer sequence lengths.
FlashTPA shows excellent scalability. While at very short sequence lengths (c.g., 21 to 2 its performance is comparable to MQA/GQA and MLA. its decoding time increases at notably slower rate with sequence length. Consequently. FlashTPA becomes significantly faster than GQA for sequences longer than approximately 2
Compared to MLA, FlashTPA is consistently among the top two performers. In many instances, particularly at sequence lengths greater than 2 or 2 FlashTPA matches or slightly surpasses MLA in speed. The logarithmic scale for time suggests that these differences can be substantial in practice for very long contexts. For example, at sequence length of 2 across various batch sizes, FlashTPA often shows visible advantage over MLA.
Figure 7 (Embedding Dimension 3072): With larger embedding dimension of 3072, the relative performance trends observed in Figure 5 largely persist.
FlashTPA (red line) remains of the most efficient decoding methods. MHA (orange line) is consistently the slowest, while MQA (purple line) and GQA (green line) offer considerable improvements over MHA.
MLA (blue line) and FlashTPA are the top two performers. FlashTPA consistently matches exceeds the speed of MLA, particularly at longer sequence lengths (e.g., beyond 2 or 26 depending on the batch size). Its advantage often becomes more pronounced at the longest sequences tested (219). For instance, in batch size 1, TPA is clearly faster than MIA for sequence lengths 216 and above. similar trend is seen across other batch sizes, where TPA maintains competitive edge becomes superior at longer contexts.
25

---

Algorithm 3 Triton FlashTPA Decoding Kernel
Require: Input Tensors: AQ( N, RQ, a H), a (B, (B N, RQ, D),
6t (B, M, , bV (B,M,,)
Require: Scaling factors: Stocal, SQ. sv; Dimensions: B, (= 1), H, RQ, D,E
Require: Kernel Block dims: Вh BR, BD. Bg: Sequence Blocking: Mylecks Mehunk
Require: Program IDs: Pidy Pidy PidM
Ensure: Partial Output O purtial ( B, NumM, N, H.H.Ey E), Log-Sum-ENy LSEpartal B, Num H)
1: 4 Pidnis Matart + Pidu Bи
2: Mistock start Pidze Mblock's Mhloock end (Pidg 1) M M)
3:
b BU, BR, Bn, Be are tile sizes for dimensions H, R, D, E respectively,
4:
> Initialize accumulators for the head block
5: Osccum f ((EBBR); Mmax f (Bx) Sexp_sum
0 BB); Cscale +- Stotal SQ SK
6:
Load query factors (fixed for this program N=1)
7: Load A (Rg (RgxBrl Bu) from Ao(b, 6, O, h.a.....
Q.
8: Load B xDx Ro)
from Bq (b, o,
> Dimensions may be transposed after loading for matmul
Q.local
9:
Iterate over Mohunk- chunks within the K/V block
10:
for Mchunk_stan start from Willock start to Millek end - step Mehunk do
11:
Mechunkeea and f- minlmehhuknaan.. Mehumk start + Mehunk _eld
12:
Meur t- Mehunkeend Mechukk start
13:
Load K/V factors for the current chunk
14:
Load achunk (Meurr chunk, Bu);
a (Meurr chunk; Bh); by / (Meurr [Meurrebuukk> shunkkssskkk D);
by
chunk E, MMeurr _chun)
b Layouts optimized for memory access and matmuls
15:
chunk f- b 8V
16:
D Core TPA Score Calculation for the chunk >- Shape: (Meunr_huunks chunk: R2)
17:
chuunnk f- MatMul UK
chunk? BQ,ocaa))
18:
S2chunk MatMul( S Slehunk AQ,local)
>> Shape: (Mear coart_ccunk- Bu
19:
S'chunk S2chunk chunk Cscale
P> Shape: (Mear charr_chuk. Bh)
20:
D Online Softmax Update for the chunk P Shape: (Bh)
21:
Mmmax Mocal_lcca local f- maxxoris=0 (S33hunk)
22:
Immax + max max, Mmnax local
23:
Prum exp(S3ehunk
Mmmax_eww [None, :])
24:
Sexp_sun hasp_sum_Jcaa local - Lozie=o (Pru))
25:
Pweighted +- (Prun/ Seaup Sesp_sum__ae Joc_l
None, :]) @ a
26:
Ochunk < MatMul(B5) brunk Pweighted_av
> Shape: (E, Bu) Update global (M-block level) accumulators
27:
28:
Sexp_sum sum prev rescaled f- Sexp_sum sum exp (Mmax
Mmax _eaxx__ew)
29:
Sexp_sum t- Sexp_sum poee_ resualed Sexp_sne sum local
30:
ratio +- Sexp sum local / /8xp_summ
D. This is Sexp sum homalocaa/// /Sxp Sexp_sss__meeeeeeeee sempp_sum__ew
31:
Oaccum (1 ratio) Osccum ratio Oehunk
32:
M2max + Mmaxx_nw
33:
end for
Store partial results for this program's (batch, head block, M block)
35: Store Oaccum into O partial (b, Pidar: O. Asturt :..:
LSE,al < (og((Sep_suun)) (0oxp_summ + mmax
37:
Store LSE,a into LSE [b, Pidza Istant...
This suggests that Flas efficiency well-maintained even as the model's embedding dimension increases.
Figure 8 (Embedding Dimension 1024): For a smaller embedding dimension of 1024, similar trends are observed:
26

---

FlashTPA (red line) is highly competitive. MHA (orange line) remains the least performant. MQA (purple line) and GQA (green line) are faster than MHA.
However, as sequence length increases, both MLA (blue line) and FPashTPA demonstrate superior scalability. FlashTPA generally matches outperforms MLA, particularly for sequences longer than 2 For example, with batch size of 16, TPA shows a clear speed advantage over MLA for sequence lengths 2 and greater.
These results across different embedding dimensions highlight the robustness of Plass decoding speed advantages, especially for long sequences where it consistently ranks one of the fastest, if not the fastest, attention mechanisms among those tested.
Suich Sec Envosdiing Dive. 3972
Betch Siee Eneessing Dir: 3072
Bucch Star Entruddingg 3012
RCC
12
(a) Batch Size=1
(b) Batch Size=2
(c) Batch Size=4
Batch Stos: Erbeeding Diny 2072
Bstcc Simt 14 Erbedding Der: 2672
(d) Batch Sizc=8
(c) Batch Sizc=16
Figure 7: Decoding time comparison of different attention mechanisms with embedding dimension of 3072 and d,, = 64.
Barch Soc. Emnoddinngg 1004
Batch Sib Erkkddingg 2024
Banch Skar. Einbedddnng 1024
(a) Batch Size=1
(b) Batch Size=2
(c) Batch Size=4
Batch Suas: 8. Erbeeding Dine 2004
Botecs Sim 15 Embedding Den 2004
Meggh)
(d) Batch Size=8
(c) Batch Sizc=16
Figure 8: Decoding time comparison of different attention mechanisms with embedding dimension of 1024 and d,, = 64.
27

---

C Higher-Order Tensor Product Attention
All prior discussions have focused TPA where the query, key, and value matrices (e.g.. Q: R ) are formed as a sum of Ro components. Each component is an outer product of two context-dependent vectors, one spanning the head dimension (R.A) and the other spanning the feature- 1 Ax(x;) Bx(x)) implies per-head dimension (Rd)), ) detailed in Section 3.1 (e.g.. Q+ Rg
Q -,=E, b, where a, are columns of AQ and b, are rows of Bq). We now generalize this by introducing additional latent factors in the construction of the feature-per pere-eeerddddde vectors, leading to what term higher-order TPA. This approach allows for more complex interactions in forming these feature vectors.
For instance, in a third-order factorization, the query tensor Qe E R for a single token is constructed as:
Ra Q2
RQ Σ a (xx) 8
@ (x:)),
where a& (x1) E Rh The term by (xx) ) Rd and the newly introduced factor c9 (x1) E Rd first form a matrix by (x+) @ c (x.) E Xod via outer product (as defined in Section 2). This matrix then vectorized by veo(.) into column vector of dimension d, = dgde. The final query Q. is formed by the sum of outer products between a, (x,) and these resulting d, di-dimensional vectors.
Analogous expansions apply to K, and V.-
The additional factor c, (xx) can be viewed as a learnable, context-dependent modulation or gating term for the features generated by by (xx).
(xx) E R
(x E ER, dn dsde-
This higher-order construction can enhance expressiveness. While introducing c, increases the parameter count for the factors, it might allow for the use of smaller base ranks (Rg. Rx RRy) to achieve comparable representational power, thus offering a different design choice. One could also explore tying sharing ce across queries, keys, and values manage parameter overhead.
From a memory perspective, during inference, higher-order TPA maintains the benefit of factorized KV caching. Only the constituent factors aœ ax(x),, ), (xxx),, CK (x:) (and similarly for values) for each past token need to be stored. A trade-off arises between model capacity and the overhead of memory and computation. Higher-order tensor decompositions can provide additional flexibility and potentially increased capacity.
C.1 RoPE Compatibility Higher-Order TPA
Rotary positional embeddings (RoPE) remain compatible with higher-order factorizations. In second- order TPA, RoPE applies rotations to the dy-diimenioonal feature vectors. This compatibility extends to higher-order TPA. Consider the case where RoPE is intended to primarily rotate feature pairs derived from the b (xx) components, while the structural influence of (xx) components the dy di-dimensional vector is preserved. More formally, RoPE acts on the da dy.dimensioonal vector vec(b, @ cf)) such that the transformation is equivalent to rotating by to by = R, (where R, is the RoPE rotation matrix for d, dimensions) and then forming vec(b, @ cf This is achieved by a specific RoPE transformation matrix T, acting on the full d, dimeenmnsll vector, as stated in the following theorem.
Theorem C.1 (RoPE Compatibility in Higher-Order TPA). Consider the higher-order (3-order) Tensor Product Attention (TPA) query factorization
Ra
1
a (x,)
vec(bb & cf (x.)) E AAxd,
RQ
r=l
28

---

where a (xx) IR,,, (x,r E R (x,x) E R with dy dyde. Define the RoPE-transformed
query Q: ;oPE; (Q:) I, where
(R
0
T. = Id. @ (Ra)
(R))==
12 xdh
(R,)
Id, the identity matrix size de des and R2 E RRdxxdo (ds Z. even) the standard RoPE block-diagonal matrix composed of 2 x rotation matrices:
(os((10) sin((01) cos((01)
sin(((1)
(00) - ((02) sin (t0g) cos(t82)
R,
cos(t0d.,2) sin (tn,2)
sin sin cos n,/2)
for / E {I,..., T and j (l,...,ds, ds/2). The transformation T, = Ia. @ (R,) operates on the da oi-dimensional vectorized features by post-multiplication. This structure of T, ensures that the rotation effectively applied to the by (x,) component (which is a column vector) corresponds to pre-muliiplication by R,, as detailed the proof (Appendix D.2). This preserves the structure induced by c, (x,) while rotating b (x.).
Under these conditions, the RoPF-transformed query RoPE, (Q:) admits higher-order TPA factor- ization the same rank RQ:
Re
1
Σ
2 8 vee (6, 60 cf(x)) Ro.F),QQ.)
(C.1)
RQ
where (x.( = R2 (x+).
Please Appendix D.2 for the proof. For fourth-order higher, this result still holds.
To assess its empirical performance, we implemented third-order TPA. Table 4 lists the evaluation results for small model. These results provide an initial indication of its viability. comprehensive comparison with second-order TPA variants of similar parameter counts or ranks would be necessary to fully evaluate the trade-offs.
Table 4: The evaluation results small models with third-order TPA pre-trained using Fine Web- Edu 100B dataset with iovalluatoonhareess Abbreviations: HellaSw. HellaSwag, W.G. WinoGrande.
Few-shot ARC-E
ARC-C
BoolQ
HellaSw.
OBQA
PIQA
W.G.
MMLU
SciQ
Avg.
0-shot
49.24
24.91
57.06
34.01
31.80
63.33
50.59
23.23
66.9
44.56
2-shot
53.37
25.34
48.78
34.00
29.20
62.79
52.33
26.41
75.3
45,28
D Proofs of Theorems
D.1 Proof of Theorem 3.1
Proof. Because RoPE is linear orthogonal transform, we write
Q = = Q. 1
Rq (A)().) (x()))) T. 1 Ax(x,) (BQ((,)) T,),
RQ
where T₃ the block-diagonal matrix encoding RoPE. This allows us to define
Bx(x,) = (x((x)) Te.
29

---

thereby obtaining
RoPE, (Q:)
Ax(x:) ( (xx(x)).
RQ
Similarly, for the key tensor K,, we have
1 (Ax(x)) Bx(xx)) T. = 1 Ax(xx) (xx(xx)
K, = K. T.
Rx
Rk
which defines
Bx(x,) = Bx((x)) To,
and thus
RoPE, (K.) = R
(x(x)) Bx((x).
Now, consider the product of the rotated queries and keys:
Q K.
RoRk (AQ(x)) (BQ (xxx
1
BR(xA))
1
RORK Ao(x,) (x B (x,) (x(x,,
Since T₁ and T, encode positional rotations, the product T.T. corresponds a relative rotation Tt-s. Therefore, we can express the above as
1
K
Aq(x,) (BQ( Bx (x,)) Ak (xx)
Ro
1
RORk (x
(Bo(x))T,_,xx (x(),)
1
Ax(x:)
(T-(())TT-a)(())) Ax(xx))
Ro
- -
1
RK Ax(x,) Bx(xx)
RQ
This shows that
RoP))_.QQ.)K; Q.
Focusing individual heads i, the above matrix equality implies:
KJ (g.:T.) (k..
where
q = RoPE)(q:,)) = qr T. R
k,, = RoPE, = K,,T, Ixds
This equality confirms that the relative positional encoding between queries and keys is preserved under TPA's factorization and RoPE's 's rotation. Thus, TPA maintains compatibility with RoPE. This completes the proof of Theorem 3.1.
D.2 Proof of Theorem C.I
Theorem C.I addresses the compatibility of RoPE with higher-order (specifically, 3rd-order) Tensor Product Attention. The theorem considers the query factorization:
Q: RQ
(xi)
vec (b, (x:) of (x:)) Axds
where (x,) R (column vector), (x,)) Rdds (column vector), a (x;) E Ted (col- umn vector), and da = dode. The term b (xx) & c, (x,) is interpreted as the matrix M,
30

---

(x, (xx)) E Rddxd. The notation v for a E R² and v E R (column vectors) implies the outer product av Thus, Q R2 D a (x (vec(M).))
Ra
The RoPE-transformed query is defined as Q2 RoPE, (Q:) = Q,Tr Crucially, for the theorem's 's conclusion to hold as intended (i.e., that the by component is transformed by pre-multiplication with the standard RoPE matrix R2), the global transformation matrix T, E Rd Xda (that post-multipies Q+) is given by:
T, It. (R.)",
where Id, is the de x de identity matrix, and R. E Rd x is the standard RoPE block-diagonal matrix that pre-multiplies dg di-dimensional column vectors (as defined explicitly in the theorem statement in Section C).
The theorem claims that, under these conditions, Q admits higher-order TPA factorization:
Ro
Q
(xx) @ vec 6. (x.) @ of ex((x.
Ro r=1
where (x( = R,b, (x:).
Proof. Let at = (x;,) bf (xx), and c c for brevity. Let M, = b,(())T Rdaxd, Let vec( ) dR be the column vector obtained by stacking the columns of M, The query tensor is Q Ro
Rio a (v.)
The RoPE transformation is Q Q, Substituting the factorization and the revised definition of Te
Q
a (vr))
(Ia, 60 (R.)
RQ
Ra Ro a ((v,) (Ia, & (R.)
Let's analyze the transformed vector part for the r-th component: (v,) (Ls @ (R,)). This row vector is the transpose of ((Ta, @ (R,) v,). Let's compute the pre-multiplying matrix:
(ILa & TT = (Ia,)" & ((R.) Ia, @ R.-
So, the column vector transformation is (Id, R,),, Substitute v, voo(M,) = roc(b,(c,))
(Id, @ R.) (c)
We the Kronecker product identity: (Bo Ao) vec(Xn) veo( Ao) To match expression (Ia, @ RR) vec(M,), we identify: Ao R2 Bo Iid Bo It. Xo M, (c Applying the identity, get:
vec (BB) = vec ((R)))
Let 6,2 = R. This is precisely the transformation for the by component as claimed in the theorem. So the transformed column vector is vec (c.)T). The corresponding row vector in the sum for Q is therefore (vec(b, (c,) (C) Substituting this back into the expression for Q::
1 Ra
a ((vec(6((cc
RO
This equivalent to the theorem's claimed factorization, using the definition a @ coc_vec = a(col_ved vec)
Ra
Q.
a x vec (b, & c,y)
RQ
31

---

where 6 = R.b.. This completes the proof, showing that RoPE can be consistently applied to higher-order TPA representations if the global RoPE transformation matrix T, (that post-multiplies Q.) is appropriately defined as La, @ (R,) ensuring that the standard RoPE matrix R, effectively pre-multiplies the by component.
E More Related Work
Transformers and Attention. As a sqquence-to-sequeer architecture, Transformer [60] intro- duced Multi-Head Attention (MHA), enabling more effective capture of long-range dependencies. Subsequent work has explored a variety of attention mechanisms aimed at improving scalability and efficiency, including sparse patterns [10, 49, 16, 30, 27, 31], kernel-based projections [11], and linearized transformers [59, 25, 44, 69, 54, 67]. To decrease memory usage and circumvent the limi- tation of memory bandwidth in training. [46] proposed Multi- Suali--ueryyaaa Attention (MQA) where multiple query heads share the same key head and value head. To tackle the issue of quality degradation and instability in training. Grouped- -eurry Attention (GQA) [2] divides queries into several groups, and each group of queries shares single key head and value head. Recently, DeepSeek-V V2 [32] applied multihead latent attention (MLA) to achieve better performance than MHA while reducing KV cache in inference time by sharing the same low-rank representation of key and value. Concurrently, [21] proposed Multi-matrix Factorization Attention (MFA), which can be simply seen as MQA with low- rank factorized Q Compared to the approaches above, TPA applied contextual tensor decompositions to represent queries, keys, and values activations compactly, achieving better reduction on the size of KV cache with improved performance.
KV Cache Optimization. During the auto-regressive inference of Transformers, key and value (KV) tensors from previous tokens are cached to avoid recomputation, a technique first proposed by [40]. This Key-Value (KV) cache, while crucial for efficiency, consumes significant memory and can introduce latency bottlenecks due to memory bandwidth limitations [1]. Consequently, various studies have explored methods to mitigate these issues. These include KV cache eviction strategies that discard less significant tokens [70, 62, 8, 1]. dynamic sparse attention mechanisms focusing on selected keys and values [42, 55, 50], offloading the KV cache to CPU memory [17, 26, 53], and quantizing the KV cache [61, 34, 19]. In contrast to these approaches, TPA focuses on reducing the intrinsic size of the KV cache by employing cecoorddecomooee key and value representations.
Low-Rank Factorizations. Low-rank approximations are widely used to compress model parameters and reduce computational complexity. Notable examples include LoRA [20], which factorizes weight updates during fine-tuning mune-tunig, and its derivatives tailored for various training scenarios such as efficient pretraining (ReLoRA [28], MoRA [22]), long-context training (LongLoRA [9], SinkLoRA [66]), and continual training (InfLoRR A [29], GS-LORA [71], I-LoRA [41]). These methods generally produce static low-rank expansions that are independent of the input context. Theoretical justifications for the
expressiveness of low-rank approximations have been provided by [38, 65]. Initialization strategies for these factorization matrices have also been explored: OLoRA [7] utilizes QR-cccomposition of pretrained weights for improved language model performance, while LoLDU [48] employs LDU-decomposition to accelerate LoRA training. Furthermore, AdaLoRA [68] uses Singular Value Decomposition (SVD) on pretrained weights and introduces parameter importance scores to dynamically adjust ranks. TPA, in contrast, constructs Q. K, and V tensors using contextually-awane factorizations, allowing for dynamic adaptation based on the input.
F
More on Attention Mechanisms
F. Multi-Query Attention (MQA)
Multi- Multi-Query Attention (MQA) [46] significantly reduces memory usage, particularly for the KV cache, by sharing a single key and value projection across all attention heads, while each head maintains a unique query projection. Given a sequence of input embeddings X R the query, shared key, and shared value tensors are computed as:
Q: XW K XWK K
V Vihared
XWV shared-
shared:
Thus, each head uses distinct query projection Qi Txdx but shares the common key K R xd, and value V wared E RTT xxd,, tensors. The weight matrices are:
W IR Amnse xd,,
W shared W.V.
IR (aaaaXda
shared
32

---

The resulting MQA operation is:
MQA(X) Concat headd head,) ) wo,
where
head, Attention Q shared: V sshared).
By sharing key and value projections, MQA substantially reduces memory demands, especially for the KV cache during autoregressive inference. However, this comes at the cost of reduced model expressivity, all heads must utilize the same key and value representations.
F.2 Grouped Query Attention (GQA)
Grouped Query Attention (GQA) [2] generalizes Multi-Head Attention (MHA) and MQA by dividing the total h attention heads into G groups. Within each group, heads share a common key and value projection, while each head maintains its own unique query projection. Formally, let g denote the group index for head E (1....,h) ,h), where g(i) (1.....G) G). The projections are:
K K0)
X w a(t)>
x(t)
X W
Q -XW
and
head,
Attention
Q K g(1), V Vote):
Here, W.K and wy are the shared weight matrices for group g. each in R and w.O Rd dn is the query weight matrix for head i. The complete output is again concatenation of all heads:
GQA(X) Concat hhead] head,, wo
By varying from 1 (equivalent to MQA) to h (equivalent to MHA), GQA offers a trade-off between memory efficiency and model capacity.
F.3 Multi-head Latent Attention (MLA)
Multi-head Latent Attention (MLA), as used in DeepSeek-V2 [32] and DeepSeek-V V3 [33], introduces low-rank compression for keys and values reduce KV caching costs during inference.
CKV = XW
Concat ( K.
K = K =
K ROOP((XWK))
K; Concat (K. K").
Concat (VC, VC,. Vg = VC CKVV
Here, DDV E IR rndee Xdc projects a compressed dimension da, WUK E (d.)( up-projects the compressed keys, WKR E dd projects to a residual key component for RoPE, and WUV E Rdex (dah) up-projects the compressed values. CVV E Tx the shared compressed KV latent (where de < d,,h). The RoPE transformation is applied to separate key embedding K E Txx, Thus, only CKV and K are eached, reducing KV memory usage while largely preserving performance compared to standard MHA [60].
MLA also compresses the queries, lowering their training-time memory footprint:
Q DQ
Concat Q- QC
Q Q cuw
Concat (Q, Q .....
Q R = Q R RoPE(C
Q = Concat (Q. Q*)
33

---

The weight matrices are W E R p E R( and W E x Here,
Cº E RR'''' (where d' < d, is the compressed query latent. The final query Q for each head, formed by concatenating Q and Q has a dimension of d,, + d
Given compressed queries, keys, and values, the final attention output for the t-th token is:
O Softmax (Vanee
V
(ds+ty
U = Concat (01, O2
Oz W
where V, is typically V as residual value component is explicitly defined, and Wo (dd is the output projection.
During inference, CKV and K are cached to accelerate decoding. In detail, if RoPE were ignored for the compressed components, the inner product k,, (where Ir,i kg, (dd) of the i-th head between t-th token query and s-th token key could be calculated using the current hidden state X2 Redaaaa and the cached latent state c E Rd for the s-th token:
ks, = [(W!9)T (WPP)T T xx] (WWKK) c
(F.I)
[,, DQ W.
UQ (
(F.2)
where w denotes the i-th head's portion of the respective weight matrix. The term
[W, W uru
[W could be pre-computed for faster decoding. However, noted by [51],
this pre- pre-computttion strategy is not directly compatible with RoPE if RoPE were applied to these
compressed representations. RoPE applies rotation matrix T2 Rd xd based position t (see
Section F.5), satisfying T.T. = T,-a (Equation F.4). If RoPE were applied to the up-projected Q and Kc:
9. k, = [T, T (W T xx] [T. (WWUK) TK
(F.3)
[ DQ (W.DOW TW. Jock
is not possible because depends the relative position (t-s) s) and thus varies for different (f, s) pairs. To maintain RoPE compatibility while benefiting from compression, MLA introduces an additional, smaller key component K (and similarly Q to which RoPE is applied, while the main compressed components K and Vc (derived from CKV) remain RoPE-free. As we will demonstrate in Section 3.2 of the main paper, TPA offers a different approach to integrate RoPE efficiently with factorized attention through its tensor product formulation.
Unlike Equation (F.2), acceleration by pre-computing the term [W. DQ
F.4 Multi-matrix Factorization Attention (MFA)
[21]] proposed Multi-matin Factorization Attention (MFA), which can be conceptualized variation of MQA where the shared key and value projections have a dimension de. and the query projection for each head is low-rank factorized:
Q: = XW DQ W
K Khhared
XWK
shared
shared
XWV
shared
where
DQ DQ E daedd de
UQ
E Rdd
WK shared
W shared
E
dunss Xd.
F.5 Rotary Position Embedding (RoPE)
Many recent LLMs use rotary position embedding (RoPE; 52) to encode positional information in the query/key vectors. Specifically, for vector at position I, RoPE applies a rotation matrix T, E Rdxd (where d is the dimension of the query/key vectors, typically d, per head). T, is a block-diagonal
cos t02) -sin(to,))
matrix composed of d/2 rotation blocks of the form
for E
sin (10,) (t0,)
The frequencies {0,} are typically defined as 02 base 2//d with common base like 10000. If qe Rd is query (or key) row vector for specific head position RoPE is applied as:
RoPE(q.) T.
34

---

A key property of RoPE is that the inner product between traasfraanf vectors depends only on their relative position. For a query q and key k,; (q.T.)(k,T.) = =q.T.T,k, = q.
This relies on the property:
T.T. T 1-81
(F.4)
which embeds relative positional information (t 5) into the attention scores.
G More on TPA
Parameter Initialization for TPA Factors. We initialize the weight matrices for TPA factors, - Who Wh and W (or their combined forms W W etc.),
such as W W using Xavier initialization [15]. Specifically, each entry of a weight matrix is drawn from uniform distribution ((-bouund, bound), where bound /6/(nm, + Plom). Here, 7mm and nout are the input and output dimensions of the respective weight matrix. This initialization strategy is chosen to help maintain the variance of activations and gradients as they propagate through the network layers, contributing to stable training.
TPA with Non-contextual B. In Section 4.1, we have introduced TPA with non-contextual A, where rreadddmenssio factors a², a a. E E are fixed. Conversely, one may fix the token-dimension factors b b,", by E Rd as learned parameters, while allowing a& (xx). a (xx), a. (x,) to adapt to the input token X,. The key tensor for token t. K, E Axd,, would then be constructed as:
RK
1 RK
K,
b
A similar formulation applies values. This configuration might be effective if the fundamental token-level features (captured by b,) are relatively stable, while their combination across heads (captured by a, (x:)) needs adapt to the context. Performance comparisons for TPA with non- contextual A factors versus non-contextual B factors on small and medium-sized models are presented in Tables 5, 6, and 8.
Table 5: Evaluation results of small models with TPA using non-contextual A or B factors, pre-trained on FineWeb-Edu 100B dataset (O-shot with Im-evaluation--aarnes) Abbreviations: HellaSw. HellaSwag, W.G. WinoGrande.
Method
ARC-E
ARC-C
BoolQ
HellaSw.
OBQA
PIQA
W.G.
MMLU
SciQ
Avg.
TPA (non-ctx-A)
$0.17
25.60
$7.95
36
31.40
64.80
49.57
24.88
64.80
45.03
TPA (non-ctx-B)
47.39
26.37
54.8
32.71
30,2
63.38
50.2
23.13
64.8
43.66
Table 6: Evaluation results of small models with TPA using non-contextual A or B factors, pre-trained on FineWeb-Edu 100B dataset (2-shot with Im-evaluatioo--aarnes) Abbreviations: HellaSw.
HellaSwag, W.G. WinoGrande.
Method
ARC-E
ARC-C
BoolQ
HellaSw.
OBQA
PIQA
W.G.
MMLU
SciQ
Avg.
TPA (non-ctx-A) 55.09
27.65
53.82
36.24
30.20
64.53
50.75
26.01
78.60
46,99
TPA (non-ctx-B)
50.8
26.96
57.65
32.4
29.4
63.22
49.57
23.96
66.4
44.48
Table Evaluation results of medium models with TPA using non-contextual A or B factors, pre- trained on Wen-WedEdu 100B dataset (O-shot with Im-evaluation - Abbreviations: HellaSw.
HellaSwag, W.G. = WinoGrande.
Method
ARC-E
ARC-C
BoolQ
HellaSw.
OBQA
PIQA
W.G.
MMLU
SciQ
Avg.
TPA (non-ctx-A)
58.96
31.48
59.76
45.07
80
69.21
53.59
25.42
76.40
50.52
TPA (non-ctx-B)
55.43
29.69
58.32
40.77
34.40
66.92
51.38
25.66
71.10
48.
TPA KV Only. A simpler variant involves using a standard linear projection for queries,
Q = Wwx, E Raxds
and factorize only the key and value tensors (K, V,). This approach, termed TPA-K TPA-KVonly. Vonly, maintains the standard query projection mechanism but still achieves significant KV cache reduction through factorized key and value representations.
35

---

Table 8: Evaluation results of medium models with TPA using non-contextual A or B factors, pre-
trained on FineWeb-Fdi 100B dataset (2-shot with Im-evaluattion-aaress)) Abbreviations: HellaSw.
HellaSwag, W.G. WinoGrande.
Method
ARC-E
ARC-C
BoolQ
HellaSw.
OBQA
PIQA
W.G.
MMLU
SciQ
Avg.
TPA (non-ctx--A
65.45
33.79
56.88
45.23
33.60
68.61
54.22
25.00
85.00
51.98
TPA (non-ctx-B)
61,20
30.20
55,93
40.45
34.40
68.23
51.78
26.11
78.10
49.60
TPA KV with Shared B. Further parameter reduction can be achieved by sharing the token- dimension factors b, between keys and values:
b (xx) = by (xx) (if contextual), or b k bb
b/"" (if non-contextual).
This sharing reduces both parameter count and the KV cache footprint. Although constrains K and V, to be constructed from the same token-level basis vectors, this variant can still offer strong performance with additional memory savings.
Nonlinear Head Factors. Instead of using purely linear transformations to derive the contextual head-dimension factors a& (xe). a. (x:). a. (x:), one can introduce element-wiss nonlinearities (e.g.. sigmoid (-) or softmax). Applying softmax, for instance, to the coefficients that generate a, (x:) could be interpreted as form of Mixture-of-Heade where the network learns to dynamically weight different head configurations based on the input context.
Discussion. These variants highlight the flexibility of the TPA framework, allowing for different trade-offs between memory efficiency, computational cost, and model expressiveness. By carefully choosing which factor components thead-dimension or lakeen-imension are contextual versus non- contextual, and by adjusting the ranks (Ro, Rx, Rv), TPA can not only unify existing mechanisms like MHA, MQA, and GQA but also significantly reduce KV cache poteenttey by an order of magnitude during autoregressive inference.
H More on Experiments
H.1 Experimental Settings
We list the main architecture cyper-parametery and training devices in Table 9. For all models, the head dimension d, is fixed at 64. Specific architectural choices include: 2 KV heads for GQA models; residual key dimension d 32 for MLA models; and ranks Rg Ry 2 and Rą 6 for TPA and TPA-K -Vonly models, unless otherwise specified. Other relevant papper-prraeterss are listed in Table 10.
Training Setup Details. We follow the nanoGPT training configuration [24]. In particular, we use the Adam W [35] optimizer with (8,, (3) = (0.9, 0.95), weight decay of 0.1, and gradient clipping at 1.0. We follow the same setting nanoGPT that the learning rate is managed by cosine annealing scheduler [36] with 2,000 warmup steps and (total) global batch size of 480. For the small, medium, large and XL models, we set maximum learning rates of 6 x 10 3 10 10- and 1 x 10 (respectively), and minimum learning rates of 3 x 5 6 x 10 - x 10th and 1 x 10 (respectively).
Table 9: The architecture hyper- syper-parameters and training devices of models. Abbreviations: BS. Batch Size, GAS. = Gradient Accumulation Steps.
MODEL SIZE
PARAMETERS
DEVICES
MICRO BS.
GAS.
#LAYERS
dscaun
SMALL
124M
4x A100 GPUs
24
5
12
768
MEDIUM
353M
8x A100 GPUs
20
3
24
1024
LARGE
772M
8x A100 GPUS
15
4
36
1280
XL
1.55B
A100 GPUS
10
48
1600
H.2
Additional Experimental Results
H.2.1
Perplexity Curves
We display the perplexity curves for medium, large, and XL size models in Figure 9.
36

---

Table 10: The architecture hypee-parameters for different models.
MODEL. SIZE
SMALL
MEDIUM
LARGE
XI.
(MHA)
12
16
20
25
(MQA)
31
39
h (GQA)
30
38
h (MLA)
12
23
34
49
h (TPA-KVONLY)
22
29
37
47
(TPA)
34
47
61
78
de (MLA)
256
$12
512
512
&e (MLA)
512
1024
1024
1024
Medium Mode, fewlleb-edulcon
Larga Model, F0000
Model, Fineleeeedd 100
VIA
V116
MI5
VCA
Wetret Neptete
MLA TRA-CVOY
TRA-C0015
THH, stvryty
TM
1P0
12
54
12
Taining zalend 391
Taaiiing Tobens S11
Trakling nikens
(a) Validation Perplexity
(b) Validation Perplexity
Validation Perplexity
Figure 9: The validation perplexity of medium-size (353M) models, large-size (773M), and XL-size (1.5B) models with different attention mechanisms on the FineWeb- Web-Edu 100B dataset.
H.2.2 Ablation Study on Different Ranks
Figure 10 illustrates the training loss, validation loss, and validation perplexity for XL-sized (1.5B parameters) TPA models with varying key/value ranks (Rx = Ry = R. as indicated in the figure legend), trained on the FineWeb- Edu 100B dataset. Corresponding O-shot evaluation results presented in Table 12 (rows for TPA-KVonly with different Æx,v). These results indicate that increasing the ranks for key and value factorizations generally improves the performance of the TPA models.
XL Model. Rnw.Web-ed... edu110e
Wonty tranc
Thuusbbly Yank Tou-sb--y Yank
Kiarly
The Kieely franc
Kienty
2.
The mwubbablyy Jank
37
2.r-
2.6-
25
25
12
25
Ynalinn Sw0ee 110
Trakiing rakens 11
Traiiing mikert
(a) Training Loss
(b) Validation Loss
Validation Perplexity
Figure 10: The training loss, validation loss and validation perplexity curves of XL-size (1.5B) TPA models with different key/value ranks (Rx = Rv = R) on the Feb-Wed-Edu 100B dataset.
H.2.3 0-shot Evaluation with Im-evaluation-harme
We present 0-shot evaluation results using the Im-evaluation-haare harnesss for small (124M parameters) and XL (1.5B parameters) models in Tables 11 and 12, respectively.
H.2.4 2-shot Evaluation with Im-evaluation-haanee
Similarly, 2-shot evaluation results are provided in Tables (Small), 14 (Medium), (Large), and 16 (XL).
37

---

Table 11: Evaluation results of small models (124M) with different attention mechanisms, pre-trained on FineWeb- 100B dataset (0-shot with ee---elluaiion-haeeeee The best scores in each column are bolded. Abbreviations: HellaSw. = HellaSwag, W.G. WinoGrande.
Method
ARC-E
ARC-C
BoolQ
HellaSw.
OBQA
PIQA
W.G.
MMLU
SciQ
Avg.
MHA
50.63
26.96
59.39
36.18
32.00
64.96
51.85
23.40
70.30
46.19
MQA
49.62
25.34
55.72
35.94
31,40
64,85
$1.30
23.37
68.70
45.14
GQA
48.70
25.68
56.15
35.58
31.40
64,91
$1.62
23.12
68.20
45.04
MIA
50.21
26.71
58.01
36.25
32,80
64.69
50.59
24.67
71.90
46.20
TPA-K Vonly
$1.05
26.54
57.25
36.77
32.60
65.02
50.91
23.64
69.70
45.94
TPA
51.26
27.39
57.00
36.68
32.80
64.47
49.72
24.61
72.00
46.21
Table 12: Evaluation results of XL models (1.5B) with different attention mechanisms, pre-trained on the FineWeb-Edu 100B dataset (0-shot with Im-evalluaioon-hanns) The best scores in each column bolded. Abbreviations: HellaSw. HellaSwag, W.G. WinoGrande. If not specified, TPA and
TPA-K Vonly models use Rx = Rv 2.
Method
ARC-E
ARC-C
BoolQ
HellaSw.
OBQA
PIQA
W.G.
MMLU
SciQ
MHA
64.81
35.41
61.90
54.32
37.20
72.74
55.80
25.44
82.80
54.49
MQA
64.10
36.01
62.26
54.38
39.00
72.58
56.43
23.70
81.90
54.48
GQA
63.68
35,92
60.46
54.17
38.40
73,56
56.27
24.77
81.70
54.33
MLA
64.14
35.92
60.12
53.60
39.20
72.25
55.17
24.71
81.60
54.08
TPA-K
65.61
36,77
63.02
54.17
37.00
73.34
54.62
25.02
81.60
54.57
TPA-K (Riyy = 4)
64.52
37.03
63.27
54.89
39.80
72.91
56.51
24.74
81.60
55.03
TPA-KVonly (Rx, 6)
65.78
35.92
61.71
54.86
38.60
72.69
57.93
25.59
82,20
55.03
TPA
66.71
36.52
61.38
54.03
40.40
72.52
56.83
24.49
82.20
55.01
Table 13: Evaluation results of small models (124M) with different attention mechanisms, pre-trained
on FineWeb-Fddu 100B dataset (2-shot with in-evaluationn-haness. The best scores in each column are bolded. Abbreviations: HellaSw. = HellaSwag, W.G. = WinoGrande.
Method ARC-E
ARC-C
BoolQ
HellaSw.
OBQA
PIQA
W.G.
MMLU
SciQ
Avg.
MHA
$7.66
28.24
57.28
36.43
29.60
64.09
51.14
26.57
82.00
48.11
MQA
53.79
26.35
44.95
34.18
28.80
62.79
52.01
25.91
78.10
45.21
GQA
55.01
25.94
55.72
35.68
31.80
65.29
51.93
25.27
77.80
47.16
MLA
54.76
27.13
58.07
36.13
31.40
65.07
51.30
25.90
78,90
47.63
TPA-KVonly 54.25
27.90
57.06
36.36
31.80
64.31
53.59
26.18
79.20
47.85
TPA
57.53
28.07
56.33
36.49
31.80
64.36
51.14
25.92
79.70
47.93
Table 14: Evaluation results of medium models (353M) with different attention mechanisms, pre- trained on Fine Web-Edu 100B dataset (2-shot with Im-evaluation -hanness, default LR x 40.4). The best scores in each column are bolded. Abbreviations: HellaSw. HellaSwag, W.G. = WinoGrande.
Method
ARC-E
ARC-C
BoolQ
HellaSw.
OBQA
PIQA
W.G.
MMLU
SciQ
Avg.
MHA
64.73
32.42
58.29
45.89
34,20
68.50
53.20
25,86
88.00
52.34
MQA
64.98
33.62
55.02
45.81
34.00
69.59
53.43
24.30
85.20
51.77
GQA
65.24
33.19
56.54
45.41
34,80
69.04
55.72
24.73
87.90
52.51
MLA
64.98
33.62
53.52
45.94
33.00
68.55
51.85
25,46
89.10
51.78
TPA-KVonly
64.69
32.34
59.48
46.23
35.40
70.08
54.06
25.64
86.30
52.69
TPA
67.97
34.56
57.22
46.87
34.60
69.91
52.01
25.07
89.90
53.12
Table 15: Evaluation results of large models (772M) with different attention mechanisms, pre-trained on the FineWeb-Edu 100B dataset (2-shot with in-evaaluation-harnes. The best scores in each column are bolded. Abbreviations: HellaSw. HellaSwag, W.G. WinoGrande.
Method
ARC-R
ARC-C
BoolQ
HellaSw.
OBQA
PIQA
W.G.
MML.U
SciQ
Avg.
MHA
67.85
36.35
59.82
50.22
35.00
70.67
53.35
23.92
91.10
54.25
MQA
68.86
36.09
53.79
50.50
37.00
70.89
54.70
25.01
88.00
53.87
GQA
69.15
36.09
58.84
50.29
36.20
70.73
54.22
26.08
90.00
54.62
MLA
70.54
38.74
61.50
51.86
36.00
70.89
54.22
25.47
92.40
55.74
TPA-KVonly
71.34
37.71
59.76
51.10
36.00
71.49
54.62
25.83
90.10
55.33
TPA
70.41
37.71
60.06
$1.30
34.00
71.06
54.54
25.79
90.30
55.02
38

---

Table 16: Evaluation results of XL models (1.5B) with different attention mechanisms, pre-trained on the Fen-We--Edu 100B dataset (2-shot with nn-evaluationn-arnesss The best scores in each column are bolded. Abbreviations: HellaSw. HellaSwag, W.G. WinoGrande. If not specified,
RK= = Ry = for TPA and Von-y models.
Method
ARC-E
ARC-C
BoolQ
HellaSw.
OBQA
PIQA
W.G.
MMLU
SciQ
Avg.
MHA
70.83
39.93
59.85
54.05
36,20
72.52
55.17
25.42
91.70
56.18
MQA
71.34
39.76
58.93
54.27
39.40
72.96
57.38
24.74
91.90
56,74
GQA
71.17
39.08
73.07
24.87
92,20
MLA
70.79
37.54
50.83
53.33
40.00
72.09
56.51
24.93
91.80
55.31
TPA-K Vonty
72.85
39.68
60.92
53.81
37.00
73.34
56.83
26.19
91.30
56.88
Ton-K (RR,v = 4)
72.98
40.27
60.15
54.88
36.80
73.29
56.43
25.50
92.10
56.93
TPA-KVonly (Rx,
73.95
39.76
58.99
54.73
36.80
72.91
$9.04
24.93
92.90
$7.11
TPA
71.76
39.16
61.25
53.74
37.80
72.80
55.49
23.86
90.70
56.28
H.3
Ablation Studies on Learning Rates
To assess sensitivity learning rates, conducted parallel experiments medium- models using a learning rate of 3 x 10 (compared to the default 6 10 used for other medium model results). The training loss, validation loss, and validation perplexity curves are shown Figure 11. Performance on standard benchmarks for these models trained with the 3 x 10-4 learning rate are reported in Tables 17 (0-shot) and 18 (2-shot). The results demonstrate that TPA and TPA-K Vonly maintain their performance advantages over other attention mechanisms even with this alternative learning rate.
32
Medium Model RheWee eou1008
1.1
Medim Model Riweleb-dd-1100
Medium Hodel Fre
20-
GGA
GDA
VLA
H.A.
va betten pee
28-
2.0
27
2.74
34-
Thenin 100
Thining Sokens 311
Thaning tokers
(a) Training Loss
(b) Validation Loss
(c) Validation Perplexity
Figure 11: The training loss, validation loss, and validation perplexity of medium-size (353M) models (learning rate 3 x 10-1) with different attention mechanisms on the Fine 100B dataset.
Table 17: The evaluation results of medium models (learning rate 3 10 10-4) with different attention mechanisms pretrained using the FineWeb-BB Web-Edu 100B dataset (0-shot with aim-evaluation-hamees The best scores in each column bolded. Abbreviations: HellaSw. HellaSwar W.G. WinoGrande.
Method
ARC-E
ARC-C
BoolQ
HellaSw.
OBQA
PIQA
W.G.
MMLU
SciQ
Avg.
МHA
56.52
29.27
58.84
44.06
35.00
68.44
51.07
25.35
76.40
49.44
MQA
55.68
28.24
60.86
44.17
35.20
68.66
52.72
25.14
72.90
49.29
GQA
54,88
29.61
56.36
43.77
35.20
68.82
52.57
25.41
74.80
49,05
MLA
59.64
29.78
60.73
45.17
34.20
68.66
52.80
25.34
75.70
50.22
TPA-KVonly
57.11
30.03
61.25
44.83
34.60
69.04
54.54
23.35
74.60
49.93
TPA
59.30
31.91
60.98
45.57
34.60
69.48
53.91
24.93
77.20
50.88
I Broader Impacts and Limitations
This work allows for the processing of much longer sequences of information with limited hardware resources by reducing the KV cache size. This could make advanced AI capabilities accessible to entities with limited computational budgets, potentially fostering improvement on downstream tasks, including in-depth document analysis, complicated-conter reasoning. and code generation, promoting innovation across various sectors in fields of scientific research, education, and software development.
39

---

Table 18: The evaluation results medium models (learning rate 10 10-4) with different attention mechanisms pre-trained using the FineWeb-Edu 100B dataset (2-shot with im-evaauutionnhhaness. The best scores in each column are bolded. Abbreviations: HellaSw. HellaSwag, W.G.
WinoGrande.
Method
ARC-E
ARC-C
BoolQ
HellaSw.
OBQA
PIQA
W.G.
MMLU
SciQ
Avg.
MHA
64.44
32.85
59.05
44.18
33.20
68.72
50.12
26.01
87.40
51.77
MQA
64.27
32.94
57.71
44.36
31.80
68.01
51.70
25.99
86.00
$1.42
GQA
61.70
32.17
52.81
43.99
33,80
68.50
53.35
24.44
86.40
50,80
MLA
65.95
31.48
50.98
44.99
32.20
68.93
51.93
88.80
$1.24
TPA-KVonly
65.99
33.70
57.49
44.47
34.20
69,53
53.28
24.23
86.50
52.15
TPA
66.54
34.47
58.96
45.35
33.00
69.21
53.99
24.51
91.30
53.04
Although our work proposes a KV-cache efficient architecture for large language models, it may contain certain limitations. For instance, generalization other modalities deserves more extensive investigation.
NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
We describe all the contributions and scope in the abstract and introduction
parts.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
We discussed the limitations in Appendix I.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be
40

---

used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms
and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification:
We list all the assumptions and proofs in Appendix D.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification:
We listed all the experiment details in Section 6 for reproduction of our work.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
41

---

(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g.. a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g.. to registered users), but it should possible for other researchers have some path reproducing verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
The code and data required to reproduce the main experimental results are provided at https: /// /anonymoouuu .open. science/rr /eeeee/rr/e---// anonymous-2000 -0025. The sup- plemental material will contain instructions for their use.
Justification:
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.. cc/ gwwidee//odeeuuuaeeeeeeeeeeeeeee
www.sssssiiiiiiiiiioiioioooiiipi oliccy for more details.
While we encourage the release of code and data, we understand that this might be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https: / //nips. htc/pullic/aae wttp///w///e/eeee/eeoecce CoodeSuumissiooPoooo
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
We just list all the training and test details in Section 6.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can provided either with the code, appendix, or as supplemental material.
42

---

7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification: The error bars are not reported because it would be too computationally expensive for repeated experiments on LLMs.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
OK to report 1-sigma error bars, but one should state The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative етгог rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We just list all the computer resources in Section 6.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: h/tppp/////////..../////// wwt.iis.....pppppeee hptpplic/ttiisssiieelllleeeee
Answer: [Yes]
Justification: Our research only explores a novel framework for large language models with better KV-Cache efficiency. Therefore, the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
43

---

If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. there a special consid- eration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: We discussed the potential positive societal impacts and negative societal impacts in Appendix
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deploymennt of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Our work proposes a novel framework of large language models. To our knowledge, this work has no direct path to any negative applications.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
44

---

Question: Are the creators or original owners of assets (e.g.. code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We add the citation all the codes (nanoGPT and Im-evaluation-hanne MIT License) and datasets (FineWeh-Bie Edd-1000: odc-by) that we used in this work. No other models are included in our work.
Guidelines:
The answer NA means that the paper does use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license dataset.
For existing datasets that re-packaged, both the original license and the license of the derived asset (if has changed) should provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: The code implementing our proposed TPA model and experimental setup is released at https: pettps//anoole //anonymous 4open. sci eeiiennn////eeee wwtee////-annnnnnoooo anonymous- sonymous-2255. This code will be documented facilitate understanding and use by other researchers. No new datasets pre-trained models are introduced beyond the code for the methods.
Guidelines:
The answer NA means that the paper does release assets.
Researchers should communicate the details the dataset/code/model as part their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: The paper only use open-source codes and datasets which do not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
45

---

Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper only use open-source codes and datasets which do not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if' applicable), such the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality the research, declaration not required.
Answer: [Yes]
Justification: This work aims at exploring more efficient architecture for large language models. Therefore, LLM architectures are well described in the main part of this paper.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy /https: //tp://nen..../eeee wwwpuripsscc/oofffeee ccc//ooneeeeccce/2255/////////// for what should or should be described.
46


Reviews:
[{"review_id": "87DcwmAL2v", "reviewer": "Reviewer_nLZi", "summary": "This paper proposes Tensor Product Attention and TPA-based transformer (T6) as an alternative to vanilla attention mechanism and transformer. TPA enjoys substantial reduction in inference-time KV cache size compared to multi-head attention and other well-known models while preserving performances. In addition, this paper theoretically proves that TPA is well adapted to RoPE and develops FlashTPA decoding, which is an efficient autoregressive inference algorithm for TPA.", "questions": "1. Typo: Line 121 should be $R_Q,R_K,R_V$?\n2. Is T6 easy to train compared to standard transformers? I think the training aspect is also worth studying.", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths:\n1. KV cache is a major bottleneck for transformers, so developing new architectures that alleviate this issue is valuable.\n2. The paper is clear and relatively easy to follow. TPA is well defined; major contributions are results are stated.\n3. TPA greatly improves memory usage from $2hd_h$ to $(R_K+R_V)(h+d_h)$, which is small when $R_K,R_V$ are small.\n4. Empirical results seem promising as T6 outperforms standard transformers in many tasks.\n\nWeakness:\n1. It is unclear to me what the motivation for TPA is. It seems like the main motivation is that $Q_t$ (and similarly $K_t,V_t$) can be written as the average of $R_Q$ many outer products of vectors such that we only need to store the vectors here. However, I do not understand why this is necessarily a good way besides that it saves memory in theory.\n2. It is mentioned that MHA etc can be expressed as TPA, but I think with pretty bad $R_Q,R_K,R_V$? It is unclear to me that TPA has better expressivity compared to standard attention, especially since the paper mentions setting $R_K,R_V$ to be 1 or 2 in line 178.\n3. It can be theoretically proved that TPA saves memory, but there is no related experiment that illustrates this. Line 265 states \"factorizing the attention activations shrinks autoregressive KV cache requirements by up to 5×–10×\", but I don't see where this comes from empirically? This makes the whole argument less convincing.", "was_revised": true, "final_justification": "I think the idea presented in this paper is interesting, and there are supporting experiments (the authors conduct a few more larger scale experiments after the submission, if that can be counted). The memory and inference speed can be improved with this technique with a reasonable performance. Thus I will recommend borderline accept.", "rebuttal": "We thank the reviewer for their thoughtful feedback and positive assessment of our work. We are encouraged that the reviewer found our paper clear, the problem valuable, and the empirical results promising. Below, we address the questions raised.\n\n> Q1: It is unclear to me what the motivation for TPA is. It seems like the main motivation is that $Q_t$ (and similarly $K_t, V_t$) can be written as the average of $R_Q$ many outer products of vectors such that we only need to store the vectors here. However, I do not understand why this is necessarily a good way besides that it saves memory in theory.\n\nA2: Thank you for this insightful question. While memory savings are a key benefit, the primary motivation for TPA is to introduce a more expressive and efficient attention mechanism through contextual low-rank factorization.  Unlike standard attention mechanisms (MHA, MQA, GQA) that use static projection matrices for all tokens, TPA dynamically computes factorized representations for queries, keys, and values based on each token's specific context ($\\mathbf{x}_t$). As shown in Line 127 of our paper, the factors (e.g., $\\mathbf{a}^Q(\\mathbf{x}_t)$, $\\mathbf{b}^Q(\\mathbf{x}_t)$) are functions of the input. This allows the model to learn input-dependent, low-rank subspaces for attention, offering greater representational flexibility.  We hypothesize that this contextual adaptation is the reason TPA not only reduces the KV cache but also achieves superior performance (e.g., lower perplexity and higher downstream accuracy), as shown in our experiments. In essence, TPA achieves a powerful combination of memory efficiency and improved model quality, which is a critical goal for scaling language models to longer contexts. We will clarify this motivation in the introduction of our revised manuscript.\n\n> Q2: It is mentioned that MHA etc can be expressed as TPA, but I think with pretty bad $R_Q,R_K,R_V$? It is unclear to me that TPA has better expressivity compared to standard attention, especially since the paper mentions setting $R_K,R_V$ to be 1 or 2 in line 178.\n\nA2: Thank you for raising this important point. Our goal in showing that MHA is a special non-contextual case of TPA was to provide a unifying perspective, not to suggest that TPA should operate with high ranks to mimic MHA.  On the contrary, our central claim is that TPA achieves strong expressivity with very low ranks. Its expressive power stems not from a high rank, but from its contextual nature. By dynamically generating the Q/K/V representations for each token, TPA can adapt its behavior to the specific input in a way that static projections cannot. Our empirical results, showing that TPA with $R_{K,V}=2$ consistently outperforms MHA, MQA, and GQA, support the idea that this contextual factorization is a more efficient and powerful inductive bias for attention.  \nTo further validate TPA's superior performance and expressivity, we are pleased to share new results from large-scale MoE models (>10B parameters) from our industry collaborators. As shown below, TPA with a KV rank of 2 significantly outperforms a strong GQA baseline across multiple standard benchmarks.  \n\n| Method | MMLU | MMLU-Pro | AGIEval | GSM8k | MATH |\n| :--- | :--- | :--- | :--- | :--- | :--- |\n| GQA-MoE | 50.9 | 20.1 | 30.7 | 19.1 | 11.2 |\n| TPA-MoE | 52.1 | 22.7 | 32.8 | 26.5 | 14.0 |\n\nThese experiments (each consuming over 10,000 GPU hours) provide strong evidence that TPA's architecture leads to substantial quality gains in large-scale settings. \n\n> Q3: It can be theoretically proved that TPA saves memory, but there is no related experiment that illustrates this. Line 265 states \"factorizing the attention activations shrinks autoregressive KV cache requirements by up to 5×–10×\", but I don't see where this comes from empirically? This makes the whole argument less convincing.\n\nA3: We appreciate you highlighting the need for a more explicit calculation. The 5x-10x memory reduction is a direct calculation based on the KV cache formulas and the hyperparameters of our models, which we will make explicit in the revised paper.  The KV cache size per token for a baseline like MHA is $2 \\cdot h \\cdot d_h$. For TPA, it is $(R_K + R_V)(h + d_h)$. The reduction factor is therefore $\\frac{2 h d_h}{(R_K + R_V)(h + d_h)}$.  \n\nLet's use our XL (1.5B) model as a concrete example: \n- MHA baseline: $h=25, d_h=64$. KV Cache size = $2 \\times 25 \\times 64 = {3200}$ floats/token. \n- TPA-KVonly ($R_{K,V}=2$): $h=47, d_h=64$. KV Cache size = $2+2)(47+64) = 4 \\times 111 = {444}$ floats/token.     \n- This represents a 7.2x reduction ($3200 / 444$) compared to the MHA baseline of a similar parameter count. \n- TPA ($R_{K,V}=2$): $h=78, d_h=64$. KV Cache size = $(2+2)(78+64) = 4 \\times 142 = {568}$ floats/token. \n-  This is a 5.6x reduction ($3200 / 568$).  The up to 10x part of our claim comes from using a rank of $R_{K,V}=1$. For instance, TPA-KVonly with $R_{K,V}=1$ would have a cache size of $(1+1)(47+64) = 222$ floats/token, yielding a 14.4x reduction. We will add a table and clear calculations to Section 3.3 to make this more transparent.\n \n> Q4:Typo: Line 121 should be $R_Q,R_K,R_V$\n\nA4: Thank you for your suggestion, and we will correct it in the revised version.\n\n> Q5: Is the T6 easy to train compared to standard transformers? I think the training aspect is also worth studying.\n\nA5: This is an excellent point. We have measured the training throughput and can confirm that T6 trains are stable and efficient. With our current implementation, throughput is already highly competitive.  The tables below show that while the current, unoptimized version of TPA has modest overhead compared to MHA/GQA, its throughput is comparable to other efficient architectures like MLA, and becomes nearly equal to GQA at a larger scale.  \n\n##### Dense Models (350M) on 8x A100 GPUs (TPS means tokens per second): \n\n| Method | MHA | MQA | GQA | TPA-KVonly | TPA | MLA |\n| :--- | :--- | :--- | :--- | :--- | :--- | :--- |\n| **TPS (M)** | 0.515 | 0.456 | 0.452 | 0.427 | 0.356 | 0.355 |\n\n##### MoE Models (>10B) on 64x GPUs (TPS means tokens per second):\n\n| Method / Hardware | GQA | TPA |\n| :--- | :--- | :--- |\n| H800 TPS (M) | 0.975 | 0.891 |\n| H20 TPS (M) | 0.407 | 0.405 |\n\nThe overhead in the current dense model benchmarks is due to a vanilla implementation that explicitly computes the tensor factors. However, this is not a fundamental limitation.  In addition, we are developing FlashTPA training and prefilling kernels. \n\nAs analyzed in Appendix C of our submission, these optimized kernels operate directly on the factorized components without materializing the full Q, K, and V tensors. This approach significantly reduces the overall computational load of the attention mechanism. Our analysis suggests this can make TPA's training and prefilling throughput even faster than the baselines, turning a minor trade-off into another advantage of our method.  Given that TPA already trains stably and at a competitive speed, with a clear path toward even faster, optimized kernels, we believe it is a highly practical architecture. We will add these throughput results and a note on the ongoing kernel development to the appendix."}, {"review_id": "9pBz8slwho", "reviewer": "Reviewer_VFNC", "summary": "This paper introduces Tensor Product Attention (TPA), a novel attention mechanism aimed at mitigating the substantial memory overhead of the KV cache during inference in large language models. The central idea is to represent the query, key, and value tensors using contextual low-rank tensor decompositions, which shrinks the KV cache size. The authors propose the T6, a model architecture based on TPA. They highlight that TPA integrates seamlessly with rope, a common feature in modern LLMs.", "questions": "- The average performance scores reported in Tables 2, 3, and 12 include the MMLU benchmark, where the models perform close to random chance. This appears to inflate the average and potentially overstate the advantage of TPA. Could you provide a discussion or table that re-calculates the average scores excluding MMLU? For the 1.5B models (Table 12), my own calculation suggests the gap between TPA-KVonly (56.96%) and MLA (56.91%) becomes negligible. This, combined with the mixed results in Tables 13-15, challenges the claim of superior performance. How do you justify the adoption of TPA over MLA if the reasoning performance is not clearly superior?\n\n- The paper provides a compelling case for TPA's inference efficiency but does not present any data on its training efficiency. Since TPA introduces additional linear projections and tensor operations compared to a standard attention block, it could negatively impact training throughput (tokens/sec). Could you provide a comparative analysis of the training speed of T6 against the baselines, particularly MHA and MLA, which have highly optimized kernels? Understanding this trade-off is crucial for assessing the overall viability of the architecture.\n\n- The primary motivation for reducing KV cache size is to enable models to process longer sequences. However, the evaluation is performed on benchmarks that typically do not require long-context reasoning or recall. To provide more direct evidence for the benefits of your approach, have you considered evaluating TPA on tasks specifically designed for long-context understanding? For example, benchmarks from the ruler framework or \"Needle-in-a-Haystack\" style evaluations would more directly demonstrate the practical value unlocked by TPA's memory efficiency.", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "**Strengths**\n\n-  The proposed solution using contextual tensor factorization for Q, K, and V activations is an original and technically interesting approach. It moves beyond static compression methods by making the factorization dependent on the input, offering a potentially more expressive way to compress attention components.\n\n- The primary strength of this work lies in the demonstrated efficiency improvements. The accompanying FlashTPA decoding algorithm is shown to be highly competitive, often outperforming optimized kernels for other attention mechanisms, especially as sequence length increases.\n\n**Weaknesses**\n-  The significance of the claimed performance improvements on downstream tasks is questionable. The evaluation includes MMLU scores for models that perform near random chance on this benchmark (e.g., ~25%). Including these scores in the average can create a misleading impression of TPA's effectiveness. When MMLU is removed from the average calculation, the performance gap between TPA and the strongest baseline, MLA, becomes marginal on common reasoning tasks. Furthermore, the results for smaller models (Tables 13-15) do not consistently show TPA as the superior method.\n\n| Method         | Avg (No MMLU) |\n| -------------- | ------------- |\n| MHA            | 56.22         |\n| MQA            | 55.48         |\n| GQA            | 55.68         |\n| MLA            | 56.91         |\n| TPA-KVonly | 56.96    |\n| TPA            | 56.69         |\n\n\n-  While the paper provides an excellent analysis of inference speed, it completely omits any evaluation of training efficiency (i.e., training throughput). The baselines compared against (MHA, MLA) often have highly optimized CUDA or Triton kernels. TPA introduces additional operations, and without a throughput comparison, it is difficult for practitioners to assess the full computational cost and trade-offs of adopting this new architecture for training.\n\n- The paper's core motivation is enabling longer context windows. However, the empirical evaluation is restricted to benchmarks from the lm-evaluation-harness, which predominantly test for general knowledge and reasoning over short contexts. The evaluation suite lacks tasks specifically designed to measure long-context recall and understanding (e.g., \"Needle-in-a-Haystack\" tests or benchmarks from the ruler library). The absence of such targeted evaluation makes it difficult to verify the practical utility of TPA's memory savings for long-context applications.", "was_revised": true, "final_justification": "I am satisfied with this work and would lean towards a borderline accept. The training throughput analysis addresses my concerns.\n\nHowever, I have a suggestion regarding the experimental comparison. The new MoE results comparing TPA-MoE against GQA-MoE are indeed impressive, showing substantial improvements across multiple benchmarks. Yet, I believe these comparisons would be significantly more meaningful if conducted against MLA-MoE rather than GQA-MoE. MLA achieves similar KV cache compression ratios to TPA through its sophisticated two-stage projection mechanism, where inputs are first projected to a latent space before expansion to multiple heads.\n\nGiven that both TPA and MLA target the same problem of KV cache reduction while maintaining model expressiveness, and both achieve comparable cache compression ratios, MLA constitutes the most relevant baseline for evaluating TPA's contributions. I encourage the authors to include MLA-MoE comparisons in their revised manuscript, especially for these large-scale experiments.", "rebuttal": "Dear reviewer, we sincerely thank you for your detailed and constructive feedback. We appreciate the positive comments on the originality and technical interest of our approach. We are grateful for the opportunity to address the your questions and suggestions.\n\n> Q1: The significance of the claimed performance improvements on downstream tasks is questionable. The evaluation includes MMLU scores for models that perform near random chance on this benchmark (e.g., ~25%). ...\n\nA1: Thank you for your feedback. We acknowledge your point that for smaller models, the performance gains on some benchmarks can be marginal, and the low MMLU scores can affect the interpretation of the average. The primary justification for adopting TPA rests on three pillars that become increasingly evident at scale: (1) superior performance on large models, (2) a smaller KV cache, and (3) faster inference decoding.  \nWhile the gains on our 1.5B models are modest, we have new results from our industry collaborators on >10B parameter MoE models that demonstrate TPA's advantages much more clearly. In these large-scale experiments, a TPA-based model (with a 24% smaller KV cache) significantly outperforms a strong GQA baseline on challenging reasoning benchmarks.  \n\n| Method | MMLU | MMLU-Pro | AGIEval | GSM8k | MATH |\n| :--- | :--- | :--- | :--- | :--- | :--- |\n| GQA-MoE | 50.9 | 20.1 | 30.7 | 19.1 | 11.2 |\n| TPA-MoE | 52.1 | 22.7 | 32.8 | 26.5 | 14.0 |\n\nThese results, obtained from experiments consuming over 10,000 GPU hours each, show that TPA's architectural benefits scale with model size, leading to substantial improvements in reasoning capabilities (e.g., +38% on GSM8k, +25% on MATH relatively).  Therefore, we justify adopting TPA over MLA and other baselines because it offers a superior package for large-scale deployment: it not only achieves better performance at scale but does so while being more memory-efficient (smaller KV cache) and faster at inference (as shown by our FlashTPA decoding results). We will add these new results to the paper to better illustrate TPA's advantages.\n\nQ2: While the paper provides an excellent analysis of inference speed, it completely omits any evaluation of training efficiency (i.e., training throughput). The baselines compared against (MHA, MLA) often have highly optimized CUDA or Triton kernels. \n\nA2: Thanks for your suggestion. We have since benchmarked the training throughput and agree this is crucial information for practitioners.  Our results show that even with a vanilla implementation, TPA's training throughput is highly competitive with other efficient attention mechanisms like MLA. On a larger scale, its throughput is nearly the same as the GQA baseline.\nDense Models (350M) on 8x A100 GPUs (TPS means tokens per second): \n\n##### Dense Models (350M) on 8x A100 GPUs (TPS means tokens per second): \n\n| Method | MHA | MQA | GQA | TPA-KVonly | TPA | MLA |\n| :--- | :--- | :--- | :--- | :--- | :--- | :--- |\n| **TPS (M)** | 0.515 | 0.456 | 0.452 | 0.427 | 0.356 | 0.355 |\n\n##### MoE Models (>10B) on 64x GPUs (TPS means tokens per second):\n\n| Method / Hardware | GQA | TPA |\n| :--- | :--- | :--- |\n| H800 TPS (M) | 0.975 | 0.891 |\n| H20 TPS (M) | 0.407 | 0.405 |\n\nFurthermore, we are actively developing FlashTPA training/prefilling kernels. As analyzed in Appendix C, an optimized kernel that avoids materializing the full Q/K/V tensors can significantly reduce the FLOPs required for the attention computation. This gives TPA a clear path to potentially exceeding the training throughput of existing baselines. \nIn summary, TPA already trains at a competitive speed for large scale model training, and we have a concrete plan for further optimization. We will add these results and discussion to the appendix.\n\n> Q3: The paper's core motivation is enabling longer context windows. However, the empirical evaluation is restricted to benchmarks from the lm-evaluation-harness, which predominantly test for general knowledge and reasoning over short contexts. The evaluation suite lacks tasks specifically designed to measure long-context recall and understanding (e.g., \"Needle-in-a-Haystack\" tests or benchmarks from the ruler library). The absence of such targeted evaluation makes it difficult to verify the practical utility of TPA's memory savings for long-context applications.\n\nA3: Thanks for your feedback. We can provide strong evidence from our new, large-scale experiments on benchmarks that require reasoning over longer inputs.  The GSM8k and MATH benchmarks, especially when evaluated with few-shot examples, require the model to process and reason over significantly longer context than typical classification tasks. \n\nOur >10B TPA-MoE model shows substantial improvements to these very tasks:  \n- GSM8k: TPA achieves 26.5 vs. GQA's 19.1 (+38% relatively) \n- MATH: TPA achieves 14.0 vs. GQA's 11.2 (+25% relatively)  \n\nThese results demonstrate that TPA's architecture leads to better performance on complex, multi-step reasoning problems that benefit from longer context.  \n\nIn addition, we plan to run \"Needle-in-a-Haystack\" (NIAH) style tests in the next steps. Regarding dedicated benchmarks like NIAH, we acknowledge their importance. However, the computational cost to train and reliably evaluate multiple large models suitable for these tests is extremely high (often exceeding 50,000 GPU-hours ($100K USD) per experiment), which we believe is beyond the scope of a typical academic conference submission. We plan to pursue these evaluations in future work. We will add a discussion emphasizing these points to the revised paper to better contextualize our results.\n\nWe believe the combination of (1) a fundamentally smaller KV cache, (2) faster decoding speeds for long sequences, and (3) superior performance on reasoning tasks like GSM8k and MATH provides compelling evidence for TPA's practical utility in long-context scenarios. We will add a discussion emphasizing these points to the paper."}, {"review_id": "xFs43a7acC", "reviewer": "Reviewer_Z9Nz", "summary": "The paper introduces Tensor Product Attention (TPA), a novel attention mechanism that factorizes query, key, and value representations into low-rank contextual tensor products. This approach significantly reduces the memory footprint of key-value caches during inference without compromising model quality. TPA integrates naturally with Rotary Position Embeddings (RoPE) and is used as the core component in a new transformer architecture called T6. The authors show that TPA can outperform or match traditional attention variants like Multi-Head Attention (MHA), Multi-Query Attention (MQA), Grouped-Query Attention (GQA), and Multi-Head Latent Attention (MLA) on language modeling benchmarks. They also propose FlashTPA, an efficient decoding algorithm tailored to TPA that enhances autoregressive inference speed, especially for long sequences. Overall, the paper offers a practical and scalable alternative to standard attention mechanisms, advancing memory and computational efficiency in large language models.", "questions": "1. Could you provide ablations on the choice of tensor ranks (RQ, RK, RV)? It's unclear how sensitive performance and efficiency are to these hyperparameters. A systematic study would clarify trade-offs and could improve the quality rating.\n\n2. The paper lacks evaluation on actual long-context tasks. How does TPA perform on problems involving 32k+ tokens, such as document QA or summarization? This is important for backing the scalability claims and could raise the significance score.\n\n3. Comparisons are missing to other recent memory-efficient attention methods like linear or state-space models. Including them would better situate TPA within the broader field and help assess its relative merit.\n\n4. The paper presents no discussion of failure modes. Are there cases where TPA underperforms or is harder to train? Clarifying this would strengthen the honesty of the evaluation and improve the quality score.", "limitations": "No. The paper does not sufficiently address limitations or potential negative societal impacts. It lacks discussion on where TPA might degrade performance, such as under very low-rank settings or on tasks beyond language modeling. Additionally, potential risks of enabling longer context generation, such as misuse in misinformation or surveillance, are not considered. The authors should explicitly discuss cases where TPA may fail or introduce trade-offs, and briefly reflect on broader implications of more efficient long-context inference.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 4, "strengths_and_weaknesses": "Quality: The paper is technically sound and well-supported by both theoretical and experimental evidence. The proposed method is carefully motivated, and the authors demonstrate a clear understanding of prior work. Experimental comparisons are fair and span several baselines and model sizes. However, the analysis could be strengthened with deeper ablations, particularly on factorization ranks and the interaction with RoPE. The paper could also more clearly discuss potential failure modes or limitations, such as performance trade-offs under extremely low-rank settings.\n\nClarity: The writing is generally clear and well-organized. Core concepts like tensor product factorization and FlashTPA decoding are explained with useful figures and mathematical definitions. Some dense theoretical sections may challenge readers unfamiliar with tensor operations, and the paper could benefit from slightly more intuitive explanations or visual summaries. Important implementation details, especially for reproducibility, are mostly deferred to the appendix.\n\nSignificance: The work addresses a critical bottleneck in scaling language models - the cost of KV caching at inference - and proposes a solution that shows compelling improvements in memory and decoding efficiency without sacrificing model quality. This has practical implications for real-world deployment and theoretical relevance for sequence modeling research. However, significance would be better supported by evaluations on more diverse, real-world long-context tasks.\n\nOriginality: The core idea of contextual tensor product factorization for attention is novel and distinct from prior work in low-rank adaptation or shared-key attention. The combination of RoPE and its efficient implementation through FlashTPA adds to the originality. While built on existing tools like RoPE and attention factorization, the paper's contributions are more than a simple synthesis - they represent a meaningful conceptual and practical advance.", "was_revised": true, "final_justification": "The clarifications provided by the author have resolved my main issues, so I have slightly increased my evaluation.", "rebuttal": "We are very grateful to the reviewer for their positive evaluation and insightful, constructive feedback. We are encouraged that they found our work technically sound, original, and significant. We address the questions and suggestions for improvement below.\n\n> Q1: \"The analysis could be strengthened with deeper ablations, particularly on factorization ranks and the interaction with RoPE... Could you provide ablations on the choice of tensor ranks ($R_Q, R_K, R_V$)?\"  \n\nThank you for these excellent suggestions to strengthen our analysis. We have performed additional ablations to address these points.  \n\n- Ablation on Ranks ($R_K, R_V$): We agree a systematic study is crucial. As you noted, we included an ablation on the key/value ranks in Appendix J.2.2 (Figure 10, Tables 12) for our 1.5B parameter models. We compared ranks of {2, 4, 6} and found that performance consistently improves with higher ranks. Importantly, even at a very low rank of 2, TPA remains highly competitive, demonstrating a graceful trade-off between performance and KV cache size rather than a sharp failure point. We will highlight these findings more prominently in the main paper.  \n\n- Interaction with RoPE: To further analyze the model, we conducted new experiments using a partial RoPE setting (applying RoPE to only a fraction of the head dimension). Our results for the Medium (353M) models show that TPA-based models remain top performers in this alternative configuration, underscoring the robustness of our approach. \n\n### Medium 0-shot Results\n\n| Medium 0-shot | ARC-E | ARC-C | BoolQ | HellaSwag | OBQA | PIQA | W.G. | MMLU | SciQ | Avg. |\n| :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- |\n| llama-gqa-partialrope | 56.06 | 29.69 | 59.60 | 44.17 | 34.20 | 68.01 | 54.38 | 24.78 | 73.50 | 49.37 |\n| llama-mha-partialrope | 58.04 | 29.44 | 58.59 | 44.47 | 34.00 | 68.28 | 52.96 | 23.59 | 76.20 | 49.51 |\n| T6_partialrope | 56.57 | 30.12 | 59.11 | 45.30 | 35.20 | 69.37 | 53.59 | 26.57 | 85.40 | **51.25** |\n| T6_kvonly_partialrope | 58.84 | 31.91 | 60.64 | 45.27 | 34.20 | 68.66 | 53.83 | 24.15 | 74.30 | 50.20 |\n| MLA | 56.65 | 29.52 | 57.83 | 46.05 | 34.60 | 69.42 | 52.80 | 24.62 | 79.70 | 50.13 |\n\n### Medium 2-shot Results\n\n| Medium 2-shot | ARC-E | ARC-C | BoolQ | HellaSwag | OBQA | PIQA | W.G. | MMLU | SciQ | Avg. |\n| :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- |\n| llama-gqa-partialrope | 64.10 | 31.83 | 57.31 | 44.14 | 35.40 | 67.90 | 53.43 | 25.84 | 87.90 | 51.98 |\n| llama-mha-partialrope | 64.81 | 33.02 | 54.74 | 44.32 | 31.80 | 68.66 | 51.22 | 25.80 | 88.30 | 51.41 |\n| T6_partialrope | 65.82 | 33.96 | 53.85 | 45.51 | 32.80 | 69.59 | 54.14 | 25.33 | 87.00 | 52.00 |\n| T6_kvonly_partialrope | 66.50 | 34.98 | 58.81 | 45.77 | 35.60 | 68.39 | 53.67 | 26.15 | 88.20 | **53.12** |\n| MLA | 64.98 | 33.62 | 53.52 | 45.94 | 33.00 | 68.55 | 51.85 | 25.46 | 89.10 | 51.78 |\n\n> Q2: \"The paper could benefit from slightly more intuitive explanations or visual summaries. Important implementation details, especially for reproducibility, are mostly deferred to the appendix.\"  \n\nA: Thank you for this valuable feedback. In the revised version, we will improve the paper's clarity and accessibility. Specifically, we will add a more intuitive visual summary to Section 3 to better illustrate the core mechanism of contextual factorization. While page limits are a constraint, we will also reorganize the main text and appendix to bring the most critical implementation details (such as our initialization strategy) into the main paper to improve readability and reproducibility.\n\n> Q3: \"The paper's significance would be better supported by evaluations on more diverse, real-world long-context tasks... How does TPA perform on problems involving 32k+ tokens, such as document QA or summarization?\"  \n\nA: We completely agree that demonstrating performance on long-context tasks is key. While running full 32k+ context benchmarks on multiple models was computationally prohibitive for this submission, we can provide strong evidence from new, large-scale (>10B parameter) MoE models on benchmarks like GSM8k and MATH, which require reasoning over longer and more complex prompts. In these experiments, our TPA model significantly outperformed a strong GQA baseline:  \n\n- GSM8k: TPA achieves 26.5 vs. GQA's 19.1 (+38%)\n- MATH: TPA achieves 14.0 vs. GQA's 11.2 (+25%)  \n\nThese results show that TPA's architectural advantages translate to superior performance on challenging reasoning tasks. We argue that the combination of (1) proven memory savings, (2) faster long-sequence decoding (via FlashTPA), and (3) superior performance on long-form reasoning tasks provides compelling evidence of TPA's utility for real-world, long-context applications. \n\n> Q4: \"Comparisons are missing to other recent memory-efficient attention methods like linear or state-space models. Including them would better situate TPA within the broader field and help assess its relative merit.\"  \n\nA: Thank you for this suggestion. We view TPA as an improvement to the standard dot-product attention mechanism, which is largely orthogonal to methods that fundamentally replace it (like linear attention or state-space models). TPA's factorization principles can be integrated with these other architectures.  To demonstrate this, we ran new experiments applying a linear attention mechanism to both a GQA baseline and our TPA architecture. Our results show that a TPA model with linearized attention outperforms the GQA model with linearized attention. For instance, in our medium model experiments, a linear GQA baseline achieved a 46.79% average 0-shot accuracy and 47.60% 2-shot accuracy, whereas a TPA variant with linear attention performed better with a 47.43% average 0-shot accuracy and 48.82% average 2-shot accuracy (see results shown in tables below). This suggests that the benefits of TPA’s contextual factorization are complementary and can enhance other efficient attention schemes. We will add this discussion and the new results to the revised paper.\n\n\n| Medium 0-shot | ARC-E | ARC-C | BoolQ | HellaSwag | OBQA | PIQA | W.G. | MMLU | SciQ | Avg. |\n| :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- |\n| linearattn-gqa-partialrope | 52.06 | 27.39 | 61.35 | 37.81 | 31.80 | 65.83 | 51.22 | 23.49 | 70.20 | 46.79 |\n| linearattn-tpa-partialrope | **53.83** | **28.24** | 56.54 | **38.55** | **32.40** | **66.32** | **53.51** | **24.24** | **73.20** | **47.43** |\n\n| Medium 2-shot | ARC-E | ARC-C | BoolQ | HellaSwag | OBQA | PIQA | W.G. | MMLU | SciQ | Avg. |\n| :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- |\n| linearattn-gqa-partialrope-decay | 56.14 | 27.30 | 59.94 | 37.21 | 31.20 | 65.56 | 50.36 | 25.37 | 75.30 | 47.60 |\n| linearattn-tpa-partialrope | **58.42** | **28.50** | 51.04 | **38.16** | **34.00** | **66.49** | 51.70 | **25.96** | **85.10** | **48.82** |\n\n\n\n> Q5: \"The paper does not sufficiently address limitations or potential negative societal impacts... The paper presents no discussion of failure modes. Are there cases where TPA underperforms or is harder to train?\"  \n\nA: We sincerely thank you for pushing for a more complete discussion of limitations and impacts. We acknowledge this was an oversight and will add a dedicated section in the revision.  \n\n1. TPA trains as stably as standard Transformers. The most critical factor we identified is parameter initialization. Similar to other factorized methods like LoRA, TPA's performance is sensitive to how the factor matrices are initialized. A potential failure mode is using a poor initialization scheme, which can hinder convergence; we found Xavier initialization to be crucial. We will also discuss that while performance is strong even at low ranks, an extreme rank setting (e.g., rank 1 for a highly complex task) would present a trade-off. Finally, we will note that our evaluation is confined to language, and performance on other modalities is an area for future work.  \n\n2. We will add a broader impact statement. We recognize that enabling more efficient, longer-context models is a dual-use technology. While beneficial for many applications, it could also lower the barrier for generating more coherent large-scale information. We will explicitly acknowledge these risks in the revised manuscript."}, {"review_id": "BsDkkDiJBo", "reviewer": "Reviewer_z6jQ", "summary": "This paper introduces Tensor Product Attention (TPA), a novel attention mechanism that factorizes query, key, and value tensors via contextual low-rank tensor products. TPA significantly reduces KV cache size during inference while maintaining or improving model performance. The authors build a new architecture called T6 (Tensor ProducT ATTenTion Transformer) and propose an efficient inference algorithm FlashTPA Decoding, demonstrating improved memory and computational efficiency over standard MHA, MQA, GQA, and MLA.", "questions": "1. Can TPA generalize beyond causal language modeling? Would it work effectively in encoder-decoder or bidirectional transformer settings?\n2. Are there any observed issues with training stability or convergence speed for TPA compared to MHA?\n3. Is there any learned redundancy between the AQ/BQ and AK/BK factors? Can the parameter count be further reduced?", "limitations": "There is no limitation mentioned in this paper.", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n1. The use of contextual tensor product factorization for Q/K/V is original and theoretically well-grounded.\n2. TPA reduces KV cache memory usage by up to 10×, enabling longer sequence lengths at inference time without sacrificing quality, highly relevant for LLM scalability.\n3. TPA is fully compatible with rotary positional embeddings.\n\nWeaknesses\n1. While the authors emphasize inference efficiency, it's unclear if the TPA factorization introduces training instability or overhead compared to MHA.\n2. The rank settings (e.g., RQ = 16, RK = 1) appear somewhat heuristic.", "was_revised": true, "final_justification": "i think the reviewers answered my questions well. So I will maintain my current score.", "rebuttal": "We are sincerely grateful to the reviewer for their positive assessment and support for our work. We appreciate the thoughtful questions, which help us clarify important aspects of TPA.\n\n> Q1: \"Can TPA generalize beyond causal language modeling? Would it work effectively in encoder-decoder or bidirectional transformer settings?\" \n\nA: Thank you for this excellent question. Yes, TPA is a general attention mechanism that can be readily applied beyond causal language modeling. TPA replaces the standard Q/K/V projection layer and is independent of the attention mask. By simply using a bidirectional attention mask (for encoders like BERT) or a combination of bidirectional and causal masks (for encoder-decoder models like T5), TPA can be seamlessly integrated into these architectures.\n\nGiven that standard Multi-Head Attention is a special, non-contextual case of TPA (as discussed in Appendix A), and MHA is the foundation for these other architectures, we are confident that TPA would also perform effectively in these settings, likely offering similar benefits in terms of expressivity and efficiency.\n\n> Q2: \"Are there any observed issues with training stability or convergence speed for TPA compared to MHA?\"  \n\nA: Thank you for asking about this crucial practical aspect. We found TPA to be as stable to train as its baseline counterparts.  \n1. Throughout our experiments across multiple model sizes, we observed no training instabilities such as exploding gradients or loss. The training and validation loss curves presented in the paper (e.g., Figures 3 and 4) are smooth and demonstrate stable convergence, similar to the MHA, GQA, and MLA baselines. \n\n2.  To quantify the training overhead, we benchmarked the throughput. The results show that TPA's training speed is highly competitive. For instance, on large-scale MoE models, its throughput is nearly identical to GQA's.     \n\n##### Dense Models (350M) on 8x A100 GPUs (TPS means tokens per second): \n\n| Method | MHA | MQA | GQA | TPA-KVonly | TPA | MLA |\n| :--- | :--- | :--- | :--- | :--- | :--- | :--- |\n| **TPS (M)** | 0.515 | 0.456 | 0.452 | 0.427 | 0.356 | 0.355 |\n\n##### MoE Models (>10B) on 64x GPUs (TPS means tokens per second):\n\n| Method / Hardware | GQA | TPA |\n| :--- | :--- | :--- |\n| H800 TPS (M) | 0.975 | 0.891 |\n| H20 TPS (M) | 0.407 | 0.405 |\n\nWe will add a full table of these throughput results to the appendix to provide a clear picture of the computational trade-offs.\n\n3. Furthermore, we are actively developing FlashTPA training/prefilling kernels. As analyzed in Appendix C, an optimized kernel that avoids materializing the full Q/K/V tensors can significantly reduce the FLOPs required for the attention computation. This gives TPA a clear path to potentially exceeding the training throughput of existing baselines. In summary, TPA already trains at a competitive speed for large scale model training, and we have a concrete plan for further optimization.\n\n> Q3: \"The rank settings (e.g., $R_Q = 16, R_K = 1$) appear somewhat heuristic.\"  \n\nA: That is an insightful observation. While the ranks are hyperparameters, our choices were not arbitrary but rather guided by hardware efficiency considerations. Modern deep learning accelerators like GPUs achieve maximum performance when computation sizes are aligned with their architectural specifics (e.g., multiples of 8 or 16 for tensor core operations).  \n\nFor our FlashTPA kernels, implemented in Triton, performance is best when block sizes are powers of two, typically starting at 16. Choosing ranks like $R_Q=16$ and $R_K=1$ or $R_K=2$ ensures that the resulting tensor operations map efficiently to the hardware, minimizing padding and maximizing utilization. Therefore, the settings were chosen as a practical balance between model expressivity and computational performance, a principle we will clarify in the paper.\n\n> Q4: \"Is there any learned redundancy between the A_Q/B_Q and A_K/B_K factors? Can the parameter count be further reduced?\"  \n\nA: This is a very insightful question. Yes, there is potential for parameter sharing to further reduce the model's parameter count. For instance, one could tie the weights that generate the head-dimension factors between query and key (i.e., use the same projection for $\\mathbf{A}_Q(\\mathbf{x}_t)$ and $\\mathbf{A}_K(\\mathbf{x}_t)$). This would create a more parameter-efficient variant of TPA, which could be advantageous in highly resource-constrained environments.  For this work, we chose the untied version to maximize the model's expressive capacity and performance. However, exploring different parameter-sharing strategies within the TPA framework is a promising direction for future research, and we thank you for highlighting it."}]

Author rebuttal:

Meta-review: This paper proses tensor product attention, an approach to shrink the KV cache by factorizing queries, keys, and values into contextual low-rank components. The approach matches or surpasses standard attention baselines including MHA, GQA, and MLA across a variety of metrics. The authors show promising scaling results, including up to 10B MoE's, with strong training throughput and quality.

Overall this is a strong paper, and I recommend it for a spotlight.
Decision: Accept (spotlight)