Произведение тензорных вниманий - все, что вам нужно.

Tensor Product Attention Is All You Need

January 11, 2025

Авторы: Yifan Zhang, Yifeng Liu, Huizhuo Yuan, Zhen Qin, Yang Yuan, Quanquan Gu, Andrew Chi-Chih Yao

cs.AI

Аннотация

Масштабирование языковых моделей для обработки более длинных последовательностей в общем случае требует больших кэшей ключ-значение (KV), что приводит к значительному избыточному использованию памяти во время вывода. В данной статье мы предлагаем Tensor Product Attention (TPA), новый механизм внимания, который использует тензорные декомпозиции для компактного представления запросов, ключей и значений, что значительно уменьшает размер кэша KV во время вывода. Факторизуя эти представления на контекстуальные низкоранговые компоненты (контекстуальная факторизация) и плавно интегрируя с RoPE, TPA достигает улучшения качества модели наряду с эффективностью использования памяти. Основываясь на TPA, мы представляем Transformer Tensor ProducT ATTenTion (T6), новую архитектуру модели для моделирования последовательностей. Проводя обширное эмпирическое исследование задач языкового моделирования, мы демонстрируем, что T6 превосходит производительность стандартных базовых моделей Transformer, включая MHA, MQA, GQA и MLA по различным метрикам, включая перплексию и ряд известных оценочных бенчмарков. Особенно стоит отметить, что эффективность использования памяти TPA позволяет обрабатывать значительно более длинные последовательности при фиксированных ресурсных ограничениях, решая критическую проблему масштабируемости в современных языковых моделях. Код доступен по ссылке 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.