PaperlayerView exact revision on arXiv

Source-linked AI summary

Simple linear attention language models balance the recall-throughput tradeoff

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

arXiv:2402.18668v2cs.CLcs.LG

TL;DR

Attention language models recall context well but use a growing KV-cache during inference, raising the question of whether efficiency can improve without sacrificing recall. The paper studies the state-size–recall tradeoff and proposes BASED, combining linear and sliding-window attention. BASED matches strong sub-quadratic models in quality, improves recall-intensive performance, and reaches up to 24× higher generation throughput than FlashAttention-2.

  • Problem

    Attention excels at recall, but its inference throughput is bottlenecked by aggressive KV-cache memory consumption.

  • Method

    The paper combines softmax-approximating linear attention with small sliding-window attention and varies their state-determining dimensions to explore the recall-memory Pareto frontier.

  • Results

    Across language modeling, DNA modeling, and LM Eval Harness evaluations up to 1.3b parameters, BASED competes with strong sub-quadratic baselines; it also outperforms Mamba by 10.36 accuracy points on recall-intensive tasks and reaches up to 24× higher generation throughput than FlashAttention-2.

  • Takeaways & Limitations

    Combining local fine-grained attention with long-range linear attention expands the recall-throughput tradeoff frontier while retaining an efficient recurrent inference view.

  • Takeaways & Limitations

    Small sliding-window attention has a limited receptive field, so matching a 4096-token window with w = 64 would require 64 BASED layers.

Abstract

from arXiv · show

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.

1. Introduction

Attention offers strong recall but incurs growing inference memory costs, motivating architectures that balance recall with throughput. BASED combines linear attention and small sliding-window attention to navigate this tradeoff.

  • Attention excels at recall but inference throughput is bottlenecked by aggressive KV-cache memory consumption.
  • A fundamental tradeoff links recurrent state size to recall capacity across broad architecture classes.
  • Linear attention alone struggles with associative recall, while sliding-window attention limits recall range by its window width.
  • BASED combines linear attention with tiny sliding-window softmax attention, whose components complement long-range modeling and precise local token operations.
  • 24× higher generation throughput than FlashAttention-2 is enabled by IO-aware algorithms for BASED's linear-attention implementation.

2. Preliminaries and Related Work

This section reviews attention and efficient alternatives that reduce inference cost through sparsity, kernelization, recurrence, or attention-free sequence mixing. These approaches target attention’s growing generation-time memory and computation, while recall quality remains a central concern.

  • Attention: Softmax attention computes outputs from query, key, and value projections of the input sequence.In causal generation, each new output attends to previously cached keys and values.
  • Attention: Attention is effective for recall and efficient to train, but generation requires operations over a growing KV-cache.The resulting memory consumption and lower throughput become more problematic for longer sequences.
  • Efficient attentions: Sparse attention reduces time and memory by attending only to selected patterns, including strided patterns and local sliding windows.Prior sparse patterns can underperform full attention when they fail to capture dense interactions.
  • Efficient attentions: Linear attention replaces exp(q^⊤k) with feature-map products ϕ(q)⊤ϕ(k), enabling associative computation and recurrent generation.The cited formulation uses O(Nd^2) time and space and permits constant-memory, O(1)-time-per-token generation.
  • Attention alternatives: Related work combines sparse and linear attention, while other attention-free mixers use state-space models, gated convolutions, or input-dependent recurrences.These approaches seek improved efficiency without sacrificing language-model quality, but some evidence suggests weaker recall or in-context learning than Transformers.

3. No Free Lunch: Memory-Recall Tradeoff

