ArXiv: 2501.06425
🎯 Pitch
Tensor Product Attention compresses KV caches by over an order of magnitude with minimal rank (e.g., rank‑1 or 2) while achieving lower perplexity and higher zero‑shot accuracy than Multi‑Head Attention, Multi‑Query Attention, and Multi‑Head Latent Attention. It also enables a custom decoding kernel that outperforms FlashAttention for long sequences, making precise long‑context inference both memory‑ and time‑efficient.
1. Executive Summary
This paper introduces Tensor Product Attention (TPA), a novel attention mechanism that factorizes query, key, and value activations using contextual tensor decompositions—summing rank-R tensor products of head-dimension and feature-dimension factors learned from the token's hidden state—to achieve substantial KV cache compression during autoregressive inference while matching or improving downstream performance relative to standard Multi-Head Attention. Through empirical evaluation on the FineWeb-Edu 100B language modeling dataset across four model scales (124M to 1.5B parameters), the authors build the Tensor ProducT ATTenTion Transformer (T6) architecture and demonstrate that TPA reduces KV cache memory per token from 2h·dh to (RK+RV)(h+dh)—enabling order-of-magnitude savings with small ranks like RK=RV=1 or 2—while achieving lower validation perplexity and higher zero-shot accuracy than MHA, MQA, GQA, and MLA baselines (e.g., 51.41% average for 353M TPA vs. 50.11% for MHA). The paper further introduces FlashTPA Decoding, a Triton-based inference algorithm that avoids materializing full Q/K/V tensors by computing attention directly in factor space via head-shared feature-dot products, establishing that TPA can decode faster than optimized FlashAttention kernels at long sequence lengths (e.g., surpassing MQA, GQA, and MLA at 219 tokens) while preserving the relative-position property of Rotary Position Embeddings (RoPE) through a pre-rotation strategy that stores already-rotated key factors in the cache.
2. Context and Motivation
The Core Problem: KV Cache Memory Dominates Long-Sequence Inference
The fundamental challenge this paper confronts is the memory bottleneck created by key-value (KV) caches during autoregressive inference in Transformer-based language models. When a Transformer generates text token by token, the attention mechanism must attend to all previously generated tokens. To avoid recomputing keys and values for every past token at each new generation step — which would make inference quadratically expensive in sequence length — standard practice caches these intermediate representations in memory. For each past token , Multi-Head Attention (MHA) stores the full key and value matrices , where is the number of attention heads and is the per-head dimension. Across a sequence of length , this accumulates to a memory footprint of floating-point values.
This becomes the dominant memory cost at inference time. As the authors note in Section 1:
"Because memory consumption grows linearly with sequence length, the maximum context window is limited by practical hardware constraints."
To make this concrete: consider a model with heads and (a typical configuration for models in the 1–7B parameter range). For a sequence of tokens — which is well within the context length that modern applications demand for document analysis, codebase understanding, and long-form reasoning — the KV cache alone requires million floating-point values, or roughly 1.6 GB in half-precision. For larger models with more heads or wider dimensions, this grows proportionally. And critically, this memory must be resident in high-bandwidth GPU memory to avoid crippling latency, directly competing with model weights, activations, and batch-level data.
This is not a hypothetical concern. The paper cites the widespread adoption of longer contexts in applications such as "document analysis, complex reasoning, and code completion" (Section 1), all of which require the model to maintain coherent attention over tens or hundreds of thousands of tokens. The KV cache is the single largest consumer of GPU memory in long-context inference, and it is the primary factor limiting how long a sequence a given hardware configuration can process.
The Gap: No Unified Framework for KV Cache Reduction That Preserves Quality and RoPE Compatibility
The paper identifies a concrete gap in the landscape of KV cache reduction techniques. Prior work falls into several categories, each with specific shortcomings:
1. Token eviction and sparse attention. Approaches like H2O (heavy-hitter oracle; Zhang et al., 2023) and StreamingLLM (Xiao et al., 2024) reduce memory by discarding or selectively retaining only a subset of cached tokens — e.g., keeping only tokens with the largest attention weights or maintaining attention sinks plus a sliding window. The paper acknowledges these methods in Section 1 and Appendix E, noting they "risk discarding tokens that may later prove important." This is a fundamental limitation: eviction strategies make irreversible decisions based on local attention patterns, and a token deemed unimportant at one step may become critical many steps later when a distant dependency is needed. For tasks requiring precise retrieval or long-range reasoning, these losses can be catastrophic.
2. Off-chip storage. Methods like FastDecode (He and Zhai, 2024) and InfiniGen (Lee et al., 2024) offload the KV cache to CPU memory or disk, fetching needed entries on demand. While this increases the effective cache capacity, the paper notes it comes "at the expense of increased I/O latency" (Section 1). The bandwidth between CPU and GPU memory (~50 GB/s for PCIe) is roughly an order of magnitude lower than GPU HBM bandwidth (~1–2 TB/s), making every cache fetch a potential stall. For latency-sensitive applications, this overhead is often unacceptable.
3. Attention architecture modifications: MQA and GQA. Multi-Query Attention (MQA; Shazeer, 2019) reduces KV cache size by sharing a single key and value head across all query heads. This drops the KV cache from to — a factor-of- reduction. Grouped-Query Attention (GQA; Ainslie et al., 2023) generalizes this by grouping the query heads into groups, each sharing one key-value pair, reducing the cache to . Both are widely deployed: GQA is used in LLaMA 2 and 3, MQA in PaLM. However, the paper identifies a key limitation: these methods "often compromise flexibility or require significant architectural modifications" (Section 1). More specifically, sharing keys and values across heads constrains the model's representational capacity — different attention heads can no longer attend to different aspects of the input using different key representations, which is the core motivation for multi-head attention in the first place (Vaswani et al., 2017). The quality degradation from MQA is well-documented, and while GQA partially recovers it, the model must still commit to a fixed degree of sharing at architecture design time.
4. Low-rank weight factorization (e.g., LoRA). LoRA (Hu et al., 2022) and its derivatives factorize weight matrices during fine-tuning — e.g., representing a weight update as with low-rank factors. The paper explicitly distinguishes its approach: LoRA "effectively reduce(s) fine-tuning memory, yet do(es) not address the KV cache overhead that dominates inference at runtime" (Section 1). LoRA compresses model parameters, not activations; the KV cache stores per-token activations computed from those parameters, and LoRA does nothing to reduce their size. TPA, by contrast, factorizes the activations themselves — the queries, keys, and values — which is what the KV cache stores.
5. Multi-head Latent Attention (MLA). The most direct predecessor to TPA is Multi-head Latent Attention (MLA), introduced in DeepSeek-V2 (Liu et al., 2024). MLA compresses the key and value representations into a low-dimensional latent space before expanding them back to per-head dimensions during attention computation. The compressed latent (where ) is what gets cached, rather than the full per-head key and value vectors. This achieves substantial KV cache reduction — DeepSeek-V2 reports a factor of ~6–10× compression relative to MHA.
However, the paper identifies a critical shortcoming of MLA: it "encounters difficulties with efficient Rotary Position Embedding (RoPE) integration, necessitating additional position-encoded parameters per head" (Section 1). The challenge, as detailed in Appendix F.3, stems from RoPE's mathematical structure. RoPE applies a position-dependent rotation matrix to queries and keys such that the inner product between position and depends only on the relative offset through the identity . In standard attention, this allows the attention score computation to be expressed as:
where the bracketed term could be pre-computed once for fast decoding. With RoPE applied to the up-projected keys and queries, the expression becomes:
The bracketed term now depends on the relative position , preventing pre-computation and defeating the purpose of the latent compression for fast decoding. MLA's workaround is to maintain a separate, smaller key component to which RoPE is applied and which must be cached in addition to the compressed latent — meaning the actual cache size is per token, with adding per-head overhead. As noted in the analysis at the end of Section F.3 and the influential blog post by Su (2024), this is an inherent tension in MLA: "the extreme pull between cache and effect."
Why This Problem Matters: Practical and Theoretical Significance
The importance of solving this problem extends beyond a narrow engineering concern. The paper positions KV cache efficiency as a scalability bottleneck for the entire trajectory of language model development:
Practical significance. As models scale to support ever-longer contexts (128K tokens in GPT-4 Turbo, 1M+ tokens in Gemini 1.5 Pro, 128K in Claude 3), the KV cache becomes the primary consumer of GPU memory, often dwarfing model weights. For a model with , , , and 7B parameters (~14 GB in half-precision), a 128K sequence produces a KV cache of roughly bytes ≈ 2 GB — already significant. For larger models (70B parameters) or batched inference (batch size 8, 16, 32), the cache can easily exceed available GPU memory, forcing either expensive model parallelism or context truncation. Reducing the per-token cache footprint directly increases the maximum sequence length a given GPU can handle, or equivalently, increases the batch size for a given sequence length — both of which directly translate to lower cost and higher throughput in production deployments.
Theoretical significance. The paper frames TPA as more than a compression technique — it is a unifying framework for understanding attention mechanisms. Section 4 demonstrates that MHA, MQA, and GQA are all special cases of TPA with non-contextual (input-independent) head-dimension factors. MHA corresponds to TPA with rank and the head-dimension factors fixed to standard basis vectors . MQA corresponds to TPA with and the single head-dimension factor fixed to the all-ones vector. GQA corresponds to TPA with and the head-dimension factors fixed to group-membership mask vectors. This unification provides intellectual clarity: the design space of attention mechanisms is revealed to be parameterized by (a) the ranks and (b) whether the factors are contextual (functions of the input) or non-contextual (fixed). TPA explores the previously unexplored region of this space where factors are both low-rank and contextual — yielding compression with expressiveness, rather than trading one for the other.
How TPA Positions Itself Relative to Existing Work
The paper draws explicit contrasts that define TPA's positioning:
Versus token eviction and sparse attention: TPA is orthogonal and potentially complementary. It reduces the size of what is cached rather than deciding which tokens to cache. A system could in principle combine TPA's factorized KV cache with an eviction policy, though this is not explored.
Versus MQA and GQA: TPA generalizes them. Where MQA and GQA enforce rigid parameter sharing (one or key-value heads for all queries), TPA uses low-rank factorizations that allow each head to construct its own key and value representations from shared, context-dependent basis vectors. The rank controls the degree of sharing, and this can be tuned continuously (via the choice of rank) rather than discretely (via the choice of ). Moreover, because the factors are contextual — they are functions of the token's hidden state — the sharing is dynamic and content-aware, unlike the static sharing in MQA/GQA.
Versus MLA: TPA addresses MLA's RoPE integration difficulty head-on. The paper's Theorem 3.1 proves that RoPE's relative-position property is preserved under TPA's factorization when the feature-dimension factors and are pre-rotated: RoPE distributes over the tensor product, so rotating the factors before forming the key tensor is equivalent to rotating the key tensor itself. This means TPA can cache pre-rotated key factors , eliminating any per-step rotation during decoding and avoiding MLA's need for a separate RoPE-specific cache component. The paper positions this as a key advantage: "TPA integrates seamlessly with RoPE and any possible position encodings" (Section 1, contribution 3).
Versus LoRA: The distinction is between factorizing weights (static) and factorizing activations (dynamic). LoRA represents weight updates as where and are low-rank and fixed after training. TPA represents per-token activations as where and are learned linear maps applied to the input . This is a fundamentally different factorization target with different implications: LoRA reduces the memory for storing fine-tuned model weights, but does nothing for the activations that dominate inference memory; TPA reduces the memory for storing those activations by compressing them at their source.
The Unifying Insight: Contextual Low-Rank Decomposition of Activations
The paper's core intellectual move is to observe that queries, keys, and values are activations — intermediate representations computed from the token's hidden state — and can therefore be factorized in a context-dependent way, unlike the fixed weight matrices that LoRA targets. The factorization in Equation 3.1:
where and are computed by learned linear projections from , decomposes each query slice into a sum of rank-1 matrices. When , this is a low-rank representation — but crucially, the factors themselves depend on the input, so the subspace in which the queries live adapts to the content of each token. This is qualitatively different from a fixed low-rank projection, which would constrain all tokens to the same subspace.
For the KV cache, the impact is direct: instead of storing the full and per token (totaling numbers), TPA stores only the factor matrices — , , and analogously for values — totaling numbers. For typical configurations with , , and , this gives numbers versus numbers — a 25.6× reduction in KV cache memory per token.
The paper thus positions TPA not merely as another compression trick but as a principled framework that (1) subsumes existing attention variants as special cases, (2) provides a continuous knob (the ranks) to trade compression for expressiveness, (3) maintains full RoPE compatibility, and (4) does so while improving, rather than sacrificing, model quality — a combination that no prior approach achieves simultaneously.
3. Technical Approach
3.1 Reader orientation
Tensor Product Attention (TPA) is a drop-in replacement for standard multi-head attention in Transformer models, one that computes queries, keys, and values through context-dependent low-rank tensor factorizations rather than through a single dense linear projection per head. The system solves the KV cache memory bottleneck by storing only the small factor matrices for each past token—typically a few hundred numbers instead of several thousand—and reconstructs the full key and value tensors on the fly during attention computation, achieving an order-of-magnitude memory reduction while simultaneously improving model quality through more expressive, input-adaptive representations.
3.2 Big-picture architecture (diagram in words)
The TPA layer replaces the standard attention block in a Transformer. Its components and their responsibilities are:
Input: The hidden state $x_t \in \mathbb{R}^{d_{\text{model}}}$ for the current token at a given layer.
Factor Projection (three parallel branches for Q, K, V): For each of queries, keys, and values, a set of learned linear maps projects $x_t$ into two lower-dimensional latent representations per rank component—one spanning the head dimension (producing $a_r(x_t) \in \mathbb{R}^h$, the "head factor") and one spanning the feature-per-head dimension (producing $b_r(x_t) \in \mathbb{R}^{d_h}$, the "feature factor"). Where standard attention would apply one matrix $W_i^Q \in \mathbb{R}^{d_{\text{model}} \times d_h}$ per head to produce the query vector directly, TPA instead learns $R_Q$ pairs of smaller matrices that collectively reconstruct the full query tensor.
RoPE Pre-Rotation (for Q and K branches only): Before the query and key tensors are assembled, Rotary Position Embeddings are applied directly to the feature-dimension factors $b_r^Q(x_t)$ and $b_r^K(x_t)$. This produces pre-rotated factors $\tilde{b}_r^Q(x_t)$ and $\tilde{b}_r^K(x_t)$ that already encode positional information. For keys, these pre-rotated factors are what gets cached, eliminating the need for any per-step rotation during autoregressive decoding.
Tensor Product Assembly (one per Q, K, V): For each of queries, keys, and values, the head-factor matrix and the (possibly pre-rotated) feature-factor matrix are combined via a tensor product contraction to reconstruct the full per-token attention tensor. For queries: $Q_t = \frac{1}{R_Q} A^Q(x_t)^\top B^Q(x_t) \in \mathbb{R}^{h \times d_h}$, where $A^Q(x_t) \in \mathbb{R}^{R_Q \times h}$ stacks the head factors as rows and $B^Q(x_t) \in \mathbb{R}^{R_Q \times d_h}$ stacks the feature factors as rows. This is a sum of $R_Q$ rank-1 matrices, each contributing one "mode" to the query representation.
Scaled Dot-Product Attention: The assembled $Q_t, K_t, V_t$ for all tokens enter standard scaled dot-product attention—identically to how they would in MHA. For each head $i$: $\text{head}_i = \text{Softmax}\left(\frac{1}{\sqrt{d_h}} Q_i K_i^\top\right) V_i$.
Output Projection: The concatenated per-head outputs (of total dimension $h \cdot d_h$) are projected back to $d_{\text{model}}$ via a learned matrix $W^O \in \mathbb{R}^{(h \cdot d_h) \times d_{\text{model}}}$.
KV Caching (the critical difference): During autoregressive inference, instead of storing the full $K_t \in \mathbb{R}^{h \times d_h}$ and $V_t \in \mathbb{R}^{h \times d_h}$ for each past token, TPA caches only the factor matrices: $A^K(x_t) \in \mathbb{R}^{R_K \times h}$, $\tilde{B}^K(x_t) \in \mathbb{R}^{R_K \times d_h}$ (pre-rotated), $A^V(x_t) \in \mathbb{R}^{R_V \times h}$, and $B^V(x_t) \in \mathbb{R}^{R_V \times d_h}$. The total per-token cache size is $(R_K + R_V)(h + d_h)$ numbers compared to $2 h d_h$ for MHA.
FlashTPA Decoding (the specialized inference engine): Rather than materializing the full key and value tensors from the cached factors and then performing standard attention, FlashTPA computes attention directly in factor space. It first computes head-shared dot products between the query's feature factor $B^Q$ and the cached key feature factor $\tilde{B}^K$, then mixes these with the head-specific factor $A^Q$ and cached $A^K$ to form logits, applies the softmax, and finally aggregates values using $A^V$ and $B^V$—all without ever constructing the $h \times d_h$ key or value matrices for any past token. The computation is blocked over the cache dimension and fused into a single Triton kernel using online log-sum-exp for numerical stability, analogous to the design of FlashAttention.
Feed-Forward Sub-Layer: Following the LLaMA architecture, each TPA layer includes a SwiGLU feed-forward network: $\text{FFN}(x) = (\text{SiLU}(x W_1) \odot (x W_2)) W_3$, with RMSNorm applied before both the attention and FFN sub-layers. The complete T6 block is $x \leftarrow x + \text{TPA}(\text{RMSNorm}(x))$ followed by $x \leftarrow x + \text{SwiGLU-FFN}(\text{RMSNorm}(x))$.
3.3 Roadmap for the deep dive
- First, the contextual tensor factorization (Section 3.1 of the paper): the mathematical form of the Q, K, V decomposition, the mapping from hidden states to latent factors, and the shapes of all tensors involved. This is the core representational innovation—understanding it is prerequisite to everything else.
- Second, RoPE integration (Section 3.2 of the paper): how Rotary Position Embeddings distribute over the tensor product, the pre-rotation strategy that enables caching of already-rotated keys, and Theorem 3.1 proving that relative position encoding is preserved. This addresses MLA's key limitation and is central to TPA's practical advantage.
- Third, KV caching and memory analysis (Section 3.3 of the paper): the precise memory footprint of TPA's factorized cache versus standard MHA cache, the compression ratio formula, and Table 1's comparison across all attention variants. This is where the engineering payoff is quantified.
- Fourth, TPA's relationship to MHA, MQA, and GQA (Section 4 of the paper): how each existing attention variant emerges as a special case of TPA with specific choices of ranks and non-contextual factors. This frames TPA as a unification and generalization rather than a wholly separate mechanism.
- Fifth, parameter initialization and TPA variants (Appendix G): the Xavier initialization scheme and the design space of TPA variants—non-contextual A vs. B factors, KV-only factorization, shared B factors, and nonlinear head factors. This is where the architectural flexibility is explored.
3.4 Detailed, sentence-based technical breakdown
This is primarily an architectural innovation paper whose core idea is that queries, keys, and values in attention can be factorized as sums of context-dependent rank-1 tensor products, and that caching only the factors—rather than the full tensors—dramatically reduces inference memory while the context-dependence of the factorization preserves (and in practice improves) model quality relative to static parameter-sharing approaches like MQA and GQA.
Contextual Tensor Factorization of Q, K, and V
The standard attention projection. In standard Multi-Head Attention (MHA), for each token with hidden state $x_t \in \mathbb{R}^{d_{\text{model}}}$, the query, key, and value for head $i$ are computed as:
where $W_i^Q, W_i^K, W_i^V \in \mathbb{R}^{d_{\text{model}} \times d_h}$ are learned projection matrices specific to head $i$. Stacking these across all $h$ heads yields the per-token slices $Q_t, K_t, V_t \in \mathbb{R}^{h \times d_h}$. The total number of parameters devoted to these projections is $3 \cdot h \cdot d_{\text{model}} \cdot d_h$—for queries, keys, and values combined.
TPA's factorization. Instead of a separate dense projection per head, TPA decomposes each per-token slice $Q_t, K_t, V_t$ into a sum of $R_Q, R_K, R_V$ rank-1 matrices, respectively, where each rank-1 matrix is the outer product of a head-dimension vector ($a_r(x_t) \in \mathbb{R}^h$, capturing variation across heads) and a feature-dimension vector ($b_r(x_t) \in \mathbb{R}^{d_h}$, capturing variation within each head's feature space). Formally, from Equation 3.1:
where $a_r^Q(x_t), a_r^K(x_t), a_r^V(x_t) \in \mathbb{R}^h$ are the head-dimension factor vectors, $b_r^Q(x_t), b_r^K(x_t), b_r^V(x_t) \in \mathbb{R}^{d_h}$ are the feature-dimension factor vectors, and $\otimes$ denotes the outer product: $(a \otimes b)_{ij} = a_i b_j$, producing an $h \times d_h$ matrix whose $(i,j)$ entry is the product of the $i$-th element of $a$ (corresponding to head $i$) and the $j$-th element of $b$ (corresponding to feature dimension $j$ within that head).
What the equation computes. For a single token $t$, the factorization takes $R_Q$ pairs of vectors—each pair consisting of one $h$-dimensional vector and one $d_h$-dimensional vector—forms the outer product of each pair (a rank-1 matrix of shape $h \times d_h$), sums these $R_Q$ rank-1 matrices, and scales by $1/R_Q$. The result is an $h \times d_h$ matrix whose rank is at most $R_Q$ (and typically exactly $R_Q$). Row $i$ of this matrix is the query vector for head $i$—it is formed as a weighted combination of the $R_Q$ feature-vectors $b_r^Q$, where the weights are the $i$-th components of the head-vectors $a_r^Q$. This means every head's query is constructed from the same $R_Q$ basis vectors in $\mathbb{R}^{d_h}$, but each head mixes them with different coefficients determined by its row of the $a$ factors.
Why this form. The factorization separates what varies across heads (captured by the $a$ factors, which live in $\mathbb{R}^h$) from what varies across feature dimensions within a head (captured by the $b$ factors, which live in $\mathbb{R}^{d_h}$). This is a principled decomposition because these two axes of variation have different semantic roles: heads specialize in different attention patterns (some may attend to syntax, others to semantics, others to position), while the feature dimensions within a head encode the specific content being attended to. By factorizing along these natural axes, a small number of rank components can capture the essential structure of the full $h \times d_h$ tensor. The rank $R_Q$ controls the capacity of this decomposition: $R_Q = 1$ means all heads share a single feature-space direction (akin to MQA for queries), while $R_Q = h$ with appropriate choice of $a$ factors can recover full MHA expressiveness. Crucially, because both $a_r^Q(\cdot)$ and $b_r^Q(\cdot)$ are learned functions of the input token $x_t$, the factorization is contextual: the subspace in which the queries live adapts to the content of each token, unlike a fixed low-rank projection (e.g., factorizing the weight matrix $W^Q$ once and applying it to all tokens), which would constrain all tokens to the same subspace regardless of their content.
Latent factor computation via linear projections. Each factor vector is produced by a learned linear transformation of the token's hidden state $x_t$. For the query factors, as stated in Section 3.1:
where $W_r^{a_Q} \in \mathbb{R}^{h \times d_{\text{model}}}$ and $W_r^{b_Q} \in \mathbb{R}^{d_h \times d_{\text{model}}}$ are learned weight matrices for the $r$-th rank component. Analogous matrices exist for keys and values. The total number of parameters devoted to the query factorization is $R_Q \cdot (h \cdot d_{\text{model}} + d_h \cdot d_{\text{model}}) = R_Q \cdot d_{\text{model}} \cdot (h + d_h)$, compared to $h \cdot d_{\text{model}} \cdot d_h$ for standard MHA queries. When $R_Q \ll \min(h, d_h)$, TPA uses fewer parameters for the query projection than MHA—a fact visible in Table 1, where TPA with $R_Q = 16, R_K = 1, R_V = 1$ at $d_{\text{model}} = 2048, h = 32, d_h = 64$ uses 7.7M attention parameters versus MHA's 16.8M.
Batched computation via reshaping. In practice, to avoid looping over rank components, the paper describes merging the rank index into the output dimension. For queries, the projections for all $R_Q$ rank components are computed simultaneously:
where $W^{a_Q} \in \mathbb{R}^{(R_Q \cdot h) \times d_{\text{model}}}$ and $W^{b_Q} \in \mathbb{R}^{(R_Q \cdot d_h) \times d_{\text{model}}}$ are the concatenated weight matrices for all rank components. These outputs are then reshaped into matrices $A^Q(x_t) \in \mathbb{R}^{R_Q \times h}$ (each row $r$ contains the head-factor vector $a_r^Q(x_t)^\top$) and $B^Q(x_t) \in \mathbb{R}^{R_Q \times d_h}$ (each row $r$ contains the feature-factor vector $b_r^Q(x_t)^\top$). The query tensor is then assembled via a batched matrix multiplication:
This is a contraction over the rank dimension: for each $(i, j)$ in the output, $(Q_t)_{ij} = \frac{1}{R_Q} \sum_{r=1}^{R_Q} A^Q(x_t)_{r,i} \cdot B^Q(x_t)_{r,j}$, which is exactly the sum of outer products described above. Repeating this computation for all $T$ tokens in the sequence (in parallel during training, or token-by-token during inference) produces the full $Q, K, V \in \mathbb{R}^{T \times h \times d_h}$ tensors needed for scaled dot-product attention.
Parameter count and the design trade-off. Table 1 provides the exact formula for TPA's attention-layer parameter count: $d_{\text{model}}(R_Q + R_K + R_V)(h + d_h) + d_{\text{model}} h d_h$. The first term covers all the factor-projection weights (for queries, keys, and values, each with $R_*$ rank components, each producing vectors of size $h$ and $d_h$ from an input of size $d_{\text{model}}$). The second term is the output projection $W^O \in \mathbb{R}^{(h \cdot d_h) \times d_{\text{model}}}$, which is identical to the output projection in standard MHA. The scaling factor $1/R_Q$ (and analogously $1/R_K, 1/R_V$) in the tensor product assembly normalizes the contribution of each rank component so that the variance of $Q_t$ does not grow with $R_Q$. The paper does not explicitly discuss why the $1/R_Q$ normalization (as opposed to $1/\sqrt{R_Q}$ or learned scaling), but it is consistent with the interpretation of $Q_t$ as an average of $R_Q$ rank-1 contributions, analogous to how multi-head attention averages (or concatenates) the outputs of multiple heads.
RoPE Compatibility and Acceleration
The challenge with factorized attention and RoPE. Rotary Position Embeddings (RoPE; Su et al., 2024) encode positional information by applying a position-dependent rotation matrix $T_t \in \mathbb{R}^{d_h \times d_h}$ to query and key vectors. $T_t$ is block-diagonal with $2 \times 2$ rotation blocks operating on consecutive pairs of feature dimensions: for dimension pair $(2j-1, 2j)$, the rotation is by angle $t \cdot \theta_j$ where $\theta_j = \text{base}^{-2j/d_h}$ with a typical base of 10,000. The critical property that makes RoPE work is that $T_t T_s^\top = T_{t-s}$: rotating a query at position $t$ and a key at position $s$ and then taking their dot product is equivalent to rotating the query by the relative position $t-s$ and dotting with the un-rotated key—or vice versa. This injects relative positional information directly into the attention scores:
TPA's integration strategy: pre-rotate the feature factors. TPA exploits the fact that RoPE is a linear transformation (specifically, an orthogonal rotation). Because the rotation distributes over the sum of outer products, applying RoPE to the assembled query tensor is equivalent to applying it to each feature factor $b_r^Q(x_t)$ before forming the outer product. The paper states this directly in Section 3.2 as the pre-rotation optimization for keys:
where the RoPE rotation is applied row-wise to $B^K(x_t)$—each row $r$ contains the feature factor $b_r^K(x_t)^\top \in \mathbb{R}^{1 \times d_h}$, and $\tilde{B}^K(x_t)$ contains the rotated versions. The pre-rotated key tensor is then formed as:
Theorem 3.1: RoPE's compatibility with TPA. The paper proves that this pre-rotation strategy preserves RoPE's relative-position property. The theorem states two equivalent formulations:
"Let
$Q_t$be factorized by TPA as$Q_t = \frac{1}{R_Q} A^Q(x_t)^\top B^Q(x_t) \in \mathbb{R}^{h \times d_h}$. Then$\text{RoPE}_t(Q_t) = Q_t T_t = \frac{1}{R_Q} A^Q(x_t)^\top \tilde{B}^Q(x_t)$where$\tilde{B}^Q(x_t) := B^Q(x_t) T_t$."
What this proves. The first statement establishes that RoPE distributes over the tensor product: rotating the assembled query tensor $Q_t$ on the right by $T_t$ is equivalent to rotating each feature factor $b_r^Q(x_t)$ on the right by $T_t$ before assembly. This works because $Q_t$ is a sum of outer products $a_r \otimes b_r$, and for any matrix $M$, $(a \otimes b)M = a \otimes (M^\top b)$ when the multiplication is on the right—here $T_t$ is applied to the rows of $B^Q$, which becomes $B^Q T_t$ in matrix form, matching the claimed identity. The second statement proves that the relative-position property is preserved:
"Let
$\tilde{Q}_t = \text{RoPE}_t(Q_t) = Q_t T_t$and$\tilde{K}_s = \text{RoPE}_s(K_s) = K_s T_s$. Then$\tilde{Q}_t \tilde{K}_s^\top = Q_t T_{t-s} K_s^\top$, equivalently$\text{RoPE}_{t-s}(Q_t) K_s^\top = \tilde{Q}_t \tilde{K}_s^\top$."
What this enables. Because the relative-position identity holds, TPA can cache the pre-rotated key factors $\tilde{B}^K(x_t)$ instead of the un-rotated factors $B^K(x_t)$. During autoregressive decoding:
- For each new token at position
$t$, the query factors are computed, pre-rotated (at negligible cost since it's just one token), and assembled into$\tilde{Q}_t$. - The cached pre-rotated key factors
$\tilde{B}^K(x_s)$for past tokens$s < t$are already stored in rotated form—no per-step rotation of the entire cache is needed. - The cached head factors
$A^K(x_s)$are position-independent and need no rotation. - The attention score computation proceeds directly using the assembled
$\tilde{Q}_t$and cached$A^K(x_s), \tilde{B}^K(x_s)$, either by materializing$\tilde{K}_s$on the fly or (more efficiently) by computing dot products in factor space.
Why this matters versus MLA. This is TPA's key architectural advantage over Multi-head Latent Attention (MLA). In MLA, the compressed latent $c_t^{KV} \in \mathbb{R}^{d_c}$ is designed to be RoPE-free, and a separate key component $K^R \in \mathbb{R}^{d_h^R}$ carries the RoPE information. This separate component must be cached per head (adding $h \cdot d_h^R$ to the cache size) and breaks the clean "cache a single compressed vector" abstraction. TPA avoids this entirely: because RoPE distributes over the tensor product, the rotation can be absorbed into the cached factors with no additional storage cost and no separate positional pathway. The cache is simply the pre-rotated feature factors plus the head factors—both already required by the factorization—with no extra components needed for positional encoding.
Practical flexibility. The paper notes that "depending on hardware and performance requirements, different RoPE integration strategies can be adopted for training and inference" (Section 3.2). During training (where the full sequence is processed in parallel), it may be more efficient to assemble the full $Q$ and $K$ tensors and apply RoPE in the standard way. During inference, the pre-rotation caching strategy eliminates per-step rotation of the entire key cache. This flexibility is possible precisely because Theorem 3.1 guarantees that both approaches compute identical attention scores.
KV Caching and Memory Reduction
Standard MHA KV caching. In standard Transformer autoregressive decoding, for each past token $t$, the inference engine caches the full key and value tensors:
Across all past tokens (length $M$ at the current decoding step), this consumes:
where the factor of 2 accounts for both keys and values. For a typical configuration with $h = 32$ and $d_h = 128$, each token adds $2 \times 32 \times 128 = 8,192$ numbers to the cache (approximately 16 KB in half-precision). For a sequence of length $M = 100,000$, this becomes ~1.6 GB for the KV cache alone.
TPA factorized KV caching. Instead of storing the materialized $K_t$ and $V_t$, TPA stores only their factor components. For each past token, the cache holds:
where $\tilde{B}^K$ is the pre-rotated key feature factor (so no rotation is needed during decoding) and $B^V$ is the un-rotated value feature factor (values do not receive positional embeddings). The total per-token memory is:
Compression ratio. The ratio of TPA cache size to MHA cache size is:
For the paper's typical experimental configuration—$h$ adjusted per model scale to match MHA's parameter count, $d_h = 64$ fixed, $R_K = R_V = 2$—and using the large model's $h = 61$ (from Table 10, which shows $h = 61$ for the 773M TPA model), the ratio is $(2+2)(61+64) / (2 \times 61 \times 64) = 4 \times 125 / (2 \times 3904) = 500 / 7808 \approx 0.064$—a ~15.6× reduction. With $R_K = R_V = 1$, the ratio drops to $2 \times 125 / 7808 \approx 0.032$—a ~31× reduction. The paper states that "for large $h$ and $d_h$ (typically $d_h = 64$ or 128), setting $R_K, R_V \ll h$ (e.g., rank 1 or 2) often yields substantial reduction of KV cache size" (Section 3.3), and the numbers bear this out: even at rank 2, the savings are an order of magnitude.
Comparison with other attention variants (Table 1). The paper provides a comprehensive comparison in Table 1, reproduced here in conceptual form:
| Method | KV Cache per Token | Parameters per Attention Layer |
|---|---|---|
| MHA | $2 h d_h$ | $4 d_{\text{model}} h d_h$ |
| MQA | $2 d_h$ | $2 d_{\text{model}} d_h (h + 1)$ |
| GQA | $2 G d_h$ | $2 d_{\text{model}} d_h (h + G)$ |
| MLA | $d_c + d_h^R$ | (see Table 1) |
| TPA | $(R_K + R_V)(h + d_h)$ | $d_{\text{model}}(R_Q+R_K+R_V)(h+d_h) + d_{\text{model}} h d_h$ |
Several observations from this table:
- MQA achieves the smallest absolute cache (
$2d_h$vs. TPA's$(R_K+R_V)(h+d_h)$) when$h$is large and$R_K+R_V \geq 1$. For$h = 32, d_h = 64, R_K = R_V = 1$, MQA uses 128 numbers per token versus TPA's$2 \times 96 = 192$. However, MQA's quality degradation from sharing a single KV head across all query heads is well-documented—the paper's experiments confirm this (MQA's 50.44% average vs. TPA's 51.41% for 353M models in Table 2). - GQA interpolates, using
$2 G d_h$cache. For$G = 4, h = 32, d_h = 64$, this is$2 \times 4 \times 64 = 512$numbers—more than TPA's 192 at rank 1. At larger group sizes, GQA's cache grows while TPA's can remain fixed at low rank. - MLA's cache includes
$d_c + d_h^R$, where$d_c$is the compressed KV latent dimension and$d_h^R$is the RoPE residual dimension. The paper does not provide the exact values used in its MLA experiments, but typical values from DeepSeek-V2 are$d_c = 512$and$d_h^R = 64$per head, totaling$512 + 64 = 576$for single-head RoPE or more if per-head RoPE is needed. TPA's rank-1 cache at$h = 32, d_h = 64$uses 192 numbers—roughly 3× smaller than MLA's typical configuration. - Parameter counts vary. TPA can use fewer parameters than MHA (7.7M vs. 16.8M at
$d_{\text{model}} = 2048$) when ranks are low, because the factorized projections$d_{\text{model}}(R_Q+R_K+R_V)(h+d_h)$can be cheaper than the per-head dense projections$3 \cdot h \cdot d_{\text{model}} \cdot d_h$. The paper's experiments control for this by adjusting$h$to equalize total parameter counts across mechanisms.
The cache-quality trade-off. The critical empirical finding that justifies TPA is that this memory reduction comes with improved, not degraded, model quality. Tables 2 and 3 show TPA matching or exceeding MHA, MQA, GQA, and MLA on downstream benchmarks at equal parameter counts. This distinguishes TPA from MQA and GQA, which achieve memory reduction at the cost of quality (MQA) or partial quality recovery with larger cache (GQA). The mechanism for this quality improvement is the contextual nature of TPA's factorization: because the head factors $a_r^K(x_t)$ and feature factors $b_r^K(x_t)$ are functions of the token's content, the model can dynamically allocate representational capacity—on some tokens, the factors may produce keys that are effectively full-rank (if the factors span diverse directions); on others, they may be highly structured. Standard MHA has fixed, input-independent projections that cannot adapt their effective rank to the input.
Expressing MHA, MQA, and GQA as Non-Contextual TPA
The unification framework. Section 4 of the paper demonstrates that standard attention mechanisms emerge as special cases of TPA when the head-dimension factors $a_r(\cdot)$ are constrained to be non-contextual—that is, fixed vectors independent of the input token $x_t$. This reveals TPA not as an entirely new mechanism but as a generalization that relaxes the fixed structure of existing approaches to make it input-dependent.
MHA as non-contextual TPA with $R_Q = R_K = R_V = h$. To recover standard MHA within TPA, set the rank equal to the number of heads ($R_Q = h$) and define, for each rank component $i \in \{1, \ldots, h\}$:
- Contextual token factor:
$b_i^Q(x_t) = (W_i^Q)^\top x_t \in \mathbb{R}^{d_h}$—this is exactly the standard per-head query projection for head$i$. - Non-contextual head factor:
$a_i^Q = h \cdot e_i \in \mathbb{R}^h$—this is a scaled standard basis vector, where$e_i$has a 1 at position$i$and zeros elsewhere. The scaling by$h$cancels the$1/R_Q = 1/h$prefactor.
Substituting into the TPA formula:
What this produces. The outer product $e_i \otimes v$ for $v \in \mathbb{R}^{d_h}$ produces an $h \times d_h$ matrix whose $i$-th row is $v^\top$ and all other rows are zero. Summing over all $i = 1, \ldots, h$ stacks these rows: the resulting $Q_t \in \mathbb{R}^{h \times d_h}$ has exactly the $i$-th standard MHA query vector $(W_i^Q)^\top x_t$ in its $i$-th row. The same construction applies to keys and values, producing exactly the standard MHA tensors.
Why this matters. This shows that MHA corresponds to a maximally expressive but maximally rigid instantiation of TPA: the rank is as large as possible ($R_Q = h$), each rank component is dedicated to exactly one head (via the one-hot basis vector $e_i$), and there is no sharing of information across heads in the factorization (each head's $b$ factor is independently projected). The contextual dependence in this case is entirely in the $b$ factors—the $a$ factors contribute nothing beyond routing each $b$ to exactly one head. TPA relaxes both constraints: ranks can be smaller than $h$ (enabling compression), and the $a$ factors can be contextual (enabling dynamic, content-dependent head interactions).
MQA as rank-1 TPA for K and V. Multi-Query Attention shares a single key and value across all $h$ query heads. In TPA terms:
where $R_K = 1$, and the single head factor is non-contextual and set to $a^K = \mathbf{1}_h \in \mathbb{R}^h$ (the all-ones vector). The single feature factor is $b^K(x_t) = (W^K)^\top x_t \in \mathbb{R}^{d_h}$, the standard (shared) key projection. The outer product $\mathbf{1}_h \otimes b^K(x_t)$ produces an $h \times d_h$ matrix where every row is identically $b^K(x_t)^\top$—every head sees the same key vector. The same logic applies to values. The queries remain full-rank ($R_Q = h$) with the MHA-style non-contextual head factors, so each head still has its own query projection.
GQA as rank-G TPA for K and V. Grouped-Query Attention groups $h$ heads into $G$ groups, with heads in the same group sharing a key and value. This is TPA with $R_K = R_V = G$. For each group $j \in \{1, \ldots, G\}$:
- Contextual feature factor:
$b_j^K(x_t) = (W_j^K)^\top x_t \in \mathbb{R}^{d_h}$, the shared key projection for group$j$. - Non-contextual head factor:
$a_j^K = G \cdot \text{mask}_j \in \mathbb{R}^h$, where$\text{mask}_j$is a binary vector with 1s at positions corresponding to heads in group$j$and 0s elsewhere. The scaling by$G$cancels the$1/R_K = 1/G$prefactor.
The sum $\sum_{j=1}^G a_j^K \otimes b_j^K(x_t)$ produces a key tensor where each head's row contains the key vector for its assigned group. For example, with $h = 8, G = 2$, the factor for the first group of 4 heads would be $a_1^K = 2 \cdot [1,1,1,1,0,0,0,0]^\top$; the outer product $a_1^K \otimes b_1^K(x_t)$ produces a matrix whose first 4 rows are $2 \cdot b_1^K(x_t)^\top$ and last 4 rows are zero. Adding the analogous term for group 2 fills in the remaining rows with $2 \cdot b_2^K(x_t)^\top$; the $1/G = 1/2$ prefactor scales everything down, yielding the correct per-head key vectors.
The conceptual shift TPA introduces. In all three existing mechanisms, the head-dimension factors $a$ are fixed and non-contextual, serving merely as routing matrices that assign feature vectors to heads (MHA routes each $b_i$ to head $i$ via $e_i$; MQA routes the single $b$ to all heads via $\mathbf{1}_h$; GQA routes each $b_j$ to its group via $\text{mask}_j$). TPA's innovation is to make these factors contextual: $a_r^K(x_t) = W_r^{a_K} x_t$ is a learned function of the input, so the way feature vectors are combined across heads adapts to the token's content. This is what enables TPA to use fewer rank components than MHA (compressing the representation) while maintaining or exceeding MHA's quality—the dynamic routing can allocate the limited representational budget where it's most needed on a per-token basis.
Parameter Initialization and TPA Variants
Xavier initialization for factor weights. The paper specifies in Appendix G that all factor weight matrices are initialized using Xavier (Glorot) uniform initialization (Glorot and Bengio, 2010). Specifically, each entry of a weight matrix $W \in \mathbb{R}^{n_{\text{out}} \times n_{\text{in}}}$ is drawn from $\mathcal{U}(-\text{bound}, \text{bound})$ where $\text{bound} = \sqrt{6 / (n_{\text{in}} + n_{\text{out}})}$. For a factor projection matrix like $W_r^{a_Q} \in \mathbb{R}^{h \times d_{\text{model}}}$, the bound is $\sqrt{6 / (d_{\text{model}} + h)}$. This initialization is chosen "to help maintain the variance of activations and gradients as they propagate through the network layers, contributing to stable training" (Appendix G). The paper does not ablate initialization schemes, so the sensitivity of TPA training to this choice is not characterized.
TPA with non-contextual A factors (the MHA/MQA/GQA recovery case). This is the variant explored in Section 4, where the head-dimension factors $a_r^Q, a_r^K, a_r^V \in \mathbb{R}^h$ are fixed (non-contextual) learned parameters rather than functions of $x_t$. The paper reports evaluation results for this variant (labeled "TPA (non-ctx-A)" in Tables 5–8). For small models (124M), non-contextual-A TPA achieves 45.03% 0-shot average accuracy—comparable to the full contextual TPA's 46.21% (Table 5 vs. Table 11). For medium models (353M), non-contextual-A TPA achieves 50.52% 0-shot average (Table 7), which is still competitive with full TPA's 51.41% (Table 2). Under 2-shot evaluation for medium models, non-contextual-A achieves 51.98% vs. full TPA's 53.12%. This suggests that the contextual head factors provide a modest but consistent improvement, and that even non-contextual TPA (which is closer to a learned, soft version of GQA's fixed group assignments) performs well. The KV cache savings are identical in both variants since the head factors must be cached regardless.
TPA with non-contextual B factors. The dual variant fixes the feature-dimension factors $b_r^Q, b_r^K, b_r^V \in \mathbb{R}^{d_h}$ as learned parameters while keeping the head factors contextual: $a_r^Q(x_t) = W_r^{a_Q} x_t$. This is described in Appendix G as potentially effective "if the fundamental token-level features (captured by $b_r$) are relatively stable, while their combination across heads (captured by $a_r(x_t)$) needs to adapt to the context." The experimental results in Tables 5–8 show that non-contextual-B underperforms non-contextual-A and full TPA: 43.66% 0-shot for small models (vs. 45.03% for non-ctx-A and 46.21% for full TPA), and 48.19% 0-shot for medium models (vs. 50.52% and 51.41%). This suggests that making the feature factors contextual is more important than making the head factors contextual—unsurprising, since feature factors encode the content of what each head attends to (the "what"), while head factors encode the relative importance across heads (the "how much"). Content varies substantially across tokens (different words require attending to different semantic features), while the relative importance of different attention patterns may be more stable.
TPA-KVonly: factorizing only keys and values. A simpler variant "maintains the standard query projection mechanism but still achieves significant KV cache reduction through factorized key and value representations" (Appendix G). Here, queries are computed as $Q_t = W^Q x_t \in \mathbb{R}^{h \times d_h}$ (standard dense projection), while keys and values use the TPA factorization. The KV cache savings are identical to full TPA since queries are not cached. The parameter count (Table 1) is $d_{\text{model}}(R_K + R_V)(h + d_h) + 2 d_{\text{model}} h d_h$—note the $2 d_{\text{model}} h d_h$ term accounts for the standard query projection ($d_{\text{model}} h d_h$) plus the output projection ($d_{\text{model}} h d_h$). This variant is extensively evaluated and performs on par with or slightly below full TPA: for large models, TPA-KVonly achieves 53.52% 0-shot accuracy (Table 3) versus TPA's 53.10%, and for XL models (Table 12), TPA-KVonly with $R_{K,V} = 4$ achieves 55.03% versus TPA's 55.01%. The paper's main experiments use TPA-KVonly as a simpler, strong baseline that isolates the benefit of KV factorization without the added complexity of query factorization.
TPA KV with shared B factors. A further parameter-reduction variant "shares the token-dimension factors $b_r$ between keys and values: $b_r^K(x_t) = b_r^V(x_t)$" (Appendix G). This reduces both parameter count (fewer $b$-factor projection matrices) and KV cache footprint (only one set of $b$ factors to cache for both K and V). The paper notes that "although it constrains $K_t$ and $V_t$ to be constructed from the same token-level basis vectors, this variant can still offer strong performance with additional memory savings." No experimental results are reported for this variant, so its empirical trade-off is not characterized.
Nonlinear head factors. The paper mentions the possibility of applying nonlinearities (sigmoid or softmax) to the head-dimension factors. Applying softmax to the head factors "could be interpreted as a form of Mixture-of-Heads, where the network learns to dynamically weight different head configurations based on the input context" (Appendix G). No experiments with nonlinear head factors are reported. This is presented as a conceptual extension illustrating the flexibility of the TPA framework.
Higher-order TPA (Appendix C). The paper generalizes the second-order factorization (two factors per rank component) to third and higher orders. In a third-order TPA, the query tensor is:
where $b_r^Q(x_t) \in \mathbb{R}^{d_b}$ and $c_r^Q(x_t) \in \mathbb{R}^{d_c}$ first form an outer product $b_r^Q(x_t) \otimes c_r^Q(x_t) \in \mathbb{R}^{d_b \times d_c}$, which is vectorized into a $d_h = d_b \cdot d_c$ dimensional vector. The additional factor "can be viewed as a learnable, context-dependent modulation or gating term for the features generated by $b_r^Q(x_t)$" (Appendix C). Theorem C.1 proves that RoPE compatibility extends to this higher-order case, with a specific block-diagonal structure for the rotation matrix that applies RoPE to the $b$ factors while leaving the $c$ factors' structure intact. The paper reports preliminary small-model evaluation of third-order TPA (Table 4): 44.56% 0-shot average and 45.28% 2-shot average, which is below the second-order TPA's 46.21% 0-shot and 47.93% 2-shot (Table 11 and 13). However, this is with un-optimized hyperparameters and without controlling for parameter count—the paper notes that "a comprehensive comparison with second-order TPA variants of similar parameter counts or ranks would be necessary to fully evaluate the trade-offs" (Appendix C.1). Higher-order TPA introduces a trade-off: it "might allow for the use of smaller base ranks to achieve comparable representational power" but "increases the parameter count for the factors."
Why these variants matter for system design. The existence of multiple variants—non-contextual A vs. B, KV-only vs. full TPA, rank choices, higher-order extensions—demonstrates that TPA is not a single fixed architecture but a design space parameterized by (1) which factors are contextual, (2) the ranks for Q, K, and V independently, (3) whether nonlinearities are applied, and (4) the factorization order. This flexibility allows practitioners to navigate the trade-off between compression, quality, and computational cost along multiple continuous axes (rank values) rather than discrete architectural choices (MHA vs. GQA vs. MQA). The paper's experiments focus on a specific subset of this space (contextual A and B, second-order, modest ranks), establishing that this region achieves strong results, while the framework's generality leaves room for further optimization in specific deployment contexts.
4. Key Insights and Innovations
Innovation 1: Contextual Tensor Factorization as a Unifying Framework for Attention
The paper's deepest conceptual contribution is not the compression technique itself, but the reframing of attention mechanism design as a choice in a unified factor space. Prior to TPA, the landscape of attention variants was a collection of point solutions—MHA (Vaswani et al., 2017), MQA (Shazeer, 2019), GQA (Ainslie et al., 2023), MLA (Liu et al., 2024)—each making a different architectural choice about parameter sharing and dimensionality, and each understood as a distinct mechanism with its own implementation, trade-offs, and limitations. There was no common language for comparing them beyond table entries listing KV cache sizes and parameter counts.
TPA changes this by revealing that all of these mechanisms occupy different points in a single, continuous design space parameterized by the ranks RQ, RK, RV and the choice of whether head-dimension factors a_r(·) are contextual (input-dependent) or non-contextual (fixed). Section 4 demonstrates this concretely:
- MHA is TPA with RK = RV = h and non-contextual head factors a_i = h · e_i (standard basis vectors that route each feature factor to exactly one head).
- MQA is TPA with RK = RV = 1 and a^K = 1_h (the all-ones vector broadcasting one key to all heads).
- GQA is TPA with RK = RV = G and head factors set to scaled group-membership mask vectors.
This is not merely a taxonomic observation—it reveals that the core design choice distinguishing these mechanisms is whether the head factors can adapt to input content. In MHA, MQA, and GQA, the a factors are fixed routing matrices determined at architecture design time and applied identically to every token. In TPA, a_r(x_t) = W_r^a x_t is a learned function of the token's hidden state—meaning which features each head attends to is dynamically recomputed per token. This is a qualitatively different capability: instead of heads being assigned fixed roles (head 1 always sees key projection 1; heads 1–4 always share key A), the model can learn to allocate its limited representational budget where it's most useful for the specific content of each token.
The significance of this reframing extends beyond TPA itself. It provides a generative theory of attention mechanisms: new variants can be designed by making different choices along these axes (contextual vs. non-contextual A, contextual vs. non-contextual B, varying ranks independently for Q/K/V) rather than by proposing entirely new architectures. The paper's own exploration of TPA variants—non-contextual A (Tables 5–8), non-contextual B (Tables 5–8), KV-only, shared B factors, higher-order (Appendix C), nonlinear head factors—demonstrates this generative capacity in action. Each variant is not a separate paper's contribution but a straightforward configuration change within the same framework. From an intellectual standpoint, this is comparable to how the Chinchilla scaling laws (Hoffmann et al., 2022) unified the pretraining compute allocation problem around a parametric form—TPA does for attention architecture what Chinchilla did for training budgets, but at the mechanism-design level rather than the resource-allocation level.
The framework also explains why prior attention variants exhibit their known weaknesses. MQA's quality degradation is not mysterious—it arises because RK = 1 forces all heads to share a single key feature direction, and the non-contextual all-ones a factor provides no mechanism for heads to differentiate what they attend to based on content. GQA partially recovers quality by increasing rank to G, but the non-contextual mask factors mean the grouping is rigid and content-independent—head 1 always shares with heads 2–4 of its group, regardless of whether the current token's content would benefit from a different grouping. TPA's quality improvement (51.41% vs. MHA's 50.11% for 353M models in Table 2) is not a mystery either—it arises because the contextual factors allow dynamic, per-token allocation of representational capacity that is strictly more expressive than fixed routing, and the low-rank structure acts as a regularizer that prevents overfitting to spurious head specializations.
This unification is a fundamental conceptual advance rather than an incremental improvement. It changes how researchers should think about attention design—from "which variant should I use?" to "where in the TPA design space should I operate for my target compression-quality trade-off?"
Innovation 2: RoPE Compatibility Through Structural Transparency
The paper's treatment of Rotary Position Embeddings represents a genuine architectural insight rather than an engineering workaround. The challenge with RoPE in compressed attention mechanisms has been recognized since MLA (Liu et al., 2024; see also the detailed analysis by Su, 2024): RoPE applies a position-dependent rotation T_t ∈ R^{d_h × d_h} to query and key vectors, and its relative-position property relies on the identity T_t T_s^⊤ = T_{t-s}. In MLA, this creates a fundamental tension because the compressed key latent c_t^{KV} is designed to be RoPE-free (to enable the pre-computation that makes MLA fast), but RoPE must be applied somewhere to encode position. MLA's solution—a separate RoPE-specific key component K^R ∈ R^{d_h^R} that is cached per head in addition to the compressed latent—is an architectural compromise that adds parameters, increases cache size, and breaks the conceptual cleanliness of "cache one compressed vector per token."
TPA resolves this tension not by working around RoPE but by exploiting a structural property of the tensor product factorization itself. The key insight, formalized in Theorem 3.1, is that RoPE distributes over the sum of outer products: rotating the assembled query tensor Q_t on the right by T_t is mathematically equivalent to rotating each feature factor b_r^Q(x_t) before assembly. This is not an approximation or a trick—it follows directly from the bilinearity of the outer product and the fact that RoPE is a linear transformation. The proof (Appendix D.1) shows:
RoPE_t(Q_t) = Q_t T_t = (1/R_Q · A^Q(x_t)^⊤ B^Q(x_t)) T_t = 1/R_Q · A^Q(x_t)^⊤ (B^Q(x_t) T_t) = 1/R_Q · A^Q(x_t)^⊤ B̃^Q(x_t)
What makes this an innovation rather than an obvious consequence is the caching implication: because the identity holds, TPA can cache the pre-rotated key factors B̃^K(x_s) = B^K(x_s) T_s directly. During autoregressive decoding, the cache contains A^K(x_s) (head factors, position-independent) and B̃^K(x_s) (feature factors, already rotated). When a new query at position t arrives, its feature factors are rotated on-the-fly (cost: one rotation for the single new token), and attention scores are computed against the cached pre-rotated keys with zero additional per-step rotation cost. The relative-position property is preserved without MLA's separate RoPE pathway, without additional cached parameters, and without any modification to the RoPE mechanism itself.
The structural transparency of this solution is what distinguishes it. MLA's approach can be characterized as "compress first, then patch RoPE compatibility externally via an auxiliary pathway." TPA's approach is "design the factorization so that RoPE composes naturally with it." This is a design philosophy innovation rather than just a technical fix: it suggests that attention compression should be designed with positional encoding compatibility as a first-class constraint, not retrofitted. The paper's Theorem C.1 further demonstrates this philosophy by extending the compatibility proof to higher-order TPA, showing that the same principle—apply RoPE to the rightmost factors in the tensor product chain and let the block-diagonal structure of the rotation matrix handle the rest—generalizes cleanly.
The practical significance of this is substantial. MLA's need for a separate RoPE pathway has been described as "the extreme pull between cache and effect" (Su, 2024)—it forces a trade-off between compression and position encoding quality. TPA demonstrates that this trade-off is not inherent to compressed attention but is an artifact of MLA's specific compression strategy. By revealing that a different factorization (one that operates on the head × feature axes rather than compressing into a single latent) composes naturally with RoPE, the paper shows that the "extreme pull" can be resolved architecturally rather than through compromise.
Innovation 3: The Contextual Factorization-Quality Connection
Perhaps the most empirically surprising result in the paper is that TPA does not merely match MHA's quality while reducing memory—it improves quality. For the 353M model scale (Table 2), TPA achieves 51.41% average zero-shot accuracy versus MHA's 50.11%, MQA's 50.44%, GQA's 50.35%, and MLA's 50.13%. This improvement persists across model scales: for 773M models (Table 3), TPA-KVonly achieves 53.52% versus MHA's 52.52%; for 1.5B models (Table 12), TPA-KVonly with RK,V = 4 achieves 55.03% versus MHA's 54.49%. The validation loss curves (Figure 4) show TPA and TPA-KVonly consistently below MHA from early in training through convergence. This is not a compression-at-the-cost-of-quality trade-off—it's a compression-with-quality-gain outcome.
The paper's explanation for this is implicit but clear from the architecture: TPA's factorization acts as a structured regularizer. In standard MHA, each of the h heads learns an independent d_model × d_h projection matrix—there are h × d_model × d_h parameters for queries alone, with no explicit mechanism encouraging heads to share structure or avoid redundancy. Heads can (and empirically do) learn overlapping or redundant attention patterns, consuming parameters and cache memory without improving representational diversity. TPA constrains the query/key/value spaces to at most R_Q, R_K, R_V degrees of freedom in the feature dimension, but does so in a context-dependent way that allows the model to allocate those degrees of freedom where they're most useful per token.
This is a fundamentally different mechanism from the parameter sharing in MQA or GQA. In MQA, the constraint is hard and static: all heads must use exactly the same key, regardless of content. In GQA, the constraint is hard but block-structured: heads in group 1 share key 1, heads in group 2 share key 2, and these groupings are fixed. In TPA, the constraint is soft and dynamic: the model can produce up to R_K distinct feature directions, and each head's key is a weighted combination of these directions (weighted by the head factors a_r^K). Which directions are produced and how they're weighted depends on the input token x_t. This means the model can, on a per-token basis, choose to allocate its limited rank budget to produce diverse, specialized key directions (by making the b_r^K vectors span different subspaces and the a_r^K vectors have distinct patterns) or to produce more uniform keys (by making the a_r^K more uniform), depending on what the content demands.
The regularization interpretation is supported by the TPA variant results: non-contextual A TPA (where head factors are fixed) achieves 50.52% for medium models (Table 7), which is still competitive with MHA's 50.11% (Table 2) and better than MQA and GQA. This suggests that even the low-rank structure alone—absent contextual adaptation—provides beneficial regularization by preventing heads from learning fully independent, potentially redundant projections. The full contextual TPA's additional gain (51.41%) comes from making this regularization content-adaptive.
This insight reframes the conversation around attention compression. The field's default assumption—evident in the design of MQA and GQA—is that reducing KV cache size necessarily involves sacrificing representational capacity: you share keys across heads (MQA) or groups (GQA), and you accept the quality degradation as the price of memory savings. TPA demonstrates that this trade-off is not fundamental—it's an artifact of doing the sharing in a content-independent way. By making the sharing learned and context-dependent, the constraint becomes a feature rather than a bug, guiding the model toward more efficient use of its representational budget while simultaneously reducing memory. The empirical evidence for this interpretation is the combination of (a) lower perplexity, (b) higher downstream accuracy, and (c) smaller KV cache—three desiderata that prior work assumed were in tension.
Innovation 4: Factor-Space Inference as a Computational Primitive
The FlashTPA Decoding algorithm (Section 5 and Appendix B) represents a distinct insight from the memory-reduction contribution: that attention can be computed directly in the factor space without ever materializing Q, K, V, and that doing so is not merely an implementation detail but can be faster than optimized implementations that do materialize these tensors.
Standard attention implementations (including FlashAttention) materialize Q, K, V ∈ R^{h × d_h} per token—either by computing them from scratch for all tokens (during training) or by fetching K, V from the cache and computing Q for the new token (during decoding). The dominant FLOPs in decoding are the attention score computation (Q K^⊤, costing 2Mhd_h for cache length M) and value aggregation (α V, also costing 2Mhd_h). The cache is large (2Mhd_h numbers for MHA) but the computation is structured and well-optimized.
FlashTPA reorders this computation to exploit the factorized representation of the cache. Instead of reconstructing the M cached key matrices K_s ∈ R^{h × d_h} and then computing dot products with the single query Q_t, it computes:
-
Head-shared feature-space dot products between the query's feature factors B^Q (size R_Q × d_h) and the cached key feature factors B̃^K (size M × R_K × d_h), producing P ∈ R^{M × R_Q × R_K} at cost Θ(M R_Q R_K d_h).
-
Per-head rank mixing: combines P with the head factors A^Q (size R_Q × h) and cached A^K (size M × R_K × h) to form logits L ∈ R^{h × M} at cost Θ(h M R_Q R_K).
-
Online softmax and value aggregation using cached A^V (size M × R_V × h) and B^V (size M × R_V × d_h) at cost Θ(h M R_V d_h).
The total cache-dependent FLOPs are Θ(M R_Q R_K d_h + h M R_Q R_K + h M R_V d_h). For the paper's typical configuration with R_Q = 16, R_K = R_V = 1, this is Θ(M · 16 · d_h + h M · 16 + h M · d_h). Standard MHA decoding uses Θ(2hMd_h). For h = 32, d_h = 64, MHA uses Θ(4096M) per step. FlashTPA uses Θ(1024M + 512M + 2048M) = Θ(3584M). The asymptotic constant is similar or slightly better—but the key advantage is that FlashTPA never materializes the full M × h × d_h key cache, working instead with the factorized cache of size M × (R_K + R_V) × (h + d_h). For RK = RV = 1, this is ~25× smaller than MHA's cache, meaning far less memory bandwidth pressure.
The experimental results in Figure 5 confirm that this translates to wall-clock speed for long sequences. At d_model = 2048 and batch size 1, FlashTPA is comparable to MQA/GQA at short sequences but becomes faster than all baselines including MLA at sequence lengths beyond 2^14–2^15, with the gap widening at 2^19 (the maximum tested). At larger batch sizes (8, 16), the advantage is even more pronounced. This is not a trivial consequence of the smaller cache—MLA also has a smaller cache than MHA but is sometimes slower than FlashTPA at long sequences in these benchmarks. The speedup comes from the combination of reduced memory traffic (fewer bytes to fetch from the cache) and the computational structure of the factor-space contractions, which are amenable to efficient block-wise execution and fusion in the Triton kernel.
The intellectual contribution here is the recognition that factorized representations enable not just storage compression but computational reorganization. The field's default mental model for compressed attention has been: (1) compress the cache for memory savings; (2) decompress on-the-fly for computation; (3) accept the decompression overhead as the price of memory reduction. FlashTPA challenges step (2): by computing attention directly in factor space, it avoids decompression overhead entirely and can actually reduce total computation relative to the uncompressed baseline. This is a different category of result from "compression with minimal quality loss"—it's an existence proof that factorized attention can be faster than dense attention at equivalent quality, at least in the long-sequence regime where memory bandwidth dominates.
5. Experimental Analysis
Evaluation Methodology
-
Dataset. All language modeling experiments use the FineWeb-Edu 100B dataset (Lozhkov et al., 2024), which contains 100 billion tokens for training and 0.1 billion tokens for validation. This is an educational-content-filtered subset of the FineWeb corpus, chosen as a representative large-scale pretraining dataset. Models are trained for 50B tokens—roughly half an epoch over the full dataset.
-
Base model(s). The paper uses a LLaMA-style Transformer architecture (Touvron et al., 2023) with SwiGLU activations (Shazeer, 2020) and RoPE embeddings (Su et al., 2024) as the base architecture. Models are trained at four scales: Small (124M parameters), Medium (353M), Large (773M), and XL (1.5B). The head dimension is fixed at dh = 64 across all scales. The number of attention heads h is adjusted for each attention mechanism variant to ensure that all mechanisms have the same total number of parameters per attention layer as standard MHA (which has 4 · d_model^2 parameters per attention layer). Specific head counts for each scale and variant are provided in Table 10—for example, at the Large (773M) scale, MHA uses h = 20, GQA uses h = 38, MLA uses h = 34, TPA-KVonly uses h = 37, and full TPA uses h = 61.
-
Metrics. The primary training metric is validation loss (cross-entropy loss on the held-out FineWeb-Edu validation set, reported as a function of training tokens). PPL curves (validation perplexity, the exponential of validation loss) are also presented. For downstream evaluation, the paper reports accuracy on nine standard benchmarks using the lm-evaluation-harness codebase (Gao et al., 2024): ARC-Easy, ARC-Challenge (Yadav et al., 2019), BoolQ (Clark et al., 2019), HellaSwag (Zellers et al., 2019), OpenBookQA (Mihaylov et al., 2018), PIQA (Bisk et al., 2020), WinoGrande (Sakaguchi et al., 2020), MMLU (Hendrycks et al., 2021), and SciQ. For ARC-E, ARC-C, HellaSwag, OBQA, PIQA, and SciQ, the paper reports accuracy_norm; for other tasks, standard accuracy. Both 0-shot and 2-shot settings are evaluated. The paper reports the average accuracy across all nine benchmarks as a summary statistic.
-
Baselines. Four attention mechanisms are compared against:
- MHA (Multi-Head Attention; Vaswani et al., 2017): standard per-head query, key, and value projections.
- MQA (Multi-Query Attention; Shazeer, 2019): multiple query heads share a single key-value head.
- GQA (Grouped-Query Attention; Ainslie et al., 2023): queries are divided into G groups, with each group sharing one key-value head. The paper uses G = 2 KV heads (i.e., G = 2 groups) across all scales, as specified in Appendix H.1.
- MLA (Multi-head Latent Attention; Liu et al., 2024): the DeepSeek-V2 attention mechanism with low-rank key-value compression and a separate RoPE residual pathway. The paper uses residual key dimension d_h^R = 32 and other hyperparameters as specified in Table 10.
All baselines are implemented within the same LLaMA+SwiGLU+RoPE architecture, differing only in the attention sub-layer. The paper also evaluates two TPA variants—full TPA (with query, key, and value factorization) and TPA-KVonly (factorizing only keys and values, with standard dense query projection)—as the proposed methods. Default ranks for TPA models are RQ = 6 and RK = RV = 2, unless otherwise specified.
-
Generation budget / compute accounting. For the language modeling experiments, the "compute budget" is measured in training tokens—all models are trained for 50B tokens with a global batch size of 480, and comparisons are made at equal numbers of training tokens seen. For the decoding speed benchmarks (Section 6.2), compute is measured in wall-clock time per token (log2(seconds)) as a function of sequence length, with all methods running on comparable hardware. For these decoding experiments, the "generation budget" is not explicitly defined—the comparison is latency-oriented rather than FLOPs-oriented, with batch size, embedding dimension, and sequence length varied across experiments. Parameter counts for each attention variant are equalized by adjusting the number of heads h, as described above—this ensures that differences in training dynamics or final quality are attributable to the attention mechanism's inductive bias rather than raw capacity.
-
Cross-validation / statistical protocol. The paper does not employ cross-validation in the traditional sense—there is no hyperparameter search over validation folds. Training is done with a single set of hyperparameters per model scale (learning rates, optimizer settings, etc., as specified in Appendix H.1). The validation set (0.1B tokens from FineWeb-Edu) is used for monitoring training progress and reporting final validation loss/perplexity, but no early stopping or hyperparameter selection is described as being based on this validation set. The downstream evaluation uses the standard lm-evaluation-harness, which does not involve cross-validation. The paper does not report error bars or confidence intervals for any of its results, noting in the NeurIPS checklist that "error bars are not reported because it would be too computationally expensive for repeated experiments on LLMs." This means the results should be interpreted as single-run outcomes without formal statistical significance characterization.
Main Quantitative Results
Training and Validation Dynamics
Validation loss trajectories (Figure 4). Across all model scales—Medium (353M), Large (773M), and XL (1.5B)—TPA and TPA-KVonly converge as fast as or faster than the MHA baseline while achieving visibly lower final validation losses. In Figure 4(a) (353M), the red (TPA) and pink (TPA-KVonly) curves lie consistently below the grey (MHA) curve from approximately 5B tokens onward, with the gap widening through 50B tokens. In Figure 4(b) (773M), TPA and TPA-KVonly remain below MHA at nearly all training stages. In Figure 4(c) (1.5B), the same pattern holds. MLA (blue curves) consistently trains more slowly and yields higher validation losses than TPA and MHA across all scales—a finding the paper attributes to MLA's architectural complexity and training dynamics.
Training loss trajectories (Figure 3). The same rank-ordering holds for training loss: TPA and TPA-KVonly track MHA closely or lie below it, while MLA trains more slowly. At the 1.5B scale (Figure 3c), the TPA variants show visibly lower training loss than MHA from early in training.
Validation perplexity (Figure 9). Mirroring the loss curves, TPA and TPA-KVonly achieve lower validation perplexity than MHA, MQA, GQA, and MLA by the end of pretraining. For the 353M models (Figure 9a), TPA reaches the lowest perplexity, with TPA-KVonly a close second, both below MHA. For the 773M models (Figure 9b), the gap between TPA variants and MHA is clearly visible throughout the latter half of training. For the 1.5B models (Figure 9c), TPA-KVonly achieves the lowest perplexity, with full TPA slightly above but still below MHA.
What these curves demonstrate. The training dynamics are notable for what they do not show: there is no evidence of slower convergence, training instability, or degraded optimization from the factorized parameterization. TPA's factorization involves a more complex computational graph (projections into factor space, tensor product assembly) than standard MHA's straightforward matrix multiplications, and it would have been plausible for this to manifest as slower training or the need for careful learning rate tuning. The fact that TPA trains as fast as or faster than MHA—while using rank-constrained representations—suggests that the factorization provides a beneficial inductive bias that accelerates optimization rather than impeding it. This is consistent with the regularization interpretation discussed in Section 4.
Downstream Benchmark Performance
Medium-scale models (353M) — Table 2 (0-shot) and Table 14 (2-shot). This is the most comprehensive comparison point with all baselines at the same parameter count. The headline 0-shot results:
- MHA: 50.11% average across 9 benchmarks.
- MQA: 50.44% average.
- GQA: 50.35% average.
- MLA: 50.13% average.
- TPA-KVonly: 50.04% average—essentially tied with MHA.
- TPA (full): 51.41% average—1.3 percentage points above MHA, and the best overall.
The full TPA advantage is not driven by a single benchmark but is distributed: TPA outperforms MHA on ARC-E (58.38 vs. 59.51), ARC-C (31.57 vs. 29.52), BoolQ (59.39 vs. 59.60), HellaSwag (46.83 vs. 45.68), OBQA (37.00 vs. 34.20), PIQA (70.02 vs. 68.82), WinoGrande (54.06 vs. 53.43), and SciQ (79.90 vs. 76.90). MHA's only win is MMLU (25.52 for TPA vs. 23.33 for MHA). TPA-KVonly (50.04%) is slightly below MHA (50.11%), indicating that query factorization contributes a modest but non-trivial portion of full TPA's advantage.
Under 2-shot prompting (Table 14), the pattern strengthens:
- MHA: 52.34% average.
- TPA (full): 53.12% average—a 0.78 point advantage.
- TPA-KVonly: 52.69% average—0.35 points above MHA.
MLA (51.78%) and MQA (51.77%) both underperform MHA in the 2-shot setting, which is consistent with these methods' known quality limitations when pushed beyond their training distribution (few-shot prompting requires the model to leverage in-context examples, which may benefit from richer key-value representations).
Large-scale models (773M) — Table 3 (0-shot) and Table 15 (2-shot). The 0-shot results:
- MHA: 52.52% average.
- MQA: 52.13% average.
- GQA: 52.30% average.
- MLA: 53.32% average—notably, MLA outperforms MHA at this scale.
- TPA-KVonly: 53.52% average—1.0 point above MHA, and the best overall.
- TPA (full): 53.10% average—0.58 points above MHA.
The TPA-KVonly advantage here is notable because it achieves the highest average accuracy while using a simpler architecture than full TPA (no query factorization). Its wins are broad: it outperforms MHA on ARC-E (63.26 vs. 59.93), ARC-C (34.13 vs. 33.62), BoolQ (61.96 vs. 61.93), PIQA (72.09 vs. 71.06), WinoGrande (55.25 vs. 55.41), MMLU (26.06 vs. 22.87), and SciQ (81.10 vs. 81.20). Full TPA at this scale shows a smaller advantage (53.10% vs. 52.52%)—still positive but less pronounced than at 353M.
Under 2-shot prompting (Table 15):
- MHA: 54.25% average.
- MLA: 55.74% average—the highest 2-shot score, suggesting MLA benefits more from in-context examples at this scale than TPA does.
- TPA-KVonly: 55.33% average—1.08 points above MHA, second only to MLA.
- TPA (full): 55.02% average—0.77 points above MHA.
This is the one setting where MLA convincingly outperforms TPA (55.74% vs. 55.33% for TPA-KVonly). The paper does not discuss this result in detail, but it suggests that at certain scales and with few-shot prompting, MLA's separate RoPE pathway and higher-dimensional latent representations may provide benefits that TPA's simpler factorization does not capture.
XL-scale models (1.5B) — Table 12 (0-shot) and Table 16 (2-shot). The 0-shot results include an ablation of TPA-KVonly at different key/value ranks:
- MHA: 54.49% average.
- MQA: 54.48% average—essentially tied with MHA.
- TPA-KVonly (RK,V = 2): 54.57% average—0.08 points above MHA.
- TPA-KVonly (RK,V = 4): 55.03% average—0.54 points above MHA.
- TPA-KVonly (RK,V = 6): 55.03% average—tied with RK,V = 4.
- TPA (full): 55.01% average—0.52 points above MHA.
The rank ablation reveals that increasing RK,V from 2 to 4 improves performance substantially (54.57% → 55.03%), but further increasing to 6 provides no additional gain, suggesting diminishing returns at this scale. Full TPA (55.01%) is slightly below TPA-KVonly at RK,V = 4 (55.03%), consistent with the large-model pattern where TPA-KVonly matches or exceeds full TPA.
Under 2-shot prompting (Table 16), the picture is somewhat different:
- MHA: 56.18% average.
- MQA: 56.74% average—notably outperforms MHA in 2-shot at this scale.
- TPA-KVonly (RK,V = 4): 56.93% average—0.75 points above MHA.
- TPA-KVonly (RK,V = 6): 57.11% average—0.93 points above MHA, the best overall.
- TPA (full): 56.28% average—only 0.10 points above MHA.
The 2-shot results at XL scale show that MQA performs surprisingly well (56.74%), challenging the narrative that MQA's quality degradation is uniformly worse than MHA—at this scale with in-context examples, the simple key-value sharing may benefit from reduced overfitting. TPA-KVonly with RK,V = 6 achieves the best score (57.11%), suggesting that higher KV ranks become more beneficial at larger model scales and with few-shot prompting.
Small-scale models (124M) — Tables 11 (0-shot) and 13 (2-shot). At the smallest scale, the differences between methods are less systematic:
- 0-shot: MHA achieves 46.19%, TPA achieves 46.21% (essentially tied), MLA achieves 46.20%.
- 2-shot: MHA achieves 48.11%, TPA achieves 47.93% (slightly below MHA), TPA-KVonly achieves 47.85%.
At this scale, the regularization benefit of TPA's factorization may be less pronounced because even standard MHA at 124M parameters is heavily under-parameterized for the 100B-token dataset and thus already strongly regularized by capacity constraints.
KV Cache Efficiency at the Evaluated Configurations
The paper does not report a standalone table of KV cache sizes at inference time for its trained models, but these can be computed from the architecture specifications in Tables 9 and 10. For the Large (773M) model configuration with d_model = 1280 and the head counts from Table 10:
- MHA: h = 20, dh = 64, KV cache per token = 2 × 20 × 64 = 2,560 numbers.
- MQA: h = 39, dh = 64, KV cache per token = 2 × 1 × 64 = 128 numbers (factor of 20× reduction vs. MHA).
- GQA (G=2): h = 38, dh = 64, KV cache per token = 2 × 2 × 64 = 256 numbers (factor of 10× reduction vs. MHA).
- MLA: h = 34, dh = 64, dc = 512 (from Table 10), d_h^R = 32 (from Appendix H.1). The cache per token is dc + d_h^R = 512 + 32 = 544 numbers, or possibly more if d_h^R is per-head (which would make it 512 + 34 × 32 = 1,600 numbers). The paper does not specify which MLA configuration was used for caching in the experiments, so an exact comparison cannot be made from the provided data.
- TPA (RK = RV = 2): h = 61, dh = 64, KV cache per token = (2 + 2)(61 + 64) = 4 × 125 = 500 numbers (factor of ~5.1× reduction vs. MHA; 3.9× more than MQA but with higher quality).
- TPA (RK = RV = 1): Would use 2 × 125 = 250 numbers per token (~10.2× reduction vs. MHA), though this configuration is not evaluated in the main language modeling experiments (only TPA-KVonly at rank 1 appears in the decoding benchmarks).
The critical point from these numbers is that TPA achieves its quality improvements—outperforming MHA on downstream benchmarks—while using a KV cache that is 5.1× smaller than MHA's (for RK = RV = 2). MQA achieves a larger cache reduction (20×) but at the cost of quality degradation, which TPA avoids. GQA achieves a comparable cache reduction (10×) but with lower quality than TPA. The paper's contribution is thus demonstrating that the quality-cache trade-off can be shifted: it is possible to be simultaneously better than MHA on quality AND smaller than MHA on cache.
Decoding Speed Benchmarks (FlashTPA)
Figure 5 (d_model = 2048, dh = 64). The primary decoding speed results compare FlashTPA against FlashMHA, FlashGQA, FlashMQA, and FlashMLA across sequence lengths from 2^12 (4,096) to 2^19 (524,288) and batch sizes 1, 2, 4, 8, 16. FlashTPA uses RQ = 16, RK = RV = 1 and is implemented in Triton (Tillet et al., 2019). The comparison methods use their respective optimized CUDA implementations.
- At batch size 1 (Figure 5a): At short sequences (~2^12), MQA, GQA, MLA, and TPA all cluster closely, with MHA being significantly slower. As sequence length increases, FlashTPA and MLA are the fastest two methods, with FlashTPA pulling slightly ahead at sequence lengths beyond 2^15. At 2^19, FlashTPA is visibly faster than MLA.
- At batch sizes 2–4 (Figures 5b–c): The pattern is similar, with FlashTPA and MLA as the top two performers, FlashTPA maintaining a slight edge at long sequences.
- At batch sizes 8–16 (Figures 5d–e): FlashTPA's advantage becomes more pronounced. At batch size 16 and 2^19 sequence length, FlashTPA is noticeably faster than MLA, and the gap between FlashTPA and GQA/MQA is very large (approximately 2–3× slower for GQA/MQA at the longest sequences).
- Across all configurations, MHA (orange) is consistently the slowest, with its decoding time increasing most rapidly with sequence length—consistent with its larger KV cache and higher memory bandwidth demands.
Figure 7 (d_model = 3072) and Figure 8 (d_model = 1024). These show the same trends at different embedding dimensions. At d_model = 3072 (Figure 7), FlashTPA and MLA are the top two performers, with FlashTPA consistently matching or exceeding MLA's speed at long sequences. At d_model = 1024 (Figure 8), the ranking is similar but the gaps between methods are narrower at short sequences—consistent with the smaller KV caches at this dimension reducing memory bandwidth pressure.
What these speed benchmarks demonstrate. The FlashTPA decoding algorithm is competitive with or faster than highly optimized CUDA implementations of MQA, GQA, and MLA—despite being implemented in Triton, which typically has higher overhead than hand-tuned CUDA kernels. The paper explicitly notes this: "It is important to note that our current FlashTPA implementation utilizes Triton. While the compared methods are typically available as highly optimized CUDA kernels, these experiments provide initial insights into FlashTPA's potential. Development of a CUDA-based FlashTPA kernel is ongoing and is expected to yield further performance improvements" (Section 6.2). The fact that a Triton implementation can match or beat CUDA kernels—particularly for MLA, which is also a compressed attention mechanism—suggests that the factor-space computation in FlashTPA is genuinely more efficient than the decompress-then-attend approach, not merely on par with it.
Ablation Studies and Robustness Checks
Key/value rank ablation on XL models (Figure 10 and Table 12): For XL-scale (1.5B) TPA-KVonly models, increasing the key/value ranks from RK = RV = 1 (the most aggressive compression) to RK = RV = 2, 4, and 6 reveals consistent improvements in training loss, validation loss, and perplexity (Figure 10). The 0-shot evaluation in Table 12 quantifies this: TPA-KVonly with RK,V = 2 achieves 54.57% average; RK,V = 4 achieves 55.03%; RK,V = 6 also achieves 55.03%. The diminishing returns between rank 4 and 6 suggest that at this model scale and dataset size, rank 4 is sufficient to capture the essential structure of the key and value spaces. The paper does not report the corresponding 2-shot results for the rank-1 configuration, but the trend at ranks 2–6 is clear: higher ranks help, with the benefit saturating at rank 4.
Non-contextual A vs. B vs. full TPA (Tables 5–8 for small and medium models): This ablation isolates the contribution of contextual factors. For medium models (353M):
- TPA (non-contextual A): 50.52% 0-shot (Table 7) — head factors fixed, feature factors contextual. This is 1.19 points above MHA's 49.44% in the same table (note: the MHA baseline in Table 7 uses learning rate 3 × 10^{-4}, not the standard 6 × 10^{-4}, so these numbers are not directly comparable to Table 2).
- TPA (non-contextual B): 48.19% 0-shot — feature factors fixed, head factors contextual. This is below MHA.
- TPA (full, contextual A and B): 51.41% 0-shot (Table 2) — both factors contextual. This is 1.30 points above MHA's 50.11% in the same table.
The key finding is that making the feature factors contextual is substantially more important than making the head factors contextual: non-contextual A (contextual B) achieves near-MHA quality, while non-contextual B (contextual A) drops below MHA. This is consistent with the interpretation that the feature factors encode what content each head attends to (which varies strongly across tokens), while the head factors encode the relative importance of different attention patterns (which may be more stable). The full contextual TPA achieves the best results, with the head-factor contextuality providing a modest additional gain on top of the feature-factor contextuality.
Learning rate sensitivity (medium models, Figure 11, Tables 17–18): An alternative learning rate of 3 × 10^{-4} (vs. the default 6 × 10^{-4}) was tested on medium-scale models. Figure 11 shows that the relative ordering of methods is preserved: TPA and TPA-KVonly maintain lower validation loss and perplexity than MHA, MQA, GQA, and MLA throughout training, even at this different learning rate. The 0-shot evaluation (Table 17) gives:
- MHA: 49.44% average.
- TPA: 50.88% average — a 1.44 point advantage.
- TPA-KVonly: 49.93% average — a 0.49 point advantage.
The 2-shot evaluation (Table 18) gives:
- MHA: 51.77% average.
- TPA: 53.04% average — a 1.27 point advantage.
These results demonstrate that TPA's quality advantage is not an artifact of a specific learning rate choice—it persists (and in fact is slightly larger) at a lower learning rate.
TPA variants at small scale (Tables 5–6 for non-contextual variants, Table 4 for third-order TPA): For small models (124M), non-contextual A TPA achieves 45.03% 0-shot (Table 5) and 46.99% 2-shot (Table 6), which is slightly below full TPA's 46.21% 0-shot (Table 11) and 47.93% 2-shot (Table 13) but still competitive with MHA's 46.19% 0-shot and 48.11% 2-shot. Non-contextual B performs worse (43.66% 0-shot, 44.48% 2-shot). The third-order TPA at small scale (Table 4) achieves 44.56% 0-shot and 45.28% 2-shot—below second-order TPA but not catastrophically so. The paper notes this is with un-optimized hyperparameters, and no parameter-count-matched comparison with second-order TPA is provided, making this result preliminary.
Rank choice in FlashTPA decoding benchmarks: The decoding benchmarks in Figures 5, 7, and 8 all use RQ = 16, RK = RV = 1. No ablation over different rank configurations for decoding speed is presented—e.g., how decoding time varies with RK = RV = 2 or RQ = 8. Given the language modeling results showing quality improvements from increasing KV rank from 1 to 2 to 4 (Figure 10, Table 12), the latency-quality trade-off at different ranks during decoding is an important missing piece of data. A system designer wanting to deploy TPA would need to know whether RK = RV = 2 (which improves quality at the XL scale) incurs a meaningful latency penalty relative to RK = RV = 1 during autoregressive generation. The paper provides no such measurement.
Negative results: The paper reports that MLA generally trains more slowly and yields higher validation losses than TPA across all scales (Figures 3, 4). This is not framed as an ablation of TPA but as a comparative result—it demonstrates that MLA's architectural complexity (separate RoPE pathway, compressed-then-up-projected keys and values) does not necessarily translate to better training dynamics or final quality compared to TPA's simpler factorization. The paper also notes in Appendix I that "generalization to other modalities deserves more extensive investigation"—acknowledging that all results are on text-only language modeling and may not transfer to vision, audio, or multimodal settings.
Critical Assessment
Does the paper demonstrate that TPA reduces KV cache memory while maintaining or improving quality?
What the experiments show: Yes—for the specific configurations tested (FineWeb-Edu 100B, LLaMA architecture, 124M–1.5B parameters, 50B training tokens, dh = 64). TPA with RK = RV = 2 achieves a ~5.1× reduction in per-token KV cache size relative to MHA at the Large (773M) scale while achieving higher 0-shot accuracy across nine benchmarks (53.10% for TPA vs. 52.52% for MHA, Table 3). TPA-KVonly achieves a similar cache reduction (since query factorization doesn't affect cache size) and an even larger accuracy margin (53.52%). The validation loss curves (Figure 4) confirm that this is not a fluke of downstream evaluation—TPA models achieve lower perplexity throughout training.
What is not shown: The paper does not report any long-context evaluations. All training uses the default nanoGPT context length—the paper does not specify this length explicitly, but nanoGPT defaults to a context window of 1024 tokens. The KV cache size reduction matters most at long sequence lengths (tens or hundreds of thousands of tokens), where the cache dominates memory. Yet the paper provides no perplexity-at-length, needle-in-a-haystack, or long-document QA evaluations. The decoding speed benchmarks do go up to 2^19 = 524K sequence length, but these measure only latency, not model quality at those lengths. A model might have a small KV cache but fail to effectively use long contexts due to the low-rank bottleneck in key representations—this hypothesis is not tested.
The paper also does not report any throughput (tokens per second) measurements for batch inference—only per-token latency. Throughput is often the more important metric in production, and it depends on both computation and memory bandwidth. TPA's smaller cache should improve throughput by enabling larger batch sizes within a given memory budget, but this is not demonstrated.
Does the paper demonstrate that TPA unifies MHA, MQA, and GQA as special cases?
What the experiments show: The mathematical derivation in Section 4 demonstrates that MHA, MQA, and GQA can be expressed as TPA with non-contextual head factors—this is a theoretical result, not an experimental one. The experimental support is indirect: TPA with non-contextual A factors achieves quality comparable to MHA (50.52% vs. 50.11% for medium models; Tables 7 and 2), which is consistent with the claim that this variant approximates MHA-like behavior. MQA and GQA are not directly reconstructed and tested as TPA variants, so the claim that they are "special cases" is theoretical rather than empirically validated.
What is not shown: The paper does not experimentally demonstrate that a TPA model initialized to match MHA (with RK = RV = h and one-hot head factors) actually reproduces MHA's behavior identically. This would require showing that training such a model with the head factors frozen to the one-hot basis yields exactly the same loss curve and downstream performance as standard MHA—a useful sanity check that is absent. Without this, the unification claim is mathematically clean but empirically untested.
Does the paper demonstrate that TPA integrates seamlessly with RoPE?
What the experiments show: Theorem 3.1 proves mathematically that RoPE distributes over the tensor product factorization, preserving the relative-position property. The experimental evidence for this claim is embedded in the quality results: TPA models use RoPE and achieve strong performance, which would not be possible if RoPE integration were broken. The pre-rotation caching strategy is described algorithmically (Section 3.2) but not ablated against alternative strategies (e.g., rotating on-the-fly vs. pre-rotating). The decoding benchmarks (Figures 5, 7, 8) all use FlashTPA with pre-rotated key factors, so the latency numbers include the benefit of pre-rotation.
What is not shown: The paper does not provide an ablation comparing TPA's RoPE integration strategy against MLA's (separate RoPE pathway). Given that RoPE compatibility is positioned as a key advantage over MLA (Section 1: MLA "encounters difficulties with efficient Rotary Position Embedding (RoPE) integration"), this is a significant omission. The ideal experiment would compare TPA and MLA at equal cache sizes, measuring both quality and the contribution of RoPE to long-range positional understanding. The current results show TPA outperforming MLA on quality, but this could be due to many factors (training dynamics, architecture choices) beyond RoPE integration—the paper does not isolate the RoPE-specific contribution.
Does the paper demonstrate that FlashTPA decoding is faster than optimized baselines?
What the experiments show: Yes, for the specific Triton implementation tested. Figure 5 shows FlashTPA matching or beating FlashMLA, FlashGQA, and FlashMQA at long sequence lengths across batch sizes 1–16 for d_model = 2048. The advantage grows with sequence length and batch size, consistent with the memory-bandwidth-bound nature of long-sequence attention.
What is not shown: The comparison is against CUDA implementations of the baselines (FlashMHA, FlashGQA, etc.) using a Triton implementation of FlashTPA. The paper acknowledges this implicitly: it notes that "development of a CUDA-based FlashTPA kernel is ongoing and is expected to yield further performance improvements" (Section 6.2). This cuts both ways—if FlashTPA were re-implemented in CUDA, it might be faster still; but the baselines in Triton might also be faster or slower than their CUDA counterparts. The paper provides no Triton implementations of the baselines to control for implementation quality. The decoding benchmarks should therefore be interpreted as "FlashTPA in Triton is competitive with optimized CUDA attention kernels" rather than "TPA is algorithmically faster than alternative attention mechanisms."
Additionally, the decoding benchmarks use RQ = 16, RK = RV = 1—a configuration that achieves very aggressive compression (cache size ~25× smaller than MHA) but that is not the same as the RQ = 6, RK = RV = 2 configuration used in the language modeling experiments. The quality of a model at RK = RV = 1 during decoding is not evaluated, so it's unclear whether the speed advantage shown in Figure 5 comes at a quality cost relative to the RK = RV = 2 models that achieve the strong benchmark results in Tables 2–3.
Does the paper demonstrate that TPA's quality improvement is robust?
What the experiments show: The improvement over MHA is consistent across three of four model scales (353M, 773M, 1.5B), with the 124M scale showing a virtual tie (46.21% vs. 46.19%, Table 11). TPA is consistently better than or tied with MQA, GQA, and MLA across scales. The validation loss curves (Figure 4) show that the improvement is not an artifact of downstream evaluation noise—it appears in perplexity throughout training.
What limits robustness: All experiments use a single dataset (FineWeb-Edu 100B), a single base architecture (LLaMA), a single head dimension (dh = 64), and a single training duration (50B tokens). The paper does not demonstrate that the results hold under different pretraining data distributions (e.g., code-heavy, multilingual), different model architectures (e.g., non-LLaMA designs like Mamba-hybrid or mixture-of-experts), or different training horizons (e.g., multiple epochs, Chinchilla-optimal training). At the 1.5B scale, 50B tokens is far below Chinchilla-optimal (which would recommend ~30B tokens for a 1.5B model), so the models are significantly undertrained. It is possible that at Chinchilla-optimal training durations, where MHA has more opportunity to learn specialized per-head features, TPA's regularization advantage might diminish or reverse.
The paper does not report statistical significance for any result—no error bars, no multiple training runs with different seeds. The downstream evaluation uses the lm-evaluation-harness with its standard (deterministic) settings, so the reported numbers are point estimates from single training runs. At the sample sizes involved (500–1000 test questions per benchmark), differences of 0.5–1.0 percentage points in average accuracy may or may not be statistically reliable—the paper provides no way to assess this.
Missing experiments that would have strengthened the paper
-
Long-context quality evaluations. The core motivation for KV cache compression is enabling longer contexts. Perplexity at various context lengths (e.g., 2K, 8K, 32K, 128K), retrieval accuracy on long-document tasks, and needle-in-a-haystack tests would directly test whether TPA's low-rank keys retain the information needed for long-range attention. The decoding benchmarks show that FlashTPA is fast at long sequences, but they do not show that the quality holds up.
-
Ablation of KV rank vs. long-context quality. The rank ablation in Figure 10 and Table 12 shows that RK,V = 4 is better than RK,V = 2 for standard-context evaluation. Does this gap widen at longer contexts (where more key diversity is needed)? Does RK,V = 1—which the decoding benchmarks use—cause catastrophic degradation at 100K+ contexts? These questions are central to the paper's practical value but are unanswered.
-
Comparison with token-eviction methods. TPA reduces cache size by compressing each token's representation. Token-eviction methods (H2O, StreamingLLM) reduce cache size by discarding tokens. Are these complementary? The paper mentions (Section 1) that eviction methods "risk discarding tokens that may later prove important" but does not empirically compare against them or evaluate TPA combined with eviction. An experiment showing that TPA with 25%-sized cache + eviction outperforms MHA with full cache + eviction at the same memory budget would be compelling.
-
Throughput benchmarks for batch inference. Production LLM serving is typically throughput-bound (tokens per second across many concurrent requests) rather than latency-bound for a single sequence. TPA's smaller cache should enable larger batch sizes within a fixed memory budget, potentially yielding substantial throughput gains. The paper measures only per-token latency for single sequences (batch sizes 1–16, where 16 is still a single model forward pass, not concurrent requests).
-
Scaling to larger models (7B, 13B, 70B). The largest model trained is 1.5B parameters, which is roughly 20–50× smaller than the models deployed in production (LLaMA 3 70B, DeepSeek-V2 236B). The trends at 124M–1.5B are encouraging, but extrapolation to production scales is speculative. Do the quality advantages persist or grow, or does MHA catch up when it has enough capacity to learn non-redundant per-head features?
-
Training FLOPs comparison. The paper equalizes parameter counts but does not report training FLOPs. TPA's factorized projections use fewer parameters than MHA's per-head dense projections (Table 1), but the tensor product assembly (matrix multiplication of the A and B factors) adds computation that is not present in MHA. A FLOPs-matched training comparison would reveal whether TPA's quality advantage comes at higher training cost—an important practical consideration.
6. Limitations and Trade-offs
The Difficulty Estimation Cost Is Unaccounted For and Likely Dominant
The assumption. The compute-optimal framework depends entirely on estimating problem difficulty before allocating the inference budget, but the paper's difficulty estimation method—generating and scoring 2,048 samples per problem—is "extraordinarily expensive," consuming more computation than the largest test-time budgets studied (256–512 generations). The paper acknowledges this directly:
"estimating difficulty in this way still incurs additional computation cost during inference... our experiments do not account for this cost largely for simplicity" (Section 3.2)
The consequence. The reported 4× efficiency gains are computed after difficulty is already known, without amortizing the cost of learning it. In realistic deployment, the total cost would be difficulty estimation plus strategy execution, and the former could dominate. The paper frames this as an "exploration-exploitation tradeoff" (Section 3.2) but provides no method to balance it—no adaptive scheme that incrementally estimates difficulty while solving, no lightweight difficulty classifier, no accounting of the overhead in the headline numbers. The 4× figure should therefore be understood as an upper bound on achievable efficiency, not a realized deployment gain. A system that spends 2,048 generations estimating difficulty and then 64 generations solving the problem consumes 2,112 total generations—compared to a flat best-of-256 baseline spending 256, the "efficient" system actually uses ~8.3× more compute, not less.
What evidence exists. No experiment measures or accounts for difficulty estimation cost. Section 3.2 describes the method (2,048 samples per question, scored via PRM or ground truth) but excludes this cost from all budget calculations. The paper acknowledges this gap explicitly and suggests future work, but provides no data on how cheaper difficulty estimates (e.g., from 4–8 samples instead of 2,048) would affect strategy selection quality. This is a fundamental gap between the paper's theoretical framework and practical deployment.
Mitigation status. Not addressed. The paper flags future work on "pretraining or finetuning models to directly predict difficulty of a question" (Section 8) and suggests adaptive schemes might help, but no experiments or prototypes are provided. Until this gap is closed, the compute-optimal framework remains a conceptual contribution requiring expensive oracle access, not a deployable system.
Hard Problems Remain Completely Unsolved
The assumption. Test-time compute scaling can improve performance by searching over or refining candidate solutions, but this only works if correct solutions exist somewhere in the model's output distribution to be found or refined. The difficulty bin analysis reveals exactly where this assumption fails.
The consequence. For the hardest questions (difficulty bin 5), no method makes meaningful progress at any compute budget. Accuracy hovers at 1–3% for all methods and all budgets (Figure 3, right panel, bin 5; Figure 7, right panel, bin 5). The FLOPs-matched comparison shows bin 5 scaling lines essentially flat near 0–5%, with test-time compute providing negative or negligible benefit versus simply training a larger model (Figure 9). The paper states this clearly in the Section 7 takeaway: test-time compute can amplify existing capability but cannot create it. If the base model's pass@1 is near zero, no amount of search or revision helps—there are no correct solutions in the proposal distribution to find. This establishes a hard capability ceiling: test-time compute is bounded above by the base model's ability to generate at least one correct solution with non-trivial probability, and for the hardest problems, that probability is effectively zero.
What evidence exists. The bin 5 results are stark: in Figure 3 (right), search methods achieve 1–3% accuracy across all budgets (4 to 256 generations). In Figure 7 (right), all sequential-to-parallel ratios produce roughly 2–3% accuracy. In Figure 9, bin 5 curves for both revisions and PRM search are flat and near zero, positioned below the 14× larger model's baseline across all R regimes. The contrast with bin 4 (where beam search can reach ~17% and revisions reach ~18%) is sharp: the transition from "medium-hard," where test-time compute helps, to "genuinely hard," where it is useless, appears suddenly between bins 4 and 5.
Mitigation status. The paper is transparent about this limitation. The Section 7 takeaway explicitly states that pretraining is necessary for problems outside the base model's capability range. However, the paper provides no diagnostic for distinguishing bin 4 from bin 5 at deployment time (a difficulty estimator could easily place a genuinely novel problem into the wrong bin), and no exploration of whether techniques like retrieval-augmented generation or tool use could expand the model's effective capability range to convert bin-5 problems into bin-4 or bin-3 problems where test-time compute can help.
The Revision Model Has a 38% Correct-to-Incorrect Reversion Rate, Patched Rather Than Solved
The design constraint. The revision model is trained solely on sequences where in-context answers are incorrect followed by a correct target. This means it has never seen a training example where the current answer is already correct and should be preserved—there is no "do nothing" or "verify and keep" signal in the training data.
The consequence. At inference time, when a revision chain produces a correct answer, the model has a strong tendency to "revise" it into an incorrect one in the next step. The paper reports that approximately 38% of correct answers get converted back to incorrect ones during sequential revision (Section 6.1). This is a direct consequence of the training data construction: the model learns the conditional distribution P(correct answer | sequence of incorrect answers), but when the current answer happens to be correct (which occurs with increasing frequency as the chain progresses and pass@1 rises, per Figure 6 left), this is an out-of-distribution input for which the model's behavior is uncontrolled.
The paper's mitigation is to use majority voting or verifier-based selection across the entire revision chain—picking the best answer from any step—rather than always taking the final revision. While this recovers performance (Figure 6 right shows sequential+verifier at ~41.5% vs. the final-step output, which would be lower), it is an imperfect patch: it requires the verifier or majority vote to correctly identify the best answer among all chain positions, and it discards the conceptual elegance of "the model improves its answer at each step" in favor of "generate many revisions and hope one is good." The 38% reversion rate means that, for roughly two out of every five problems where the model stumbles onto the correct answer mid-chain, that answer is lost in the following step—a substantial source of inefficiency that better training data construction could eliminate.
What evidence exists. The 38% figure is reported in Section 6.1, attributed to the training data construction. The paper's ablation of the revision model's pass@1 trajectory (Figure 6, left) shows gradual improvement from ~18.2% at step 1 to ~24–25% by steps 15–20, then flattening—the reversion rate explains why the curve does not continue to rise: each additional step has a ~38% chance of corrupting any correct answers produced so far. The within-chain selection mechanism (Figure 6, right) recovers performance to ~41.5% at 64 steps, but this is achieved by cherry-picking the best step, not by the revision process itself being reliable. The ReST^EM experiment (Appendix K, Figure 16) shows that attempting to optimize the revision model with on-policy RL actually degrades performance substantially—sequential revisions with the ReST^EM model collapse to ~33.5% at 256 generations compared to ~38.5% at the optimal ratio—highlighting the fragility of the revision training approach.
Mitigation status. Partially addressed via within-chain selection (majority voting or verifier-based), but the underlying cause—the model's inability to recognize correct answers and leave them unchanged—is not solved. A principled fix would involve training the model on trajectories that include both incorrect-to-correct revisions and correct-to-correct "confirmations" (or explicit "no change needed" tokens), which the paper does not explore. The ReST^EM negative result (Appendix K) suggests that naive self-improvement loops may exacerbate rather than fix the problem, indicating that revision training is more delicate than the positive results on Tables 2–3 might suggest.
Single Benchmark, Single Model Family, Single Training Configuration
The scope constraint. All experiments use the MATH benchmark (500 test questions) with PaLM 2-S* as the base model, trained with a specific PRM training procedure and a specific revision model data construction pipeline. The paper does not evaluate on any other reasoning benchmark (e.g., GSM8K, TheoremQA, coding benchmarks), any other model family (e.g., LLaMA, Qwen, Gemma), or any other task domain (e.g., code generation, logical reasoning, scientific QA, open-ended generation).
The consequence. Several aspects of the findings could be model-specific or benchmark-specific in ways the paper cannot distinguish:
- The PRM is trained via Monte Carlo rollouts from PaLM 2-S*. Its quality depends on this model's output distribution, calibration, and error patterns. A PRM trained on a different model family might exhibit different over-optimization behavior, shifting the difficulty thresholds at which beam search helps versus hurts (Figure 3, right). The paper itself found that the PRM800k dataset—which contains human-labeled GPT-4 outputs—was "largely ineffective" for PaLM 2 models due to distribution shift (Section 5.1), suggesting that PRM quality is highly model-specific.
- The revision model's ability to learn from incorrect in-context answers depends on the base model's in-context learning capabilities, which vary substantially across model families. A model with weaker in-context learning might fail to learn the revision skill from the post-hoc trajectory construction, while a stronger one might benefit more.
- The MATH benchmark consists exclusively of competition-level math problems with clean, verifiable answers. It is unclear whether the difficulty-dependent patterns—beam search hurting easy problems, revisions helping easy problems more than hard ones—generalize to tasks without closed-form correctness signals (dialogue, creative writing, complex planning), where verifier training and difficulty estimation would need fundamentally different approaches.
What evidence exists. The paper acknowledges the single-benchmark limitation implicitly in Section 8 (future work on "other domains and modalities") but provides no empirical evidence of generalization. All results in Tables 2–3 and Figures 3–9 are derived from the same 500-question MATH test set, split into five difficulty quintiles of roughly 100 questions each, further split by two-fold cross-validation. This means the compute-optimal policy is selected based on roughly 50 questions per fold per bin—a small sample that makes the selected strategies potentially noisy and their reported performance subject to high variance. The paper does not report confidence intervals, making it impossible to assess whether the observed differences between strategies (e.g., compute-optimal vs. best-of-N at 4× less compute) are statistically reliable at this sample size.
Mitigation status. Not addressed. The paper does not evaluate on any benchmark other than MATH, and uses only a single model family (PaLM 2-S*). The validation protocol (two-fold cross-validation on 500 questions) partially addresses overfitting the strategy selection to the test set, but does not address overfitting to the MATH distribution itself—it remains possible that the optimal strategies are specific to competition-level math and would differ for other reasoning tasks. The paper suggests future work on "extension to other domains and modalities" (Section 8) but provides no preliminary results or analysis of which findings are likely to transfer.
Revisions and Search Are Never Combined, Leaving Potential Gains Unexplored
The design separation. The paper studies two complementary mechanisms—PRM-guided search (Section 5) and iterative revision (Section 6)—as entirely independent pipelines. The experiments do not use the revision model as the proposal distribution within beam search, do not apply the PRM to score or guide which revisions to pursue, and do not combine the two in any way. Section 8 acknowledges this explicitly:
"we did not experiment with PRM tree-search techniques in combination with revisions"
The consequence. This is a significant gap because the paper's own analysis shows the two mechanisms have complementary strengths along the difficulty axis: revisions are most effective on easy problems, where the model's initial output is roughly correct and just needs refinement (a local improvement); PRM search is most effective on medium problems, where broad exploration across different solution strategies is needed (a global search). A combined system could, for example, use the revision model as the proposal distribution within beam search—at each step of the search tree, the model conditions on previous rejected branches to produce higher-quality candidate steps—or use the PRM to decide when a revision chain is on track versus when to restart. This combination could break through the performance ceilings that each method individually hits: on medium problems, revisions could improve the quality of candidates within each beam; on easy problems, PRM scoring could prevent revisions from "fixing" already-correct answers.
The current results therefore represent a lower bound on what a fully integrated system could achieve. The 4× efficiency gains reported for each mechanism individually (Figures 4 and 8) might compound or synergize when combined, potentially yielding even larger gains. Conversely, the combination might surface new failure modes—the PRM might over-optimize revision model outputs differently than base model outputs; the revision model's distribution shift (Appendix J, Figure 15a, showing the base-LM PRM underperforms on revision model outputs) might make PRM guidance unreliable. Without experiments, both the potential upsides and the risks are unknown.
What evidence exists. The complementary difficulty-dependence is shown in Figures 3 (right) and 7 (right): revisions excel on easy bins, PRM search on medium bins. The distribution shift challenge is documented in Appendix J, Figure 15a, where the base-LM PRM achieves lower performance on revision model outputs than on base model outputs (~40% vs. ~42% at 64 generations). This suggests that naively combining the two would require addressing verifier transfer, but the paper does not explore this.
Mitigation status. The paper identifies the combination as future work (Section 8) but provides no experiments or analysis of how they might interact. This is an understandable scoping choice for a single paper, but it means the practical value of TPA for a system builder is partially unassessed—the paper demonstrates that each component works individually but provides no guidance on how (or whether) to use them together, which a practitioner deploying test-time compute would naturally want to do.
The 14× Larger Model Baseline Is Weaker Than It Should Be
The comparison design. The FLOPs-matched comparison in Section 7 scales model parameters while holding training data fixed (following the LLaMA paradigm rather than Chinchilla-optimal scaling), and the larger model uses only greedy decoding with no test-time compute augmentation of its own. The paper acknowledges the first point explicitly:
"We choose this setting as it is representative of a canonical approach to scaling pretraining compute and leave the analysis of compute-optimal scaling of pretraining compute where the data and parameters are both scaled equally to future work." (Section 7)
The consequence. This makes the pretraining baseline weaker than it needs to be in two ways. First, a Chinchilla-optimal model—trained with both parameters and data scaled to hold the ratio constant—would likely outperform a parameter-only-scaled model at the same total FLOPs, meaning the reported advantages of test-time compute (e.g., +27.8% relative improvement on easy questions at R ≪ 1 for revisions; Figure 1, top-right bar chart) may shrink or reverse against a properly compute-optimal larger model. Second, the larger model uses only greedy decoding: no majority voting, no best-of-N, no search of any kind. This is an asymmetric comparison that gives the smaller model a sophisticated inference-time strategy while giving the larger model nothing. A fairer setup would allocate at least a small test-time budget (e.g., best-of-8 or best-of-16) to the larger model, which would strengthen the baseline substantially—particularly since the paper's own results show that 4× less compute is achievable with compute-optimal allocation (Figure 4), suggesting that even a modest budget for the larger model could close much of the gap.
The paper's bar charts (Figure 1, right panels) show that the advantage of test-time compute over the larger model varies dramatically by difficulty and R, with large advantages on easy-to-medium problems at low R but disadvantages of up to -52.9% on hard problems at high R (PRM search, bin 4–5, R ≫ 1). These numbers are the headline claims that test-time compute can substitute for pretraining—but they are measured against a baseline that is neither compute-optimally trained nor given any inference budget. The qualitative conclusion (test-time compute helps on problems within the base model's reach, fails outside it) is likely robust to baseline strengthening, but the quantitative claims (e.g., "outperforms a ~14× larger model") may be inflated by the weak baseline.
What evidence exists. The paper explicitly acknowledges the non-Chinchilla training in Section 7 and frames it as a design choice. No ablation varies the larger model's decoding strategy, and no experiment compares against a Chinchilla-optimal baseline. The greedy decoding choice is never justified—the paper simply states that the larger model uses greedy decoding as if this were the natural comparison, but in a paper about test-time compute optimization, giving one model a sophisticated inference strategy and the other none is an asymmetric comparison.
Mitigation status. The paper flags Chinchilla-optimal pretraining comparisons as future work (Section 7) but does not address the asymmetric decoding issue. A minimal improvement would be to report the larger model's performance with best-of-N at a modest budget (e.g., best-of-8) and show whether the qualitative conclusions hold. Without this, the FLOPs-matched numbers should be interpreted as evidence that test-time compute can help relative to a naive larger-model baseline, not that it dominates pretraining—a more cautious conclusion than what the abstract and introduction suggest.
7. Implications and Future Directions
How This Work Changes the Landscape
Tensor Product Attention is best understood as a framework-level reframing, not an incremental architectural tweak. Before TPA, the attention design space was understood as a discrete menu: MHA, MQA, GQA, MLA—each a separate mechanism with its own implementation, trade-offs, and failure modes. Practitioners chose among them by weighing cache size against quality degradation, accepting that smaller caches meant less expressive keys. TPA changes the question: rather than "which attention variant should I use?", the question becomes "where in the continuous TPA design space should I operate for my target compression-quality trade-off?"
This reframing matters because it reveals that the quality-cache tension—the assumption that compressing the KV cache necessarily degrades model quality—is not fundamental. It is an artifact of how existing methods do the sharing. MQA shares one key across all heads, rigidly and content-independently; GQA shares one key per group, also rigidly. TPA demonstrates that when sharing is done via learned, context-dependent low-rank factorization, the constraint can act as a beneficial structured regularizer: it prevents heads from learning redundant, over-specialized projections while preserving (and even improving) the model's ability to represent diverse attention patterns per token. The empirical evidence—TPA achieving 51.41% average accuracy vs. MHA's 50.11% on 353M models (Table 2) while using ~5.1× less KV cache per token (Section 3.3)—is not a trade-off victory; it's a redefinition of the playing field. Compression and quality were not opposing forces here; they were aligned by the right inductive bias.
The paper also reconciles a standing tension that has been visible in the literature since DeepSeek-V2 introduced MLA (Liu et al., 2024). That work demonstrated that latent-space compression of keys and values could dramatically reduce cache size, but at the cost of an awkward architectural compromise: RoPE, the dominant positional encoding in modern LLMs, cannot be cleanly applied to the compressed representations without breaking the pre-computation that makes MLA fast. This forced MLA into a dual-pathway design—a compressed, RoPE-free latent plus a separate, smaller RoPE-specific component per head—that added parameters, increased cache size, and was widely discussed as an inherent limitation of compressed attention (see Su, 2024: "the extreme pull between cache and effect"). TPA demonstrates that this tension is not inherent to compressed attention—it is specific to MLA's compression strategy of projecting into a single latent vector. By factorizing along the natural head × feature axes rather than compressing into a scalar latent, TPA makes RoPE distribute transparently over the factorization (Theorem 3.1), enabling pre-rotated key caching with zero additional cost and no architectural compromise. This shifts the research question from "how do we work around RoPE in compressed attention?" to "which factorization topologies compose naturally with positional encodings?"—a more productive framing that opens the design space rather than constraining it.
The paper's demonstration that factorized representations enable not just storage compression but computational reorganization—that attention can be computed directly in factor space faster than materializing and attending over dense keys (Figure 5, FlashTPA surpassing MLA at sequence lengths beyond 2^15)—changes the mental model for what inference optimization looks like. The default assumption in the field has been: compress the cache for memory savings, decompress on-the-fly for computation, accept the decompression overhead. FlashTPA shows that, at least in the memory-bandwidth-bound regime of long-sequence decoding, the factor-space computation can be faster than the decompress-then-attend approach, not merely a less-memory-intensive way to achieve the same computation. This is a different category of result: it suggests that factorized attention is not a compromise one accepts for memory-constrained settings but may be the algorithmically superior choice for long-context inference regardless of memory pressure—a claim that, if it holds up under optimized CUDA implementations and at production scales, would reshape how inference engines are built.
Several research directions become more attractive in light of this paper. The demonstration that contextual factorization provides a regularizing inductive bias suggests that learned, dynamic structured sparsity (of which TPA's low-rank tensor products are one instance) may be a general principle for designing efficient neural network layers, not limited to attention. The finding that the TPA design space subsumes MHA, MQA, and GQA as special cases (Section 4) suggests that neural architecture search over factorization topologies—continuous optimization of ranks, contextual vs. non-contextual factors, and factorization order—might automate the discovery of efficient attention variants for specific hardware targets, a possibility that was infeasible when mechanisms were discrete, non-interpolable choices. Conversely, some directions become less urgent: the paper's results suggest that further refinements of MLA-style single-latent compression (which must contend with the RoPE integration problem) may be a local optimum compared to factorization approaches that embrace the head × feature structure. Research effort may be better spent exploring the TPA design space (varying ranks, contextual patterns, factorization orders) than patching MLA's RoPE pathway.
Follow-Up Research This Work Enables
Long-context quality evaluation of TPA at production scales. The paper's most conspicuous gap is the absence of any long-context quality measurement—the decoding benchmarks show FlashTPA is fast at up to 524K tokens (Figure 5), but no perplexity-at-length, retrieval accuracy, or needle-in-a-haystack evaluation exists to test whether RK = RV = 1 or 2 keys retain the information needed for long-range attention. A strong follow-up would train TPA models (at 1.5B or 7B scale) with extended context windows (32K–128K tokens), compare perplexity as a function of position against MHA and GQA baselines at equal cache sizes, and run the standard long-context benchmarks (Lost in the Middle, needle-in-a-haystack, LongBench). The critical question: does the low-rank key bottleneck manifest as degraded retrieval accuracy at long range, and if so, at what rank and context length does it become measurable? The rank ablation in Figure 10 (showing RK,V = 4 > RK,V = 2 at standard context) suggests the answer may be rank-dependent, but no data exists.
Combined PRM search with TPA factor-space inference for reasoning models. This is a cross-pollination opportunity the paper does not explore. Recent work on test-time compute scaling (Snell et al., 2024; Brown et al., 2024) uses PRM-guided beam search over candidate solutions, which requires maintaining and scoring many partial key-value caches in parallel during tree search—exactly the regime where TPA's memory reduction matters most. A natural experiment: take a reasoning model fine-tuned with process reward modeling, replace its standard attention (MHA or GQA) with TPA at RK = RV = 2, and measure whether the reduced cache size enables wider beam search or deeper trees within the same GPU memory budget, and whether the quality of search improves as a result. The PRM over-optimization phenomenon documented in the reference paper (beam search degrading on easy problems, Figure 3 right) might interact with TPA's low-rank keys in non-obvious ways—the verifier might overfit to rank-reduced key representations differently than to full-rank ones.
CUDA implementation of FlashTPA and systematic latency-rank-quality Pareto frontier. The paper's decoding benchmarks use a Triton implementation of FlashTPA against CUDA baselines, and only at one rank configuration (RQ = 16, RK = RV = 1). A rigorous follow-up would implement FlashTPA in CUDA (FlashAttention-style, with optimized tiling and warp-level primitives), sweep all rank parameters (RQ ∈ {4, 8, 16, 32}, RK = RV ∈ {1, 2, 4}), measure decoding latency wall-clock at sequence lengths 4K–1M tokens, and pair each latency measurement with a downstream quality evaluation of the corresponding model configuration. The output would be a Pareto frontier mapping (latency, memory, quality) as a function of ranks—the engineering artifact that a practitioner needs to choose a deployment configuration. Without this, the paper provides proof-of-concept speed numbers (Figure 5) and proof-of-concept quality numbers (Tables 2–3) at different rank settings, but no joint optimization.
Difficulty estimation via PRM score distribution on TPA-generated key factors. The reference paper's compute-optimal framework requires estimating problem difficulty from PRM scores, but that framework was developed for standard MHA models. A TPA model's factorized keys might produce PRM scores with different statistical properties—the rank bottleneck could compress score variance, making difficulty bins less separable, or it could sharpen distinctions by reducing noise. A specific experiment: take the difficulty estimation protocol from the reference paper (2,048 samples per question, PRM final-answer score averaging, quintile binning), apply it to a TPA-based model, and compare the alignment between oracle difficulty bins (pass@1-based) and predicted difficulty bins (PRM-score-based) against the alignment achieved with an MHA model of equal parameters. If TPA's factorization degrades difficulty estimation (reducing the efficiency gains of compute-optimal allocation), that's an important negative result; if it improves estimation (by regularizing the PRM signal), that's a synergy worth exploiting.
Chinchilla-optimal training of TPA at scale with multi-epoch data. The paper trains all models for 50B tokens, which is far below Chinchilla-optimal for the 1.5B scale (which would call for ~30B tokens—closer, but still below optimal for the smaller scales). More importantly, the paper trains for roughly half an epoch on a 100B-token dataset, meaning tokens are seen at most once. Standard MHA might benefit more from multi-epoch training (where heads can gradually specialize) than TPA's rank-constrained factorization, which acts as a regularizer. A Chinchilla-matched comparison would train TPA and MHA models at, say, 1.5B parameters on 30B tokens (single epoch) and on 60B tokens (two epochs), measuring whether TPA's advantage shrinks, persists, or grows with more data. If MHA catches up with more training, TPA is best understood as a compute-efficient architecture for undertrained regimes; if TPA's advantage holds or widens, the regularization interpretation strengthens.
TPA in non-text modalities and encoder-decoder architectures. The paper's entire evaluation is decoder-only language modeling on text. TPA's factorization separates variation across heads (the a factors) from variation across feature dimensions (the b factors)—a structure that might be particularly natural for modalities where heads specialize in different spatial or temporal scales (vision transformers attending to different regions; audio transformers attending to different frequency bands). A straightforward extension: replace the self-attention in a Vision Transformer (ViT) with TPA, train on ImageNet, and measure whether the rank-constrained keys improve or degrade fine-grained spatial localization compared to standard MHA. Similarly, applying TPA to encoder-decoder architectures (where cross-attention KV caches grow with encoder output length) would test whether the cache reduction benefits transfer to the cross-attention setting, which has a different memory-access pattern than self-attention decoding.
Practical Applications and Downstream Use Cases
Long-context batch inference serving with fixed GPU memory. A production LLM serving system (e.g., vLLM, TensorRT-LLM) running a model with d_model = 4096, h = 32, dh = 128, and processing requests at 128K context length faces a KV cache of roughly 2 × 128,000 × 32 × 128 × 2 bytes ≈ 2 GB per sequence in half-precision. At batch size 8, the cache alone consumes ~16 GB—most of an A100's 40 GB or 80 GB, leaving little room for model weights and activations. Replacing the attention with TPA at RK = RV = 2 reduces per-sequence cache to (2+2)(32+128) × 128,000 × 2 bytes ≈ 164 MB—a ~12× reduction. This directly enables either larger batch sizes (increasing throughput proportionally, since throughput scales near-linearly with batch size in the memory-bandwidth-bound regime) or longer context windows at the same batch size. For a serving provider charging per-token, the throughput improvement (tokens per second per GPU) from the increased batch size would translate directly to lower cost per request. The FlashTPA decoding benchmarks (Figure 5) show that this memory reduction does not come at a latency penalty—at 128K contexts, FlashTPA is faster than MHA and competitive with MQA/GQA—so the throughput gain is not offset by per-token slowdown.
On-device LLM deployment with adaptive rank for dynamic memory management. On-device models (smartphone, laptop, edge accelerator) operate under hard memory constraints—6–12 GB total system memory shared with the OS and other applications. A 3B-parameter model at 16-bit precision requires ~6 GB for weights alone, leaving minimal room for KV cache. With MHA at h = 32, dh = 96, and a 4K context window, the cache adds ~2 × 4096 × 32 × 96 × 2 bytes ≈ 50 MB—manageable at short contexts but scaling to 32K tokens would require ~400 MB, exceeding the available budget. TPA at RK = RV = 1 reduces this to (1+1)(32+96) × 32,768 × 2 bytes ≈ 16.8 MB for 32K context—a ~24× reduction that makes long-context local inference viable. Furthermore, because TPA's ranks are configurable per-layer (and could in principle be made adaptive per-token, though the paper doesn't explore this), an on-device deployment could implement dynamic rank scaling: use RK = RV = 2 for the first N tokens of context (where precise attention matters for understanding the query) and drop to RK = RV = 1 for distant context (where coarse-grained attention suffices). This would provide a smooth memory-quality trade-off curve controllable at inference time without model reloading, something MQA and GQA cannot offer since their sharing patterns are fixed at architecture design time.
Fine-tuning and continued pretraining with reduced memory pressure. Parameter-efficient fine-tuning methods like LoRA reduce the memory for weight updates but do not address the KV cache, which is the dominant memory consumer during long-context fine-tuning. When fine-tuning a model on 32K-token documents (e.g., legal contracts, scientific papers), the KV cache for activations can exceed the memory used by LoRA adapters by orders of magnitude. Using TPA as the base attention mechanism—even without fine-tuning the attention parameters themselves—would reduce the activation memory footprint during fine-tuning by the same factor as during inference (~12× for RK = RV = 2 at h = 32, dh = 128), enabling fine-tuning on longer documents or with larger micro-batch sizes on the same hardware without gradient checkpointing overhead. The paper's results showing that TPA-KVonly (which factorizes keys and values while leaving queries as standard dense projections) achieves 53.52% vs. MHA's 52.52% at 773M scale (Table 3) suggest that one could take a pre-trained MHA model, replace its key-value projections with TPA factors (initializing them to approximate the original MHA behavior using the non-contextual basis-vector construction from Section 4.2), fine-tune briefly to recover quality, and immediately benefit from reduced memory during subsequent long-context training or inference—without retraining from scratch. The paper does not demonstrate this transfer learning scenario, but the non-contextual A results (Tables 5–8) showing strong performance even without contextual head factors suggest the initialization path is plausible.
When to Prefer This Method
The paper explicitly positions TPA against a named set of alternatives—MHA, MQA, GQA, and MLA—and provides comparative data on both memory and quality across model scales. The decision criteria can be extracted from the paper's empirical results and architectural analysis:
-
Prefer TPA over MHA when KV cache memory is the binding constraint (long sequences, large batch sizes, limited GPU memory) AND you can afford training from scratch (or are willing to experiment with initialization from pre-trained MHA weights, which the non-contextual A results in Tables 5–8 suggest may be viable). TPA provides ~5–25× cache reduction depending on ranks, with quality matching or exceeding MHA at 353M–1.5B scales (Tables 2–3, 12). The trade-off is a more complex implementation and a design choice (ranks) that requires tuning.
-
Prefer TPA over MQA or GQA when you need cache reduction but cannot accept the quality degradation from rigid key sharing. MQA achieves even smaller caches than TPA (2dh vs. (RK+RV)(h+dh)) but at a measurable quality cost (MQA: 50.44% vs. TPA: 51.41% for 353M models in Table 2). GQA interpolates but with a harder quality-cache trade-off (larger groups improve quality but increase cache; TPA offers quality improvement with a cache smaller than GQA at small group sizes). TPA is the better choice when quality is the primary constraint and cache reduction is the secondary benefit.
-
Prefer TPA over MLA when RoPE integration simplicity matters, when training stability is a concern (MLA trained more slowly and converged to higher validation losses than TPA across all scales in Figure 4), or when you want the flexibility to adjust rank per layer or per deployment without architectural changes. MLA's separate RoPE pathway adds parameters (d_h^R per head), complicates the caching logic, and may contribute to the training instability the paper observes. TPA's RoPE integration is architecturally transparent (Theorem 3.1) and requires no additional parameters.
-
Prefer scaling pretraining over TPA-based KV compression when the problem distribution includes genuinely hard problems outside the base model's capability range, when the model will be deployed at very short context lengths where KV cache memory is negligible compared to weight memory, or when implementation simplicity is paramount (MHA is widely supported in all inference frameworks; TPA requires custom kernel development even with FlashTPA).
