L'attention de produit tensoriel est tout ce dont vous avez besoin.

Tensor Product Attention Is All You Need

January 11, 2025

Auteurs: Yifan Zhang, Yifeng Liu, Huizhuo Yuan, Zhen Qin, Yang Yuan, Quanquan Gu, Andrew Chi-Chih Yao

cs.AI

Résumé

L'adaptation des modèles de langage pour traiter des séquences d'entrée plus longues nécessite généralement de grandes caches clé-valeur (KV), entraînant une surcharge mémoire importante lors de l'inférence. Dans cet article, nous proposons l'Attention par Produit de Tenseurs (TPA), un nouveau mécanisme d'attention qui utilise des décompositions tensorielles pour représenter de manière compacte les requêtes, les clés et les valeurs, réduisant significativement la taille de la cache KV au moment de l'inférence. En factorisant ces représentations en composantes de rang faible contextuelles (factorisation contextuelle) et en les intégrant de manière transparente avec RoPE, TPA permet d'améliorer la qualité du modèle tout en étant efficace en termes de mémoire. Basé sur TPA, nous introduisons le Transformateur d'Attention par Produit de Tenseurs (T6), une nouvelle architecture de modèle pour la modélisation de séquences. À travers une évaluation empirique approfondie des tâches de modélisation de langage, nous démontrons que T6 dépasse les performances des modèles Transformer standards, y compris MHA, MQA, GQA et MLA, sur diverses mesures, y compris la perplexité et une gamme de benchmarks d'évaluation renommés. Notamment, l'efficacité mémoire de TPA permet le traitement de séquences significativement plus longues dans le cadre de contraintes de ressources fixes, répondant à un défi critique de scalabilité dans les modèles de langage modernes. Le code est disponible sur https://github.com/tensorgi/T6.

English

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.