PaperlayerView exact revision on arXiv

Source-linked AI summary

Zoology: Measuring and Improving Recall in Efficient Language Models

Simran Arora, Sabri Eyuboglu, Aman Timalsina, Isys Johnson, Michael Poli, James Zou, Atri Rudra, Christopher Ré

arXiv:2312.04927v1cs.CLcs.LG

TL;DR

The paper asks why efficient gated-convolution language models still lag attention despite strong performance and favorable scaling. It evaluates 17 models, introduces multi-query associative recall, and analyzes the gap theoretically and empirically. The study attributes most of the gap to associative recall and finds that input-dependent sparse attention hybrids can substantially close it while retaining sub-quadratic scaling.

  • Problem

    Gated-convolution models can match attention on prior synthetic associative-recall tests yet still show a language-modeling quality gap on real data.

  • Method

    The paper pretrains and analyzes attention and gated-convolution models, formalizes multi-query associative recall, and studies it with theory, experiments, and convolution-attention hybrids.

  • Results

    The study finds that associative-recall errors explain 82% of the perplexity gap, while input-dependent sparse attention hybrids close much of the gap to attention.

  • Takeaways & Limitations

    Efficient sequence mixers need input-dependent sequence mixing to handle the variable, multi-query recalls found in language.

  • Takeaways & Limitations

    Input-dependent filters may not easily satisfy causality, and applying an O(N log N) filter per gap can be expensive when gaps contain few bigrams.

Abstract

from arXiv · show

Attention-free language models that combine gating and convolutions are growing in popularity due to their efficiency and increasingly competitive performance. To better understand these architectures, we pretrain a suite of 17 attention and "gated-convolution" language models, finding that SoTA gated-convolution architectures still underperform attention by up to 2.1 perplexity points on the Pile. In fine-grained analysis, we find 82% of the gap is explained by each model's ability to recall information that is previously mentioned in-context, e.g. "Hakuna Matata means no worries Hakuna Matata it means no" $\rightarrow$ "??". On this task, termed "associative recall", we find that attention outperforms gated-convolutions by a large margin: a 70M parameter attention model outperforms a 1.4 billion parameter gated-convolution model on associative recall. This is surprising because prior work shows gated convolutions can perfectly solve synthetic tests for AR capability. To close the gap between synthetics and real language, we develop a new formalization of the task called multi-query associative recall (MQAR) that better reflects actual language. We perform an empirical and theoretical study of MQAR that elucidates differences in the parameter-efficiency of attention and gated-convolution recall. Informed by our analysis, we evaluate simple convolution-attention hybrids and show that hybrids with input-dependent sparse attention patterns can close 97.4% of the gap to attention, while maintaining sub-quadratic scaling. Our code is accessible at: https://github.com/HazyResearch/zoology.

1 Introduction

The paper finds that gated-convolution language models remain behind attention largely because they struggle with associative recall: retrieving previously seen associations in context. It introduces multi-query associative recall to better reflect language and uses theory and experiments to explain and reduce this gap.

  • Up to 2.1 perplexity points separate state-of-the-art convolution-based architectures from strong Transformer baselines on Pile language modeling.
  • Associative recall errors account for 82% of the perplexity gap to attention on average.
  • A 70M-parameter Transformer predicts associative-recall hits better than a 1.4B-parameter Hyena gated-convolution model.The Transformer is 20× smaller than the Hyena model.
  • Multi-query associative recall: Multi-query associative recall models multiple recalls at varying positions with large-vocabulary tokens, unlike prior synthetic formulations with one fixed-position query and small vocabularies.
  • Architectural analysis: Fixed convolution filters make variable-distance token interactions parameter- and FLOPs-inefficient compared with attention's input-dependent mixing.
  • Architectural analysis: Theoretical and empirical results show that gated-convolution model dimension must grow with sequence length for MQAR, whereas attention can solve it independently of sequence length.

2 Background and Preliminaries

The paper frames language models as autoregressive next-token predictors whose sequence mixers aggregate token representations. It contrasts attention's input-dependent all-to-all mixing with convolutional alternatives that offer more efficient sequence-length scaling.

  • Language modeling: Autoregressive language models predict each token from the preceding tokens and output a probability distribution over the vocabulary.
  • Sequence mixers: Sequence mixers aggregate token embeddings through weighted sums, with weights determined by the mixer.
  • Attention: Attention computes input-dependent token interactions using query, key, and value projections, with O(N^2d) time complexity.
  • Gated convolutions: Convolutions combine token representations with a learnable kernel and can be computed in O(Nd log N) time using the FFT.
  • Gated convolutions: Gating combines elementwise multiplication with long convolutional filters to enable interactions between distant tokens.

