モバイル向け言語モデル「MobileLLM」、On-Deviceに向けた10億未満のスケールの大規模言語モデルのアーキテクチャ最適化
3つの要点
✔️ 効率的な小規模モデルの重要性：モバイルデバイス向けに、クラウドコストとレイテンシーを考慮した10億未満のパラメータを持つ大規模言語モデルの開発に注力し、データとパラメータの量よりもモデルアーキテクチャの設計が性能向上に重要であることを提示。
✔️ モデル改良と精度向上：SwiGLU FFN、深くて細いアーキテクチャ、埋め込み共有、グループ化されたクエリアテンションといった4つのモデル設計技術を用いて、強力な基盤モデル「MobileLLM」を開発。
✔️ オンデバイス利用の最適化：MobileLLMは、チャットやAPI呼び出しなどのオンデバイスタスクで、既存の小規模モデルよりも顕著に優れた性能を示し、効率的なタスク処理の可能性を示唆。
MobileLLM: Optimizing Sub-billion Parameter Language Models for On-Device Use Cases
written by Zechun Liu, Changsheng Zhao, Forrest Iandola, Chen Lai, Yuandong Tian, Igor Fedorov, Yunyang Xiong, Ernie Chang, Yangyang Shi, Raghuraman Krishnamoorthi, Liangzhen Lai, Vikas Chandra
(Submitted on 22 Feb 2024)
Comments: Published on arxiv.
Subjects: Machine Learning (cs.LG); Artificial Intelligence (cs.AI); Computation and Language (cs.CL)
code：
本記事で使用している画像は論文中のもの、紹介スライドのもの、またはそれを参考に作成したものを使用しております。
概要
大規模言語モデルは、コミュニケーション、仕事、さらには日々のエンターテインメントに至るまで、私たちの生活に深く浸透しています。ChatGPTやPerplexity AIのような、クラウドベースで運用される大規模言語モデルの製品は、その最前線です。特に、ChatGPT4のようなモデルは1兆を超えるパラメータと言われています。しかし、将来、大規模言語モデルが会話インターフェースから推薦システムまで、日常のあらゆる面で不可欠になっていくと想定すると、人々の日常時間の約5％がこれらに依存することになります。このような未来を想像すると、GPT-4を使用するだけで約1億台のH100 GPUが必要になり、エネルギー消費と二酸化炭素排出は膨大な環境負荷をもたらします。
この課題に対する解決策は、よりコンパクトなモデルへの移行にあります。たとえば、トークンあたりわずか0.035Jを消費する350M 8ビットモデルを使用すれば、iPhoneで一日中会話アプリを使用することが可能です。さらに、最新のiPhoneアプリがLLaMA 7Bモデルで3〜6トークン/秒の処理速度を達成しているのに対し、125Mモデルなら50トークン/秒の速度で動作することが可能です。このような観点から、10億未満のパラメーターの大規模言語モデルの設計と実装を目指した研究が必要とされています。
モバイル技術の進展とともに、DRAMの容量制限やエネルギー消費の問題が浮上しています。例えば、現在のiPhoneやGoogle Pixelなどのデバイスは、高速アプリケーションの実行に必要なDRAMを6GBから12GBと限定しており、アプリケーションはこの容量の10％以内を使用するべきとされています。これらの制約を考慮し、省エネかつ効率的な言語モデルの展開が、今後の研究開発の重要な方向性になると考えられます。
10億未満のスケールの大規模言語モデルの設計改良
10億未満のパラメータを持つ言語モデルの設計と改良に関する研究を紹介しています。具体的には、125Mおよび350Mのモデルを対象に、限られたリソースの中で如何にして効率的にモデルの性能を向上させるかを研究しています。特にOn-Deviceの使用例では、モデルサイズが重大な制約となるため、重みパラメータの効果的な割り当てがこれまで以上に重要になります。下図は、基準となる10億未満のパラメーターモデルから最新モデルへの進化の過程を表しています。
この論文では、4つのモデル設計技術を用いて、強力な基盤モデルである「MobileLLM」を開発しています。SwiGLU FFN、深くて細いアーキテクチャ、埋め込み共有の再考、そしてグループ化されたクエリアテンションの活用が導入されています。これらの改善によって、モデルの性能を一貫して向上させています。さらに、即時ブロック単位のレイヤー共有手法（Immediate Block-Wiselayer Sharing Method）を開発しています。これによって、追加のメモリオーバーヘッドを発生させることなく、わずかな遅延で精度を向上させています。これにより、メモリ制約があるデバイス上でも高いパフォーマンスを実現できるモデル「MobileLLM-LS」を構築しています。
また、この論文では、32のA100 GPUを使用し、バッチサイズ32で120kイテレーションの実験を行い、その後、1Tトークンで480kイテレーションの学習を実施しています。これによって、ARC-easy、ARC-challenge、BoolQ、PIQA、SIQA、HellaSwag、OBQA、WinoGrandeなどのゼロショットの常識推論タスクや、TQA、RACEデータセットを使用した質問応答および読解タスクでのモデルの性能を評価しています。
また、この論文では、モデル設計に導入された技術について概観されています。
フィードフォワードネットワークの進化（SwiGLUの利用）：フィードフォワードネットワーク（FFN）における活性化関数としてSwiGLUを採用することで、小規模モデルの性能を向上させています。SwiGLUを使用することで、ゼロショットの推論タスクにおける125Mパラメータモデルの性能が顕著に向上させています。これは、従来のバニラFFNをSwiGLUに置き換えるだけで、平均パフォーマンスを改善させることができます。
アーキテクチャの深さ対幅（深さが重要）：トランスフォーマーモデルのアーキテクチャに関する一般的な見解に挑戦し、特に小規模モデルではアーキテクチャの深さが重要であることを明らかにしています。深くて細いモデルが、浅くて広いモデルよりも一貫して優れていることを実証しています。この発見は、小規模モデルのパフォーマンス向上においてアーキテクチャの設計が重要であることを示しています。
下図（a）および（b）は、ARC-easy、ARC-challenge、PIQA、HellaSwag、OBQA、WinoGrandeを含むほとんどのゼロショット推論タスクで、深いネットワークの優れたパフォーマンスを示しています。特に、この傾向はTQAおよびRACEデータセットでさらに顕著であり、下図（c）-（f）に示されています。この論文では、125Mのトランスフォーマーモデルにおいて、12層のものよりも30層や42層のモデルが著しく優れていることを示唆しています。
埋め込みの共有（効率と性能のバランス）：埋め込み層がパラメータ数の大きな割合を占めるため、入力-出力埋め込みの共有を再考しています。このアプローチにより、モデルのパラメータ数を大幅に削減しつつ、性能を維持できるようにしています。特に、30層の125Mモデルでは、パラメータ数を約11.8%削減できる一方で、平均精度のわずかな低下は、追加の層に割り当てることで容易に回復させています。
ヘッド数とKVヘッド数の最適化（グループ化されたクエリアテンション）：小規模トランスフォーマーモデルにおけるヘッドサイズの最適化を調査しています。グループ化されたクエリアテンション（GQA）の採用により、キー・バリュー（KV）ヘッドの数を効果的に削減し、モデルサイズを維持しながら精度を向上させています。
下図は、125Mおよび350Mモデルで望ましいヘッドサイズを決定する実験を行った結果です。16のクエリヘッドを使用すると最良の結果が得られることが示されています。さらに、KVヘッドの数を16から4に減らすと、125Mモデルで比較可能な精度が得られ、350Mモデルではほぼ10%のモデルサイズ削減で精度が0.2ポイント低下しています。これらの結果から、モデルアーキテクチャ設計のガイドラインとして機能しています。
この論文では、小規模言語モデルの効率と性能をさらに向上させる方法として、「レイヤー共有」の概念を検証しています。これは、前述の研究から導き出された深いレイヤーの利点を活用し、モデルストレージコストを増加させることなく、隠れたレイヤーの数を効果的に増やす戦略です。このアプローチは、特にストレージ容量が限られるデバイス上でのモデルの利用において重要な意味を持ちます。
追加のトランスフォーマーブロックを複製することで、アーキテクチャの変更やモデルサイズの増加なしに、モデルの精度を向上させています。この結果は、レイヤー共有が小規模モデルの設計において有効な手段であることを示しています。さらに、下図のようにベースラインと3つの異なる重み共有戦略を検証し、それぞれのパフォーマンスを比較しています。
なお、下図（a）はレイヤー共有を行わないベースラインモデル（baseline）、（b）は即時ブロック単位の共有（Immediate block-wise share）、（c）は全体繰り返しの共有（Repeat-all-over share）、（d）逆共有（Reverse share）の設計を表しています。トランスフォーマーブロックには、マルチヘッド自己注意（MHSA）とフィードフォワードネットワーク（FFN）が含まれます。全体繰り返し共有（Repeat-all-over share）は若干性能が高いものの、即時ブロック単位の共有はキャッシュを最も有効に活用しています。なぜなら、共有された重みがキャッシュに留まり、直ちに2回計算されるからです。
この比較から、レイヤー共有戦略が他の戦略よりも優れたパフォーマンスを提供することが明らかになっています。しかし、モデルの計算に用いられるSRAMの容量が通常約20MBに限られていることを考慮する必要があります。この制限に対処するために、共有重みをキャッシュに配置し、重みの転送を最小限に抑えることで、自己回帰推論の全体的な実行速度を向上させる戦略を採用しています。この研究の結果を踏まえ、即時ブロック単位の共有戦略を採用した新しいモデル「MobileLLM-LS」を提案しています。
実験
この論文では、ゼロショット常識推論タスク、質問応答と読解タスクの性能を比較しています。ベースラインメソッドの結果は、一貫した評価手順を確保にするため、オープンソースのHugging Faceモデルを使用して評価されています。
ゼロショット常識推論タスクの結果は、下表のようになっています。MobileLLMは、OPT、BLOOM、Galactica、Cerebras、GPT-neo、Pythia、RWKVといった初期のオープンソースの大規模言語モデルや最新のリリースを含む10億未満のパラメータモデルと比較して、顕著な改善を達成しています。
特に、125MモデルサイズのMobileLLMは、同じサイズの既存モデルを大幅に上回り、さらに、MobileLLM-125Mは、それぞれ22%、26%小さいながらも、Pythia-160MとRWKV-169Mよりも高い精度を達成しています。また、MobileLLM-LS125Mにレイヤー共有を取り入れることで、さらに精度を0.7ポイント向上させています。特筆すべきは、MobileLLM-LS-125Mが、350Mモデルの多くよりも同等あるいはそれ以上の成果を出したことです。350Mモデルサイズにおいては、MobileLLMは以前の最先端モデルを4ポイント以上上回る成果を示しています。
質問応答と読解タスクの結果は、下表のようになっています。TQAとRACEのベンチマークを用いて事前学習モデルの評価を行い、MobileLLM-125MはTQAベンチマークで前モデルを4.3ポイント以上上回る改善を示しています。
さらに、MobileLLM-350Mモデルは、他の350Mサイズのモデルと比較して約10ポイントの大幅なパフォーマンス向上を達成しています。読解タスクにおいても、MobileLLMモデルファミリーは、以前のモデルよりも顕著に高いスコアを示しています。
また、この論文では、10億未満のパラメータを持つ言語モデルが、特にチャットとAPI呼び出しなどのオンデバイスアプリケーションにおける有用性を検証しています。
チャットタスク：MobileLLMモデルは、チャットベースのタスクにおいて、一貫した条件下で既存の最先端モデルと比較して評価されています。結果は、下表の通りです。この評価は、AlpacaEvalとMT-Benchの2つのベンチマークを使用し、MobileLLMが10億パラメータを持つモデルをも上回る優れた性能を示すことが確認されています。特に、MobileLLM-LS-350Mは、GPT-3と比較して48.2%の高い勝率を達成しています。この結果は、MobileLLMの応答品質の高さを示しています。
API呼び出しタスク：API呼び出しでは、MobileLLMを使用して自然言語入力からAPIコマンドへの変換を行っています。例えば、「7時30分にアラームを設定して」というリクエストに対して、モデルは適切なAPIコールを生成し、ユーザーに確認の応答を提供します。5000サンプルのトレーニングセットと2500サンプルのテストセットで構成される合成データセットを用いたファインチューニングにより、MobileLLM-350Mは意図と構造の正確さにおいて優れたパフォーマンスを示しています。
また、0.25Tトークンで学習された125Mおよび350Mモデルサイズを持つMobileLLMとMobileLLM-LSモデルの両方に対して、トークンごとの最小最大ポストトレーニング量子化（PTQ）実験を実施しています。下図は、W8A8 PTQを採用することで、精度の低下が0.5ポイント未満であり、レイヤーシ共有と互換性があることを示しています。これは、レイヤー共有との互換性を保ちながら、効率的なモデル運用が可能であることを示しています。
さらに、125Mおよび350Mモデルの知識蒸留を試みましたが、この結果は、知識蒸留が必ずしも最適なアプローチではないことを示唆しています。他にもiPhone 13上でのMobileLLMモデルのレイテンシを評価しており、重み共有とレイヤー数の増加により、ローディングと初期化時間にわずかな増加が見られましたが、実行時間のオーバーヘッドは僅かであり、効率的なオンデバイス実行が可能であることが示しています。
まとめ
この論文では、クラウドコストの増大とレイテンシの問題を背景に、モバイルデバイス上で効果的に動作する大規模言語モデルへのニーズに応えるための研究を行っています。特に、10億パラメータ未満というモバイル展開に適した高品質な大規模言語モデルの開発に注力しています。
従来のデータ量とパラメータ量の多さがモデル品質を左右するという見解に挑戦し、10億未満のパラメータを持つ大規模言語モデルにおけるモデルアーキテクチャの重要性を明らかにしています。深さと細さを重視したアーキテクチャ、埋め込み共有、グループ化クエリアテンションメカニズムを組み合わせることで、MobileLLMと名付けられた強固なベースラインネットワークを構築し、これまでの最先端モデルに比べて顕著な精度向上を実現しています。
さらに、モデルサイズを増やさずに微小なレイテンシ増加で実現可能な即時ブロック単位重み共有手法を導入し、その結果として得られたMobileLLM-LSモデルはさらなる精度の向上を示しています。MobileLLMモデル群は、チャットベンチマークにおいて既存の10億未満のパラメータを持つ大規模言語モデルを大きく上回る成果を示し、API呼び出しタスクではLLaMA-v2 7Bと同等の性能を持つことを証明し、一般的なオンデバイス使用例での小規模モデルの優れた潜在能力を示しています。
このように、チャットとAPI呼び出しなどの典型的なオンデバイスにおけるケースシナリオにおいて、MobileLLMモデルの効果を実証し、これらのタスクを見事にこなす能力があることを示しています。これらの成果は、小規模モデルがいかにして高度なタスクを効率的に処理できるか、そしてそのポテンシャルの広がりを示しています。
この記事に関するカテゴリー
