Tensor Product Attention Is All You Need

Thank you for your thoughtful review and constructive feedback on our submission. Your comments have provided valuable insights that will help us improve the clarity and impact of our work. We appreciate your positive assessment and recommendation. Below, we address each of your points in detail and outline our planned revisions.

1. Regarding the Presentation of Results:

Q1: "I think it would be nice if the results for 0-shot and 2-shot and small, medium, large, and XL models could be included in the main body of the paper somehow... The results in the main body currently seem a little cherry-picked."

A1: We appreciate your suggestion to enhance the presentation of our experimental results. We agree that a more comprehensive summary in the main text would provide a clearer and less selective view of TPA’s performance, directly addressing the concern about potential cherry-picking. To address this, we plan the following changes:

• We will replace the current Tables 2 and 3 with a new, consolidated table. This table will summarize the average performance (e.g., average perplexity and key downstream task accuracy) across small, medium, large, and XL model sizes for both 0-shot and 2-shot evaluations, offering a concise overview of TPA’s effectiveness compared to baselines like MHA, MQA, GQA, and MLA.
• The full, detailed results for each model size, task, and shot configuration, currently spread across the main text and appendix (Tables 2, 3, 6-11), will be consistently located in the appendix. This allows interested readers to explore specifics while keeping the main text focused.

2. Regarding the Claim of Superiority vs. Comparability:

Q2: "It seems to me that it's very fair to say that TPA is comparable to existing methods... but I don't know that the results support it being superior. I think the paper overstates this... I think it would be better to focus on the decrease in KV cache cost."

A2: We value your perspective on the tone of our performance claims and agree that a more nuanced presentation, emphasizing the significant memory efficiency alongside competitive performance, better reflects our contribution. We will revise the paper as follows:

• We will adjust the language throughout the paper—particularly in the abstract(), introduction, and conclusion—to state that TPA achieves competitive or comparable performance relative to baselines, rather than asserting outright superiority. For instance, the sentence you highlighted will be rephrased similarly to: "T6 achieves competitive performance compared to standard Transformer baselines... across various metrics, while offering significant memory savings."
• We will strengthen the focus on TPA’s primary practical advantage: KV cache reduction. In Section 3.3, we will further highlight and clearly present the quantification showing that TPA can reduce the KV cache size by approximately 10x compared to MHA for typical configurations, explicitly linking this to the ability to handle much longer sequences under memory constraints.

3. Regarding Typos and Mathematical Complexity (Other Comments):

Q3: "Last sentence on page 1 'The recently introduced...' has a typo. I found the math section a bit needlessly complex at times... I think it could be shortened and improved."

A3: Thank you for catching the typo and for your feedback on the mathematical presentation.

• We will correct the typo on page 1 and perform a thorough proofread of the entire manuscript.
• We will revise Sections 2 (Background) and 3 (Tensor Product Attention) to improve clarity and accessibility. This will involve consolidating notation where possible, ensuring all symbols are clearly defined upon first use, and refining explanations. We will also consider moving particularly lengthy derivations or detailed proofs (such as aspects of the FLOP analysis currently in Appendix A) to the appendix to maintain a smoother narrative flow in the main text, while ensuring the core methodology remains rigorously presented.

Thank you once again for your time and valuable suggestions. We are confident that incorporating these revisions will significantly strengthen the paper.

We sincerely appreciate your insightful feedback and constructive comments. We have carefully considered all concerns and we provide detailed responses to your questions below.

Q1: The paper describes a few different variants of their basic idea, such as making part of the QKV computation non-contextual, i.e. independent of token embedding and the experiments seem promising, albeit incomplete.

A1: Thank you for your attention to the non-contextual TPA. According to the experiment results of TPA small and medium models with non-contextual B shown below (compared with Tables 2, 6, 8, and 9 in the paper), the non-contextual versions are slightly worse than the current version of TPA. Therefore, to achieve better performance, the original TPA, TPA-KV only and TPA with non-contextual A are recommended.

Q2: The paper seems to only focus on order 2 tensor products and the equations in 3.1/3.2/3.3 seem to essentially implement low rank decompositions of Q, K, V matrices. Have you experimented with higher order tensors at all that are mentioned in Appendix B?

A2: We just implemented experiments on small-size models with third-order TPA and that with only KV factorized. The performances on small models are shown below. It is worse than other TPA series models. Therefore, for small models, we recommend the second-order TPA but the high-order TPA is still potential for much larger models.

Q3: At the end of Appendix A there is a brief comparison of TPA vs MLA in terms of the computation required and it’s claimed that d_rope + d_c can inflate the dot product cost by roughly 4.5x to 9x compared to MQA. The numbers used to arrive at this conclusion seem to contradict the conclusions in DeepSeek v3 where MLA only causes an increase of 2.25x vs MQA. Could you please elaborate on this part a little more.

A3: Thank you for your careful consideration of the comparison of computation between these attention mechanisms. Here's a clearer explanation:

• As described in Appendix A.8, the dot product cost is determined by the hyperparameters, including number of heads ($h$) and dimension for each head ($d_h$). Different sizes of these dimensions may result in differences in dot product computation cost. For MHA, MQA, GQA, the cost to compute $Q_i(x_T)K_i^\top$ is $\mathcal{O}(hd_hT)$ and for MLA, the cost is$\mathcal{O}(h(d_{\text{rope}}+d_c)T)$ if the current sequence length is $T$.