3 Identifying the associative recall problem

The paper finds that associative recall explains much of the language-modeling gap between gated convolutions and attention, despite gated convolutions succeeding on prior synthetic tests. It therefore introduces multi-query associative recall (MQAR), which better reflects repeated recalls at varying sequence positions in real language.

  • Fine-grained analysis: Attention outperforms gated convolutions by at least a third of a perplexity point across the 70M, 160M, and 360M scales.The minimum gaps are +2.14, +0.59, and +0.35 PPL at those scales, respectively.
  • Fine-grained analysis: The largest gap occurs on rare associative-recall hits, while the other-token slice shows no gap between attention and gated convolutions.AR Hits comprise 6.4% of tokens and are identified using repeated n-grams adjusted for training frequency.
  • Fine-grained analysis: 82% of the average perplexity gap between gated convolutions and attention is attributed to the associative-recall slice.The slice measures the fraction of the overall gap that would close if a model matched attention on these tokens.
  • From AR to MQAR: Prior synthetic associative-recall tasks use one fixed-position query, whereas real language requires multiple recalls at varying positions in one forward pass.The paper names this broader setting multi-query associative recall (MQAR) and formally defines it over sequences containing multiple queries.

4 Explaining the associative recall problem

The paper explains the associative-recall gap through theoretical and empirical analysis of MQAR using BaseConv, a minimal gated-convolution model. Attention solves MQAR with sequence-length-independent model dimension, whereas gated-convolution solutions require scaling resources; input-dependent filters improve this scaling but introduce practical costs.

  • Analytical framework: BaseConv provides a minimal gated-convolution framework for analyzing the complexity of MQAR solutions and simulates a broad class of gating-and-convolution architectures.The analysis uses the correspondence between gating and convolution primitives and arithmetic circuits.
  • Theoretical capacity: Attention solves MQAR with O(c^2) parameters, O(Nc^2 + N^2c) time complexity, and O(1) layers.This construction uses model parameters independent of sequence length.
  • Theoretical capacity: Data-independent convolution filters solve MQAR with approximately O(N log c) parameters, time complexity, and layers up to poly-logarithmic factors.The theorem assumes one-hot or distinct embeddings for the relevant token vocabulary.
  • Input-dependent filters: Input-dependent filters reduce the layer-depth issue for a subclass of inputs by using filters that respond to matching-key interaction distances.The construction defines a filter that spikes at the distance separating matching keys.
  • Empirical scaling: Empirically, attention solves MQAR with constant dimension 64, while gated convolutions require d ≥ N to exceed 0.9 accuracy.The comparison covers BaseConv and popular gated-convolution architectures including RWKV, H3, and Hyena.
  • Empirical scaling: Programmatic input-dependent filters achieve near-constant dimension scaling, while autocorrelation filters improve scaling over input-independent filters.The two approaches differ in whether matching is hard-coded from token ids or learned through autocorrelation.
  • Practical limitations: Input-dependent filters have practical limitations because causality is difficult to enforce and O(N log N) filtering can be expensive for sparse bigram interactions.The paper motivates adding a small amount of attention to BaseConv as a potentially simpler alternative.

5 Closing the Associative Recall Gap

The paper evaluates hybrid BaseConv-attention models that add full or selectively sparse attention to gated convolutions. Input-dependent sparse attention closes much of the associative-recall and overall quality gap while retaining efficient scaling.

  • Input-dependent sparsity can close most of the gap between gated convolutions and attention.The hybrids selectively apply attention according to token-dependent selection patterns.
  • The downstream evaluation tests Hyena and BaseConv hybrids with hybridization and selective look-up at 160M and 355M parameter scales.
  • At 360M parameters, three attention layers let BaseConv outperform the Transformer with fewer FLOPs.
  • Full-attention BaseConv hybrids outperform attention-only models by 0.85 perplexity points while reducing total FLOPs by 18%.These hybrids still use quadratic attention.
  • Programmatic selection closes 85% of the AR gap, while learned selection closes 72% using k = 256 sub-quadratic attention positions per example.

6 Discussion and Conclusion

