Einfache lineare Attention-Sprachmodelle balancieren den Kompromiss zwischen Recall und Durchsatz.
Simple linear attention language models balance the recall-throughput tradeoff
February 28, 2024
Autoren: Simran Arora, Sabri Eyuboglu, Michael Zhang, Aman Timalsina, Silas Alberti, Dylan Zinsley, James Zou, Atri Rudra, Christopher Ré
cs.AI
Zusammenfassung
Aktuelle Arbeiten haben gezeigt, dass auf Aufmerksamkeit basierende Sprachmodelle hervorragende Fähigkeiten beim Abruf besitzen, also der Fähigkeit, Erzeugungen in zuvor im Kontext gesehenen Tokens zu verankern. Allerdings wird die Effizienz von auf Aufmerksamkeit basierenden Modellen während des Inferenzprozesses durch den aggressiven Speicherverbrauch des KV-Caches eingeschränkt. In dieser Arbeit untersuchen wir, ob wir die Effizienz von Sprachmodellen verbessern können (z. B. durch Reduzierung des Speicherverbrauchs), ohne dabei den Abruf zu beeinträchtigen. Durch Experimente und theoretische Analysen an einer breiten Palette von Architekturen identifizieren wir einen zentralen Kompromiss zwischen der Zustandsgröße eines Modells und seiner Abruffähigkeit. Wir zeigen, dass effiziente Alternativen zur Aufmerksamkeit (z. B. H3, Mamba, RWKV) einen festen rekurrenten Zustand beibehalten, jedoch Schwierigkeiten beim Abruf haben. Wir schlagen BASED vor, eine einfache Architektur, die lineare und gleitende Fensteraufmerksamkeit kombiniert. Durch die Variation der Fenstergröße von BASED und der Feature-Dimension der linearen Aufmerksamkeit können wir die Zustandsgröße anpassen und die Pareto-Grenze der Abruf-Speicher-Kompromisskurve durchlaufen, wobei wir auf der einen Seite die volle Qualität der Aufmerksamkeit und auf der anderen Seite die geringe Zustandsgröße von Aufmerksamkeitsalternativen erreichen. Wir trainieren Sprachmodelle mit bis zu 1,3 Milliarden Parametern und zeigen, dass BASED die stärksten subquadratischen Modelle (z. B. Mamba) in Bezug auf die Perplexität erreicht und sie bei realen, abrufintensiven Aufgaben um 6,22 Genauigkeitspunkte übertrifft. Implementierungen der linearen Aufmerksamkeit sind oft weniger effizient als optimierte Standardimplementierungen der Aufmerksamkeit. Um BASED wettbewerbsfähig zu machen, entwickeln wir IO-bewusste Algorithmen, die einen 24-mal höheren Durchsatz bei der Sprachgeneration ermöglichen als FlashAttention-2, wenn 1024 Tokens mit 1,3-Milliarden-Parameter-Modellen generiert werden. Der Code für diese Arbeit ist unter https://github.com/HazyResearch/based verfügbar.
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.