LLMの効率化: Llama 3.1 8BからLlama-3.1-Minitron 4Bへのプルーニングと蒸留

Posted at 2024-08-17

はじめに：LLMの課題と小型化の重要性

大規模言語モデル（LLM）は自然言語処理の分野で革命を起こしていますが、その巨大なサイズゆえに課題も抱えています。Llama 3.1 405BやNVIDIA Nemotron-4 340Bのような大規模モデルは優れた性能を示す一方で、デプロイに多大なリソースを必要とします。

この課題に対応するため、業界では小規模言語モデル（SLM）の開発が進んでいます。SLMは多くの言語タスクで十分な性能を発揮しつつ、より少ないリソースで大規模展開が可能です。

本記事では、NVIDIAの研究者たちが開発した、構造化重みプルーニングと知識蒸留を組み合わせた効果的な手法を紹介します。この手法を用いて、Llama 3.1 8Bモデルから新たにLlama-3.1-Minitron 4Bモデルを作成する過程を詳しく解説していきます。

プルーニングは、モデルをより小さく効率的にする手法です。具体的には以下の2つの方法があります：

プルーニング後は通常、精度回復のための再トレーニングが行われます。

知識蒸留は、大規模な教師モデルから小規模な生徒モデルへ知識を転移する技術です。目的は、元のモデルの予測能力を維持しつつ、より効率的で軽量なモデルを作成することです。

知識蒸留には主に2つのアプローチがあります：

本記事では主に古典的知識蒸留アプローチに焦点を当てています。

Llama-3.1-Minitron 4Bの開発は、以下のステップで行われました：

まず、元のLlama 3.1 8Bモデルを94Bトークンのデータセットでファインチューニングしました。これにより、元のトレーニングデータとの分布の違いを補正し、蒸留時により適切な指導を行えるようにしました。

8Bから4Bへの圧縮では、16レイヤー（50%）を削減しました。各レイヤーの重要度を評価し、最終的に16〜31番目のレイヤーを削除することで最適な結果が得られました。

幅方向では以下の圧縮を行いました：

興味深いことに、プルーニング直後は幅プルーニングの方が言語モデル損失が高くなりましたが、短期間の再トレーニング後にはこの傾向が逆転しました。

Llama-3.1-Minitron 4Bの性能を、元のLlama 3.1 8Bや同サイズの他のモデル（Phi-2 2.7B、Gemma2 2.6B、Qwen2-1.5Bなど）と比較しました。

結果として、幅プルーニング戦略の有効性が再確認されました。Llama-3.1-Minitron 4B（幅プルーニング版）は、多くのベンチマークで元の8Bモデルに迫る性能を示し、同サイズの他のモデルと比較しても競争力のある結果を残しました。

NVIDIA TensorRT-LLMを用いて最適化したモデルのパフォーマンスを測定しました。

FP8精度での展開により、BF16と比較して全モデルで約1.3倍のパフォーマンス向上が見られました。

プルーニングと古典的知識蒸留を組み合わせた手法は、大規模言語モデルを段階的に小型化する上で非常に効果的であることが示されました。この方法は、ゼロからのトレーニングと比較して、より高い精度と効率性を実現します。

Llama-3.1-Minitron 4Bの開発は、最先端のオープンソースモデルであるLlama 3.1ファミリーを基にした初めての取り組みです。今後、この技術がさらに発展し、より効率的で高性能な小規模言語モデルの開発につながることが期待されます。

NVIDIAは、この研究成果を広く共有し、コミュニティの発展に貢献していく方針です。詳細な情報やモデルは、NVIDIAのGitHubリポジトリやHugging Faceで公開されています。

大規模言語モデルの効率化は、AI技術の民主化と持続可能な発展に向けた重要な一歩です。今後も、モデルの小型化と性能向上の両立に注目が集まるでしょう。