← 기사 목록
日本語https://zenn.dev/topics/llm/feed

Mamba3をゼロから実装してピュアMamba3言語モデルを作ってみた

원문 사이트로 이동 ↗

추출된 키워드

43개
言語モデル·5Mamba3·5SSM·4Transformer·4事前学習·4SFT·4ピュアMamba3モデル·4状態空間モデル·4karpathy/climbmix-400b-shuffle·3hotchpotch/fineweb-2-edu-japanese·3BPEトークナイザ·3B_bias·3C_bias·3AdamW·3Weight Decay·3GPT2·3MMLU·3GSM8K·3HumanEval·3chunk_size·3MIMO rank·3nanochat·3Mamba2·3Mamba·3Attention·3ICLR·3val_bpb·2JCommonsenseQA·2SpellingBee·2ARC-Challenge·2ARC-Easy·2Nemotron·2Jamba·2Magpie JA·2OASST2 JA·2Dolly JA·2SmolTalk·2RTX3090·2Falcon Mamba·2gradient accumulation·2PLaMo·2wandb·2ubuntu 24.04·2

원문

10,786자
Mamba3をゼロから実装してピュアMamba3言語モデルを作ってみた

Mamba3をゼロから実装してピュアMamba3言語モデルを作ってみた

はじめに

みなさんMamba使ってますか？
このMamba、今年は最新バージョンのMamba3の論文がICLRに採択されていたりと盛り上がってますね！ https://iclr.cc/virtual/2026/poster/10010352

ただなんか...やっぱりまだTransformerと比較するとそこまで浸透していないのか、メジャーなLLMだとNemotron、Jamba、Falcon Mambaや国内だとPLaMoくらいにしか採用されていないのかな～って感じですね。
あと大体がTransformerとハイブリッドのアーキテクチャになっていて、mambaだけの言語モデルってほとんどないですよね～（Falcon Mambaくらい？）

そんなことを考えていたら純粋にmambaだけで構成されたLMを作ってみたくなったので、せっかくならと最新版アーキテクチャのMamba3をゼロから実装して小さなLMを構築して、そのまま事前学習からSFTまで一気通貫でやってみました！
結構軽い気持ちで始めたんですが、事前学習がほんとうに大変で何度かくじけそうになってました...

そもそもMamba3とは

まずMambaとは、状態空間モデル（State Space Model: SSM）を核にしたアーキテクチャで、Transformerのような全トークン間のAttentionとは異なり、各ブロックで固定サイズの隠れ状態に情報を圧縮しながら処理するため、隠れ状態が固定サイズのため入力のシーケンス長に対してO(1)のメモリ・計算量で動く点が特徴。
Mamba3はその系譜の最新版で、Mamba, Mamba2からの改善によってより精度が向上している感じです。

モデルの実装

Mamba3 論文、本家のレポジトリや有志の方の再現実装の情報をもとに実装を進めました。
モデル以外の学習用フレームワークやトークナイザーに関しては、 nanochatをforkして使用させていただきました！
いや～nanochatがあって本当にめっちゃ助かりました！
これがなかったらこのプロジェクトはそもそも動いてないかも...

できたモデルはこんな感じの構成で、全12層がMamba3ブロックで構成されたピュアMamba3モデルになってます！
いやー昨今のLLMと比べるとめっちゃ小さい...

パラメータ値
アーキテクチャMamba3
層数12
d_model768
MIMO rank2
chunk_size32
総パラメータ数113M
vocab size32,768

コードやモデルはリンク先において置くので、試したい方がおりましたら是非！https://github.com/RyotaroNumata/nanomamba3 https://huggingface.co/kikyo0114/nanochat-mamba3-mimo-r2

検証環境

本プロジェクトは非常に低予算でお送りしているため、筆者のおうちにあるGPUマシンを使って何とか検証を進めました...
使っているマシンの構成は以下の感じ。

  • CPU: i9-9900k RAM 64GB
  • GPU: RTX3090 24GB
  • OS: ubuntu 24.04

事前学習

