Tensor Product Attention Is All You Need

• URL: http://arxiv.org/abs/2501.06425v2
• Date: Fri, 07 Feb 2025 22:00:48 GMT
• Title: Tensor Product Attention Is All You Need
• Authors: Yifan Zhang, Yifeng Liu, Huizhuo Yuan, Zhen Qin, Yang Yuan, Quanquan Gu, Andrew Chi-Chih Yao,
• Abstract summary: Product Attention (TPA) is a novel attention mechanism that uses tensor decompositions to represent queries, keys, and values compactly. TPA achieves improved model quality alongside memory efficiency. We introduce the ProducT ATTion Transformer (T6), a new model architecture for sequence modeling.
• Score: 54.40495407154611
• License:
• 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, TPA's 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.

Related papers

• FlexiGPT: Pruning and Extending Large Language Models with Low-Rank Weight Sharing [59.12511498024836]
  We present a method to prune large language models (LLMs) that selectively prunes model blocks based on an importance score. We propose a principled metric to replace each pruned block using a weight-sharing mechanism. Empirical evaluations demonstrate substantial performance gains over existing methods.
  arXiv  Detail & Related papers  (2025-01-24T18:46:37Z)
• CSR:Achieving 1 Bit Key-Value Cache via Sparse Representation [63.65323577445951]
  We propose a novel approach called Cache Sparse Representation (CSR) CSR transforms the dense Key-Value cache tensor into sparse indexes and weights, offering a more memory-efficient representation during LLM inference. Our experiments demonstrate CSR achieves performance comparable to state-of-the-art KV cache quantization algorithms.
  arXiv  Detail & Related papers  (2024-12-16T13:01:53Z)
• Symmetric Dot-Product Attention for Efficient Training of BERT Language Models [5.838117137253223]
  We propose an alternative compatibility function for the self-attention mechanism introduced by the Transformer architecture. When applied to the pre-training of BERT-like models, this new symmetric attention mechanism reaches a score of 79.36 on the GLUE benchmark against 78.74 for the traditional implementation.
  arXiv  Detail & Related papers  (2024-06-10T15:24:15Z)
• Layer-Condensed KV Cache for Efficient Inference of Large Language Models [44.24593677113768]
  We propose a novel method that only computes and caches the KVs of a small number of layers. Our method achieves up to 26$times$ higher throughput than standard transformers.
  arXiv  Detail & Related papers  (2024-05-17T08:59:46Z)
• CORM: Cache Optimization with Recent Message for Large Language Model Inference [57.109354287786154]
  We introduce an innovative method for optimizing the KV cache, which considerably minimizes its memory footprint. CORM, a KV cache eviction policy, dynamically retains essential key-value pairs for inference without the need for model fine-tuning. Our validation shows that CORM reduces the inference memory usage of KV cache by up to 70% with negligible performance degradation across six tasks in LongBench.
  arXiv  Detail & Related papers  (2024-04-24T16:11:54Z)
• Keyformer: KV Cache Reduction through Key Tokens Selection for Efficient Generative Inference [2.8241099113277666]
  "Keyformer" is an innovative inference-time approach to mitigate the challenges associated with KV cache size and memory bandwidth utilization. We evaluate Keyformer's performance across three foundational models: GPT-J, Cerebras-GPT, and MPT.
  arXiv  Detail & Related papers  (2024-03-14T02:42:42Z)
• Repeat After Me: Transformers are Better than State Space Models at Copying [53.47717661441142]
  We show that while generalized state space models are promising in terms of inference-time efficiency, they are limited compared to transformer models on tasks that require copying from the input context.
  arXiv  Detail & Related papers  (2024-02-01T21:44:11Z)
• CoTFormer: A Chain-of-Thought Driven Architecture with Budget-Adaptive Computation Cost at Inference [36.753384415107774]
  Scaling language models to larger and deeper sizes has led to significant boosts in performance. We propose CoTFormer, a novel architecture which closely mimics Chain-of-Thought (CoT) at the token level. We show that it is possible to reduce the computation cost significantly without any reduction in accuracy.
  arXiv  Detail & Related papers  (2023-10-16T21:37:34Z)
• Direction is what you need: Improving Word Embedding Compression in Large Language Models [7.736463504706344]
  This paper presents a novel loss objective to compress token embeddings in Transformer-based models by leveraging an AutoEncoder architecture. Our method significantly outperforms the commonly used SVD-based matrix-factorization approach in terms of initial language model Perplexity.
  arXiv  Detail & Related papers  (2021-06-15T14:28:00Z)
• Improve Variational Autoencoder for Text Generationwith Discrete Latent Bottleneck [52.08901549360262]
  Variational autoencoders (VAEs) are essential tools in end-to-end representation learning. VAEs tend to ignore latent variables with a strong auto-regressive decoder. We propose a principled approach to enforce an implicit latent feature matching in a more compact latent space.
  arXiv  Detail & Related papers  (2020-04-22T14:41:37Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.