PaperlayerView exact revision on arXiv

Source-linked AI summary

Repeat After Me: Transformers are Better than State Space Models at Copying

Samy Jelassi, David Brandfonbrener, Sham M. Kakade, Eran Malach

arXiv:2402.01032v2cs.LGcs.AIcs.CL

TL;DR

The paper asks what GSSMs sacrifice for fixed-size-memory efficiency, focusing on their ability to retrieve and repeat input context. Through theoretical analysis and experiments, it compares GSSMs with transformers on copying and retrieval tasks. Transformers copy exponentially long strings in the construction and outperform GSSMs in synthetic efficiency, length generalization, and pretrained context-retrieval evaluations.

  • Problem

    The paper investigates what GSSMs sacrifice for improved fixed-size-memory efficiency, specifically whether they lose the ability to retrieve and repeat arbitrary parts of the input context.

  • Method

    The paper combines a theoretical copying analysis with synthetic experiments and pretrained-model evaluations comparing transformers and GSSMs.

  • Results

    Transformers copy strings of exponential length in the theoretical construction and outperform GSSMs on copying efficiency, length generalization, and pretrained context retrieval.

  • Takeaways & Limitations

    The results indicate a fundamental capability gap between transformers and GSSMs on tasks requiring access to arbitrary parts of the input context.

  • Takeaways & Limitations

    GSSMs retain advantages for long inputs because their memory and computational complexity do not increase with input length, and they can be better at tracking state variables across long sequences.

Abstract

from arXiv · show

Transformers are the dominant architecture for sequence modeling, but there is growing interest in models that use a fixed-size latent state that does not depend on the sequence length, which we refer to as "generalized state space models" (GSSMs). In this paper we show that while GSSMs are promising in terms of inference-time efficiency, they are limited compared to transformer models on tasks that require copying from the input context. We start with a theoretical analysis of the simple task of string copying and prove that a two layer transformer can copy strings of exponential length while GSSMs are fundamentally limited by their fixed-size latent state. Empirically, we find that transformers outperform GSSMs in terms of efficiency and generalization on synthetic tasks that require copying the context. Finally, we evaluate pretrained large language models and find that transformer models dramatically outperform state space models at copying and retrieving information from context. Taken together, these results suggest a fundamental gap between transformers and GSSMs on tasks of practical interest.

1. Introduction

The paper argues that fixed-size-memory GSSMs trade away input-context retrieval capabilities that transformers retain, establishing this gap through theory and experiments on copying and lookup.

  • Motivation: GSSMs use fixed-size memory to reduce per-token inference memory, while transformers require memory that grows with sequence length in naive implementations.The paper groups state space models, RNNs, linear attention, and parallel RNNs under GSSMs.
  • Research gap: The paper identifies retrieving and repeating arbitrary parts of the input context as a capability sacrificed by GSSMs for improved efficiency.It connects this limitation to transformers’ stronger performance on tasks requiring context access.
  • Theory: A small transformer can copy strings of exponential length in its number of heads by storing and retrieving n-grams that track where to copy.The theoretical construction contrasts with GSSMs, whose fixed latent state limits accurate copying when inputs contain more bits than that state.
  • Theory: Any state space model fails to solve the copy task unless its latent state grows linearly with sequence length.This formalizes the fixed-state bottleneck for copying.
  • Experiments: Synthetic experiments find that transformers learn copying much more efficiently and generalize better to longer inputs than GSSMs.The experiments use models of approximately 160M parameters, addressing whether theoretically representable solutions are actually learned.
  • Pretrained models: Pretrained Pythia transformers substantially outperform similarly sized Mamba GSSMs on copying and context retrieval, despite Mamba achieving lower Pile perplexity.The result indicates that language-modeling perplexity does not capture all downstream memory capabilities tested here.

2. Theory: Representational Capacity

The theory compares transformers and GSSMs on copying: a depth-2 transformer can copy very long sequences, while GSSMs require latent-state capacity that grows with sequence length. The transformer construction uses n-gram hashing and retrieval, with errors controlled by repeated n-grams.

  • Main theoretical comparison: A depth-2 transformer can copy sequences of exponential length in its size, whereas GSSMs need latent-state capacity that grows linearly with sequence length.The GSSM lower bound applies even when parameter memory and state-update runtime are left unconstrained.
  • Copy task: The copy task asks a model to reproduce the input sequence after beginning-of-sequence and copy tokens.The formal setup samples strings over non-special dictionary tokens and expects the generated prefix to equal the input.
  • Transformer construction: The transformer construction hashes n-grams, attends to the previous occurrence of the latest n-gram, and outputs its succeeding token.Hard-ALiBi supplies local positional masking and supports applying the hash mechanism across the input.
  • Transformer guarantee: For any n, a depth-2 transformer of dimension O(n log(D)) has copy error below the probability of repeated n-grams when 2n ≤ L ≤ D^n.The construction copies perfectly whenever the input contains no repeated n-gram pattern.
  • Transformer guarantee: For uniformly sampled inputs, the probability of repeated n-grams is less than L^2D^-n, yielding transformer parameters that depend only logarithmically on sequence length.The paper states that finite-precision versions require O(log(log(L))) bits for the construction.
  • GSSM lower bound: Every GSSM with mem(S) < L log(D) − 1 has error greater than 1/2 on the uniform copy distribution.Thus, for copying, transformer input-dependent memory is optimal up to logarithmic factors despite being less memory-efficient than GSSMs in general.

