Tensor Product Attention Is All You Need

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.

Q1: "Can TPA generalize beyond causal language modeling? Would it work effectively in encoder-decoder or bidirectional transformer settings?"

A: 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.

Given 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.

Q2: "Are there any observed issues with training stability or convergence speed for TPA compared to MHA?"

A: Thank you for asking about this crucial practical aspect. We found TPA to be as stable to train as its baseline counterparts.

1. 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.

2. 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.

Dense Models (350M) on 8x A100 GPUs (TPS means tokens per second):

MoE Models (>10B) on 64x GPUs (TPS means tokens per second):

We will add a full table of these throughput results to the appendix to provide a clear picture of the computational trade-offs.

3. 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.

Q3: "The rank settings (e.g., $R_Q = 16, R_K = 1$) appear somewhat heuristic."

A: 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).

For 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.

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?"

A: 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.

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.

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$)?"

Thank you for these excellent suggestions to strengthen our analysis. We have performed additional ablations to address these points.

• 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.

• 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.

Medium 0-shot Results

Medium 2-shot Results

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."

A: 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.

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?"

A: 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:

• GSM8k: TPA achieves 26.5 vs. GQA's 19.1 (+38%)
• MATH: TPA achieves 14.0 vs. GQA's 11.2 (+25%)

These 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.

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."

A: 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.

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?"

A: 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.

1. 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.

2. 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.

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.

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%). ...

A1: 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.
While 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.

These 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.

Q2: 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.

A2: 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. Dense Models (350M) on 8x A100 GPUs (TPS means tokens per second):

Dense Models (350M) on 8x A100 GPUs (TPS means tokens per second):

MoE Models (>10B) on 64x GPUs (TPS means tokens per second):

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. We will add these results and discussion to the appendix.

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.

A3: 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.

Our >10B TPA-MoE model shows substantial improvements to these very tasks:

• GSM8k: TPA achieves 26.5 vs. GQA's 19.1 (+38% relatively)
• MATH: TPA achieves 14.0 vs. GQA's 11.2 (+25% relatively)

These results demonstrate that TPA's architecture leads to better performance on complex, multi-step reasoning problems that benefit from longer context.

In 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.

We 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.

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.

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.

A2: 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.

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.

A2: 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.
To 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.

These experiments (each consuming over 10,000 GPU hours) provide strong evidence that TPA's architecture leads to substantial quality gains in large-scale settings.

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.

A3: 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)}$.

Let's use our XL (1.5B) model as a concrete example:

• MHA baseline: $h=25, d_h=64$. KV Cache size = $2 \times 25 \times 64 = {3200}$ floats/token.
• TPA-KVonly ($R_{K,V}=2$): $h=47, d_h=64$. KV Cache size = $2+2)(47+64) = 4 \times 111 = {444}$ floats/token.
• This represents a 7.2x reduction ($3200 / 444$) compared to the MHA baseline of a similar parameter count.
• TPA ($R_{K,V}=2$): $h=78, d_h=64$. KV Cache size = $(2+2)(78+64) = 4 \times 142 = {568}$ floats/token.
• 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.

Q4:Typo: Line 121 should be $R_Q,R_K,R_V$

A4: Thank you for your suggestion, and we will correct it in the revised version.

Q5: Is the T6 easy to train compared to standard transformers? I think the training aspect is also worth studying.

A5: 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.

Dense Models (350M) on 8x A100 GPUs (TPS means tokens per second):

MoE Models (>10B) on 64x GPUs (TPS means tokens per second):

The 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.

As 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.

We thank the reviewer for their time and their thoughtful follow-up comment. We appreciate the opportunity to provide further clarification.

Q: It seems like the argument is: even though TPA might not have as strong expressivity as MHA when $R_Q, R_K, R_V$ are small, empirically it still shows descent performance compared to MHA?

A: Thank you for giving us the chance to clarify this key point. Our argument is not that TPA has weaker expressivity, but rather that it possesses a more effective and efficient inductive bias.

The key difference is that TPA's projections are contextual—they adapt to each input token—whereas MHA's are static. This dynamic, input-specific factorization provides a different, and we argue more powerful, form of expressivity that is highly effective even at very low ranks. The strong empirical results, particularly the significant performance gains on larger models, support this view: TPA achieves superior results by learning a more efficient representational subspace for attention, rather than by replicating the high-rank, non-contextual structure of MHA.

Q: The TPS table still shows that TPA is slower? I hope there will be more discussions on this.

A: You are correct that our current benchmarks show a modest training overhead for TPA in some settings, but is becomes negligible under large-scale setting. We believe this is a highly favorable trade-off, as this slight increase in training computation delivers substantial improvements in (1) final model quality, (2) inference memory efficiency (KV cache), and (3) inference decoding speed (with FlashTPA).

Furthermore, this overhead is an artifact of our current, unoptimized implementation. As we mentioned, we are developing optimized FlashTPA training kernels that operate directly on the factorized components (as analyzed in Appendix C). These kernels will significantly reduce the computational load during training, and we are confident they will make TPA's throughput highly competitive with, or even faster than, the baselines.

---

We sincerely appreciate your valuable feedback throughout this process. We hope that our clarifications regarding TPA's contextual expressivity and the training speed, along with the supporting data, might provide sufficient grounds for re-evaluation. Please do not hesitate to let us know if any further questions arise; we would be delighted to provide additional details.