Matryoshka、3種類のモデルを単一の重みに内包し、学習計算量を36％削減

コーネル大学の研究者らは、異なるサイズの言語モデルを、互いに無関係な3つのcheckpointではなく、層ごとに拡張される単一の重みセットとして構成するMatryoshka Language Model Suitesを提案した。実験アーキテクチャでは、まず24層・500Mパラメータのサブモデルを実行し、続いて幅の広い10層を追加して1.5Bモデルを形成し、最後にさらに5層を加えて3Bモデルを構成する。幅の異なる部分は、パラメータを増やさないjunctionで接続される。具体的には、前段サブモデルの出力ノルムを新しいembeddingに合わせた後、追加チャネルを連結する。

この設計により、スイート全体のパラメータ数は3.2Bに抑えられ、3つの独立モデルの合計5.2Bと比べて38％少ない。1回のforward passで各出口のlogitsが同時に生成されるため、最大モデルは各ステップで小規模モデルを直接蒸留でき、教師出力を別途保存したり、教師モデルを再実行したりする必要がない。チームはFineWeb-Eduの350億tokenを用いてゼロから学習した。7つのzero-shot多肢選択タスクでは、各サイズの平均スコアと、同じtoken数で学習した独立モデルとの差が0.5ポイント以内に収まり、総学習FLOPsは36％減少した。1.5Bおよび3Bの出口では、out-of-domain perplexityもわずかに低下した。

推論時の要点は、単に読み込む重みが少ないことではなく、draftモデルとverifierモデルが前段の層とそのKVキャッシュを共有する点にある。従来の独立モデルによる500M／3B構成では、draftモデルが大きすぎるため生成速度が低下するが、入れ子型では自身の標準デコーディングと比べて約20％から40％高速化できる。テスト設定全体での総合的なスループット向上率は14％から26％だった。著者らはTransformers互換の重みと複数の学習段階を公開しているが、これは英語版FineWeb-Eduで学習されたアラインメント未実施の基盤モデルであり、`trust_remote_code`を有効にする必要がある。また、数十Bから数百B規模のMoE、long context、量子化デプロイメントでも同程度の削減効果が得られるかは、まだ実証されていない。エンジニアリングチームは次の段階として、vLLMやSGLangが共有層を用いた投機的デコーディングをネイティブ実装するか、またマルチテナントサービスでメモリとレイテンシの改善効果を実際に回収できるかを見極める必要がある。