[Paper Review] Simple linear attention language models balance the recall-throughput tradeoff

Simran Arora, Sabri Eyuboglu+7 more|arXiv (Cornell University)|Feb 28, 2024

Topic Modeling4 citations

TL;DR

This paper proposes Based, a hybrid language model architecture that combines linear attention and sliding window attention to balance recall and inference efficiency. By tuning window size and feature dimension, Based achieves full attention-level recall with significantly reduced memory usage and 24× higher throughput than FlashAttention-2, outperforming SOTA models like Mamba on recall-intensive tasks by 6.22 accuracy points.

ABSTRACT

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.

Motivation & Objective

• To address the tradeoff between recall performance and inference efficiency in autoregressive language models.
• To identify and traverse the Pareto frontier between model state size and recall ability across different architectures.
• To design a model that maintains high recall like standard attention while reducing memory consumption like efficient alternatives (e.g., Mamba, RWKV).
• To develop IO-aware optimizations that make the proposed architecture competitive in real-world throughput benchmarks.
• To demonstrate that combining linear and sliding window attention enables superior tradeoff performance compared to either component alone.

Proposed method

• Introduce a hybrid architecture, Based, that combines linear attention and sliding window attention within the same model layers.
• Use a configurable sliding window size (e.g., 64, 128) and linear attention feature dimension to control state size and tradeoff behavior.
• Apply IO-aware algorithmic optimizations to reduce memory access overhead and improve throughput during generation.
• Train models up to 1.3B parameters using standard training recipes with Adam, cosine decay, and BFloat16 precision.
• Use rotary positional encodings for sliding window components and no positional encodings for linear components.
• Implement a block-exclusive polynomial analysis to theoretically justify the need for sufficient model depth to solve associative recall tasks.

Figure 1: Based overview. Combining linear attention with tiny sliding window softmax attention (e.g., 64 or 128 tokens in width) enables improved recall accuracy with limited efficiency overhead vs. smaller tile sizes. ( Left ) Time to execute Cutlass GEMMs ( $y$ ) vs. sliding window attention size

Experimental results

Research questions

• RQ1 Can a model achieve the recall performance of standard attention while maintaining the small state size of efficient alternatives like Mamba?
• RQ2 What is the fundamental tradeoff between model state size and recall ability across different architecture classes?
• RQ3 Can combining linear attention and sliding window attention create a Pareto-optimal solution that outperforms both individual components?
• RQ4 How can I/O-efficient algorithms be designed to make hybrid attention models competitive with FlashAttention in real-world throughput?
• RQ5 Does the combination of linear and local attention mechanisms enable better long-range and precise token interactions than either alone?

Key findings

• Based achieves 6.22 higher accuracy than Mamba on real-world recall-intensive tasks, demonstrating superior recall capability despite smaller state size.
• With 1.3B parameters, Based achieves 24× higher throughput than FlashAttention-2 during autoregressive generation of 1024 tokens.
• The hybrid architecture successfully traverses the Pareto frontier of the recall-memory tradeoff, recovering full attention-level recall at the cost of only a modest efficiency overhead.
• Linear attention alone fails on associative recall tasks due to lack of precision in local token comparisons, while sliding window attention is limited by window size.
• Theoretical analysis shows that a data-independent BaseConv model requires at least ⌊log(2p)⌋ layers to solve associative recall, validating the need for sufficient depth in such models.
• IO-aware optimizations enable Based to match or exceed the throughput of optimized attention implementations, even with hybrid attention components.

Figure 2: Throughput (memory) - recall tradeoff. $x$ -axis shows state size (bytes) during generation; $y$ -axis shows accuracy on the MQAR recall task [ 1 ] . For each architecture, we train several models varying hyperparameters that affect the recurrent state size ( e.g. model dimension). The plo