マルチトークン予測（MTP、モデルに後続の複数のトークンを追加で予測させる手法）は、下流タスクの性能向上と推測デコーディングのドラフトヘッドとしての活用という両方の利点を持つ、一般的な事前学習のレシピとなっている。著者は、MTPの交差エントロピー勾配を正確に2つの部分に分解した。1つは、MTPヘッドをモデル自身の次トークン分布へ引き寄せるKL自己蒸留項（有用であり、推測デコーディングを改善する）であり、もう1つは、信号対雑音比が非常に低く、学習を妨げ、規模が大きくなるにつれて深刻化する残差項である。Soft-MTPは「ストレースルー勾配ルーティング」を用いてKL項のみを保持し、既存の交差エントロピー計算を再利用するため、語彙全体のロジット行列は不要であり、融合交差エントロピーカーネルとも互換性がある。
いいね 9 · コードのスター数 6
理由
- 変更が極めて小さい：標準的なMTPに対しては1行のコード変更で済み、パラメータは追加されず、著者によれば測定可能な追加オーバーヘッドはない。また、DeepSeek式およびMeta式のMTPの両方に対応している。
- 著者によれば、167Mから1.2Bまでのすべての規模において、次トークンのパープレキシティが低く、推測デコーディングが高速化され、ゼロショット精度が向上している。1.2Bモデルをファインチューニングすると、数学およびコード推論の性能が向上し、KLの直接最小化と比較しても、ほとんどの指標で劣らない。
- 再現可能：4つの規模 × 3つの目的関数（NTP、MTP、Soft-MTP）で合計12個のチェックポイントをHugging Faceに公開しており、訓練設定や推測デコーディングの評価スクリプトも含まれている。tests/ディレクトリでは、float64環境下で勾配が明示的なKLと一致することを検証できる。ライセンスはMIT。
使い方
- コード：git clone github.com/JaydenTeoh/Soft-MTP → pip install -r requirements.txt → python -m src.fetch_data --shards 5（事前トークン化されたFineWeb-Eduを使用。各設定で31.5Bトークンを訓練し、少なくとも4個のシャードが必要。アセンブリ時に約120 GBのディスク容量が必要）。MuonオプティマイザーにはPyTorch 2.9以上が必要。
- 訓練：torchrun --standalone --nproc_per_node=4 -m src.train --config configs/167m_fineweb.yaml --objective soft_mtp（比較対照としてhard_mtp / direct_kl / ntpに変更可能）。
- 既存のモデルを確認：AutoModelForCausalLM.from_pretrained("JaydenTeoh/soft-mtp-689m-model", trust_remote_code=True)。トークナイザーは元のgpt2を使用。推測デコーディングの評価にはpython -m src.eval_spec --ckpts …を使用。
- 自身の訓練コードへ移植する場合は、src/models/loss.py内のcombined_mtp_loss()を参照。
リスク
- 公開されているのは1.2B未満の、FineWeb-Eduでのみ訓練された研究用小規模モデルであり、そのまま対話モデルとして使用できるものではない。より大規模なモデルでも同様の効果が得られるかは検証されていない。
- 価値は主に事前学習段階にあり、恩恵を受けるにはゼロから訓練する必要がある。訓練および評価はデフォルトでWeights & Biasesに接続される。
- 要約では推測デコーディングの具体的な加速倍率が示されておらず、効果は著者による報告に基づく。READMEでは、ウォールクロックの高速化はGPUに依存するため、比較は同一マシン上で行う必要があると注意喚起されている。
開く · コード · Hugging Face論文ページ