{"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-hant","path":"/content/papers/zh-hant/mamba/repeat-after-me.345ed541cc48.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-hant","path":"/content/papers/zh-hant/mamba/jamba.0e951e0c4113.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}}}