Tensor Product Attention Is All You Need
Authors: name: Yifan Zhang; name: Yifeng Liu; name: Huizhuo Yuan; name: Zhen Qin; name: Yang Yuan; name: Quanquan Gu; name: Andrew Chi-Chih Yao
Published: 2025-01-11 03:37:10

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