[Paper Review] Repeat After Me: Transformers are Better than State Space Models at Copying

Samy Jelassi, David Brandfonbrener+2 more|arXiv (Cornell University)|Feb 1, 2024

Simulation Techniques and ApplicationsDecision Sciences3 citations

TL;DR

This paper demonstrates that Transformers outperform Generalized State Space Models (GSSMs) in copying and retrieving information from long input sequences due to their ability to store and retrieve n-grams via attention mechanisms. Theoretically, a two-layer Transformer can copy exponentially long sequences, while GSSMs are fundamentally limited by their fixed-size latent state, which cannot encode more bits than its dimensionality.

ABSTRACT

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.

Motivation & Objective

• To investigate whether GSSMs, despite their efficiency, can match Transformers in copying and retrieving information from input context.
• To analyze the theoretical limitations of GSSMs in representing long sequences due to fixed-size latent states.
• To evaluate the generalization and efficiency of Transformers versus GSSMs on synthetic copying tasks.
• To test the performance of pretrained models (e.g., Pythia vs. Mamba) on context retrieval and copying in zero-shot settings.
• To understand the practical implications of architectural differences on downstream language model capabilities beyond perplexity.

Proposed method

• Theoretical analysis of the string copying task using formal definitions of GSSMs and Transformers, proving that GSSMs cannot copy sequences longer than their latent state size allows.
• Construction of a two-layer Transformer that copies sequences of length exponential in the number of attention heads, using n-gram storage and retrieval via attention.
• Empirical evaluation on synthetic datasets where models are trained to copy sequences of increasing length, measuring both training efficiency and length generalization.
• Use of a Hard-Alibi positional encoding to improve generalization beyond the training context window in Transformers.
• Evaluation of pretrained models (Pythia and Mamba) on a context-based lookup task, measuring F1 score across varying paragraph lengths.
• Analysis of learned attention patterns in trained Transformers to confirm reliance on n-gram lookup mechanisms for copying.

(a) Copying: training efficiency.

Experimental results

Research questions

• RQ1 Can a Transformer model copy sequences of length exponential in its size, and if so, how does it achieve this?
• RQ2 What are the fundamental representational limitations of GSSMs in copying tasks due to their fixed-size latent state?
• RQ3 How do Transformers and GSSMs compare in terms of training efficiency and generalization to longer sequences on synthetic copying tasks?
• RQ4 Do pretrained Transformers outperform pretrained GSSMs in zero-shot context retrieval and copying tasks?
• RQ5 To what extent does the fixed memory capacity of GSSMs degrade performance on natural language tasks with long input contexts?

Key findings

• A two-layer Transformer can copy sequences of length exponential in the number of attention heads, demonstrating a powerful mechanism for context retrieval via n-gram storage.
• GSSMs are fundamentally limited in copying tasks: they cannot accurately copy sequences with more bits than the size of their latent state.
• On synthetic copying tasks, Transformers train significantly faster and generalize much better to longer sequences than GSSMs, with performance degrading rapidly for Mamba beyond the training length.
• Pretrained Pythia (Transformer-based) models outperform Mamba (GSSM-based) models in zero-shot context lookup tasks, with F1 scores declining more sharply for Mamba as paragraph length increases.
• The learned copying mechanism in Transformers relies on attention-based n-gram lookup, confirming the theoretical construction.
• Despite Mamba achieving lower perplexity on the Pile than Pythia, it underperforms on context-intensive tasks, indicating that training loss alone does not predict performance on retrieval and copying.

(b) Copying: length generalization