I modelli linguistici con semplice attenzione lineare bilanciano il compromesso tra richiamo e velocità di elaborazione.

Simple linear attention language models balance the recall-throughput tradeoff

February 28, 2024

Autori: Simran Arora, Sabri Eyuboglu, Michael Zhang, Aman Timalsina, Silas Alberti, Dylan Zinsley, James Zou, Atri Rudra, Christopher Ré

cs.AI

Abstract

Recenti studi hanno dimostrato che i modelli linguistici basati su meccanismi di attenzione eccellono nel richiamo, ovvero nella capacità di ancorare le generazioni a token precedentemente osservati nel contesto. Tuttavia, l'efficienza dei modelli basati su attenzione è limitata durante l'inferenza dal consumo aggressivo di memoria del KV-cache. In questo lavoro, esploriamo se è possibile migliorare l'efficienza dei modelli linguistici (ad esempio riducendo il consumo di memoria) senza compromettere il richiamo. Applicando esperimenti e teoria a un ampio insieme di architetture, identifichiamo un compromesso fondamentale tra la dimensione dello stato di un modello e la sua capacità di richiamo. Mostriamo che alternative efficienti all'attenzione (ad esempio H3, Mamba, RWKV) mantengono uno stato ricorrente di dimensione fissa, ma faticano nel richiamo. Proponiamo BASED, una semplice architettura che combina attenzione lineare e attenzione a finestra scorrevole. Variando la dimensione della finestra di BASED e la dimensione delle feature dell'attenzione lineare, possiamo regolare la dimensione dello stato e attraversare la frontiera di Pareto della curva di compromesso tra richiamo e memoria, recuperando la piena qualità dell'attenzione da un lato e la piccola dimensione dello stato delle alternative all'attenzione dall'altro. Addestriamo modelli linguistici fino a 1,3 miliardi di parametri e mostriamo che BASED eguaglia i modelli sub-quadratici più forti (ad esempio Mamba) in termini di perplessità e li supera in compiti reali intensivi di richiamo di 6,22 punti di accuratezza. Le implementazioni dell'attenzione lineare sono spesso meno efficienti delle implementazioni ottimizzate dell'attenzione standard. Per rendere BASED competitivo, sviluppiamo algoritmi consapevoli delle operazioni di I/O che consentono un throughput 24 volte superiore nella generazione linguistica rispetto a FlashAttention-2, quando si generano 1024 token utilizzando modelli da 1,3 miliardi di parametri. Il codice per questo lavoro è disponibile all'indirizzo: https://github.com/HazyResearch/based.

English

Recent work has shown that attention-based language models excel at recall, the ability to ground generations in tokens previously seen in context. However, the efficiency of attention-based models is bottle-necked during inference by the KV-cache's aggressive memory consumption. In this work, we explore whether we can improve language model efficiency (e.g. by reducing memory consumption) without compromising on recall. By applying experiments and theory to a broad set of architectures, we identify a key tradeoff between a model's state size and recall ability. We show that efficient alternatives to attention (e.g. H3, Mamba, RWKV) maintain a fixed-size recurrent state, but struggle at recall. We propose BASED a simple architecture combining linear and sliding window attention. By varying BASED window size and linear attention feature dimension, we can dial the state size and traverse the pareto frontier of the recall-memory tradeoff curve, recovering the full quality of attention on one end and the small state size of attention-alternatives on the other. We train language models up to 1.3b parameters and show that BASED matches the strongest sub-quadratic models (e.g. Mamba) in perplexity and outperforms them on real-world recall-intensive tasks by 6.22 accuracy points. Implementations of linear attention are often less efficient than optimized standard attention implementations. To make BASED competitive, we develop IO-aware algorithms that enable 24x higher throughput on language generation than FlashAttention-2, when generating 1024 tokens using 1.3b parameter models. Code for this work is provided at: https://github.com/HazyResearch/based.