{"notes":[{"id":"mamba/repeat-after-me","category":"mamba","legacyPath":"","legacySlug":"","legacyPaths":[],"title":"Repeat After Me","subtitle":"The Phone Book Test","author":"Z. Yuan","paperDate":"2024-02-01","publishedOn":"2025-12-23","paperTitle":"Repeat After Me: Transformers are Better than State Space Models at Copying","paperUrl":"https://arxiv.org/abs/2402.01032","bibliography":{"title":"Repeat After Me: Transformers are Better than State Space Models at Copying","authors":["Samy Jelassi","David Brandfonbrener","Sham M. Kakade","Eran Malach"],"url":"https://arxiv.org/abs/2402.01032","pdfUrl":"https://arxiv.org/pdf/2402.01032v2","codeUrl":"https://github.com/sjelassi/transformers_ssm_copy","preprint":{"id":"2402.01032","submittedOn":"2024-02-01","url":"https://arxiv.org/abs/2402.01032"},"publication":{"venue":"ICML","year":2024,"kind":"conference","url":"https://proceedings.mlr.press/v235/jelassi24a.html"},"citation":"@misc{mambarepeatafterme2024,\n  title = {{Repeat After Me: Transformers are Better than State Space Models at Copying}},\n  author = {Samy Jelassi and David Brandfonbrener and Sham M. Kakade and Eran Malach},\n  year = {2024},\n  eprint = {2402.01032},\n  archivePrefix = {arXiv},\n  url = {https://arxiv.org/abs/2402.01032}\n}\n","citationPath":"/content/papers/citations/mamba/repeat-after-me.6f1018ade730.bib"},"thumbnail":{"src":"/content/papers/assets/6d45fd8ab61191e1263c400e01bf76bdd9ef3534ecc4485835332074ea553443.png","width":1378,"height":416},"locale":"en","path":"/content/papers/en/mamba/repeat-after-me.a596955e7451.md","headings":[{"id":"defining-the-problem","text":"Defining the Problem","depth":2},{"id":"solving-the-problem","text":"Solving the Problem","depth":2},{"id":"the-copy-task","text":"The Copy Task","depth":3},{"id":"too-small-a-state-cant-copy-it-all","text":"Too Small a State Can't Copy It All","depth":3},{"id":"transformers-find-their-place-with-n-grams","text":"Transformers Find Their Place with n-grams","depth":3},{"id":"discussion","text":"Discussion","depth":2},{"id":"synthetic-data-learning-fast-generalizing-far","text":"Synthetic Data: Learning Fast, Generalizing Far","depth":3},{"id":"tell-it-what-to-look-for-first-and-its-fine","text":"Tell It What to Look For First, and It's Fine","depth":3},{"id":"pretrained-models-the-phone-book","text":"Pretrained Models: The Phone Book","depth":3},{"id":"what-this-paper-leaves-open","text":"What This Paper Leaves Open","depth":3},{"id":"conclusion","text":"Conclusion","depth":2}],"anchorAliases":{},"guide":null,"relatedReading":[{"id":"mamba/jamba","relation":"referenced","topic":""}]},{"id":"mamba/jamba","category":"mamba","legacyPath":"","legacySlug":"","legacyPaths":[],"title":"Jamba","subtitle":"An Eighth of Attention","author":"Z. Yuan","paperDate":"2024-03-28","publishedOn":"2026-01-10","world":"external/jamba","paperTitle":"Jamba: A Hybrid Transformer-Mamba Language Model","paperUrl":"https://arxiv.org/abs/2403.19887","bibliography":{"title":"Jamba: A Hybrid Transformer-Mamba Language Model","authors":["Opher Lieber","Barak Lenz","Hofit Bata","Gal Cohen","Jhonathan Osin","Itay Dalmedigos","Erez Safahi","Shaked Meirom","Yonatan Belinkov","Shai Shalev-Shwartz","Omri Abend","Raz Alon","Tomer Asida","Amir Bergman","Roman Glozman","Michael Gokhman","Avashalom Manevich","Nir Ratner","Noam Rozen","Erez Shwartz","Mor Zusman","Yoav Shoham"],"url":"https://arxiv.org/abs/2403.19887","pdfUrl":"https://arxiv.org/pdf/2403.19887v2","codeUrl":null,"preprint":{"id":"2403.19887","submittedOn":"2024-03-28","url":"https://arxiv.org/abs/2403.19887"},"publication":null,"citation":"@misc{mambajamba2024,\n  title = {{Jamba: A Hybrid Transformer-Mamba Language Model}},\n  author = {Opher Lieber and Barak Lenz and Hofit Bata and Gal Cohen and Jhonathan Osin and Itay Dalmedigos and Erez Safahi and Shaked Meirom and Yonatan Belinkov and Shai Shalev-Shwartz and Omri Abend and Raz Alon and Tomer Asida and Amir Bergman and Roman Glozman and Michael Gokhman and Avashalom Manevich and Nir Ratner and Noam Rozen and Erez Shwartz and Mor Zusman and Yoav Shoham},\n  year = {2024},\n  eprint = {2403.19887},\n  archivePrefix = {arXiv},\n  url = {https://arxiv.org/abs/2403.19887}\n}\n","citationPath":"/content/papers/citations/mamba/jamba.11fef67c0c54.bib"},"thumbnail":{"src":"/content/papers/assets/fa186580641946194f81adf74136e29431c3f6dd3fa63e2ecf8a1c0c29d3e5a0.png","width":1056,"height":1000},"locale":"en","path":"/content/papers/en/mamba/jamba.129c94adb5f2.md","headings":[{"id":"defining-the-problem","text":"Defining the Problem","depth":2},{"id":"solving-the-problem","text":"Solving the Problem","depth":2},{"id":"the-jamba-block","text":"The Jamba Block","depth":3},{"id":"working-out-the-kv-cache-by-hand","text":"Working Out the KV Cache by Hand","depth":3},{"id":"discussion","text":"Discussion","depth":2},{"id":"how-it-stacks-up-against-mixtral","text":"How It Stacks Up Against Mixtral","depth":3},{"id":"one-to-three-or-one-to-seven","text":"One to Three, or One to Seven?","depth":3},{"id":"where-does-pure-mamba-get-stuck","text":"Where Does Pure Mamba Get Stuck?","depth":3},{"id":"three-more-findings","text":"Three More Findings","depth":3},{"id":"conclusion","text":"Conclusion","depth":2}],"anchorAliases":{},"guide":null,"relatedReading":[{"id":"mamba/repeat-after-me","relation":"referenced","topic":""},{"id":"mamba/mamba-2-hybrid","relation":"referenced","topic":""}]}],"text":"Two months after Mamba came out, Samy Jelassi, David Brandfonbrener, Sham M. Kakade, and Eran Malach at Harvard uploaded this paper, which later appeared at ICML 2024. The title states the conclusion outright: when it comes to copying, Transformers beat state space models.\n\nCopying sounds like a dull skill. But when a reader hands a model a document and asks it to find a date, a name, or a line of code, what the model is doing is lifting a short stretch of the original text out of its context. If it can't do that well, a low perplexity won't help.\n\n## Defining the Problem\n\nThe authors discuss a whole class of models together: S4, Mamba, traditional RNNs, and linear attention and parallel RNNs, both of which can be trained in parallel. What they have in common is that, no matter how many tokens they have processed, they rely only on a state that is **fixed in size and independent of sequence length**. The authors call this class **GSSMs** (generalized state space models).\n\nThe selling point of GSSMs is clear: predicting each token takes only $O(1)$ memory. A Transformer, by contrast, has to keep every token it has seen in the KV cache, which grows with length.\n\nThe authors want to know: does saving that memory come at a price? If so, what is it?\n\nTo find out, they chose the simplest task they could: copying.\n\n## Solving the Problem\n\n### The Copy Task\n\nThe task format is:\n\n```text\n<BOS> x1 x2 ... xL <COPY>\n```\n\nAfter seeing `<COPY>`, the model must output $x_1$ through $x_L$ verbatim. Scoring uses string-level accuracy: a string counts as correct only if all of it is right.\n\n### Too Small a State Can't Copy It All\n\nStart with the GSSM side. The intuition is simple: to reproduce a whole string unchanged, you first have to remember the whole string. Where? The only place available is that fixed-size state.\n\nThe authors write the number of distinct values the state can take as $|\\mathcal{S}|$; its memory capacity in bits is $\\text{mem}(\\mathcal{S})=\\log|\\mathcal{S}|$. For strings whose characters are drawn uniformly from $D$ options and whose length is $L$, there are $D^L$ possible inputs, but the state can distinguish at most $|\\mathcal{S}|$ of them. The error rate here is the probability that the whole string is not copied exactly. The lower bound given by Theorem 2.7 is:\n\n$$\n\\text{err}(H) \\ge 1-\\frac{|\\mathcal{S}|}{D^L}\n$$\n\n(The main text of the paper prints a strict inequality, while the counting proof in Appendix D yields a non-strict one (greater than or equal to); we follow the proof here.)\n\nPut in a form that's easier to remember (Corollary 2.8): as long as the state has fewer than $L\\log_2 D-1$ bits, the error rate exceeds one half.\n\nLet's compute a few numbers to get a feel for it:\n\n```python\nimport math\n\nfor D, L in [(26, 300), (26, 1000), (50_000, 2_000)]:\n    print(f\"Vocab {D:>6}, length {L:>5}: about {L * math.log2(D):,.0f} bits\")\n\nstate_bits = 1000 * 32           # 1,000 fp32 values\nprint(f\"1,000 fp32 = {state_bits:,} bits, enough for at most about {state_bits / math.log2(50_000):,.0f} tokens from a 50K vocabulary\")\n```\n\n```text\nVocab     26, length   300: about 1,410 bits\nVocab     26, length  1000: about 4,700 bits\nVocab  50000, length  2000: about 31,219 bits\n1,000 fp32 = 32,000 bits, enough for at most about 2,050 tokens from a 50K vocabulary\n```\n\nThe example in the last line comes from the authors themselves: 1,000 32-bit floating-point numbers carry enough information to hold at least 2,000 tokens.\n\nSo this theoretical lower bound is actually quite loose. A real Mamba state is far larger than 1,000 numbers, and the 1,410 bits needed to copy 300 English letters are no problem at all. The authors acknowledge this: the theory only shows that “a state that can't hold the string can't copy all of it.” Whether a model with enough room will learn to store the string in its state is a question for experiments.\n\n### Transformers Find Their Place with n-grams\n\nOn the other side, the authors construct a two-block Transformer and prove that the length of the strings it can copy can grow **exponentially** with the model dimension (Theorem 2.3; the introduction phrases it in terms of the number of attention heads). The error is governed by the probability that the string contains a repeated n-gram, which is small for uniformly random strings.\n\nIts method is a lot like how a person copies a long passage:\n\n1. Look at the last few characters you just wrote, say the last 3.\n2. Go back to the original and find where those 3 characters appear.\n3. Write down the next character after that position.\n\nThis is an n-gram lookup. Attention is a natural fit for “comparing the current fragment against every earlier position.” To make an attention head look only at the most recent $m$ tokens and assemble an n-gram hash from them, the authors design a hard-cutoff version of the ALiBi positional bias, called **Hard-ALiBi**: beyond $m$ positions, the bias is set to $-\\infty$, and within them, to 0.\n\nThe failure case is just as clear: if the original contains two identical n-grams followed by different characters, the model can't tell which one to continue copying from. In uniformly random strings, the probability of such a repeat falls exponentially with $n$ (Lemma 2.4).\n\n## Discussion\n\n### Synthetic Data: Learning Fast, Generalizing Far\n\n[![Three line charts. (a) The x-axis is the number of training examples (log scale) and the y-axis is copy accuracy: Transformers with four kinds of positional encoding reach nearly 100% within 100,000 examples, Mamba needs about 2 million, and LSTM stays near 0. (b) The x-axis is string length: every model is at 100% up to the training length of 50; beyond it, LSTM and the RoPE Transformer drop straight to 0, Mamba holds out to about 70, NoPE to about 90, ALiBi to about 130, and Hard-ALiBi is still near 100% at 1,000. (c) The x-axis is the number of phone-book entries: Pythia 1.4B and 2.8B are still at about 90% with 200 entries, 410M falls to about 50%, and all three Mamba models are below 20% at 100 entries and near 0 at 200](/content/papers/assets/6d45fd8ab61191e1263c400e01bf76bdd9ef3534ecc4485835332074ea553443.png)](/content/papers/assets/6d45fd8ab61191e1263c400e01bf76bdd9ef3534ecc4485835332074ea553443.png)\n\n*Figure 1 from the paper, excerpted from [**page 2 of the v2 PDF**](https://arxiv.org/pdf/2402.01032v2#page=2), CC BY 4.0. (a) and (b) show models trained from scratch, with about 160 million parameters for the Transformers and Mamba and about 40 million for the LSTM; (c) shows off-the-shelf pretrained models.*\n\nStart with the left panel. All models are trained from scratch: the Transformers and Mamba have about 160 million parameters, while the LSTM, which trained unstably at larger sizes, uses about 40 million. Strings are at most 300 characters long, drawn from the 26 letters of the English alphabet.\n\nAll four Transformers learn the task within 100,000 examples; Mamba needs around two million. In the authors' words, the Transformers need about 100 times fewer examples. The LSTM doesn't learn it at all within this number of examples.\n\nThe middle panel tests length generalization: train only on strings up to length 50, then test on longer ones. As soon as they pass the training length, Mamba and LSTM quickly drop to 0. For the Transformers, it depends on the positional encoding: the RoPE version also falls off fast, NoPE and ALiBi hold out longer, and the Hard-ALiBi version is still nearly perfect at 1,000.\n\nThe authors also confirm that the Transformers really are doing n-gram lookups. When two identical n-grams are deliberately planted in the string, repeats of length 4 or less have almost no effect, while from length 5 on the accuracy starts to fall, matching the behavior of a perfect 5-gram model.\n\n### Tell It What to Look For First, and It's Fine\n\nOne controlled comparison is especially worth a look. The task: given a string, plus one of its n-grams as a key, output the $k$ tokens that follow the key. The authors built two versions:\n\nBoth versions train on strings up to length 30 and are tested up to length 100:\n\n- **Key at the end** (suffix key): the model only learns what to look for after reading the whole string. The Transformers still do reasonably well out to length 100, while GSSMs fail as soon as they leave the training length.\n- **Key at the start** (prefix key): the model knows from the outset what to look for. While reading, it only has to watch for that key and, once it appears, note the next few characters. This time the GSSMs stay perfect all the way, even beating the NoPE and ALiBi Transformers.\n\nThe only difference is when the question shows up.\n\nWhat a fixed-size state does well is “summarizing as it reads.” If it knows in advance what to keep, it can throw away whatever is irrelevant; if the question only appears after reading, it has to hold on to everything that might be needed. Once the information exceeds the state's capacity, there's no way to guarantee it can all be recovered.\n\n### Pretrained Models: The Phone Book\n\nThe experiments so far all used small models trained from scratch. Next, the authors compare off-the-shelf pretrained models: Pythia (Transformer) from 410M to 2.8B, against Mamba from 360M to 2.8B. Both were trained on the Pile with the same tokenizer, so the main difference is the architecture. Also of note: at the same scale, Mamba has slightly lower perplexity on the Pile.\n\nThe best-known result is the phone-book test in the right panel. The prompt contains a randomly generated phone book, with lines like this:\n\n```text\nJohn Powell: 609-323-7777\n```\n\nTwo example questions with answers follow, and then the model is asked for one person's number. The longer the phone book, the faster Mamba's accuracy falls; at around 70 entries, the smallest Pythia, 410M, already beats the largest Mamba, 2.8B.\n\n[![Three line charts. (a) Copying natural-language strings from C4: the three Pythia models are still at about 80% to nearly 100% at 500 tokens, while the three Mamba models decline with length, the 2.8B already below 10% at 500 tokens. (b) Copying after shuffling the word order: Pythia 1.4B and 2.8B drop slightly, 410M falls to about 50%, and all three Mamba models reach 0 by 300 tokens at the latest. (c) F1 on SQuAD question answering: the two are close on short passages, and Mamba falls further behind as passages get longer](/content/papers/assets/8d14006fc09d6d17409ad80cc3ab8e3574c16a5a2b26e22539149dda61ed153d.png)](/content/papers/assets/8d14006fc09d6d17409ad80cc3ab8e3574c16a5a2b26e22539149dda61ed153d.png)\n\n*Figure 7 from the paper, excerpted from [**page 8 of the v2 PDF**](https://arxiv.org/pdf/2402.01032v2#page=8), CC BY 4.0.*\n\nThree other tests point the same way:\n\n1. **Copying natural language**: take a passage from C4, put it in the prompt twice, then add the first word of a third copy and ask the model to continue. Even the smallest Pythia beats the largest Mamba by a wide margin.\n2. **Copying shuffled words**: natural language can be compressed, so shuffling the word order makes copying harder. Both sides get worse, Mamba more so, and the largest Mamba is already at 0 by 300 tokens. This fits the inference that “a fixed state can only cope by compressing.”\n3. **SQuAD question answering**: 2.8B against 2.8B, the two are about even on short passages; the longer the passage, the faster Mamba's F1 falls.\n\n### What This Paper Leaves Open\n\nIn the related-work section, the authors also mention examples pointing the other way: some tasks only require tracking a very small state. In flip-flop language modeling, for example, LSTMs generalize easily while Transformers struggle; for very long sequences like DNA, a fixed state is also more practical.\n\nSo this paper is about **memory-intensive** tasks. It doesn't claim GSSMs are necessarily worse at other tasks, nor does it test what happens when the two kinds of layers are mixed.\n\n## Conclusion\n\nThis paper turns “a fixed-size state has a cost” from an intuition into a measurable result:\n\n1. **In theory**, the length a two-block Transformer can copy can grow exponentially with model dimension; when a GSSM's state has fewer than $L\\log_2 D-1$ bits, the probability of failing to copy a uniformly random string exactly exceeds one half.\n2. **In practice**, the limit often shows up earlier. Even when the state is big enough, Mamba has a hard time learning to store the whole string: it needs 100 times more training examples and fails once the length changes.\n3. **Pretrained models behave the same way**. Despite slightly lower perplexity, Mamba loses on phone-book lookup, copying, and long-passage question answering to Pythia models of the same size, and even to much smaller ones.\n\nThe “know what to look for first, and it's fine” comparison is probably the most important takeaway from this paper. It shows what a fixed state is good at and when it falls short, and it gives a very direct reason for “pairing state models with attention that can look back and find things.”\n\n\n> [!TIP]\n> \n> Jamba, which AI21 released two months later, keeps a small number of attention layers between its Mamba layers; in their ablations, the authors also observed that pure Mamba often strayed from the answer format set by the few-shot examples:\n> \n> - **[[24.03] Jamba: An Eighth of Attention](/blog/papers/mamba/jamba)**\n","imageSizes":{"/content/papers/assets/6d45fd8ab61191e1263c400e01bf76bdd9ef3534ecc4485835332074ea553443.png":{"width":1378,"height":416},"/content/papers/assets/8d14006fc09d6d17409ad80cc3ab8e3574c16a5a2b26e22539149dda61ed153d.png":{"width":1378,"height":389}}}