Metaのマルチトークン予測によってAIモデルの速度が最大3倍向上

最近の研究で、Meta、エコール・デ・ポン・パリテック、パリサクレー大学の研究者たちは、AIの大規模言語モデル（LLM）の精度と速度を向上させる新しいアプローチを提案しました。このアプローチでは、モデルが複数のトークンを同時に予測できるようになります。この革新は、一度に一つのトークンを予測する従来の自己回帰モデルデザインに挑戦しています。

マルチトークン予測の利点

マルチトークン予測はすべてのLLMや言語タスクに適しているわけではありませんが、特定のシナリオにおいては重要な利点を提供します。たとえば、生成タスクを従来の方法よりも最大3倍速く加速することが可能です。この技術にはまだ改善の余地がありますが、一部のLLMアプリケーションにおいて強力なツールとなる潜在性があります。

次トークン予測の課題

従来のLLM訓練方法は「次トークン予測」と呼ばれ、この自己教師あり学習技術では、モデルにトークンのシーケンスを提示し、次のトークンを予測させます。その後、このトークンが次の入力に追加され、さらなる予測を行います。この反復的なプロセスを広範なテキストコーパスに適用することで、モデルは一貫性のあるテキストを生成する能力を学びます。

しかし、研究者たちは次トークン予測アプローチが言語処理、知識獲得、および推論能力の発展において限界があることを指摘しています。モデルが一度に一つのトークンにしか焦点を当てないため、局所的なパターンに対して過度に敏感になり、より広いコンテキストが必要な推論を見落とす可能性があるのです。さらに、次トークン予測は、流暢さを実現するために膨大なデータセットを必要とします。

Metaの最近の研究では、「複数の未来のトークンを同時に予測することは、サンプル効率を高める」と述べています。

マルチトークン予測の探究

対照的に、マルチトークン予測では、LLMがトレーニングデータの各位置で複数の未来のトークンを同時に予測します。研究者たちは、追加のトレーニング時間やメモリ要件を課さないシンプルなマルチトークン予測アーキテクチャを提案しています。

このモデルは、ほとんどのLLMの基礎となるTransformerアーキテクチャを基にしていますが、修正が加えられています。単一の出力を生成するのではなく、各トークン予測のために複数の独立した出力ヘッドを含めています。

マルチトークン予測の実装

推論時、モデルは通常の次トークン予測方法を各出力ヘッドに適用し、追加のヘッドを利用してデコーディングプロセスを効率化します。このフレームワークは、分野の先行研究を活用しています。「コスト効果が高くシンプルでありながら、マルチトークン予測は、高速で強力なTransformerモデルの訓練を大幅に向上させる」と研究者たちは述べています。

結果と観察

研究チームは、3億パラメータから130億パラメータのモデルでマルチトークン予測戦略を試しました。彼らの発見は顕著なパターンを示しています：小型モデルはマルチトークン予測の恩恵をあまり受けず、モデルサイズが大きくなるにつれてその効果が増大します。例えば、4トークン予測に訓練されたモデルは、MBPPコードベンチマークにおいて単一トークン予測に比べて数パーセントのパフォーマンス向上を示しました。

研究者たちは、「同じ計算リソースを使用して、大規模言語モデルのパフォーマンスを向上させることが可能です」と結論しています。また、マルチトークン予測は推論速度も向上させ、さまざまなバッチサイズにおいてモデルを最大3倍速くすることができます。「マルチトークン予測で事前訓練を行うと、次トークンモデルの単なるファインチューニングに比べて追加ヘッドの精度が向上し、自己推測デコードの可能性を最大限に引き出します」と彼らは説明しています。

研究はまた、マルチトークン予測がモデルに長期的なパターンを学習させることを促進し、特に「バイトレベルのトークン化」に関する実験では顕著でした。この場合、各バイトは単一のトークンとして扱われ、マルチバイト予測はベースラインの単一バイトモデルに対して大きく優れた性能を示しました。この結果は、事前に定義された語彙がないアプリケーションにおいて重要です。

今後の研究方向

マルチトークン予測は利点がある一方で、課題も存在します。予測するトークンの最適な数は、タスクやモデルサイズによって異なります。研究者たちは、最適なトークン数を自動的に特定する技術や語彙サイズとマルチトークン戦略の相互作用についての今後の研究の道を探っています。

この研究は、企業向けアプリケーションにおいても期待を持たせ、生成タスク（コード補完など）においてより迅速な推論速度と精度の向上を実現する可能性があります。既存のLLMアーキテクチャに大きな変更を加えなくても、Transformerフレームワーク内の他の最適化技術との互換性を確保しながら実現できるでしょう。