The section establishes a fundamental tradeoff between recurrent-state size and recall capacity through experiments and theory. Larger states generally improve recall, while lower-bound results show that memory requirements are intrinsic rather than architectural accidents.

  • Overview: Experiments and theory demonstrate a fundamental tradeoff between inference memory consumption, measured by recurrent-state size, and recall capacity.The study combines synthetic-data experiments with theoretical analysis across a broad class of architectures.
  • Empirical tradeoff: Larger recurrent states almost always improve MQAR accuracy within an architecture class, although architectures differ at the same state size.The empirical study varies hyperparameters controlling state size across models including Mamba and Hyena.
  • Experimental setup: The MQAR task presents key-value pairs followed by queries, requiring models to recall the matching earlier pair to predict the next token.Models train on sequences of length 256 and evaluate on sequences of length 1,024 with broader numbers of key-value pairs.
  • Empirical results: Attention achieves perfect recall but requires recurrent state that grows with sequence length, whereas Mamba makes the best use of a limited fixed-memory budget among prior architectures.H3 and Hyena fall below the Pareto frontier, while BASED expands it by varying state-controlling hyperparameters.
  • Lower bounds: Any causal recurrent model solving MQAR requires Ω(N)-bits in its state, establishing a memory lower bound for exact recall.The result supports the conclusion that the observed memory-recall tradeoff is fundamental rather than an artifact of architectural choices.
  • Gated convolutions: BaseConv cannot solve MQAR in constant-many layers, requiring log(2d) layers under the stated encoding and dimensionality conditions.For broader p-hot encodings, input-dependent BaseConv requires at least floor(log(2p)) layers.
  • Gated convolutions: For certain settings, the lower bound for BaseConv is tight: O(max(log log c, log log N)) layers suffice to solve MQAR.This completes the stated lower-bound characterization for those settings.
  • Theoretical comparison: Linear attention can be simulated by BaseConv with a poly-logarithmic layer blowup, indicating relative efficiency for linear attention over gated convolutions.The result concerns the theoretical simulation relationship rather than the empirical MQAR curve alone.

4. The BASED Architecture

BASED combines global softmax-approximating linear attention with local exact sliding-window attention to navigate the recall–memory tradeoff. Its state size and recall can be adjusted through feature and window dimensions, while small windows preserve efficient local interactions.

  • Motivation: Neither linear attention nor sliding-window attention alone suffices to navigate the recall–memory Pareto frontier.Linear attention lacks precision for local token shifts and comparisons, whereas sliding-window attention limits associative recall range.
  • BASED design: BASED combines global linear attention with exact local softmax attention in small sliding windows.The linear component preserves global token interactions, while sliding-window attention supplies precise local shifts and comparisons.
  • BASED design: Small windows such as 64–128 tokens recover 90.8% of full softmax attention’s recall accuracy at 1e-5× its latency.The global linear-attention component addresses the long-range modeling challenge introduced by small windows.
  • Taylor Linear Attention: The Taylor feature map uses a 2nd-order approximation with projected query and key dimension d′, allowing recurrent-state size to vary with d′.The cited construction gives feature dimension d′ = 16 as an example and identifies d′ as a control on recurrent-state size.
  • Local Exact Attention: Sliding-window attention gives each query access to the preceding w keys, with O(Nw) time and space and a w-sized KV-cache.BASED uses small windows, while larger windows increase state size and affect speed and recall.
  • Additional mixers: Short gated convolutions can further improve local precise shifts because they operate over the full sequence.These convolutions complement the local sliding-window layers for token comparisons.

5. Efficient Implementation

The paper develops hardware-aware kernels for BASED because naïve linear-attention implementations incur costly memory movement. These algorithms fuse feature-map and causal-dot-product computation and keep recurrent state in faster GPU memory where possible.

  • Efficiency motivation: Naïve BASED is slower than optimized standard attention because it requires substantial high-latency memory movement.The implementation challenge concerns real-world wall-clock efficiency despite linear attention’s theoretical complexity.
  • GPU execution: The algorithms exploit GPU tensor cores and the hierarchy from slow HBM to faster SRAM and registers.Tensor cores provide much higher throughput than non-matrix-multiply cores, motivating IO-aware placement and movement reduction.
  • Execution stages: The implementation targets both prompt processing and next-token prediction.These correspond to prefill during generation or training and autoregressive generation, respectively.
  • Baseline implementation: The baseline implementation computes feature maps in Python while using CUDA only for the causal dot product.This is inefficient because feature-map computation itself requires substantial work.
  • Linear-attention kernel: The hardware-aware linear-attention kernel computes the feature map and causal dot product together.For causality, output tiles are computed by splitting the dot-product computation into components handled in parallel.
  • IO analysis: The optimized procedure reduces HBM-to-SRAM movement by O(2BHND) bytes and avoids O(BHNDd) bytes of SRAM-to-register movement.It also stores the KV-state in registers when possible.

