{"notes":[{"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":""}]},{"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/mamba-2-hybrid","category":"mamba","legacyPath":"","legacySlug":"","legacyPaths":[],"title":"Mamba-2-Hybrid","subtitle":"おぼろげな記憶","author":"Z. Yuan","paperDate":"2024-06-12","publishedOn":"2026-01-22","paperTitle":"An Empirical Study of Mamba-based Language Models","paperUrl":"https://arxiv.org/abs/2406.07887","bibliography":{"title":"An Empirical Study of Mamba-based Language Models","authors":["Roger Waleffe","Wonmin Byeon","Duncan Riach","Brandon Norick","Vijay Korthikanti","Tri Dao","Albert Gu","Ali Hatamizadeh","Sudhakar Singh","Deepak Narayanan","Garvit Kulshreshtha","Vartika Singh","Jared Casper","Jan Kautz","Mohammad Shoeybi","Bryan Catanzaro"],"url":"https://arxiv.org/abs/2406.07887","pdfUrl":"https://arxiv.org/pdf/2406.07887v1","codeUrl":"https://github.com/NVIDIA/Megatron-LM","preprint":{"id":"2406.07887","submittedOn":"2024-06-12","url":"https://arxiv.org/abs/2406.07887"},"publication":null,"citation":"@misc{mambamamba2hybrid2024,\n  title = {{An Empirical Study of Mamba-based Language Models}},\n  author = {Roger Waleffe and Wonmin Byeon and Duncan Riach and Brandon Norick and Vijay Korthikanti and Tri Dao and Albert Gu and Ali Hatamizadeh and Sudhakar Singh and Deepak Narayanan and Garvit Kulshreshtha and Vartika Singh and Jared Casper and Jan Kautz and Mohammad Shoeybi and Bryan Catanzaro},\n  year = {2024},\n  eprint = {2406.07887},\n  archivePrefix = {arXiv},\n  url = {https://arxiv.org/abs/2406.07887}\n}\n","citationPath":"/content/papers/citations/mamba/mamba-2-hybrid.2add9f814f3a.bib"},"thumbnail":{"src":"/content/papers/assets/f2e0e285e0d5337a787c2e32537b37089640a036c9339c3574c147a3cbaad4eb.png","width":1245,"height":322},"locale":"ja","path":"/content/papers/ja/mamba/mamba-2-hybrid.2be6a0d9ad19.md","headings":[{"id":"問題の定義","text":"問題の定義","depth":2},{"id":"問題の解決","text":"問題の解決","depth":2},{"id":"まず純粋な-ssm-の成績を見る","text":"まず純粋な SSM の成績を見る","depth":3},{"id":"電話帳とおぼろげな記憶","text":"電話帳とおぼろげな記憶","depth":3},{"id":"アテンションをどれだけ戻すか","text":"アテンションをどれだけ戻すか？","depth":3},{"id":"メモリを手計算する","text":"メモリを手計算する","depth":3},{"id":"考察","text":"考察","depth":2},{"id":"12項目の標準タスク","text":"12項目の標準タスク","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","relation":"referenced","topic":""}]}],"text":"2023年12月、Mamba は線形時間の選択的状態空間モデルで、言語モデリングにおいて同規模の Transformer に追いつきました。ただ、論文で最大のモデルは 2.8B です。誰もが同じ疑問を抱いていました。製品レベルまで大きくしても、持ちこたえられるのか？\n\n3か月後、イスラエルの AI21 Labs が、52B パラメータで Apache 2.0 ライセンスのモデルを出してきました。名前は Jamba。確かに Mamba 層を大量に使っていますが、8層ごとに1層はアテンションを残しています。\n\nHacker News の発表スレッドでは、Mamba があるのに、なぜアテンション層を入れるのか、とずばり尋ねる人がいました。返信は短いものでした。「Mamba は長いコンテキストが苦手だ。電話番号を覚えられない」。[Hacker News、2024-03-28](https://news.ycombinator.com/item?id=39853958)\n\nその返信がリンクしていたのは、長い電話帳のように原文を正確にコピーする必要があるタスクを測った研究で、こちらは別にノートを書いています。その論文には、状態モデルが得意な場面も挙げられています。Jamba の論文は、自らのアブレーション実験から、もう一つのより具体的な理由を示します。\n\n## 問題の定義\n\nまず、長い文書を扱うモデルを思い浮かべてみてください。払う勘定は2つあります。\n\n一つはメモリです。アテンション層は読んだトークンごとに key と value を残すので、文書が長いほどこのキャッシュは大きくなり、最終的には GPU 1枚にどれだけ長いコンテキストを載せられるかを、これが決めます。もう一つは時間です。新しいトークンを生成するたびに、コンテキスト全体ともう一度照合しなければならず、長文ではスループットが落ちます。\n\nMamba のようなモデルは、固定サイズの状態を更新するだけなので、どちらの勘定も安く済みます。しかし AI21 の評価では、同じ規模なら、純粋な Mamba の言語モデルは Transformer にやや及びませんでした。\n\n2種類の層を一つのモデルに入れる試みは当時すでにありましたが、どれも規模は大きくありませんでした。Mamba とアテンションを交互に並べた実験は 1.3B 前後どまりで、7B まで作った StripedHyena も、アテンションだけの Mistral-7B に負けていました。\n\nですから、この論文が本当に答えようとしているのは、規模を大きくしてみないと答えられない問いです。2種類の層をどんな比率で混ぜるべきか？ MoE を加えると衝突しないか？ 数百億パラメータまで大きくしたとき、学習は安定するのか？\n\n## 問題の解決\n\n### Jamba block\n\n[![Jamba のアーキテクチャ図。左は一つの Jamba block で、下から順に Mamba、Mamba MoE、Mamba、Mamba MoE、Transformer、Mamba MoE、Mamba、Mamba MoE の計8層。右は4種類の層の内部構造で、いずれも RMSNorm と Attention または Mamba、続いて RMSNorm と MLP または MoE からなり、2つの部分それぞれに残差接続がある](/content/papers/assets/fa186580641946194f81adf74136e29431c3f6dd3fa63e2ecf8a1c0c29d3e5a0.png)](/content/papers/assets/fa186580641946194f81adf74136e29431c3f6dd3fa63e2ecf8a1c0c29d3e5a0.png)\n\n*原論文 Figure 1。[**v2 PDF の3ページ**](https://arxiv.org/pdf/2403.19887v2#page=3)から抜粋、CC BY-SA 4.0。左は一つの Jamba block、右は4種類の層の内部構造です。*\n\nまず右の図から。各層は2つの部分に分かれます。前半は系列方向に情報を混ぜる部分で、Attention か Mamba を使います。後半はトークンごとの特徴変換で、通常の MLP か MoE を使います。どちらの部分も前に RMSNorm があり、それぞれ残差接続で包まれています。\n\nこれを組み合わせると、右の図の4種類の層になります。そのうち Attention MoE は、Jamba では実際には使われていません。\n\n次に左の図です。一つの Jamba block は8層で、下から数えて5層目が唯一の Transformer 層、残りの7層はすべて Mamba です。MoE は1層おきに現れます。\n\n著者らは、この設計空間を5つのパラメータで表します。\n\n- $l$：1つの block の層数。\n- $a:m$：アテンション層と Mamba 層の比率。\n- $e$：何層ごとに MLP を MoE に置き換えるか。\n- $n$：各 MoE 層のエキスパート数。\n- $K$：各トークンが選ぶエキスパート数。\n\n公開版は $l=8$、$a:m=1:7$、$e=2$、$n=16$、$K=2$ で、4つの block を積み重ねた計32層のうち、アテンションは4層だけです。総パラメータは 52B で、トークンごとにアクティブになるのは 12B。int8 の重みなら、80GB の GPU 1枚に収まります。\n\nそのほかの細部は、おおむね当時の標準的なレシピを踏襲しています。アテンションは GQA、活性化関数は SwiGLU、語彙は 64K で、数字は1桁ずつ独立したトークンに分けます。特徴的なのは2点で、後の考察でも触れます。Mamba 層の内部に RMSNorm を加えていること、そしてモデル全体で RoPE のような位置エンコーディングを使っていないことです。\n\n### KV キャッシュを手計算する\n\nアテンションを4層だけにすると、実際どれだけ節約できるのでしょう。論文の Table 1 に、256K コンテキスト、16ビットでの KV キャッシュが載っています。\n\n| モデル | 総パラメータ | アクティブパラメータ | KV キャッシュ（256K） |\n| --- | ---: | ---: | ---: |\n| Llama-2 | 6.7B | 6.7B | 128GB |\n| Mistral | 7.2B | 7.2B | 32GB |\n| Mixtral | 46.7B | 12.9B | 32GB |\n| Jamba | 52B | 12B | 4GB |\n\n*表：原論文 Table 1。*\n\n自分たちでも計算してみましょう。アテンション層ごと、トークンごとに、key と value を1つずつ保存します。Mixtral も Jamba も KV ヘッドは8個、各ヘッド128次元で（両者の Hugging Face 上の設定ファイルによる）、bf16 なら1つの値が2 bytes です。\n\n```python\ndef kv_bytes(attn_layers, kv_heads, head_dim, tokens, bytes_per=2):\n    return 2 * attn_layers * kv_heads * head_dim * bytes_per * tokens\n\nGiB = 2**30\nfor name, layers in [(\"Mixtral（アテンション32層）\", 32), (\"Jamba（アテンション4層）\", 4)]:\n    per_token = kv_bytes(layers, 8, 128, 1)\n    print(f\"{name}: トークンあたり {per_token // 1024} KiB、256K トークンで計 {kv_bytes(layers, 8, 128, 262144) / GiB:.0f} GiB\")\n\n# Jamba の Mamba 28層：内部幅 2×4096、チャネルごとの状態次元 16\nmamba_layers, d_inner, d_state = 28, 2 * 4096, 16\nssm = mamba_layers * d_inner * d_state\nprint(f\"28層の Mamba の SSM 状態：{ssm:,} 個の値、bf16 で約 {ssm * 2 / 2**20:.0f} MiB\")\n```\n\n```text\nMixtral（アテンション32層）: トークンあたり 128 KiB、256K トークンで計 32 GiB\nJamba（アテンション4層）: トークンあたり 16 KiB、256K トークンで計 4 GiB\n28層の Mamba の SSM 状態：3,670,016 個の値、bf16 で約 7 MiB\n```\n\nどちらの数字も Table 1 と一致します（原表は GB 表記ですが、ここでは2のべき乗で数える GiB にそろえています）。8倍の差は、そのままアテンション32層と4層の差です。\n\n最後の行は、Mamba 側が持ち歩く状態です。コンテキストが 1K でも 256K でも、大きさは変わりません。（ここで数えたのは SSM の状態だけです。畳み込みの一時バッファは別勘定ですが、これもごく小さな量です。）\n\n\n> [!NOTE]\n> \n> Llama-2 の行が4倍大きいのは、その 7B 版が GQA を使っていないからです。32個のアテンションヘッドが、それぞれ自分の key と value を持ちます。Mistral は KV ヘッドを8組に減らしたので、KV キャッシュが4分の1になりました。\n\n\n## 考察\n\n### Mixtral と比べてみる\n\nまず標準的なベンチマークから。著者らが主に比べているのは Mixtral 8x7B で、総パラメータもアクティブパラメータも Jamba に近いモデルです。いくつか抜き出してみます。\n\n| ベンチマーク | Mixtral | Jamba |\n| --- | ---: | ---: |\n| HellaSwag（10-shot） | 86.7 | 87.1 |\n| WinoGrande（5-shot） | 81.2 | 82.5 |\n| GSM8K（3-shot CoT） | 60.4 | 59.9 |\n| HumanEval（pass@1） | 34.8 | 29.3 |\n| MMLU（5-shot） | 70.6 | 67.4 |\n| BBH（3-shot） | 50.3 | 45.4 |\n\n*表：原論文 Table 2 より抜粋。*\n\n著者らは「ほとんどのタスクで同等の性能」としています。項目ごとに見ると、上の表の HellaSwag と WinoGrande はやや高く、完全な表にある ARC-E、ARC-C、そして MMLU、BBH、HumanEval では数ポイント劣ります。\n\nJamba の強みは、主に長いコンテキストでのコストにあります。論文の Figure 2 によると、A100 80GB 1枚に載せられるコンテキスト長は、Jamba が Mixtral の約2倍、Llama-2-70B の約7倍にあたる約140K トークンです。Figure 3 はエンコードとデコードを合わせたエンドツーエンドのスループットを測ったもので、どちらの設定でも512トークンを生成します。A100 1枚、int8 量子化、8K コンテキストでは、Jamba はバッチ16が載りますが Mixtral は載らず、スループットに3倍ほどの差がつきます。A100 4枚、量子化なし、バッチ1、128K コンテキストでも、Jamba が約3倍速くなりました。\n\n著者らは、これらのスループットの数字は相対比較にしか使えないとも注意しています。純粋な Transformer にはコミュニティによる6年分の推論最適化がありますが、当時の Jamba にはまだありませんでした。\n\n長いコンテキストの能力については、2種類のテストをしています。干し草の山から針を探す needle-in-a-haystack では、256K までほぼすべて見つけられました。著者らは、これがアテンション4層だけで実現していることを特に強調しています。もう一つは、L-Eval のうち最も長い5つのデータセットを 3-shot 形式に直したもので、F1 の平均は 0.44、Mixtral は 0.43 でした。\n\n### 1対3か、1対7か？\n\n比率の選び方は、1.3B パラメータ、250B トークンの小さなモデルで試しています。\n\n| モデル | OLLM | HellaSwag | WinoGrande | NQ |\n| --- | ---: | ---: | ---: | ---: |\n| アテンションのみ | 36.4 | 62.4 | 59.6 | 14.5 |\n| Mamba のみ | 36.1 | 62.6 | 59.4 | 14.5 |\n| ハイブリッド 1:3 | 37.2 | 65.1 | 61.7 | 16.5 |\n| ハイブリッド 1:7 | 37.2 | 65.1 | 61.7 | 16.0 |\n\n*表：原論文 Table 4 より抜粋。いずれも MoE なし。OLLM は、著者らが再現した Hugging Face Open LLM リーダーボードの総合スコアです。*\n\n2つの混合比率はほとんど同じで、どちらも単一のアーキテクチャより優れています。著者らは、計算量が少なくて済むという理由で 1:7 を選びました。7B、50B トークンに大きくしてもう一度比べると、Mamba のみの OLLM は 35.3 で、アテンションのみの 36.1 をやや下回り、ハイブリッドモデルが 36.6 で最高、学習損失も終始いちばん低いままでした。\n\nここで気をつけたいのは、2組のアブレーションの学習量がそれぞれ 250B と 50B トークンだという点です。公開モデルの総学習量は、論文では公表されていません。比率についての結論は、ひとまずこれらのアブレーションの設定に限られます。\n\n### 純粋な Mamba はどこでつまずくのか？\n\nここが、この論文でいちばん面白いところです。\n\n純粋な Mamba は、パープレキシティも一般的なベンチマークも悪くないのに、3つのデータセットで目に見えて成績が落ちます。\n\n| モデル | IMDB | QuAC | NarrativeQA |\n| --- | ---: | ---: | ---: |\n| アテンションのみ | 84.1 | 27.9 | 45.8 |\n| Mamba のみ | 48.8 | 20.2 | 27.7 |\n| ハイブリッド | 90.9 | 26.6 | 43.7 |\n\n*表：原論文 Table 6。1.3B モデル、250B トークン。*\n\nIMDB は映画レビューの感情分類で、答えは Positive と Negative の2つしかありません。純粋な Mamba の出力を調べた著者らは、それがしばしば形式どおりに答えず、「Very Good」「Funny」「Bad」、さらには「3/10」などと書いていることに気づきました。\n\n感情の向きとしては合っている答えもあるかもしれません。ただ、前に示した例の形式で答えていないのです。\n\nfew-shot のプロンプトでは、どの例も「レビュー＋Sentiment: ラベル」という形です。この形式を踏襲する方法の一つは、最後のコロンのところで、前の例のラベルがどう書かれていたかを振り返ることです。こうした「さかのぼって見つけ、そのとおり書く」能力を、これまでの研究は Transformer の **induction head** と結びつけてきました。前にある同じパターンを見つけ、その次のトークンをコピーするアテンションヘッドです。\n\n著者らの推測では、純粋な Mamba にはアテンションがないため、このような仕組みが学習の中で自然に育ちにくいのです。1:7 のハイブリッドモデル（この 1.3B モデルではアテンションは計3層）に替えると、問題は消えました。（ここでいう「純粋な Mamba」は、実際には Mamba 層と MLP を交互に並べたもので、アテンションがないだけです。）\n\n[![IMDB の few-shot プロンプトの一部。3つの例の後にそれぞれ Sentiment ラベル Negative、Negative、Positive があり、ラベルが濃淡の異なる青で示されている。プロンプトは最後に Sentiment のコロンで止まっている](/content/papers/assets/ab6a0dbf1cacc4a6045d01ccebacef40614338c0765fdb2052bd1fa19ffcd408.png)](/content/papers/assets/ab6a0dbf1cacc4a6045d01ccebacef40614338c0765fdb2052bd1fa19ffcd408.png)\n\n*原論文 Figure 8。[**v2 PDF の12ページ**](https://arxiv.org/pdf/2403.19887v2#page=12)から抜粋、CC BY-SA 4.0。ハイブリッドモデルの最初のアテンション層にある1つのヘッドが、最後のコロンから見ているところです。色が濃いほどアテンションが強いことを表します。*\n\n図の例は、純粋な Mamba が間違え、ハイブリッドモデルが正解したものです。最後のコロンから見ると、アテンションはほぼすべて前の例のラベル、Negative、Negative、Positive に向いています。著者らは 1.3B のハイブリッドモデルの3つのアテンション層で、このようなヘッドを12個見つけました。\n\n著者ら自身も、これは逸話的な証拠だと述べています。ハイブリッドモデルに induction head に似たヘッドが現れたことはわかりますが、SSM が in-context learning を決して習得できないと証明するには足りません。それでもエンジニアリングの観点からは、結論はすでに十分実用的です。**少数のアテンション層で、この穴は埋められる**。\n\n\n> [!TIP]\n> \n> 「Mamba は電話番号を覚えられない」という言葉の出どころは、同じ年の2月に出た別の研究です。理論と実験の両面から、固定サイズの状態ではなぜ長い文字列を写しきれないのかを説明しています。\n> \n> - **[[24.02] Repeat After Me：電話帳テスト](/blog/papers/mamba/repeat-after-me)**\n\n\n### ほかの3つの発見\n\n1. **MoE は効く**。7B、50B トークンのハイブリッドモデルに MoE（16個から2個を選択、2層に1回）を加えると、OLLM は 36.6 から 38.1 に、NQ は 15.4 から 18.9 に上がりました。\n2. **Mamba は大きくすると追加の正規化が必要**。1.3B 以下では学習は安定していましたが、本番のモデルまで大きくすると大きなロススパイクが起きました。原因は Mamba 層内部の活性化の値が大きすぎることで、内部に RMSNorm を加えると安定しました。\n3. **位置エンコーディングなしでもよい**。アテンション層に RoPE を加えても加えなくても、結果はほぼ同じでした（1.3B、250B トークン）。著者らの推測では、アテンション層の前に置かれた Mamba 層が、暗黙の位置情報をすでにモデルに与えています。\n\n## 結論\n\nAI21 は Jamba を、製品規模に達した初のアテンションと Mamba のハイブリッドモデルと位置づけています。2024年3月に公開されたのは、指示チューニングもアライメントもまだ行っていない、事前学習済みのベースモデルです。この論文から持ち帰れるのは、おおよそ次の点です。\n\n1. **比率はかなり低くてよい**。1:3 と 1:7 の混合はほぼ同じ効果で、どちらもアテンションのみ、Mamba のみを上回ります。32層のうちアテンションを4層だけにすると、256K コンテキストの KV キャッシュは 32GB から 4GB に減ります。\n2. **純粋な Mamba の弱点は形式とコピー**。一般的なベンチマークでは差がほとんど見えず、IMDB のように few-shot の形式どおりに答えるタスクで初めて表に出ます。1:7 のハイブリッドモデルなら、これを取り戻せます。\n3. **性能は Mixtral に近く、長文のコストはずっと低い**。一部のベンチマークではまだ Mixtral に劣り、スループットの優位は主に長いコンテキストで現れます。\n\n最初の問いに戻りましょう。Mamba が来たのに、なぜアテンションを残すのか？ この論文の答えは実務的です。効果が最も高く、コストも許容できるのは、大部分の層に Mamba を使い、少数の層をアテンションに残すやり方です。3か月後、NVIDIA が 8B モデルでより大規模な比較を行い、結論もほぼ同じでした。\n\n\n> [!TIP]\n> \n> NVIDIA が同じデータで Mamba、Mamba-2、Transformer を比較した研究は、こちらを参照してください。\n> \n> - **[[24.06] Mamba-2-Hybrid：おぼろげな記憶](/blog/papers/mamba/mamba-2-hybrid)**\n","imageSizes":{"/content/papers/assets/fa186580641946194f81adf74136e29431c3f6dd3fa63e2ecf8a1c0c29d3e5a0.png":{"width":1056,"height":1000},"/content/papers/assets/ab6a0dbf1cacc4a6045d01ccebacef40614338c0765fdb2052bd1fa19ffcd408.png":{"width":1128,"height":778}}}