Skip to content
The methodOpen data
Back to the collection

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

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

2024 · arXiv

Language modelsMachine learning
Read the paper
HN score
0
Mentions
6
Readers
4
Months
4
What’s in a score?

From the 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.

Metadata: arXiv. Abstracts retain their original authors’ rights.

In the conversation

Original HN comments, with labels from simple text rules. A label describes the whole comment and may miss context or irony. Read the discussion to make up your own mind. About the labels

6 selected excerpts · 6 total mentions

trott
Neutral mention · rule
My feeling is that the answer is "no", in the sense that these RNNs wouldn't be able to universally replace Transformers in LLMs, even though they might be good enough in some cases and beat them in others. Here's why. A user of an LLM might give the model some long text and then say "Translate this into German please". A Transformer can look back at its whole history. But what is an RNN to do? While the length of its context is unlimited, the amount of information the model retains about it is bounded by whatever is in its hidden state at any given time. Relevant: https://arxiv.org/abs/2402.01032
Original comment Were RNNs all we needed?
cs702
Neutral mention · rule
Dao and Gu show that if you simplify Mamba so its state-space layer uses a diagonal matrix A_t that is a scalar times the identity matrix, i.e., A_t = a_t I, the state-space transformation can be expressed as a form of causal linear attention[a] by compounding coefficients a_1 ... a_t at each time step t. The equivalence of the simplified state-space layer and causal linear attention constitute the duality the authors refer to in the title. By taking advantage of this duality, Mamba-2 can be trained more efficiently, i.e., faster than original Mamba on GPUs. Theoretical stuff aside, Mamba-2's performance seems to scale slightly better than original Mamba:…
Read the full comment
https://tridao.me/assets/img/2024-05-31-mamba-2/pile_8k_mamb... Here's the code implementing Mamba-2: https://github.com/state-spaces/mamba/blob/main/mamba_ssm/mo... Great work by Tri Dao (of FlashAttention fame) and Albert Gu, as usual. The key question, for me and many others, is whether Mamba, Mamba-2, RWKV, and other linear RNN / linear attention models will ever match the performance of standard Softmax attention. My understanding and experience is that all the linear attention models out there [b] still underperform Softmax attention on things like recall tasks.[c] --- [a] https://arxiv.org/abs/2006.16236 [b] https://github.com/topics/linear-attention / https://github.com/topics/linear-attention-model -- this list is by no means complete! [c] https://arxiv.org/abs/2402.01032
Original comment Mamba-2 – State Space Duality
hansonw
Neutral mention · rule
Indeed: https://arxiv.org/pdf/2402.01032.pdf Perhaps future iterations of SSMs will accommodate dynamically sized (but still non-linearly-growing) hidden states / memories!
Original comment Mamba: The Easy Way
thatguysaguy
Neutral mention · rule
Too new is definitely one thing. Someone is going to have to make a gamble to actually paying for a serious pretraining run with this architecture before we know how it really stacks up against transformers. There are some papers suggesting that transformers are better than SSMs in fundamental ways (e.g. They cannot do arbitrary key-based recall from their context: https://arxiv.org/abs/2402.01032). This means it's not just a no-brainer to switch over.
Original comment Mamba: The Easy Way
cs702
Neutral mention · rule
The two issues I see with this work: * There's no mention of model performance on recall tasks. Models with attention do well on recall tasks, but models without it, like this one, tend to do poorly.[a] What is the performance of this model on recall tasks? * As others here have pointed out, the github link is dead. In addition, the pretrained weights don't seem to be available on HF or anywhere else. It's hard to check any claims if we have neither code nor weights! --- [a] https://arxiv.org/abs/2402.01032
Original comment Megalodon: Efficient LLM Pretraining and Inference with Unlimited Context Length

Keep following the thread.

More in Language models

More from the same topics.