6. Results

BASED matches strong Transformer++ and Mamba baselines across language-modeling evaluations, while improving recall-intensive performance and generation efficiency.

  • Language modeling recall: 10.36 accuracy points separate BASED from Mamba on recall-intensive tasks at 1.3b parameters and 50b training tokens.The tasks include information extraction and question answering, and BASED closes the gap to Transformer++.
  • Generation throughput: 40–60% speedups over FlashAttention-2 and Mamba occur during 4k-token prefill with the IO-aware implementation.At 1.3b parameters, BASED is 56% faster than FlashAttention-2 and 44% faster than Mamba.
  • Generation throughput: The custom kernel, rather than the baseline Fast Transformers kernel, unlocks BASED’s reported efficiency.Throughput was benchmarked with and without the proposed IO-aware algorithms.

7. Conclusion

The paper identifies a recall-throughput tradeoff and proposes BASED, which combines local attention with long-range linear attention. Its IO-aware algorithms improve generation throughput while its recall quality exceeds prior sub-quadratic architectures.

  • 7. Conclusion: The paper identifies a fundamental tradeoff between recall and throughput through theory and experiments.Attention recalls perfectly but requires a KV cache that grows with sequence length.
  • 7. Conclusion: BASED combines local fine-grained attention with long-range Taylor-approximated linear attention for sub-quadratic training and recurrent inference.The architecture uses exact softmax attention locally and a softmax approximation globally.
  • 7. Conclusion: 24× faster generation than FlashAttention-2 is reported for 1.3b-parameter models generating 1024 tokens at batch size 128.The result comes from IO-aware algorithms for Taylor linear-attention inference.
  • 7. Conclusion: BASED outperforms prior sub-quadratic architectures in recall quality by 10.36 accuracy points on average.The paper concludes that BASED extends the Pareto frontier of the recall-throughput tradeoff.

Impact Statement

The paper aims to advance machine learning by reducing its costs and unlocking new capabilities, while directing readers to existing literature for broader LLM risk discussions.

  • Impact Statement: The authors intend BASED to help reduce machine-learning costs and unlock new capabilities.The statement frames these as intended benefits rather than established outcomes.
  • Impact Statement: The paper refers readers to existing discussions of risks associated with developing and using large language models.It does not identify a specific societal consequence requiring emphasis here.

A.1. Efficient Language Modeling Architectures

Efficient language-modeling architectures replace or restrict standard attention using sparse attention, linear attention, state-space models, and systems-level optimizations. BASED combines local exact attention with global linear attention to target recall and efficiency together.

  • Efficient Language Modeling Architectures: Transformer efficiency work targets quadratic input processing and KV-cache costs during autoregressive decoding.These costs motivate alternatives that retain modeling quality while improving sequence-processing or generation efficiency.
  • Efficient Language Modeling Architectures: Efficient sequence models include attention approximations, structured sparse attention, and state-space models with subquadratic processing.Linear attention and SSMs also support better-than-O(NM) decoding through recurrent or associative formulations.
  • Efficient Language Modeling Architectures: Sliding-window attention reduces computation by attending only to a local window of past tokens.A window of w tokens gives processing cost proportional to Nw for sequences of length N.
  • Efficient Language Modeling Architectures: BASED uses small sliding windows and a hardware-guided implementation to capitalize on sparse-attention efficiency.Its design combines globally applied softmax-approximating linear attention with locally applied exact softmax attention.
  • Efficient Language Modeling Architectures: Linear attention replaces softmax with feature-map dot products and uses associativity to reduce computation.The cited formulation computes attention in O(Nd^2) time and space rather than standard quadratic scaling.
  • Efficient Language Modeling Architectures: Pure linear SSMs struggle with associative recall, motivating multiplicative interactions and input-dependent sequence mixers.Associative recall requires shifting prior information to interact with current tokens for next-token prediction.
  • Efficient Language Modeling Architectures: IO-aware systems techniques improve speed and memory use by fusing operations and coordinating data movement across GPU memory levels.FlashAttention is presented as an exact attention implementation using fused kernels and HBM/SRAM-aware computation.