The paper identifies associative recall as a central weakness of efficient gated-convolution models and introduces MQAR to analyze why attention is more parameter-efficient. It argues that insufficient data-dependent mixing explains the gap and motivates minimal hybrid fixes.

  • The study identifies a persistent quality gap between efficient convolution models and attention, largely due to associative recall.
  • MQAR is introduced as an analysis tool that correlates with downstream associative-recall quality.
  • The authors position MQAR and real-language-linked tasks as useful proxies for studying language-model architectures.
  • The paper includes theoretical proofs, synthetic experiments, scaling analyses, and implementation and experiment details in appendices.

A.3 Input-Dependence in Sequence Modeling Architectures

Input-dependent sequence models control token aggregation using the data, unlike fixed-filter convolutions. The reviewed architectures vary in how they introduce input dependence, but recall and scaling remain important trade-offs.

  • Input-dependent sequence models control how tokens are aggregated across a sequence using the data rather than only model parameters.
  • Attention obtains input-dependent mixing because its query, key, and value projections are linear combinations of the input.
  • Attention has O(N^2) sequence-length compute, whereas convolutions have O(N log N) scaling.
  • Gated convolutions introduce input dependence through element-wise multiplication, such as y = σ(Wx) ⊙x.
  • The required hidden-state dimensionality for gated-convolution architectures to solve MQAR grows with sequence length, while attention does not require this dependence.
  • The learned-selection module uses auxiliary sparsity loss and top-k selection to remain sub-quadratic, unlike methods that may use full attention in the worst case.

C.1 Measuring the Mqar Gap on Real Language Data

The protocol measures associative recall on rare repeated bigrams, then estimates how much of the quality gap between a convolution model and attention is attributable to recalling those prior occurrences.

  • Recall-hit construction: The analysis identifies recurring bigrams whose prior training frequency is below 1,250 occurrences, reducing the influence of memorization.Only these bigrams are used to measure associative-recall log-probabilities.
  • Recall-hit construction: An AR Hit is the true completion of a repeated bigram, evaluated by the log-probability assigned to that completion.The protocol assumes the model can produce the completion by recalling the prior bigram occurrence.
  • Gap attribution: The estimated AR contribution compares average AR-Hit and overall log-probabilities for the evaluated model and an attention reference.The calculation also uses the proportion of AR-Hit tokens in validation data.
  • Gap attribution: If the evaluated model is better overall but attention is better on AR, the protocol attributes 100% of the gap to associative recall.This is the protocol’s stated boundary case for gap attribution.
  • Scope and measurement: The evaluation measures explicit token-identified bigrams, not more abstract recall between related concepts or synonyms.Approximately 6.4% of validation tokens are AR Hits under the 1,250-occurrence threshold.

C.2 Gated Convolution Downstream Architectures

The study evaluates gated-convolution and attention variants under a common training setup, then tests whether input-dependent sparse attention can address the recall gap without full quadratic computation.

  • Architecture suite: The experiments evaluate previously proposed architectures alongside BaseConv, a theoretically canonical convolution baseline.The supplied passage introduces the architecture comparison but does not enumerate all proposed models.
  • Architecture suite: The study trains attention, Hyena, H3, RWKV, pure long convolution, and BaseConv using architecture-specific specifications and reference implementations.These implementations provide the configured comparison models and a non-gated convolution reference.
  • Hybrid architectures: Hybrid models retain the pure gated-convolution layer count and evenly intersperse replacement non-gated convolution layers.The hybrid construction changes selected layer types while preserving the overall layer count.
  • Input-dependent attention: The sparse-attention variants include full attention, random selection, programmatic selection on AR hits, and learned top-k selection.Programmatic selection detects repeated raw-token bigrams causally, while learned selection predicts positions for attention.
  • Input-dependent attention: The protocols test whether input-dependence suffices to address the associative-recall gap and provide multiple ways to incorporate it.The stated objective is to validate input-dependent architectural choices rather than one specific selection mechanism.
  • Data analysis: The analysis extends associative-recall measurements from the Pile to RedPajama sources including ArXiv and StackOverflow.The broader source analysis examines how MQAR appears across language distributions.

D.1.1 Pile Examples

