[論文レビュー] Sheared LLaMA: Accelerating Language Model Pre-training via Structured Pruning

Mengzhou Xia, Tianyu Gao+2 more|arXiv (Cornell University)|Oct 10, 2023

Topic Modeling被引用数 20

ひとこと要約

論文は LLM-shearing を紹介する。これは事前学習済みのより大きなモデルをターゲットを絞った構造的プルーニングと継続的な事前学習で小型ながら競争力のある言語モデルを生み出す方法で、動的バッチロードを用いてデータをドメイン間でバランスさせる。

ABSTRACT

The popularity of LLaMA (Touvron et al., 2023a;b) and other recently emerged moderate-sized large language models (LLMs) highlights the potential of building smaller yet powerful LLMs. Regardless, the cost of training such models from scratch on trillions of tokens remains high. In this work, we study structured pruning as an effective means to develop smaller LLMs from pre-trained, larger models. Our approach employs two key techniques: (1) targeted structured pruning, which prunes a larger model to a specified target shape by removing layers, heads, and intermediate and hidden dimensions in an end-to-end manner, and (2) dynamic batch loading, which dynamically updates the composition of sampled data in each training batch based on varying losses across different domains. We demonstrate the efficacy of our approach by presenting the Sheared-LLaMA series, pruning the LLaMA2-7B model down to 1.3B and 2.7B parameters. Sheared-LLaMA models outperform state-of-the-art open-source models of equivalent sizes, such as Pythia, INCITE, OpenLLaMA and the concurrent TinyLlama models, on a wide range of downstream and instruction tuning evaluations, while requiring only 3% of compute compared to training such models from scratch. This work provides compelling evidence that leveraging existing LLMs with structured pruning is a far more cost-effective approach for building competitive small-scale LLMs

研究の動機と目的

• 既存の大規模事前学習モデルから、コスト効率の高い強力な中規模LLMを作成する動機づけ。
• 実証済みの事前学習構成と整合するターゲットアーキテクチャを生み出すプルーニング法の開発。
• 継続的な事前学習中にデータドメイン間の知識保持の不均衡に対処。
• ロス削減率に基づいてドメイン別にデータを割り当てる動的バッチロードを導入。

提案手法

• 層・隠れ次元・ヘッド・中間次元でのプルーニングマスクを用いて、ソースモデルを特定のターゲットアーキテクチャへ変形させるターゲット指向の構造プルーニングを提案。
• ハードコンプリート分布と最小-最大目的関数を用いた制約付き最適化により、ターゲット形状を満たすプルーニングマスクを学習。
• プルーニング後のモデルで継続的な事前学習を実施して性能を回復または上回る。
• ドメイン固有のロス削減率に基づき、トレーニング中のドメインデータ比率を調整する動的バッチロードを導入。
• ドメイン志向のデータ戦略に基づいて、プルーニングと継続的事前学習の二段階プロセスを用いる。

実験結果

リサーチクエスチョン

• RQ1 大規模な事前学習済みLLMを指定のターゲットアーキテクチャへプルーニングすることで、より少ない計算量で競争力のある小型LLMを得られるか？
• RQ2 プルーニング後の継続的事前学習は、単なるプルーニングより性能をより回復させるか？
• RQ3 動的バッチロードはドメイン別のロス削減をバランスさせ、全体的な下流性能を改善できるか？

主な発見

• Sheared-LLaMAモデル（1.3Bおよび2.7B）は、同等サイズの最先端のオープンソースモデルを、11個の下流タスクおよび指示調整ベンチマークで上回る。
• LLaMA2-7Bを1.3B/2.7Bへプルーニングする際、プルーニングと継続的事前学習には約50Bトークンしか使わず、スクラッチからの学習に必要な計算の約3%程度で競争力のある結果を達成。
• 動的バッチロードはドメイン間のロス削減を整合させ、より難易度の高いドメインのデータ使用量を増やすことで下流の性能を向上させる。
• プルーニングされたモデルは、同等サイズの既存モデルから開始するのよりも、継続的事前学習の初期化としてより良い場合がある。
• ターゲット指向の構造プルーニングは、同程度のスパース性において非一様プルーニングより高い推論スループットをもたらす。