{"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":"ja","path":"/content/papers/ja/mamba/repeat-after-me.1c006f53a866.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":"8分の1のアテンション","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":"ja","path":"/content/papers/ja/mamba/jamba.008c0136e88f.md","headings":[{"id":"問題の定義","text":"問題の定義","depth":2},{"id":"問題の解決","text":"問題の解決","depth":2},{"id":"jamba-block","text":"Jamba block","depth":3},{"id":"kv-キャッシュを手計算する","text":"KV キャッシュを手計算する","depth":3},{"id":"考察","text":"考察","depth":2},{"id":"mixtral-と比べてみる","text":"Mixtral と比べてみる","depth":3},{"id":"1対3か1対7か","text":"1対3か、1対7か？","depth":3},{"id":"純粋な-mamba-はどこでつまずくのか","text":"純粋な Mamba はどこでつまずくのか？","depth":3},{"id":"ほかの3つの発見","text":"ほかの3つの発見","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 の発表から2か月後、ハーバード大学の Samy Jelassi、David Brandfonbrener、Sham M. Kakade、Eran Malach がこの論文を公開しました。のちに ICML 2024 で発表されています。タイトルが結論をそのまま述べています。「コピー」に関しては、Transformer のほうが状態空間モデルより優れている、と。\n\nコピーと聞くと、ずいぶん退屈な能力に思えます。でも、文書をモデルに渡して、ある日付や名前、あるいはコードの一行を探させるとき、モデルがしているのは、コンテキストから原文の一部をそのまま取り出してくることです。これがうまくできなければ、パープレキシティがいくら低くても役に立ちません。\n\n## 問題の定義\n\n著者らは、幅広いモデルをひとまとめにして議論します。S4、Mamba、従来の RNN、そして並列に学習できる線形アテンションや並列 RNN です。共通点は、何番目のトークンを処理していても、頼れるのは**サイズが固定され、系列長によって変わらない**一つの状態だけ、ということです。著者らはこの種のモデルを **GSSM**（generalized state space models、一般化状態空間モデル）と総称します。\n\nGSSM の売りははっきりしています。トークンを一つ予測するのに必要なメモリは $O(1)$ だけです。一方 Transformer は、これまでに見たトークンをすべて KV キャッシュに残す必要があり、長くなるほど大きくなります。\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           # fp32 を1,000個\nprint(f\"1,000個の fp32 = {state_bits:,} bits、収まるトークンは最大で約 {state_bits / math.log2(50_000):,.0f} 個（語彙5万）\")\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万）\n```\n\n最後の行は、著者ら自身が挙げた例です。32ビット浮動小数点数が1,000個あれば、少なくとも2,000トークン分の情報を収められます。\n\nつまり、この理論的な下界はかなり緩いのです。実際の Mamba の状態は1,000個の数よりはるかに大きく、英字300文字のコピーに必要な1,410 bits はまったく問題になりません。著者らもこの点を認めています。理論が示すのは「入りきらなければ必ず写しきれない」ということだけです。「入りきる」モデルが文字列を状態に保存するよう学習できるかどうかは、実験で答えるしかありません。\n\n### Transformer は n-gram で位置を探す\n\nもう一方の側では、著者らは2ブロックの Transformer を構成し、コピーできる文字列の長さがモデルの次元に対して**指数的に伸びる**ことを証明しました（Theorem 2.3。序論では、アテンションヘッド数に対して、という言い方です）。誤差は文字列中に同じ n-gram が繰り返し現れる確率で抑えられ、一様ランダムな文字列ならこの確率は小さくなります。\n\nその方法は、人が長い文章を書き写すやり方によく似ています。\n\n1. 自分が直前に書いた最後の数文字、たとえば最後の3文字を見る。\n2. 原文に戻って、その3文字が現れる位置を探す。\n3. その位置の次の文字を書き写す。\n\nこれは n-gram による表引きです。「今の断片を、前にあるすべての位置と照合する」作業は、アテンションの得意分野です。アテンションヘッドが直近 $m$ 個のトークンだけを見て n-gram のハッシュを組み立てられるように、著者らは ALiBi の位置バイアスを一定距離で完全に打ち切る版を設計し、**Hard-ALiBi** と名付けました。$m$ 個より離れた位置のバイアスは $-\\infty$、それ以内は 0 にします。\n\nこの方法がしくじる場面もはっきりしています。原文に同じ n-gram が2か所あり、その後に続く文字が違えば、モデルはどちらから続けて写せばよいかわかりません。一様ランダムな文字列では、こうした重複の確率は $n$ に対して指数的に下がります（Lemma 2.4）。\n\n## 考察\n\n### 合成データ：覚えが早く、遠くまで通用する\n\n[![3枚の折れ線グラフ。(a) 横軸は学習サンプル数（対数目盛）、縦軸はコピーの正解率。4種類の位置エンコーディングの Transformer は10万サンプル以内でほぼ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%に下がり、3つの 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\n4種類の Transformer はどれも10万サンプル以内で習得しましたが、Mamba は200万サンプルほどかかりました。著者らによれば、Transformer に必要なサンプルは約100分の1です。LSTM は、このサンプル数ではまったく習得できません。\n\n真ん中の図は、長さの汎化を測っています。長さ50以内の文字列だけで学習し、より長い文字列でテストします。Mamba と LSTM は学習長を超えると、すぐ0に落ちます。Transformer は位置エンコーディング次第です。RoPE 版もすぐ落ちますが、NoPE と ALiBi はもう少し先まで持ちこたえ、Hard-ALiBi 版は1,000でもほぼ全問正解です。\n\n著者らは、Transformer が本当に n-gram の表引きを使っていることも確かめています。文字列にわざと同じ n-gram を2か所入れると、重複の長さが4以内ならほとんど影響がなく、5以上で正解率が落ち始めます。完璧な 5-gram モデルの振る舞いと一致します。\n\n### 何を探すか先に教えてくれれば、問題ない\n\nぜひ見ておきたい対照実験があります。タスクは、文字列を与え、その中のある n-gram をキーとして渡し、キーの後ろに続く $k$ 個のトークンを出力させるものです。著者らは2つの版を用意しました。\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 で学習し、同じトークナイザーを使っているので、主な違いはアーキテクチャです。なお、同じ規模では、Pile でのパープレキシティは Mamba のほうがやや低くなっています。\n\nいちばん有名なのが、右の図の電話帳テストです。プロンプトにランダムに生成した電話帳を入れます。各行はこんな形です。\n\n```text\nJohn Powell: 609-323-7777\n```\n\nその後に質問と回答の例を2つ付け、ある人の電話番号を尋ねます。電話帳が長くなるほど Mamba の成績は早く落ち、70件ほどになると、最小の Pythia 410M が最大の Mamba 2.8B を上回ります。\n\n[![3枚の折れ線グラフ。(a) C4 の自然言語の文字列をコピー：3つの Pythia は500トークンでも約80%からほぼ100%を保ち、3つの Mamba は長さとともに下がり、2.8B は500トークンで10%未満。(b) 単語の順序をシャッフルしてからコピー：Pythia 1.4B と 2.8B はわずかに下がり、410M は約50%まで下がり、3つの Mamba は遅くとも300トークンより前に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ほかの3つのテストも、同じ傾向を示しています。\n\n1. **自然言語のコピー**：C4 から文章を抜き出してプロンプトに2回入れ、3回目の最初の単語を続けて、モデルに続きを書かせます。最小の Pythia でも、最大の Mamba に大差で勝ちます。\n2. **単語をシャッフルしてコピー**：自然言語は圧縮できますが、順序をシャッフルすると難しくなります。どちらも成績が下がり、Mamba の落ち方がより大きく、最大の Mamba でも300トークンで0になります。「固定された状態は圧縮で持ちこたえるしかない」という推論と一致します。\n3. **SQuAD の質問応答**：2.8B 同士の比較です。短い段落では両者ほぼ同じですが、段落が長くなるほど Mamba の F1 が早く落ちます。\n\n### この論文がまだ答えていないこと\n\n著者らは関連研究の節で、逆方向の例にも触れています。flip-flop 言語モデリングのように、ごく小さな状態を追跡するだけでよいタスクでは、LSTM が楽に汎化し、Transformer のほうがうまくいきません。DNA のような超長系列でも、固定された状態のほうが現実的です。\n\nつまり、この論文が扱っているのは**記憶集約的**なタスクです。GSSM がほかのタスクでも必ず劣るとは述べていませんし、2種類の層を混ぜたらどうなるかも試していません。\n\n## 結論\n\nこの論文は「固定サイズの状態には代償がある」ということを、直感から測定できる結果に変えました。\n\n1. **理論面**：2ブロックの Transformer がコピーできる長さは、モデルの次元に対して指数的に伸びます。GSSM は状態のビット数が $L\\log_2 D-1$ より少なければ、一様ランダムな文字列で文字列全体をコピーし損なう確率が半分を超えます。\n2. **実際面**：制約はたいてい、もっと早く表れます。状態が十分に大きくても、Mamba は文字列全体を保存することをなかなか学習できません。学習サンプルは100倍必要で、長さが変わると通用しなくなります。\n3. **事前学習済みモデルでも同じ**：パープレキシティがやや低い Mamba が、電話帳、コピー、長い段落の質問応答で、同規模の、さらにはずっと小さい Pythia に負けます。\n\n「何を探すか先にわかっていれば問題ない」という対照実験は、この論文から持ち帰るべきいちばんのポイントでしょう。固定された状態が何に向き、どんなときに不利になるかを示し、「状態モデルに、さかのぼって探せるアテンションを組み合わせる」ことへの、とても直接的な理由も与えています。\n\n\n> [!TIP]\n> \n> 2か月後に AI21 が発表した Jamba は、まさに Mamba 層の間に少数のアテンション層を残した構成です。AI21 はアブレーション実験で、純粋な Mamba が few-shot で指定した解答形式からたびたび外れることも観察しています。\n> \n> - **[[24.03] Jamba：8分の1のアテンション](/blog/papers/mamba/jamba)**\n","imageSizes":{"/content/papers/assets/6d45fd8ab61191e1263c400e01bf76bdd9ef3534ecc4485835332074ea553443.png":{"width":1378,"height":416},"/content/papers/assets/8d14006fc09d6d17409ad80cc3ab8e3574c16a5a2b26e22539149dda61ed153d.png":{"width":1378,"height":389}}}