• For example, in DeepSeek V3, the hidden dimension is 7168 with 128 heads. Compared with MQA with dimension of 7168/128=56 for each head, MLA has $d_{\text{rope}} + d_c=192$. Therefore, according to the analysis in Appendix A.8, it has 192/56≈3.43x increase in computation cost but only causes an increase of 2.25x KV caches (as described in DeepSeek V2 technical report) and has better performance than MQA. For smaller models, larger RoPE and compressed representation dimensions are needed to keep the superior performance of MLA, so a larger multiplier of computation is expected.

• Moreover, calculating $QK^\top$ is only one part of the attention mechanism, and we still need to calculate $\text{softmax}(QK^\top)$. Loading cached KV from memory needs time, so the speed is IO-aware. In modern GPU, the decoding phase's IO speed is crucial, so the less KV cache size, the faster the decoding speed is given the same computational FLOPs.

Thank you again for bringing this to our attention. We will revise this section and provide a more detailed analysis of the computational costs associated with different attention mechanisms based on your suggestions. We look forward to your further evaluation and are happy to provide any additional explanations if needed.

Thank you for your thorough review and constructive feedback on our submission. We appreciate your recognition of the novelty and clarity, and your detailed questions help us refine our manuscript.

Q1: Despite reducing the number of parameters per token, TPA does not decrease the GPU memory usage. The matrices Q, K, V still have the shape [T, h, d_h], ...

A1: We appreciate your question about memory usage and performance.

• Memory Savings: TPA's primary memory gain is the inference KV cache. MHA caches full $K_t, V_t$ ($2 T h d_h$ memory). TPA caches only low-rank factors ($A_K(x_t), \tilde{B}_K(x_t)$), reducing per-token memory to $(R_K+R_V)(h+d_h)$. With typical low ranks ( $R_K= R_V=2$ ) and $d_h=12$, this yields substantial savings (>10x), enabling longer sequences. Furthermore, Appendix A details algorithms computing TPA attention without materializing full Q, K, V tensors. By working directly with factors, these methods can reduce computational cost (flops) and peak memory usage during the forward pass, complementing KV cache savings.

• Performance Source: TPA's improved performance stems from its Contextual Factorization. Unlike MHA's fixed projections, TPA dynamically constructs Q, K, V factors based on the token's hidden state $x_t$. This contextual adaptation provides an inductive bias, enhancing representational capacity. This is validated by TPA's consistently lower validation losses and perplexities (Figures 2, 3, 4). Appendix A further explores potential computational savings via factorized computation.

We will revise the paper to better distinguish computation memory (materialized vs. non-materialized) from inference KV cache benefits and clarify how contextual factorization drives performance.

Q2: There is a lack of fundamental theories to explain why TPA is superior to MHA and MQA. Does this stem from TPA's stronger expressive power?...

A2: We appreciate your concern regarding theoretical backing for TPA's superiority.

• Our core theoretical argument rests on Contextual Factorization. Section 3.4 demonstrates that MHA, MQA, and GQA are special cases of TPA where factors are restricted to be non-contextual (fixed). TPA generalizes these by allowing factors to depend dynamically on the input $x_t$. This context-dependency is proposed as the source of enhanced expressive power.
• Furthermore, TPA's seamless integration with RoPE (Theorem 1) ensures effective use of relative positional information, a crucial aspect where other efficient attention mechanisms sometimes struggle.
• While formal proofs of superior expressive power are future work, the theoretical framing via contextuality, combined with strong empirical validation across multiple scales and tasks (Section 4), provides support for TPA's design.

We will revise to expand on these theoretical foundations.

Q3: The attempt to unify MHA and other attention mechanisms in Section 3.4 seems rather forced. The fact that TPA can be expressed in the form of a tensor product...

A3: We appreciate your feedback that the unification felt "forced" and questioned its implication for expressive power at low ranks ($R\ll h$).

• The unification aims to show TPA's flexibility as a framework encompassing MHA/MQA/GQA as specific instances with non-contextual factors and particular ranks. TPA's innovation lies precisely in making these factors trainable and contextual.

• We acknowledge the expressiveness trade-off with rank R. However, even with $R\ll h$, our experiments consistently demonstrate that the contextual nature of TPA's factors provides performance advantages over the fixed, non-contextual factors of MHA/MQA/GQA, alongside significant memory savings.

We will revise Section 3.4 to better clarify that contextuality is the key differentiator and link it more explicitly to the empirical results.

Q4: Experimentally, it is unclear when to use TPA-KVonly and when to use TPA. Why does TPA-KVonly perform better on medium-sized models, while TPA is superior on large models?

A4: Thank you for feedback on variant choice and performance dynamics.

• Variant Performance: TPA factorizes Q/K/V; TPA-KVonly factorizes only K/V. Our results show TPA slightly outperforms TPA-KVonly on medium models (353M), while TPA-KVonly slightly leads on large/XL models (773M, 1.5B). (Note: This corrects the potential misreading in the review). The reasons need further study, possibly related to optimization/capacity trade-offs at scale.

• Significance: While gains over MHA may seem modest, they are consistent and meaningful for LLMs. Crucially, TPA offers a dual advantage: competitive performance PLUS substantial (>10x potential) inference KV cache reduction. This memory efficiency enables much longer sequences, addressing a critical scalability bottleneck. This practical benefit underscores TPA's value.

We will enhance the manuscript to clarify these points.