プルーニングと知識蒸留によるコンパクト言語モデル

Compact Language Models via Pruning and Knowledge Distillation

July 19, 2024

著者: Saurav Muralidharan, Sharath Turuvekere Sreenivas, Raviraj Joshi, Marcin Chochowski, Mostofa Patwary, Mohammad Shoeybi, Bryan Catanzaro, Jan Kautz, Pavlo Molchanov

cs.AI

要旨

異なる展開規模やサイズを対象とした大規模言語モデル（LLM）は、現在、各バリアントをゼロからトレーニングすることで作成されていますが、これは非常に計算集約的です。本論文では、既存のLLMをプルーニングし、元のトレーニングデータの一部（3%未満）で再トレーニングすることが、繰り返しの完全な再トレーニングに代わる適切な代替手段となり得るかどうかを調査します。この目的のために、深さ、幅、アテンション、およびMLPのプルーニングを知識蒸留ベースの再トレーニングと組み合わせた、LLMのための実用的で効果的な圧縮ベストプラクティスを開発します。これらのベストプラクティスは、各軸に対するプルーニング戦略、軸を組み合わせる方法、蒸留戦略、および最適な圧縮アーキテクチャに到達するための検索技術に関する詳細な実証的探求を通じて導き出されました。このガイドを使用して、Nemotron-4ファミリーのLLMを2～4倍に圧縮し、さまざまな言語モデリングタスクにおいて、同サイズのモデルとその性能を比較します。既に事前トレーニングされた15Bモデルから、我々のアプローチを使用して8Bおよび4Bモデルを導出する場合、ゼロからトレーニングする場合と比較して、モデルごとに最大40倍少ないトレーニングトークンで済みます。これにより、モデルファミリー全体（15B、8B、および4B）のトレーニングにおいて、計算コストが1.8倍節約されます。Minitronモデルは、ゼロからトレーニングする場合と比較してMMLUスコアが最大16%向上し、Mistral 7B、Gemma 7B、Llama-3 8Bなどの他のコミュニティモデルと同等の性能を発揮し、文献にある最先端の圧縮技術を上回ります。我々は、Minitronモデルの重みをHuggingfaceでオープンソース化し、GitHubにはサンプルコードを含む補足資料を公開しています。

English

Large language models (LLMs) targeting different deployment scales and sizes are currently produced by training each variant from scratch; this is extremely compute-intensive. In this paper, we investigate if pruning an existing LLM and then re-training it with a fraction (<3%) of the original training data can be a suitable alternative to repeated, full retraining. To this end, we develop a set of practical and effective compression best practices for LLMs that combine depth, width, attention and MLP pruning with knowledge distillation-based retraining; we arrive at these best practices through a detailed empirical exploration of pruning strategies for each axis, methods to combine axes, distillation strategies, and search techniques for arriving at optimal compressed architectures. We use this guide to compress the Nemotron-4 family of LLMs by a factor of 2-4x, and compare their performance to similarly-sized models on a variety of language modeling tasks. Deriving 8B and 4B models from an already pretrained 15B model using our approach requires up to 40x fewer training tokens per model compared to training from scratch; this results in compute cost savings of 1.8x for training the full model family (15B, 8B, and 4B). Minitron models exhibit up to a 16% improvement in MMLU scores compared to training from scratch, perform comparably to other community models such as Mistral 7B, Gemma 7B and Llama-3 8B, and outperform state-of-the-art compression techniques from the literature. We have open-sourced Minitron model weights on Huggingface, with corresponding supplementary material including example code available on GitHub.