Simple linear attention language models balance the recall-throughput tradeoff (2402.18668v1)

Abstract: Recent work has shown that attention-based LLMs 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 LLM 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 LLMs 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.

Simple Linear Attention LLMs Balance the Recall-Throughput Tradeoff

The paper "Simple linear attention LLMs balance the recall-throughput tradeoff" investigates the recall efficiency of attention-based LLMs and proposes a novel architecture named Based to enhance performance metrics by addressing the inherent memory consumption tradeoffs during inference.

Introduction and Problem Statement

Attention-based LLMs are well-documented for their superior recall abilities, effectively grounding generations in prior context tokens. However, they suffer from significant inefficiencies regarding memory usage during inference. The paper centralizes the query: Can LLM efficiency be improved, particularly in terms of memory consumption, without deteriorating recall ability?

Empirical Analysis and Tradeoffs

Empirical evaluations demonstrate a fundamental tradeoff between a LLM’s recurrent state size (memory consumption during inference) and its recall functionality. Through a series of synthetic Multi-Query Associative Recall (MQAR) tasks and theoretical analyses, the authors highlight the inefficiencies of several model architectures related to recall abilities and their corresponding memory footprints.

Results:

1. Recall-Memory Tradeoff:
   • Attention-based models achieve perfect recall but at the cost of a recurrent state size growing linearly with the sequence length.
   • Efficient alternatives such as H3 and Mamba, which maintain a fixed-size recurrent state, exhibit poor recall performance.
2. Architecture Specifics:
   • Linear attention and sliding window attention alone fail to provide a satisfactory balance between memory and recall.
   • The combined architecture Based, incorporating both linear and sliding window attention, successfully traverses the Pareto frontier of the recall-memory tradeoff.

The Based Architecture

Elements and Design:

1. Linear Attention:
   • A significant component is softmax-approximating linear attention using a 2nd-order Taylor series feature map. This approximation maintains global token interactions with fixed recurrent state size, modulated by feature dimension (d').
2. Sliding Window Attention:
   • A combination with sliding window attention allows for efficient modeling of local interactions, where small-sized windows (e.g., 64 tokens) mitigate the inefficiencies of larger window sizes.

The resultant model traverses the recall-memory tradeoff effectively, aligning near attention-based models for recall, while sustaining memory use low — a crucial efficiency gain.

Theoretical Foundations

The authors establish lower bounds on recall-related memory consumption for recurrent models, underscoring the fundamental tradeoffs. Applying results from communication complexity theory, they ascertain the equivalence and simulation capabilities between Based and canonical gated-convolution architectures (BaseConv) in recall tasks. Additionally, the model's memory space complexity and the number of required layers for exact recall functions are theoretically bounded.

Experimental Results

Evaluations on LLMs trained on up to 1.3 billion parameters illustrate that Based:

1. Matches the strongest sub-quadratic architectures like Mamba in perplexity scores.
2. Outperforms on associative recall-intensive tasks by 6.22 accuracy points.
3. Demonstrates up to 24 times higher throughput on language generation tasks compared to FlashAttention-2.

Implications and Future Directions

Practical Implications:

• The findings emphasize the practical viability of implementing Based in real-world scenarios, leveraging high recall accuracy with vastly improved throughput rates.
• Applications span from information extraction, reading comprehension, to code generation where recall precision could significantly enhance task outcomes.

Theoretical Contributions:

• The development and analysis of Based contribute to the theoretical frameworks understanding the memory-recall tradeoff.
• Future research could explore optimizing model parameters specific to the needs of various downstream tasks, investigating simplified feature maps, or extending architectures to encompass broader input dependencies.

Speculative Future Trends in AI:

• As the field moves towards models balancing extensive context with efficient processing, architectures like Based set a precedent for future innovations targeting computational sustainability and accuracy.
• Development in AI hardware and specialized accelerators could further mitigate current limitations specific to certain approximation methods.

By establishing a frontier in the recall-throughput tradeoff, this paper effectively progresses the broader discourse on model efficiency and opens avenues for optimizing performance in complex, real-world language generation tasks.