
MaskLLM
NVIDIA Research · コーディング · その他
MaskLLMは、N:Mスパースのマスクを学習して大規模言語モデルを小さくするNVIDIA Researchのオープンソースなプルーニング手法だ。推論コストを抑えつつ、精度を元モデルの近くに保つ。

MaskLLM について
What Is MaskLLM
MaskLLMは大規模言語モデルを圧縮するための研究手法だ。モデルに不要なパラメータを取り除くプルーニングツールの一種で、半構造化スパース(N:Mスパースとも呼ばれる)を使う。考え方は単純で、M個の連続した重みのブロックごとにN個だけ残す。2:4のパターンなら4つの重みのうち2つを保つので、ハードウェアが動かすデータ量は半分になる。
MaskLLMの違いは、どの重みを残すかの決め方にある。従来の手法は固定のルールに頼っていた。小さな重みを捨てる、統計で推定する、といった具合だ。MaskLLMはマスクの選択そのものをモデルが学ぶ対象として扱う。N:Mパターンを確率分布としてモデル化し、Gumbel Softmaxというサンプリング手法で訓練するので、大規模なデータセット上でより良いマスクを見つけられる。マスクを学ぶ間、重みは凍結したままだ。この点が、特定タスクに対して2:4スパースを「ロスレス」にする理由である。
実用的な魅力は数字にある。NVIDIAによれば、密なLLaMA-2 7BモデルはWikitextでperplexity 5.12を記録する。従来のプルーニング手法はこれを10以上に押し上げた。MaskLLMは重みを凍結したままマスクだけを学び、6.72から6.78まで下げる。プロジェクトはNVIDIAライセンスの下でオープンソース化されており、研究者もエンジニアも自分で動かせる。ただし本物のGPUハードウェアと機械学習の知識は要る。
Getting Started
- Pythonパッケージを入れ、計算済みのマスクをHugging Faceから取得するスクリプトを実行する
- LLaMA-2 7BやLLaMA-3 8Bなど目的のモデルをPyTorchで読み込む
- 取得したマスクを適用し、付属のperplexityスクリプトで精度を測る
- 自分で訓練する場合はMegatron-LM環境を用意し、データを整えて自領域のマスクを新たに訓練する
AIツール情報
MaskLLMの料金、対応プラットフォーム、性能を簡単に確認できます。
こんな人におすすめ
このツールが最も力を発揮するユーザー、タスク、シーン。
ユーザー
- 機械学習の研究者
- モデルを配備するMLエンジニア
- 大学院生と研究者
- フレームワークやツールの開発者
タスク
- 大規模言語モデルのプルーニング
- 一発圧縮
- 領域特化のスパース
- スパースの転移
- プルーニング手法の比較
シーン
- 少ないGPUで7Bモデルを動かすチーム
- 論文の数字を再現する必要がある研究
- 全体は良いがニッチ領域で崩れるモデル
- モデルファミリー向けの圧縮パイプライン構築
- ライセンスやプライバシーの制約を検討
主な機能
学習可能なN:Mマスク
MaskLLMは各重みを手作業で採点する代わりに、マスクそのものを学ぶ。有効なN:Mパターンの集合を分布として扱い、訓練中にそこからサンプリングするので、モデルはデータセットに最適なマスクへ収束していく。この学習ステップが、マスクを一度選んで終わる一発方式との根本的な違いだ。
Gumbel Softmaxサンプリング
この手法は離散サンプリングの技法を借り、マスクの選択をエンドツーエンドで訓練可能にする。Gumbel Softmaxは、本来なら白黒はっきりした硬い選択に勾配を通す。だからこそ、N:Mマスク上の分布を推定ではなく実データで最適化できる。
Hugging Faceモデル向けの計算済みマスク
プロジェクトには既成のマスクが用意され、訓練なしで試せる。LLaMA-2 7B、LLaMA-3 8B、LLaMA-3.1 8B向けのマスクが公開され、それぞれC4のサブセットで訓練されている。スクリプトがHugging Faceから取得し、評価用に公式モデルへ適用する。
半構造化2:4スパース
MaskLLMが狙うのはN:Mスパース、多くの場合は4つの連続した重みに2つの非ゼロ値を残す2:4パターンだ。この形式が重要なのは、ハードウェア加速に対応するからである。小さくなったモデルは場所を取らないだけでなく、適切なGPUなら高速にも動く。
疎パターンのマスク転移
マスクを学ぶ同じ確率モデルが、マスクをある領域から別の領域へ運ぶ。このマスク転移は、マスク分布を一度訓練すれば、新しいタスクやデータセットへ移すときの手間がゼロから始めるより少ない。著者はこれを転移性と呼ぶ。モデルを複数の場面で鋭く保ちたいときに役立つ。
Megatron-LMの訓練パイプライン
本格的な訓練には、LLaMAモデルを複数GPUに分散するMegatron-LMのワークフローが付属する。リポジトリには、チェックポイントの変換、データのトークン化、SparseGPTやWanda、マグニチュードプルーニングによる事前マスクの生成、マスクの訓練、公開前のチェックポイント削減までをこなすスクリプトがある。
メリットとデメリット
メリット
- 固定の重要度ルールに頼らずデータからマスクを学ぶため、同じスパース率で精度が上がる
- コード、重み、計算済みマスクをオープンライセンスで公開し、結果を再現できる
- 転移学習に対応し、1つのマスク分布を複数のタスクと領域で使い回せる
- Hugging FaceとMegatron-LMを通じ、LLaMA-2やLLaMA-3、LLaMA-3.1など主要モデルで動く
- 重みを凍結するため、学んだ2:4スパースは調整したタスクに対してロスレスになる
デメリット
- 新しいマスクの訓練にはマルチGPU環境が要り、リポジトリは本格実行で1GPUあたり約40GBのメモリを示す
- 研究用の成果物であり製品ではないため、サポート窓口もホスティングサービスもAPIもない
- 既成マスクは少数のLLaMAモデル向けで、他のアーキテクチャは自分で訓練する必要がある
- 効果が出るのは2:4加速に対応したハードウェアで、旧型やNVIDIA以外のGPUでは恩恵が限られる
よくある質問
MaskLLMは、大規模言語モデルにN:Mスパースを加える学習型プルーニング手法だ。各重みを採点する代わりに、マスクを確率分布としてモデル化し、残す重みを学ぶ。NVIDIA ResearchがNeurIPS 2024で発表し、MaskLLMのコードはGitHubで公開されている。
関連コンテンツ
MaskLLMに関連するツール、スキル、記事を探す。
MaskLLMの代替ツール
Base Code
Base44 (Wix) · コーディング · 生産性向上Base CodeはBase44の共有クラウド開発環境だ。既存のGitHubリポジトリに対する変更をチャットで伝え、プレビューで確認し、プルリクエストとして返せる。
Google EmbeddingGemma 2
Google DeepMind · ライティング · 画像 · コーディングGoogle EmbeddingGemma 2 は 7 億 4000 万パラメータのオープンなモデルで、テキスト、コード、画像、音声、動画フレームを同じ 768 次元空間のベクトルに変え、スマートフォンやノート PC 上で動く。
GitGlow
GitGlow · コーディング · 生産性向上GitGlow は自分のコードと AI エージェントの変更を一つのレビューにまとめるデスクトップ Git クライアントだ。出荷前に、マシンから何が出ていくのかを確かめられる。
