より優れた高速な大規模言語モデル：マルチトークン予測によるアプローチ

Better & Faster Large Language Models via Multi-token Prediction

April 30, 2024

著者: Fabian Gloeckle, Badr Youbi Idrissi, Baptiste Rozière, David Lopez-Paz, Gabriel Synnaeve

cs.AI

要旨

GPTやLlamaなどの大規模言語モデルは、次のトークンを予測する損失関数を用いて学習されます。本研究では、言語モデルに複数の将来トークンを一度に予測させることで、サンプル効率が向上することを提案します。具体的には、学習コーパスの各位置において、共有されたモデルのトランク上で動作するn個の独立した出力ヘッドを使用して、続くn個のトークンを予測するようモデルに求めます。マルチトークン予測を補助的な学習タスクとして考慮することで、コードモデルと自然言語モデルの両方において、学習時間のオーバーヘッドなしに下流タスクの性能が向上することを確認しました。この手法は、モデルサイズが大きくなるほど有用性が増し、複数エポックにわたる学習においてもその魅力を維持します。特に、コーディングのような生成ベンチマークでは、我々のモデルが強力なベースラインを数パーセントポイント上回る一貫した性能を示します。13Bパラメータのモデルでは、HumanEvalで12%、MBPPで17%多くの問題を解決しました。小規模なアルゴリズムタスクでの実験では、マルチトークン予測が帰納ヘッドの開発とアルゴリズム的推論能力の向上に有利であることが示されました。追加の利点として、4トークン予測で学習されたモデルは、大規模なバッチサイズでも推論速度が最大3倍速くなりました。

English

Large language models such as GPT and Llama are trained with a next-token prediction loss. In this work, we suggest that training language models to predict multiple future tokens at once results in higher sample efficiency. More specifically, at each position in the training corpus, we ask the model to predict the following n tokens using n independent output heads, operating on top of a shared model trunk. Considering multi-token prediction as an auxiliary training task, we measure improved downstream capabilities with no overhead in training time for both code and natural language models. The method is increasingly useful for larger model sizes, and keeps its appeal when training for multiple epochs. Gains are especially pronounced on generative benchmarks like coding, where our models consistently outperform strong baselines by several percentage points. Our 13B parameter models solves 12 % more problems on HumanEval and 17 % more on MBPP than comparable next-token models. Experiments on small algorithmic tasks demonstrate that multi-token prediction is favorable for the development of induction heads and algorithmic reasoning capabilities. As an additional benefit, models trained with 4-token prediction are up to 3 times faster at inference, even with large batch sizes.