Real-language examples show attention retaining associative recall across long gaps, while convolutional models can fail on names, phrases, and other repeated associations; hit distances are mostly short but heavy-tailed.

  • Long-range recall examples: A 350M attention model and a 350M RWKV model are compared on recalling the 4-gram “Ma-ison Ber-gey” after a gap exceeding 700 tokens.The example focuses on predicting “Bergey” from the earlier occurrence.
  • Long-range recall examples: A 350M attention model and a 350M RWKV model are compared on recalling “Fest” in the repeated bigram “Pixar Fest”.The cited example identifies the repeated phrase as the associative-recall target.
  • Multi-token recall: Additional examples require recalling multi-token names or phrases including a middle and last name, “sequential centrifugation,” and “Missouri Southern.”These examples illustrate that associative-recall targets can span several tokens.
  • Long-range recall examples: A 350M attention model and a 355M Hyena model both recall “SubSpace” across short gaps, but only attention succeeds when the gap exceeds 1,000 tokens.The comparison directly links the differing outcome to the distance between repeated tokens.
  • Hit-distance distribution: Most AR hits occur within 100 token positions of the prior bigram, while a long tail requires long-range interactions.The hit-distance distribution follows a power law across the Pile training data.
  • Hit-distance distribution: The distance pattern motivates windowed or input-dependent interaction schemes rather than computing every N^2 token interaction for every sequence.The passage specifically identifies sliding-window attention as one possible efficiency strategy.
  • Cross-source variation: MQAR-hit prevalence is relatively high in ArXiv, GitHub, and Stack Exchange, where repeated terminology, citations, commands, equations, function names, and variables are common.C4, Common Crawl, Books, and Wikipedia contain more variable bigrams in the cited analysis.

E Synthetic Mqar Experimental Details

The synthetic MQAR procedure is designed to expose differences between sequence-modeling architectures by varying key-value-map size and the number of token-interaction gaps required for recall.

  • Purpose and rationale: MQAR is proposed as a tool for explaining language-modeling gaps among Transformers or Sparse Transformers, convolutions, and recurrences.The paper also provides a procedure for generating synthetic MQAR data for architecture development.
  • Purpose and rationale: The synthetic analysis compares attention, gated convolutions, and other sequence-modeling families to explain behavior observed on real language data.Synthetic recall tasks are connected to prior studies of Hyena, H3, and Transformer in-context learning.
  • Architectural analysis: Attention can perform recall with a bound independent of sequence length, whereas gated-convolution recall requires dimensionality that grows with sequence length.The analysis attributes the convolutional challenge to input-independent filters and many token-interaction distances.
  • Synthetic controls: The MQAR generator varies key-value-map size and the number of unique token-interaction gaps required by each example.Sequence length N, map size D, and power-law parameter α control these properties.
  • Synthetic controls: Prior synthetic recall work uses vocabulary sizes of at most 40 tokens, whereas language modeling uses vocabularies of roughly 30,000–50,000 tokens.This contrast motivates testing larger and more language-like synthetic settings.
  • Synthetic construction: The synthetic procedure splits a vocabulary into keys and values, pairs keys with random values, selects D pairs, places them initially, and repeats each pair at power-law-sampled distances.The second occurrences are placed between positions 2D and N.

E.2 Training Details

The experiments compare attention, gated convolutions, gated recurrences, and sparse attention on MQAR, using standardized synthetic-training protocols.

  • Architectures: The study evaluates four architecture categories: attention, gated convolutions, gated recurrences, and sparse attention.RetNet represents gated recurrences, while sliding-window and blocked-window attention represent sparse attention.
  • Architectures: RetNet combines chunked attention with recurrence and is evaluated with chunk sizes of 32 and 8.
  • Training protocol: Synthetic runs use weight decay 0.1, AdamW, 10% linear warmup, four learning rates, and 64 training epochs.
  • Training protocol: All synthetic experiments use exactly two layers, each containing one sequence mixer and one MLP.

F Extended Results on Mqar Across Architectures

Extended MQAR experiments test how model dimension affects architectures as the number of key-value pairs per example changes, including RetNet and efficient attention variants.

  • Experimental scope: The experiments vary model dimension for synthetic datasets containing different numbers of key-value pairs per example.
  • Gated recurrence: RetNet is evaluated with input-dependent gated recurrence and chunk sizes of 8 and 32, alongside Hyena and attention references.
  • Sparse attention: Sliding-window and blocked-window attention are evaluated with window sizes of 8 and 32.

F.1 Synthetic Experiments

