Soft-MTP、マルチトークン予測に潜むノイズ勾配を除去

導入

マルチトークン予測（MTP）は、言語モデルの事前学習で広く使われる手法になっている。通常の次トークンだけでなく、より先にあるトークンも予測させることで、下流タスクの性能向上や、投機的デコーディング向けの強いドラフト予測を期待できる。一方で、予測ヘッドを追加するだけでは十分ではない。共有されるモデル本体に、どのような勾配を返すかが学習の質を左右する。

標準MTPの課題

今回の研究は、一般的なMTPのクロスエントロピー勾配に隠れた問題を示した。勾配は厳密に二つへ分解できる。一つはKLダイバージェンスに相当する自己蒸留信号で、各MTPヘッドをモデル自身の次トークン予測分布へ近づける。これはドラフト側と本体モデルの分布を整合させるため、投機的デコーディングに役立つ。もう一つは、信号対雑音比の低い残差である。

この残差は単なる余分な項ではない。共有トランクへノイズを伝え、次トークン学習とマルチトークン学習の双方に干渉する可能性がある。さらに、モデル規模が大きくなるほど、その干渉は強くなると研究は説明している。標準MTPの損失をそのまま加えるだけでは、有用な監督信号を十分に活用できない可能性がある。

Soft-MTPの仕組み

研究者は「KL勾配復元」手法を提案した。NTPとMTPは線形出力ヘッドを共有するため、既存のクロスエントロピー勾配を組み合わせれば、KL項に対応する勾配を正確に復元できる。直通勾配ルーティングを使うことで、語彙サイズのロジット行列を明示的に作らずに、自己蒸留信号だけを残せる。

Soft-MTPは標準MTPへの一行程度の変更として実装できる。既存のNTP・MTPクロスエントロピー計算を再利用し、融合カーネルにも対応する。追加パラメータはなく、測定可能な計算オーバーヘッドも生じない。DeepSeek型とMeta型のMTPに対応し、167Mから1.2Bパラメータまでの評価では標準MTPを上回ったと報告されている。

意義と今後の見方

この研究の重要な示唆は、補助的な学習目標では容量を増やすだけでなく、勾配を選別することが有効かもしれないという点にある。MTPの主な役割が、将来トークン用ヘッドを本体モデルの分布へ整合させることなら、損失の全成分を無差別に逆伝播するより、KL信号に集中する方が合理的だ。

実装面でも導入しやすい。新しいアーキテクチャや追加パラメータ、最適化済みクロスエントロピー演算子の置き換えを必要としないからだ。既存のMTP事前学習パイプラインに対する軽量な目的関数の修正として検討できる。ただし、提供された素材には完全な実験表が含まれていないため、より大きなモデル、多数の予測ヘッド、追加タスクでの効果は論文本文で確認する必要がある。

出典：Hugging Face Daily Papers