B. IO Aware Implementations

The section develops IO-aware kernels for BASED’s Taylor linear attention and sliding-window components, targeting efficient forward computation and next-token prediction.

  • Forward / Generation Prefill: The BASED kernel computes Taylor feature maps and causal dot products together rather than leaving feature-map computation in Python.The algorithm processes tiled q, k, and v inputs while maintaining cumulative Taylor-order state.
  • Forward / Generation Prefill: The microbenchmarks compare Pure PyTorch, Fast Transformers, and BASED kernels for Taylor linear-attention forward passes.They vary batch size or sequence length while fixing the complementary dimension, using median timing on one NVIDIA H100 GPU.
  • Forward / Generation Prefill: The implementation uses second-order Taylor features, tiled GPU execution, SRAM buffers, registers, and parallel batch-by-head computations.The stated tile size is 16, with eight warps per block in BASED.
  • Forward / Generation Prefill: Algorithm 1 combines quadratic on-diagonal interactions with linear off-diagonal cumulative-state computation to produce y = T0 + T1 + T2.Local causal masking handles within-tile interactions, while cumulative KV state handles prior tiles.
  • Next Token Prediction: The work contributes IO-aware algorithms for Taylor linear-attention KV-state updates and sliding-window attention during next-token generation.These kernels target the expensive recurrent-state and windowed-attention computations in inference.

C. Extended Architecture Details

This section adds gated convolution and coarse input-dependent decay as optional architectural details, while identifying Taylor linear attention plus TCWINDOW as the core BASED design.

  • Architectural components: Taylor linear attention and TCWINDOW alone come within 0.1 perplexity points of the best models using the additional components.The comparison is reported in the same section as the architectural additions.
  • Convolution: Replacing some linear-attention and TCWINDOW layers with gated convolutions provides small language-modeling improvements while keeping convolutions short.The described BaseConv layers use short convolutions and a SiLU nonlinearity.
  • Convolution: The gated-convolution formulation uses a Hadamard product and three learned linear projections with specified weight and bias dimensions.The projected input and learned filter are combined through elementwise multiplication before projection back to model dimension.
  • Decay: The decay mechanism controls the relative emphasis on recent versus early tokens, with prior input-dependent approaches requiring a parallel scan.BASED instead explores a coarser input-dependent decay that avoids the parallel scan.
  • Decay: The main experiments use no decay at 50b and 30b tokens, while ablations find decay can help small models without changing BASED’s overall trends.The authors report this as an ablation-level effect rather than a change to the main conclusions.
  • Recall-memory tradeoff: Figure 8 plots MQAR accuracy against generation-time state size and shows a recall-capacity tradeoff across architecture classes.Models vary hyperparameters such as model dimension to change recurrent state size.

D.2. Downstream Language Results