3. Learning to Copy

Synthetic experiments test whether transformers and GSSMs can learn copying and generalize it beyond training lengths. Transformers learn the task more efficiently, generalize better to longer inputs, and support retrieval patterns that depend on storing the full context, while GSSMs perform better when only a small summary is needed.

  • Data efficiency on the copy task: Transformers need 100x less samples than the best GSSMs to learn in-distribution copying on strings of length up to 300.Models were trained and evaluated on randomly sampled strings using string-level accuracy.
  • Length generalization on the copy task: When trained on sequences of at most 50 tokens, GSSM performance drops to zero almost immediately on longer inputs, whereas transformer accuracy decays more gradually up to 1000 tokens.This evaluates out-of-distribution length generalization using string-level accuracy.
  • Length generalization on the copy task: ALiBi and NoPE transformers outperform RoPE on longer inputs, indicating that positional encoding affects length generalization.The paper attributes this likely difference to RoPE’s sinusoidal embeddings changing more dramatically at longer lengths.
  • Length generalization on the copy task: Hard-ALiBi transformers trained on sequences of at most 50 tokens achieve almost perfect length generalization on sequences up to 1000 tokens.The tested lengths exceed the context length encountered during training.
  • Transformers learn to use n-gram hashing: On duplicated n-gram inputs, transformer accuracy remains roughly unchanged for n ≤4 but drops when duplicate sequences contain 5 or more tokens.The result suggests reliance on approximately 5-gram retrieval, matching the performance of a perfect 5-gram model.
  • GSSMs cannot arbitrarily retrieve from context: GSSMs perform poorly on suffix-key lookup but achieve perfect length generalization on prefix-key lookup, where the key can be stored and matched online.The prefix-key variant requires less memory because the model need not retain the entire input context.

4. Pre-trained Models

The paper evaluates similarly sized pretrained Pythia transformers and Mamba GSSMs on copying and retrieval tasks across increasing context lengths. Transformers outperform GSSMs on memory-intensive tasks, including natural-language copying and long-context question answering, even though Mamba often has lower language-model perplexity.

  • Pre-trained Models: Pretrained transformers outperform pretrained GSSMs on memory-intensive tasks including long-string copying, retrieval, and few-shot question answering.The comparison uses models of similar scale and tests whether the limitation observed earlier transfers to pretrained models.
  • Pre-trained Models: Pythia transformers outperform Mamba GSSMs on copying and context retrieval despite Mamba models generally having slightly lower perplexity at a given size.The models were pretrained on the Pile with the same tokenizer and compared across similar parameter counts.
  • Copying the input context: Even the smallest transformer dramatically outperforms the largest GSSM when copying natural-language strings sampled from C4.The prompt contains two copies of a sampled string and asks the model to complete a third copy.
  • Copying the input context: Shuffling word order worsens copying for both model classes, with a more pronounced degradation for GSSMs; the largest GSSM reaches zero accuracy at length 300.The shuffle test reduces the compressibility and natural-language structure of the copied strings.
  • Retrieval from the input context: On SQuAD, Pythia and Mamba perform comparably on short paragraphs, but Mamba degrades faster as paragraph length increases while Pythia maintains similar accuracy.Performance is reported using F1 score binned by context-paragraph length.

5. Related Work

Prior work studies transformer and GSSM representation capacity, retrieval mechanisms, and length generalization, but comparable evaluations of transformers and GSSMs on context access were missing. This paper relates its copying results to n-gram retrieval and positional encoding while acknowledging tasks favoring GSSMs.

  • 5. Related Work: Prior studies analyze the representational capacity of GSSMs and transformers through several formal complexity frameworks.These include threshold circuits, RASP language, and first-order logic.
  • 5. Related Work: Transformers have been shown to use the full input context for retrieval, question answering, and in-context learning.
  • 5. Related Work: Related work identifies induction heads as transformer mechanisms performing retrieval operations similar to those used for copying.
  • 5. Related Work: This paper reports that comparable-quality comparisons between transformers and GSSMs on these context-access tasks were previously absent.
  • 5. Related Work: GSSMs can outperform transformers on flip-flop language modeling and can handle long DNA contexts that may be impractical for transformers.
  • 5. Related Work: The paper connects its findings to n-gram heads, limited GSSM state size, and Hard-ALiBi positional encoding for n-gram matching.

6. Discussion