Synthetic MQAR experiments show that recall demands expose parameter-efficiency differences: gated recurrences require more dimensionality as stored key-value pairs increase, while sliding attention is effective within its window.

  • Key-value capacity: As key-value pairs exceed the window size, gated recurrence requires increased dimensionality relative to full O(N^2) attention.The experiments vary the number of key-value pairs per example while changing model dimension.
  • Key-value capacity: Gated recurrences use Ω(N) bits to solve MQAR when d ≤ N.The analysis attributes this requirement to storing many key-value pairs in low-dimensional hidden states.
  • Windowed attention: Sliding-window attention closes the MQAR gap when the key-value pairs fit within the window, but degrades beyond that range.With window size 32, sliding attention performs well up to 16 key-value pairs and worsens beyond this point relative to attention.
  • Downstream validation: The observed MQAR trends for RetNet and sparse attention also carry over to downstream language-modeling results.

G Mqar Perplexity Gap and Model Size

Scaling experiments examine MQAR and associative-recall gaps at billion-parameter scales, while theoretical results characterize gated-convolution constructions and their circuit expressivity.

  • Scaling results: A 70M attention model is one full perplexity point better on associative recall than a 1.4B-parameter Hyena model.
  • Scaling results: At 7B parameters, RWKV remains behind attention, with the gap increasing as each sequence requires more recalls.
  • Rare-bigram evaluation: For bigrams seen less than 1,250 times during pretraining, the 350M-scale Hyena and RWKV gaps to attention are 1.85 and 1.84 perplexity points.
  • Rare-bigram evaluation: For bigrams seen once during pretraining, the corresponding gaps are 12.0 and 13.2 perplexity points.
  • Theory: The theoretical appendix defines gated-convolution models and shows BaseConv can simulate arithmetic circuits with polylogarithmic depth overhead.It also establishes primitives such as sequence shifting and equivalence results for linear arithmetic circuits.

H.5 General Arithmetic Circuits

The appendix establishes that arithmetic circuits can be simulated by BaseConv models, then applies this framework to recall-related constructions and MQAR algorithms. These reductions quantify the resulting layer, parameter, and complexity costs.

  • Circuit-to-BaseConv equivalence: Theorem H.21 gives an equivalent BaseConv for any arithmetic circuit, with N = n, depth O(∆log w), and internal dimension d.The construction first layers the circuit into addition and multiplication gates, then simulates each layer.
  • Circuit-to-BaseConv equivalence: Linear circuit layers are represented by matrix multiplication, while multiplication layers use sparse rearrangement, remembered values, and Hadamard products.The proof preserves entries needed by later layers while computing the current multiplication outputs.
  • Circuit-to-BaseConv equivalence: The resulting BaseConv simulation uses O(∆log w) layers while keeping internal dimension w fixed.This bound is stated for the composed simulation of the layered arithmetic circuit.
  • Applications to recall: A RetNet model with O(d^2) parameters and N layers has an equivalent BaseConv using O(Nd) parameters and O(N log d) layers.The reduction follows because each output can be expressed as a polynomial in the input coordinates and then simulated by an arithmetic circuit.
  • Applications to recall: Attention solves MQAR with O(c^2) parameters, O(Nc^2 + N^2c) time, and O(1) layers, while a BaseConv route solves it with ˜O(1) layers and ˜O(Nd) parameters.A parallel arithmetic-circuit construction also achieves O(Nd · log^2 N) work and O(d · log^2 N) time for MQAR.

H.7.7 The Resulting BaseConv Model

The resulting BaseConv model solves multi-query associative recall using input-dependent convolution kernels derived from autocorrelation. Its construction identifies relevant token-interaction distances, shifts keys, gates by queries, and retrieves matching values.

  • Problem setup: The formal MQAR setup represents each input token as a key, value, or query and returns the value whose key matches a preceding query.The interaction distance is the separation between the query and its matching key, with at most t distinct distances assumed.
  • Input-dependent kernels: Input-dependent kernels identify the top t interaction distances through autocorrelation and construct convolution filters for those shifts.The selected distances correspond to highly overlapping query and key values.
  • The resulting BaseConv model: BaseConv solves multi-query associative recall with O(1) layers and O(t · Nc) parameters when inputs have at most t distinct interaction distances.The theorem applies to inputs in {0, 1}^3N×c.
  • Recall computation: For a single interaction distance, convolution and gating produce a match exactly when a query equals the key at the shifted position.The retrieved output is the value following the matched key when that position exists.
  • Recall computation: The model shifts keys by each selected distance, gates them with queries, and uses the resulting match mask to retrieve the corresponding value.This procedure solves recall across the selected interaction distances and combines the resulting outputs.
Loading 2312.04927v1…