Tensor-Produkt-Aufmerksamkeit ist alles, was Sie brauchen.

Tensor Product Attention Is All You Need

January 11, 2025

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

cs.AI

Zusammenfassung

Die Skalierung von Sprachmodellen, um längere Eingabesequenzen zu verarbeiten, erfordert in der Regel große Schlüssel-Wert (KV)-Caches, was zu erheblichem Speicheroverhead während der Inferenz führt. In diesem Paper schlagen wir Tensor-Produkt-Aufmerksamkeit (TPA) vor, einen neuartigen Aufmerksamkeitsmechanismus, der Tensorzerlegungen verwendet, um Anfragen, Schlüssel und Werte kompakt darzustellen und somit die Größe des KV-Caches zur Inferenzzeit signifikant zu reduzieren. Durch die Faktorisierung dieser Darstellungen in kontextuelle niederdimensionale Komponenten (kontextuelle Faktorisierung) und die nahtlose Integration mit RoPE erzielt TPA eine verbesserte Modellqualität neben Speichereffizienz. Basierend auf TPA stellen wir den Tensor-Produkt-Aufmerksamkeits-Transformer (T6) vor, eine neue Modellarchitektur für Sequenzmodellierung. Durch umfangreiche empirische Evaluationen von Sprachmodellierungsaufgaben zeigen wir, dass T6 die Leistung von Standard-Transformer-Baselines wie MHA, MQA, GQA und MLA in verschiedenen Metriken, einschließlich Perplexität und einer Reihe renommierter Evaluierungsbenchmarks, übertrifft. Bemerkenswert ist, dass die Speichereffizienz von TPA die Verarbeitung von deutlich längeren Sequenzen unter festen Ressourcenbeschränkungen ermöglicht und somit eine wichtige Skalierbarkeitsherausforderung bei modernen Sprachmodellen angeht. Der Code ist verfügbar unter 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.