La Atención de Producto Tensorial es Todo lo que Necesitas

Tensor Product Attention Is All You Need

January 11, 2025

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

cs.AI

Resumen

Escalar modelos de lenguaje para manejar secuencias de entrada más largas típicamente requiere grandes cachés de clave-valor (KV), lo que resulta en una sobrecarga de memoria sustancial durante la inferencia. En este documento, proponemos Atención de Producto Tensorial (TPA), un mecanismo de atención novedoso que utiliza descomposiciones tensoriales para representar de manera compacta consultas, claves y valores, reduciendo significativamente el tamaño del caché KV en tiempo de inferencia. Al factorizar estas representaciones en componentes contextuales de rango bajo (factorización contextual) e integrarse perfectamente con RoPE, TPA logra una calidad de modelo mejorada junto con eficiencia de memoria. Basándonos en TPA, presentamos el Transformador de Atención de Producto Tensorial (T6), una nueva arquitectura de modelo para modelado de secuencias. A través de una extensa evaluación empírica de tareas de modelado de lenguaje, demostramos que T6 supera el rendimiento de los baselines estándar del Transformador, incluyendo MHA, MQA, GQA y MLA en varios métricas, incluyendo perplejidad y una variedad de reconocidas evaluaciones de referencia. Notablemente, la eficiencia de memoria de TPA permite el procesamiento de secuencias significativamente más largas bajo restricciones de recursos fijos, abordando un desafío crítico de escalabilidad en los modelos de lenguaje modernos. El código está disponible en 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.