Atenção por Produto Tensor É Tudo o que Você Precisa

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

Resumo

Aumentar o tamanho dos modelos de linguagem para lidar com sequências de entrada mais longas geralmente requer caches de chave-valor (KV) grandes, resultando em uma sobrecarga significativa de memória durante a inferência. Neste artigo, propomos a Atenção de Produto de Tensor (TPA), um mecanismo de atenção inovador que utiliza decomposições tensoriais para representar consultas, chaves e valores de forma compacta, reduzindo significativamente o tamanho do cache KV durante a inferência. Ao fatorizar essas representações em componentes de baixa ordem contextual (fatorização contextual) e integrar-se perfeitamente com RoPE, a TPA alcança uma qualidade de modelo aprimorada juntamente com eficiência de memória. Com base na TPA, introduzimos o Transformador de Atenção de Produto de Tensor (T6), uma nova arquitetura de modelo para modelagem de sequências. Através de uma extensa avaliação empírica de tarefas de modelagem de linguagem, demonstramos que o T6 supera o desempenho de baselines de Transformadores padrão, incluindo MHA, MQA, GQA e MLA, em várias métricas, incluindo perplexidade e uma variedade de benchmarks de avaliação renomados. Notavelmente, a eficiência de memória da TPA permite o processamento de sequências significativamente mais longas sob restrições de recursos fixos, abordando um desafio crítico de escalabilidade em modelos de linguagem modernos. O código está disponível em 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.