Across downstream language, few-shot, DNA modeling, and DNA classification evaluations, BASED is reported as competitive with recent strong architectures, while recall-related differences appear in few-shot behavior.

  • Evaluation setup: The downstream evaluations use zero-shot LM-Eval Harness settings for standard tasks and few-shot prompting for SuperGLUE.The standard-task metrics include perplexity, normalized accuracy, and accuracy across the listed benchmarks.
  • Standard language tasks: BASED performs competitively with recent and state-of-the-art architectures, including Mamba and Transformer++, for both 360m and 1.3b parameter models.The reported LM-Eval Harness averages cover six tasks excluding LAMBADA perplexity.
  • SuperGLUE few-shot results: Transformer++ and BASED improve monotonically as SuperGLUE demonstrations increase, whereas Mamba performs worse at five shots than at zero shots.The authors suggest Mamba’s limited recall ability may also affect few-shot abilities.
  • DNA modeling: BASED remains competitive with state-of-the-art architectures on HG38 DNA modeling at 1k, 4k, and 8k sequence lengths.The models are pretrained from scratch for 10Bn tokens at the evaluated lengths.
  • Downstream DNA classification: DNA classification results show similar performance across tasks, and the quality matching observed during pretraining transfers to downstream classification.The comparison uses GenomicBenchmarks top-1 classification accuracy.

D.4. BASED Quality Ablations

BASED quality ablations show that feature-map capacity, local sequence mixers, and decay each affect performance, while the architecture’s state size is tuned through model and feature dimensions.

  • Feature dimension: Larger feature dimensions improve quality, with diminishing returns from 24 to 32 dimensions.The ablation fixes the Taylor feature map while varying feature dimension.
  • Local sequence mixers: 2.29 AR Ppl. with neither local mixer improves to 2.09 with sliding window or 2.11 with convolutions.Adding either local component helps the associative-recall slice relative to using neither.
  • Local sequence mixers: Increasing the sliding-window size from 0 to 64 helps quality, while increasing it from 64 to 128 yields only marginal improvement.Both window sizes are described as efficient design points.
  • Input-dependent decay: Removing input-dependent decay causes no change to the overall trends, although decay can provide a small quality boost.The main 360m and 1.3b models use no input-dependent decay at 30b and 50b training tokens, respectively.
  • State size: BASED recurrent state size is determined by model dimension d and post-feature-map dimension d′, with experiments varying d ∈{48, 64, 128} and d′ ∈{8, 16, 24}.The state-size definition includes an additional K-state for computing the denominator.

E.3. Language Model Evaluation

The evaluation measures language modeling and recall-oriented downstream behavior using Pile slices and zero-shot extraction and question-answering tasks, alongside theoretical model definitions and results.

  • Pile language modeling: Pile evaluation reports overall perplexity and separate associative-recall and other-token slices.Associative-recall tokens are final bigram tokens previously seen in context but occurring no more than 1250 times in training.
  • Information extraction: SWDE evaluates extraction of semi-structured relations from raw HTML websites, while FDA evaluates key-value extraction from FDA PDFs.FDA documents are split into 1,920-token chunks, and accuracy checks whether generated text contains the value.
  • Question answering: SQUAD supplies a passage and a question whose answer appears in that passage, but the models struggle with direct questions because they are small and not instruction fine-tuned.GPT-4 reformats questions into next-token-prediction-style statements; 2,984 of 5,000 validation questions remain after filtering.
  • Theory: The theoretical section defines BASED using BaseConv and LinearAttention layers and studies their equivalence and MQAR lower bounds.The stated theory includes an equivalence result and lower bounds for associative recall.
  • Theory: Theorem F.9 states that one-way randomized communication complexity for the index problem on an n-bit string is Ω(n).This theorem is presented among the theoretical foundations for the paper’s lower-bound analysis.

F.4.2. LOWER BOUND FOR RECURRENT MODELS