The discussion concludes that transformers outperform GSSMs at copying from context, while GSSMs retain advantages in efficiency and state tracking. It therefore points toward hybrid architectures that add attention-like retrieval to state space models.

  • 6. Discussion: Transformers outperform GSSMs at copying from the input context, whereas GSSMs offer memory and computational complexity independent of input length.
  • 6. Discussion: GSSMs such as RNNs are better at tracking state variables across long sequences, which may help generate long consistent text.
  • 6. Discussion: The paper proposes hybrid architectures that give state space models an attention-like mechanism for retrieving relevant input text.
  • 6. Discussion: The paper states that its societal consequences do not require specific highlighting.

A. Experimental setup

The experimental setup compares transformers using NoPE, ALiBi, and Hard-ALiBi positional encodings. These schemes differ in whether and how attention incorporates positional or local-token information.

  • A. Experimental setup: Figure 8 illustrates NoPE, ALiBi, and Hard-ALiBi using three masked heads whose accessible context expands from the current token to two preceding tokens.
  • A. Experimental setup: The experiments compare NoPE, ALiBi, and Hard-ALiBi positional encoding schemes.
  • A. Experimental setup: NoPE adds no positional information to attention scores and improves length generalization on multiple tasks, including copying.
  • A. Experimental setup: ALiBi biases attention scores with a distance-proportional penalty controlled by a head-specific slope fixed before training.
  • A. Experimental setup: Hard-ALiBi uses M masked heads attending to progressively broader preceding-token windows, while the remaining H −M heads use NoPE attention.

A.2. Pretraining and evaluation details

The appendix describes the software, transformer architecture, training regimen, compute budget, evaluation and decoding procedures, and additional copy-task experiments. The setup uses comparable transformer and Mamba training runs while requiring substantially more epochs for LSTMs.

  • A.2. Pretraining and evaluation details: Training uses PyTorch, HuggingFace, and the Mamba GitHub repository.
  • A.2. Pretraining and evaluation details: The transformer backbone is GPT-NeoX with 12 layers, hidden size 1024, and H = 16 attention heads.
  • A.2. Pretraining and evaluation details: Transformers and Mamba train for 15 epochs, whereas LSTMs train for 300 epochs, using AdamW.
  • A.2. Pretraining and evaluation details: Approximately 600 GPU hours were required for the final training run on RTX8000 GPUs.
  • A.2. Pretraining and evaluation details: Evaluation uses 10 batches of 64 examples for most tasks and 50 questions for question answering, with greedy decoding at temperature 0.
  • A.2. Pretraining and evaluation details: On synthetic copying experiments, GSSMs require many more samples than transformers, and pretrained-model gaps widen for uniformly random strings.

B.1. Data efficiency on the copy task

The copy task becomes harder for recurrent models as string length grows, while transformers learn it more efficiently and generalize better. For short strings, the data-efficiency gap narrows but remains substantial.

  • Copying long strings: Transformers train much faster than GSSMs when copying strings of length 300.At the character level, Mamba has a smoother learning curve, while an LSTM does not learn the task within the available samples.
  • Copying long strings: An LSTM cannot learn the copy task on length-300 strings even at the character level.

B.2. Pre-trained models on the uniform copy task

On the uniform-string copy task, pretrained Pythia transformers clearly outperform pretrained Mamba models, while Hard-ALiBi transformers approach perfect 5-gram performance.

  • Experimental setup: The experiment samples uniform alphabet strings of fixed length and asks pretrained models to reproduce them using the paper's copy-task prompt.
  • Pretrained model comparison: The smallest Pythia model outperforms the largest Mamba model on totally random strings, with Pythia 1.4B achieving almost 100% accuracy.
  • N-gram comparison: Transformers with Hard-ALiBi achieve performance close to a perfect 5-gram model on the copy task.
  • Pretrained model comparison: Uniformly sampled strings degrade Mamba performance, whereas Pythia maintains high performance.

C.2. Proof of Theorem 2.3

The proof constructs a two-block transformer that copies by hashing preceding n-grams, matching each current query to a previous key, and returning the associated value.

  • Transformer construction: The first transformer block builds keys, queries, and values from preceding n-grams and the current n-gram.Keys encode the n-gram before each token, values store the token, and queries encode the current n-gram.
  • Transformer construction: The copy algorithm matches each current query with a previous key and retrieves the matching value.
  • Attention lookup: A second attention layer concentrates on the position whose key matches the current query, producing the token to copy.
  • Correctness: Under the construction's assumptions and temperature choice, the output map returns x_i−L+1 for every generation step after the copied prefix.
  • Assumptions: The construction assumes that the input n-grams are unique and that 2L + 2 ≤ K for K = D^n.

D. Proofs - Lower Bound

The lower-bound proof shows that a fixed-state sequence-to-sequence model's copied output is determined by its state after reading the input, enabling an error bound based on state capacity.

  • State representation: For any fixed-state sequence-to-sequence model, the entire output sequence is a function G of the state after reading the input.
  • Recursive construction: The proof constructs G recursively by applying the model's readout and update functions to generate each subsequent token and state.
  • Inductive step: The induction establishes that each recursively generated token equals the model's corresponding autoregressive output.
  • Error bound: The proof then bounds copying error by comparing the number of possible latent states with the number of possible inputs.
Loading 2402.01032v2…