データ

今回は日本語と英語の両方で会話できるバイリンガルなモデルを学習させるため、下記の２種のコーパスを使用しました。

  • 英語:karpathy/climbmix-400b-shuffle(nanochatのデフォルト)
  • 日本語:hotchpotch/fineweb-2-edu-japanese、30%の比率でミックス
    ※リソースと電気代にビビッて、一旦5.24Bトークンで事前学習を切り上げてます。
    モデルのサイズが0.1B程度なので、Chinchilla則の2倍くらいは学習できているはず。

トークナイザ

トークナイザはnanochatに実装されているBPEトークナイザを、日本語と英語で学習させて使用しました。

事前学習用訓練レシピの探索

いやーここがほんとにきつかった。
実はこのプロジェクト始めたのが3月の終わりくらいで、そこからゴールデンウィークくらいまではほとんど問題が発生しないで学習できるレシピの探索をしていたと言っても過言ではないくらいパラメータ調整してました。

安定的な訓練レシピを得るまでの苦悩

そもそもMamba3の学習に取り組んでいるグループが少なく、web上にある学習の知見がとても少ない状況でした。また、公式の論文にも学習に関わる情報があまり記載されていなかったため、学習に関わる設定は探索的に実験を繰り返す必要がありました…
wandbのログを先ほど確認したんですが、まともにlossが下がるようになるまで70回ほどの事前実験を要しました。

モデルの初期化方法

公式の論文や実装ではB_bias、C_biasを1.0で初期化していたので、何も考えず同じ設定にしていたのですが、それだとどうも学習が安定しませんでした。いろいろ試した結果0.02で初期化すると安定させることができました。

オプティマイザ、学習率等の設定でとんでもなく苦戦

オプティマイザに関しては公式の論文の設定に従いAdamWで進めることにしました。
が、全く学習が進まない...
LRを変えて1日中回してみたりしましたが、lossがすぐにプラトーになってしまいまったく学習が進まない感じでした。
こんな感じでこの後のstepも一向にlossが下がらない...

流石になんかおかしいので、SSMグループのパラメータの統計情報を確認すると、なんとSSMグループパラメータがほとんど同一の出力になっており、全く機能していないことがわかりました。色々調査した結果原因はおそらくWeight Decayの設定で、設定が0以上の場合だとSSMパラメータが0に引っ張られるような正則化がかかり、結果として学習が進まないようでした。

やっとlossが下がり始めるパラメータを見つけた！よかった～と思って寝て起きたらlossがめっちゃ大きくなっていたこともありました。記録していたSSMパラメータの統計情報を確認すると、lossが上昇し始めたあたりから
500stepくらいで寝た。寝て起きたらlossもりもり増えてて悲しかった...

試行錯誤の結果得られた安定的な訓練レシピ

学習設定

今回は3090一枚というカツカツ環境で実施しており、バッチサイズを大きくとれないのでgradient accumulationを大きくとっています。

パラメータ値
device_batch_size4
grad_accum64
warmup_steps40
warmdown_ratio0.2
grad_clip1.0
オプティマイザ設定

ここで重要なのはSSMパラメータにWeight Decayを使用しないことです。

グループLRWD
embedding0.0010.001
lm_head0.0080.01
ssm（in_proj, out_proj 等）0.0010.0
ssm_dyn（D, dt_bias）0.0001（ssm_lr × 0.1）0.0
mlp0.0030.05

事前学習の結果

事前学習後のスコアはこんな感じで、同規模のd12 Transformer（GPT2）と比較するとやや弱い結果でした。

Mamba3Transformer（GPT2)
val_bpb0.9120.857

学習中のlossカーブはこんな感じで発散することもなく長時間学習させることができました！
ここまで長かった~

学習成功時のlossカーブ

では日本語と英語がちゃんと生成できるか試してみましょう。

