L'attenzione del prodotto tensore è tutto ciò di cui hai bisogno

Tensor Product Attention Is All You Need

January 11, 2025

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

cs.AI

Abstract

L'adattamento dei modelli linguistici per gestire sequenze di input più lunghe richiede tipicamente ampi cache chiave-valore (KV), con conseguente notevole sovraccarico di memoria durante l'inferezza. In questo articolo, proponiamo l'Attenzione al Prodotto Tensorico (TPA), un nuovo meccanismo di attenzione che utilizza decomposizioni tensoriali per rappresentare in modo compatto le query, le chiavi e i valori, riducendo significativamente le dimensioni della cache KV durante l'inferezza. Mediante la fattorizzazione di tali rappresentazioni in componenti contestuali a basso rango (fattorizzazione contestuale) e l'integrazione senza soluzione di continuità con RoPE, TPA raggiunge una qualità del modello migliorata insieme a un'efficienza di memoria. Basandoci su TPA, introduciamo il Transformer di Attenzione al Prodotto Tensorico (T6), una nuova architettura di modello per il modellamento delle sequenze. Attraverso un'ampia valutazione empirica di compiti di modellazione del linguaggio, dimostriamo che T6 supera le prestazioni dei baselines standard del Transformer, inclusi MHA, MQA, GQA e MLA, su vari metriche, tra cui perplessità e una serie di rinomati benchmark di valutazione. In particolare, l'efficienza di memoria di TPA consente l'elaborazione di sequenze significativamente più lunghe in presenza di vincoli di risorse fissi, affrontando una sfida critica di scalabilità nei moderni modelli linguistici. Il codice è disponibile su 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.