MiniMind の1層は6つの決定でできている

一言でいうと

MiniMind-3は、Qwen3と同じ骨組みのデコーダーです。1つの層(block)は、Pre-Norm RMSNorm → アテンション(GQA・RoPE・QK-norm) → 残差 → RMSNorm → SwiGLU FFN → 残差で、埋め込みと出力層は重みを共有します。このモジュールでは、その設定を縮小して(768次元・8層 → 128次元・4層)約100万パラメーターのモデルを作り、それぞれの決定が数字として何をしているかを確認します。

なぜ必要なのか

最初に登場したトランスフォーマー(2017)の決定の多くが変わりました。学習が不安定で(層の後ろのLayerNorm)、推論のときにメモリを食いすぎ(ヘッドごとのK・V)、長さを延ばしにくかったからです(足し込む位置埋め込み)。最近の公開モデルは、ほとんど同じ処方を使っていて、MiniMindは、その処方を数百行のPythonで見せてくれます。アテンションそのものを手で計算することは、Transformerのコース(Transformer — アテンションを手で計算する)で行ったので、ここでは、MiniMindの実際のコードが、その処方をどう実装したのかと、その数字を見ます。

どう動くのか

MiniMindConfigのデフォルト値と、その意味です。

設定	MiniMind-3	このコース	意味
hidden_size	768	128	トークン1つを表すベクトルの長さ
num_hidden_layers	8	4	層の数
num_attention_heads / num_key_value_heads	8 / 4	4 / 2	Qヘッド / K・Vヘッド(GQA)
vocab_size	6,400	1,024	語彙
intermediate_size	ceil(768·π/64)·64 = 2,432	448	FFNの中間の幅
rope_theta	1e6	1e6	RoPEの回転周期の底

RMSNormでは、LayerNormが平均を引いて標準偏差で割るのに対し、RMSNormの論文は、平均を引く再中心化がなくてもよいとみて、二乗平均平方根(RMS)だけで割ります。MiniMindのコードも、x * rsqrt(mean(x²) + eps)に重みを掛けるだけです。そのため、出力のRMSは1に戻りますが、平均は0になりません。層の前で正規化するPre-Normなので、残差の経路は正規化を通らずにまっすぐつながり、深くなっても勾配がよく流れます。

GQAでは、Qヘッドが4個、K・Vヘッドが2個です。k_projの出力が2 × head_dimでQの半分で、計算するときにrepeat_kvでK・Vを2回ずつ複製して、Qヘッドの数に合わせます。推論のときにキャッシュに残るのは複製前のK・Vなので、トークンあたりのKVキャッシュが半分になります。GQAの論文は、K・Vヘッドを1個に減らしたMQAが品質を失うことを報告し、その中間の値を使うと、品質はMHAに近く、速度はMQAに近いと報告しました。

RoPEとQK-normでは、位置をベクトルに足さずに、qとkを位置の分だけ回転させます。2つのベクトルをそれぞれの位置で回してから内積を取ると、結果には2つの位置の差だけが残ります。(3, 7)と(103, 107)の内積が同じになるということです。MiniMindは、回転の前に、qとkをそれぞれRMSNormでもう一度正規化します(q_norm・k_norm)。Qwen3が使っている方式で、アテンションスコアが大きくなりすぎるのを防ぎます。

SwiGLUでは、FFNはdown(silu(gate(x)) * up(x))です。ゲートのあるGLU系は行列が3つなので、MiniMindは、中間の幅をhiddenのおよそπ倍(64の倍数に切り上げ)にします。GLUの変種の論文は、このようなゲート付きのFFNが、ReLU・GELUのFFNより品質が高いと報告しました。

埋め込みの共有では、tie_word_embeddings=Trueなので、lm_head.weightがそのままembed_tokens.weightです。小さなモデルで、語彙が占める分を半分に減らす仕組みです。

MoEの変種では、use_moe=Trueのとき、FFNがエキスパート4個になり、トークンごとに1個だけを使います。MiniMind-3-MoEが「198M-A64M」である理由です。パラメーターは3倍を超えて増えますが、1トークンが使う計算は、密なモデルとほぼ同じです。

現場での姿

モデルカードの設定ファイル(config.json)1枚から、パラメーター数・KVキャッシュの大きさ・サービング用のメモリを計算できなければなりません。「KVヘッド8、head_dim 128、32層」なら、トークン1つが何KBになるかがすぐに出てきて、その数字が同時ユーザー数を決めます。逆に、式で数えたパラメーターがモデルと違うなら、共有された重みや正規化の重みを数え落としています。学習を始める前に、損失がln(語彙)の近くかを見るのも、よくある確認です。初期化が壊れていれば、最初の学習ステップからわかります。

MiniMindのREADMEは、小さなモデルで幅と深さをどう配分するかについて、MobileLLMの論文を引用しています。同じパラメーター数なら、浅くて広いものより、狭くて深いほうが概して優れていますが、次元が512を下回ると、表現のボトルネックがはっきりしてくると書いています。このコースの128次元は、その境界よりはるかに下です。2つのCPUで数分のうちに学習するためにあえて選んだ大きさで、その代償は、最後のモジュールで測ります。

次のラボですること

小さな設定でMiniMindのモデルを作り、パラメーターを式とモデルの両方で数えます。GQAの射影の大きさとトークンあたりのKVキャッシュ、RMSNormの出力、RoPEの相対位置の性質、学習前の損失、MoEに変えたときの全体・アクティブなパラメーターを、順に確認します。