Research map / Architectures and efficiency / Attention and transformers
Tensor Product Attention Is All You Need
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Tensor Product Attention (TPA), which factorizes Q, K, V activations into contextual low-rank tensor products to shrink the inference KV cache, and the T6 Transformer built on it, evaluated by pretraining at 124M-1.5B scales on FineWeb-Edu.
Problem
Long-context language model inference is limited by the memory footprint of the key-value cache, which grows linearly with sequence length, and existing methods (MQA, GQA, MLA) trade off quality, flexibility, or RoPE compatibility.
Approach
Factorize per-token query, key, and value activations (not static weights) into sums of rank-R tensor products of linear latent factors, cache only the factors, and pre-rotate the RoPE-applied key factors so the cache is small while RoPE's relative-position property is preserved; MHA, MQA, and GQA are shown to be non-contextual special cases.
Key results
- TPA achieves 51.41% zero-shot average on medium (353M) models versus 50.11% for MHA and 50.13% for MLA.
- At 773M (large), TPA-KVonly reaches 53.52% zero-shot average versus 52.52% for MHA.
- TPA reduces KV cache per token to (R_K+R_V)(h+d_h) versus 2*h*d_h for MHA, claimed to decrease memory by an order of magnitude.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| replace MHA, MQA, GQA, or MLA with TPA | Llama architecture, medium (353M) · FineWeb-Edu 100B zero-shot evaluation average | accuracy: 50.11 -> 51.41; 50.44 -> 51.41; 50.35 -> 51.41; 50.13 -> 51.41 | Helped |
| replace MHA, MQA, or GQA with TPA | Llama architecture, large (773M) · FineWeb-Edu 100B zero-shot evaluation average | accuracy: 52.52 -> 53.10; 52.13 -> 53.10; 52.30 -> 53.10 | Helped |
| replace MLA with TPA | Llama architecture, large (773M) · FineWeb-Edu 100B zero-shot evaluation average | accuracy: 53.32 -> 53.10 | Hurt |
| factorize only keys and values instead of queries, keys, and values (ablation) | TPA, medium (353M) · FineWeb-Edu 100B zero-shot evaluation average | accuracy: 51.41 -> 50.04 | Hurt |
| factorize only keys and values instead of queries, keys, and values (ablation) | TPA, large (773M) · FineWeb-Edu 100B zero-shot evaluation average | accuracy: 53.10 -> 53.52 | Helped |
Surprising findings
- MHA, MQA, and GQA can be written as non-contextual special cases of TPA.
- On the large model, MLA scores highest on average (53.32%) among baselines, above MHA, but TPA-KVonly edges it at 53.52%.
Limitations
- FlashTPA decoding kernel is implemented in Triton rather than optimized CUDA, so speed comparisons are preliminary.
- Downstream gains are small (often within about 1 point) and models are at most 1.5B parameters.
- Results come from single runs without variance estimates.
Experimental setup
Datasets: FineWeb-Edu 100B
Benchmarks: ARC-E ARC-C BoolQ HellaSwag OBQA PIQA WinoGrande MMLU SciQ
Metrics: validation loss perplexity accuracy decoding time
Baselines: MHA MQA GQA MLA FlashMHA FlashMLA
Builds on: Multi-Head Attention Multi-Query Attention Grouped-Query Attention Multi-Head Latent Attention RoPE
Closest papers by meaning
- Efficient Low Rank Attention for Long-Context Inference in Large Language Models (NeurIPS 2025 · similarity 0.89)
- A Training-Free Sub-quadratic Cost Transformer Model Serving Framework with Hierarchically Pruned Attention (ICLR 2025 · similarity 0.89)
- Loki: Low-rank Keys for Efficient Sparse Attention (NeurIPS 2024 · similarity 0.88)
- RocketKV: Accelerating Long-Context LLM Inference via Two-Stage KV Cache Compression (ICML 2025 · similarity 0.88)
- QUEST: Query-Aware Sparsity for Efficient Long-Context LLM Inference (ICML 2024 · similarity 0.88)
- Reducing Transformer Key-Value Cache Size with Cross-Layer Attention (NeurIPS 2024 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.