## 日本語
[Prompt] むかしむかしあるところに
[Output] <|bos|>むかしむかしあるところに絵があり、その絵は「画伯」と呼ばれる事がありました。 「画伯」といっても「画伯」とは絵師のことではありません。絵師の名前は「画伯」ではなく、 絵師の形で描かれたものなのです。 絵師は「画伯」ではなく「画伯」ですから、 「画伯」の形で描かれているのです。
「画伯」の絵の多くは、 「画伯」の形で描かれているものが多いのです。 絵師は「画伯」の形で描かれているのです。 描かれている絵は 絵師が 絵師を名乗る「画伯」の形で描かれています。 絵師の 絵師に 絵師が これだけ描かれていると、 あまり目に入ってこないかもしれませんが、 絵師は 「画伯」の形で 描かれているのです。 「画伯」の形で描かれている絵は 描かれているはずの 絵師が 現実に 絵師が 描いているのです。 絵師は 画伯が 描かれている絵を 描くのです。 それが 「画伯」の形で描かれているのです。 絵師が 絵師を名乗る「画伯」の形で描かれているのです。
絵師は 「

## 英語
[Prompt] Once upon a time
[Output] <|bos|>Once upon a time, there was a little boy named Benny who loved exploring his surroundings. He had always been fascinated by animals and people, but he was always curious about what they were doing – how they were doing it, why they were doing it, and how they wanted to do it! One day, while exploring the woods at the park, he stumbled upon a group called the Forest Friends of the Forest.
Benny wanted to learn more about these amazing animals and how they were using their amazing abilities to survive in different environments. He asked his friends, "How do you think there are these amazing animals in the forest?" Benny replied, "Well, when we think about new places, there are things they can do too!"
Benny thought for a moment and then thought back. "That sounds so real sometimes, Benny," he explained. "That is so cool, Benny! During our exploration, we learned a lot about the environment and how animals interact with each other and their environment. We even learned how to recognize different animals in our environment."
As Benny continued his exploration, he noticed something unexpected – Benny was frequently using his cell phone while he was taking pictures. This was due to his phone being used to share important messages with

日本語、英語ともに出力できてる！SFT前は繰り返しが出やすいので、ベースモデルの挙動としては問題なさそう！

SFT

日英対応のベースモデルができたので、チャットで会話できるようにSFTを実施しました。
以前の実験で、今回のMamba3モデルとほぼ同規模の日英バイリンガルのTransformerモデル(GPT2)を作成していたので、そちらとのスコアとも比較してみました。

学習データ

SFTではnanochatで使用されていたものに加えて、日本語のInstructionデータを追加して実施しました。

  • SmolTalk
  • Dolly JA / OASST2 JA / Magpie JA
  • MMLU
  • GSM8K

SFTでも事前学習と同様のパラメータ設定で訓練を実施しました。訓練が進まないこともなく特段問題なく完了できていました。

SFT結果

いやーめっちゃ残念ながら全部のメトリクスでTransformerが上回る結果になってしまいました...
特に数学・コーディング（GSM8K・HumanEval）では結構差をつけられてしまった。コーディングや数学能力では入力情報を正確に参照しながら推論する必要があるため、Mambaアーキテクチャの固定サイズの隠れ状態への情報圧縮による損失が精度に影響しているのかもしれないです。
一応MMLU・JCommonsenseQAはランダム比を上回っており、日英ともに知識の吸収はできてはいそうなのでそれなりに会話できる...はず。

タスクTransformerMamba3
ARC-Easy36.45% 33.50%
ARC-Challenge33.28% 28.84%
MMLU31.89% 30.25%
GSM8K5.00% 1.06%
HumanEval9.15% 0.61%
SpellingBee99.22% 82.81%
JCommonsenseQA35.48% 33.24%

最後に会話してみましょうか...
肩こりの相談をしたら病院行け的にな会話にはなってる...？
113Mのモデルにしては会話できているような気がします！

まとめ

といった感じでMamba3は訓練時のパラメータ設定がかなりシビアでかなり苦戦しましたが、なんとかピュアなMamba3言語モデルをつくることができました！
この記事の情報がMamba3を使ったモデルを作りたいかたの一助になれば幸いです！