{"notes":[{"id":"mamba/repeat-after-me","category":"mamba","legacyPath":"","legacySlug":"","legacyPaths":[],"title":"Repeat After Me","subtitle":"电话簿测试","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":"zh-hans","path":"/content/papers/zh-hans/mamba/repeat-after-me.82bba41fa9d9.md","headings":[{"id":"定义问题","text":"定义问题","depth":2},{"id":"解决问题","text":"解决问题","depth":2},{"id":"抄写任务","text":"抄写任务","depth":3},{"id":"状态不够大就抄不完","text":"状态不够大，就抄不完","depth":3},{"id":"transformer-靠-n-gram-找位置","text":"Transformer 靠 n-gram 找位置","depth":3},{"id":"讨论","text":"讨论","depth":2},{"id":"合成数据学得快也推得远","text":"合成数据：学得快，也推得远","depth":3},{"id":"先告诉我要找什么就没问题","text":"先告诉我要找什么，就没问题","depth":3},{"id":"预训练模型电话簿","text":"预训练模型：电话簿","depth":3},{"id":"这篇还没回答的问题","text":"这篇还没回答的问题","depth":3},{"id":"结论","text":"结论","depth":2}],"anchorAliases":{},"guide":null,"relatedReading":[{"id":"mamba/jamba","relation":"referenced","topic":""}]},{"id":"mamba/jamba","category":"mamba","legacyPath":"","legacySlug":"","legacyPaths":[],"title":"Jamba","subtitle":"八分之一的注意力","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":"zh-hans","path":"/content/papers/zh-hans/mamba/jamba.cd7e942194ed.md","headings":[{"id":"定义问题","text":"定义问题","depth":2},{"id":"解决问题","text":"解决问题","depth":2},{"id":"jamba-block","text":"Jamba block","depth":3},{"id":"手算一次-kv-cache","text":"手算一次 KV cache","depth":3},{"id":"讨论","text":"讨论","depth":2},{"id":"和-mixtral-比一比","text":"和 Mixtral 比一比","depth":3},{"id":"一比三还是一比七","text":"一比三，还是一比七？","depth":3},{"id":"纯-mamba-卡在哪里","text":"纯 Mamba 卡在哪里？","depth":3},{"id":"另外三个发现","text":"另外三个发现","depth":3},{"id":"结论","text":"结论","depth":2}],"anchorAliases":{},"guide":null,"relatedReading":[{"id":"mamba/repeat-after-me","relation":"referenced","topic":""},{"id":"mamba/mamba-2-hybrid","relation":"referenced","topic":""}]}],"text":"Mamba 发表两个月后，哈佛大学的 Samy Jelassi、David Brandfonbrener、Sham M. Kakade 与 Eran Malach 上传了这篇论文，后来发表于 ICML 2024。标题直接把结论写了出来：在「照抄」这件事上，Transformer 比状态空间模型更强。\n\n照抄听起来是个很无聊的能力。可是读者丢给模型一份文档，让它找出某个日期、某个名字、某一行代码，模型做的就是从上下文里把一小段原文搬出来。这件事做不好，困惑度再低也没用。\n\n## 定义问题\n\n作者把一大类模型放在一起讨论：S4、Mamba、传统 RNN，以及可以并行训练的线性注意力与并行 RNN。它们的共同点是，不管处理到第几个 token，都只依靠一个**大小固定、不随序列长度变化**的状态。作者把这类模型统称为 **GSSM**（generalized state space models，广义状态空间模型）。\n\nGSSM 的卖点很清楚：每预测一个 token 只需要 $O(1)$ 的内存。Transformer 则要把所有看过的 token 都留在 KV cache 里，序列越长，它就越大。\n\n作者想问的是：省下这些内存，有没有付出代价？如果有，代价是什么？\n\n他们选了一个最简单的任务来检验：照抄。\n\n## 解决问题\n\n### 抄写任务\n\n任务格式如下：\n\n```text\n<BOS> x1 x2 ... xL <COPY>\n```\n\n模型看到 `<COPY>` 之后，要一字不差地输出 $x_1$ 到 $x_L$。评分用字符串级别的准确率，整串全对才算对。\n\n### 状态不够大，就抄不完\n\n先看 GSSM 这一边。直觉很简单：要原封不动地抄出整串，就得先把整串记住；记在哪里？只有那个固定大小的状态。\n\n作者把状态能表示的不同取值个数记为 $|\\mathcal{S}|$，以比特计的记忆容量为 $\\text{mem}(\\mathcal{S})=\\log|\\mathcal{S}|$。对于从 $D$ 个字符中均匀抽取、长度为 $L$ 的字符串，可能的输入有 $D^L$ 种，而状态最多只能区分 $|\\mathcal{S}|$ 种。这里的错误率，指的是整串没能一字不差抄对的概率。Theorem 2.7 给出的下界是：\n\n$$\n\\text{err}(H) \\ge 1-\\frac{|\\mathcal{S}|}{D^L}\n$$\n\n（论文正文印的是严格大于，附录 D 的计数证明推出的是大于或等于，这里以证明为准。）\n\n换个好记的说法（Corollary 2.8）：只要状态的比特数小于 $L\\log_2 D-1$，错误率就超过一半。\n\n我们算几组数字感受一下：\n\n```python\nimport math\n\nfor D, L in [(26, 300), (26, 1000), (50_000, 2_000)]:\n    print(f\"词表 {D:>6}、长度 {L:>5}：信息量约 {L * math.log2(D):,.0f} bits\")\n\nstate_bits = 1000 * 32           # 1,000 个 fp32\nprint(f\"1,000 个 fp32 = {state_bits:,} bits，最多约能装 {state_bits / math.log2(50_000):,.0f} 个 5 万词表的 token\")\n```\n\n```text\n词表     26、长度   300：信息量约 1,410 bits\n词表     26、长度  1000：信息量约 4,700 bits\n词表  50000、长度  2000：信息量约 31,219 bits\n1,000 个 fp32 = 32,000 bits，最多约能装 2,050 个 5 万词表的 token\n```\n\n最后一行的例子是作者自己举的：1,000 个 32 位浮点数，信息量足够装下至少 2,000 个 token。\n\n所以这个理论下界其实很宽松。实际的 Mamba 状态远大于 1,000 个数，抄 300 个英文字母所需的 1,410 bits 根本不成问题。作者也承认这一点：理论只说明「装不下就一定抄不完」；至于「装得下」的模型会不会学着把字符串存进状态里，要靠实验来回答。\n\n### Transformer 靠 n-gram 找位置\n\n另一边，作者构造了一个由两个块组成的 Transformer，证明它能照抄的字符串长度可以随模型维度**指数增长**（Theorem 2.3；引言里的说法是随注意力头数增长）；误差则取决于字符串里出现重复 n-gram 的概率，对均匀随机字符串来说，这个概率很小。\n\n它的做法很像人在抄一篇长文：\n\n1. 看看自己刚写下的最后几个字，比如最后 3 个。\n2. 回到原文里，找这 3 个字出现的位置。\n3. 把那个位置后面的下一个字写下来。\n\n这就是一次 n-gram 查表。注意力天生适合做「拿当前片段去和前面每个位置比对」这件事；为了让注意力头只看最近 $m$ 个 token 来组成 n-gram 的哈希，作者设计了一种硬截断版的 ALiBi 位置偏置，称为 **Hard-ALiBi**：$m$ 个位置以外的偏置设为 $-\\infty$，以内设为 0。\n\n这个方法会失手的情况也很明确：原文里有两段相同的 n-gram，后面接的字却不一样，模型就不知道该从哪一段接着抄。在均匀随机字符串里，这种重复的概率随 $n$ 指数下降（Lemma 2.4）。\n\n## 讨论\n\n### 合成数据：学得快，也推得远\n\n[![三张折线图。(a) 横轴是训练样本数（对数刻度），纵轴是抄写准确率：四种位置编码的 Transformer 在十万个样本内达到接近 100%，Mamba 要到约 200 万个样本，LSTM 始终接近 0。(b) 横轴是字符串长度：所有模型在训练长度 50 以内都是 100%，超过之后 LSTM 与 RoPE 版 Transformer 立刻掉到 0，Mamba 撑到约 70，NoPE 约 90，ALiBi 约 130，Hard-ALiBi 到 1,000 仍接近 100%。(c) 横轴是电话簿条目数：Pythia 1.4B 与 2.8B 在 200 条时仍有约 90%，410M 降到约 50%，三个 Mamba 模型在 100 条时都低于 20%，200 条时接近 0](/content/papers/assets/6d45fd8ab61191e1263c400e01bf76bdd9ef3534ecc4485835332074ea553443.png)](/content/papers/assets/6d45fd8ab61191e1263c400e01bf76bdd9ef3534ecc4485835332074ea553443.png)\n\n*原论文 Figure 1，取自 [**v2 PDF 第 2 页**](https://arxiv.org/pdf/2402.01032v2#page=2)，CC BY 4.0。(a)、(b) 是从头训练的模型，Transformer 与 Mamba 约 1.6 亿参数，LSTM 约 4,000 万；(c) 是现成的预训练模型。*\n\n先看左图。所有模型都从头训练，Transformer 与 Mamba 约 1.6 亿参数；LSTM 规模太大时训练不稳定，用的是约 4,000 万参数的版本。字符串长度最多 300，字符是 26 个英文字母。\n\n四种 Transformer 都在十万个样本以内学会了；Mamba 要到两百万个左右。按作者的说法，Transformer 需要的样本约少 100 倍。LSTM 在这个样本量内完全学不会。\n\n中间那张图测的是长度泛化：只用长度 50 以内的字符串训练，再拿更长的字符串测试。Mamba 和 LSTM 一超过训练长度，很快就掉到 0。Transformer 则要看位置编码：RoPE 版本也掉得很快，NoPE 和 ALiBi 撑得比较远，Hard-ALiBi 版本到 1,000 都还接近全对。\n\n作者也确认了 Transformer 确实在用 n-gram 查表：故意在字符串里放两段相同的 n-gram，重复长度在 4 以内几乎没有影响，到 5 以上准确率就开始下降，和一个完美 5-gram 模型的表现吻合。\n\n### 先告诉我要找什么，就没问题\n\n有一组对照实验很值得一看。任务是：给一段字符串，再给出其中某个 n-gram 当作钥匙，让模型输出钥匙后面的 $k$ 个 token。作者做了两种版本：\n\n两种版本都用长度 30 以内的字符串训练，再测到长度 100：\n\n- **钥匙放在最后**（suffix key）：模型读完整段字符串才知道要找什么。Transformer 推到长度 100 还不错，GSSM 一离开训练长度就不行了。\n- **钥匙放在最前面**（prefix key）：模型一开始就知道要找什么，读的时候只要盯住那把钥匙，看到了就记下后面几个字。这一次 GSSM 一路保持全对，还赢过了 NoPE 和 ALiBi 版的 Transformer。\n\n差别只在于问题出现在什么时候。\n\n固定大小的状态，擅长的是「边读边摘要」。如果事先知道要留什么，它可以把无关的内容丢掉；如果问题在读完之后才出现，它就得把所有可能用到的东西都先留着。信息量一旦超过状态容量，就没法保证完整找回来。\n\n### 预训练模型：电话簿\n\n前面的实验都是从头训练的小模型。作者接着拿现成的预训练模型来比：Pythia（Transformer）410M 到 2.8B，对上 Mamba 360M 到 2.8B。两边都在 Pile 上训练、用同一套分词器，主要差别就在架构。值得一提的是，同等规模下，Mamba 在 Pile 上的困惑度略低。\n\n最有名的就是右图的电话簿测试。提示词里放一本随机生成的电话簿，每行像这样：\n\n```text\nJohn Powell: 609-323-7777\n```\n\n后面附上两个问答示例，再问某个人的电话。电话簿越长，Mamba 掉得越快；到 70 条左右，最小的 Pythia 410M 已经赢过最大的 Mamba 2.8B。\n\n[![三张折线图。(a) 抄写 C4 自然语言字符串：三个 Pythia 在 500 个 token 时仍有约 80% 到接近 100%，三个 Mamba 随长度下降，2.8B 在 500 个 token 时已低于 10%。(b) 打乱词序后再抄：Pythia 1.4B 与 2.8B 小幅下降，410M 降到约 50%，三个 Mamba 最晚在 300 个 token 之前掉到 0。(c) SQuAD 问答的 F1：短段落时两者接近，段落越长，Mamba 掉得越多](/content/papers/assets/8d14006fc09d6d17409ad80cc3ab8e3574c16a5a2b26e22539149dda61ed153d.png)](/content/papers/assets/8d14006fc09d6d17409ad80cc3ab8e3574c16a5a2b26e22539149dda61ed153d.png)\n\n*原论文 Figure 7，取自 [**v2 PDF 第 8 页**](https://arxiv.org/pdf/2402.01032v2#page=8)，CC BY 4.0。*\n\n另外三个测试的方向一致：\n\n1. **抄自然语言**：从 C4 抽一段文本，在提示词里放两遍，再接上第三遍的第一个单词，让模型接着写。最小的 Pythia 也大幅领先最大的 Mamba。\n2. **打乱词序再抄**：自然语言可以压缩，打乱顺序后就难了。两边都变差，Mamba 掉得更多，最大的 Mamba 在 300 个 token 时已经是 0。这和「固定状态只能靠压缩撑着」的推论一致。\n3. **SQuAD 问答**：2.8B 对 2.8B，短段落时两者差不多；段落越长，Mamba 的 F1 掉得越快。\n\n### 这篇还没回答的问题\n\n作者在相关工作一节也提到了反方向的例子：有些任务只需要追踪一个很小的状态，例如 flip-flop 语言建模，LSTM 能轻松泛化，Transformer 反而不行；对 DNA 这类超长序列，固定状态也更实际。\n\n所以这篇讨论的是**记忆密集型**任务。它没有说 GSSM 在其他任务上一定更差，也没有测试把两种层混在一起会怎样。\n\n## 结论\n\n这篇论文把「固定大小的状态是有代价的」这件事，从直觉变成了可以量化的结果：\n\n1. **理论上**，两个块的 Transformer 能照抄的长度可以随模型维度指数增长；GSSM 的状态比特数少于 $L\\log_2 D-1$ 时，均匀随机字符串整串抄错的概率就超过一半。\n2. **实际上**，限制往往更早出现。即便状态大到装得下，Mamba 也很难学会把整串存进去：训练样本要多 100 倍，换个长度就失效。\n3. **预训练模型也一样**。困惑度略低的 Mamba，在电话簿、抄写与长段落问答上，都输给了同等规模、甚至小得多的 Pythia。\n\n「先知道要找什么就没问题」那组对照，大概是读这篇最该带走的一点。它说明了固定状态适合做什么、什么时候会吃亏，也为「让状态模型搭配能回头查找的注意力」提供了一个很直接的理由。\n\n\n> [!TIP]\n> \n> AI21 两个月后发布的 Jamba，就在 Mamba 层之间保留了少量注意力层；他们在消融实验里也观察到，纯 Mamba 经常偏离 few-shot 示例指定的作答格式：\n> \n> - **[[24.03] Jamba：八分之一的注意力](/blog/papers/mamba/jamba)**\n","imageSizes":{"/content/papers/assets/6d45fd8ab61191e1263c400e01bf76bdd9ef3534ecc4485835332074ea553443.png":{"width":1378,"height":416},"/content/papers/assets/8d14006fc09d6d17409ad80cc3ab8e3574c16a5a2b26e22539149dda61ed153d.png":{"width":1378,"height":389}}}