Tensorproductaandacht is alles wat je nodig hebt.

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

Samenvatting

Het opschalen van taalmodellen om langere invoerreeksen te verwerken vereist doorgaans grote key-value (KV) caches, wat resulteert in aanzienlijke geheugenoverhead tijdens inferentie. In dit artikel stellen we Tensor Product Attention (TPA) voor, een nieuw aandachtsmechanisme dat tensordecomposities gebruikt om queries, keys en values compact voor te stellen, waardoor de KV-cachegrootte aanzienlijk wordt verkleind tijdens inferentie. Door deze voorstellingen te factoriseren in contextuele laag-rangcomponenten (contextuele factorisatie) en naadloos te integreren met RoPE, bereikt TPA verbeterde modelkwaliteit naast geheugenefficiëntie. Op basis van TPA introduceren we de Tensor ProducT ATTenTion Transformer (T6), een nieuw modelarchitectuur voor sequentiemodellering. Door uitgebreide empirische evaluatie van taalmodelleringstaken tonen we aan dat T6 de prestaties van standaard Transformer-baselines inclusief MHA, MQA, GQA en MLA overtreft op verschillende metrieken, waaronder perplexiteit en een reeks gerenommeerde evaluatiebenchmarks. Opmerkelijk is dat de geheugenefficiëntie van TPA de verwerking van aanzienlijk langere reeksen onder vaste resourcebeperkingen mogelijk maakt, waarmee een kritieke schaalbaarheidsuitdaging in moderne taalmodellen wordt aangepakt. De code is beschikbaar op 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.