The section defines recurrent models as systems whose state summarizes input history through causal state-update functions. A communication-complexity reduction then shows that recurrent models solving AR require state memory growing linearly with sequence length, including causal Mamba under bounded-precision assumptions.

  • Definition: A recurrent model’s state at position i depends only on preceding inputs and evolves from the previous state plus the current input.The state represents accumulated input-dependent information and is distinct from static learned parameters.
  • Definition: Recurrent models exclude architectures whose computation at a state inherently requires the entire input sequence, such as non-causal full-sequence convolutions.
  • Lower bound: Ω(N) bits are necessary for any recurrent model that solves AR, by reducing the task to one-way communication for the index problem.Alice sends the model state to Bob; sublinear state communication would contradict the index-problem lower bound.
  • Lower bound: Causal Mamba also requires Ω(N) bits for AR when all computation entries use O(1) bits, implying a corresponding lower bound on its state dimensions.The proof establishes that causal Mamba is recurrent because its input-dependent parameters are causal.

F.4.3. LOWER BOUND ON THE NUMBER OF LAYERS FOR AR

The section proves layer lower bounds for data-independent BaseConv on AR and contrasts them with constant-depth attention for a binary-token MQAR construction. These bounds follow because BaseConv computes bounded-degree polynomials, whereas exact recall requires higher degree.

  • AR lower bound: Ω(ϵ log log N) layers are required for data-independent BaseConv to solve AR when d ≤ 2(log N)^(1−ϵ) and parameters use O(log N) bits.The lower bound is obtained by combining the polynomial coefficient communication bound with the index-problem lower bound.
  • Polynomial degree: A BaseConv model with L layers computes a polynomial over the input of degree at most 2L, with a multilinear equivalent on Boolean inputs.This degree bound follows by composing single-layer degree-2 transformations.
  • Attention comparison: Attention with linear biases and ReLU plus two MLP layers solves MQAR for d = log2(c) in O(1) layers.The construction matches queries to keys and returns the value associated with the first matching key.
  • BaseConv lower bound: A data-independent BaseConv model needs log(2d) layers to solve MQAR for d = log2(c), because its degree grows only exponentially with depth.MQAR in this setting requires a multilinear polynomial of degree 2d + 1.

F.5.1. THE EQUALITY PROBLEM

The equality problem tests whether two token encodings match and is used to derive layer lower bounds for BaseConv on p-hot representations. Equality has block-exclusive polynomial degree 2p, while MQAR inherits its lower bounds.

  • Problem and encodings: The equality problem asks whether two vocabulary-token encodings are identical, and p-hot encodings represent token digits as one-hot blocks.Almost p-hot encodings truncate the one-hot representation’s final bit.
  • Reduction: Any model solving MQAR also solves equality with the same number of layers, so equality lower bounds transfer directly to MQAR.
  • Polynomial representation: Block-exclusive polynomials forbid monomials containing products of two variables from the same encoding block, ensuring multilinearity under the defined representation.
  • Equality representation: The equality problem in almost p-hot encoding has a block-exclusive polynomial representation of degree 2p, and this representation is unique.Each block contributes a degree-2 agreement test, whose product yields equality.
  • Layer lower bound: A data-independent BaseConv model needs at least ⌊log(2p)⌋ layers to solve MQAR in the p-hot setting, where d = p · √c.The reduction uses equality as a lower-bound problem and the fact that L BaseConv layers yield degree at most 2L.

F.6.2. PROOF OF THEOREM F.30

The proof constructs an O(log log N)-layer BaseConv algorithm for one-hot MQAR with at most one key match per query. It converts binary index information into one-hot selection and then computes the selected value through staged primitives.

  • Index encoding: The construction first replaces one-hot key positions with binary indices and forms cumulative representations of preceding keys.This produces the encoded information needed to identify causal query-key matches.
  • Matching: It repeats query and key columns, computes their matching information, and stores each query’s matching key index in binary form.The unique-match assumption ensures that at most one key index needs to be represented per query.
  • One-hot conversion: The binary matching indices are converted into one-hot rows using a BaseConv stage requiring O(⌈log log N⌉) layers.The resulting matrix is L = C ⊙(QK⊤), the causal query-key match matrix.
  • Value aggregation: Repeated and permuted value blocks are masked, summed, and redistributed so each query receives the value associated with its matched key.The sequence includes repeat, permutation, diagonal masking, row summation, and copying operations.
Loading 2402.18668v2…