【LLM論文を読む＆試す】LLMの刈り込み(Pruning)：深いレイヤーの不合理な非効率性

 

huggingface.co

  • このモデルを実際に触ってみましたが、InstructモデルではなくBaseモデルを刈り込んだものなので正直なところ良し悪しはよくわかりませんでした。
    • 投稿によれば、Pruned-Llama3のMMLUスコアは76-77（4bitでテスト）で、元の79から若干低下しているものの、それでもMixtral 8x22Bと同等！の高いスコアを維持しているそうです。
  • ちょっと面白そうなので、このモデル作成手順の元ネタとされている論文を斜め読みしたうえで、モデル刈り込みを試してみたいと思います。

arxiv.org

  • これは2024年3月にMeta FAIR関係の研究者が投稿したarXiv論文です。

要旨

我々は、オープンウェイトの事前学習済みLLMのための単純なレイヤー刈り込み戦略を経験的に研究した。その結果、モデルのレイヤーの大部分（最大で半分）を除去した後でも、様々な質問応答ベンチマークでの性能の劣化を最小限に留められることを発見した。これらのモデルに刈り込みを行うにあたり、レイヤー間の類似性を考慮することによって除去するレイヤーの最適なブロックを特定した。この際パラメータ効率の良いファインチューニング（PEFT）手法、QLoRAを使用し、各実験を単一のA100 GPUで実行している。実用的な観点からは、これらの結果は、レイヤー刈り込み手法が他のPEFT戦略を補完し、一方ではファインチューニングの計算資源をさらに削減でき、他方では推論のメモリとレイテンシを改善できることを示唆している。科学的な観点からは、これらのLLMが層の削除に対して頑健であることは、現在の事前学習手法がネットワークの深いレイヤーのパラメータを適切に活用していないか、浅いレイヤーが知識の保存において重要な役割を果たしていることを示唆している。

メモ

  • Llama-2-70Bのようなベースモデルでは、最大で半分くらいのレイヤーを削除してもその性能の大部分を維持することができる。
  • というのも、刈り込むレイヤーを増やしていくと（単調に性能が低下するのではなく）一定の水準でMMLUスコアが急低下する現象が観察される。

  • これは、モデル性能にとって重要なパラメータが一部のレイヤーに偏っている（＝それ以外のレイヤーは削っても性能に影響しにくい）ことを示唆する。
  • 具体的に検証すると、浅いレイヤーではレイヤー間の類似性が低く（無駄がなく）、より深いレイヤーで類似性が高く（冗長になる）ことがわかった。
  • したがって、単純に刈り込む戦略をとるよりも、深いレイヤーを優先して刈り込む戦略をとることで、モデルの性能低下を抑えながらサイズを縮小する道が開ける。
  • なおレイヤーを選択的に刈り込む場合も、モデルの損傷を修繕するための追加学習（Healing）を施すことで術後？の性能がより安定する。

PruneMe：論文の非公式実装

  • さて、MergeKitを手がけるarcee-aiが「PruneMe」というこの論文の非公式な実装を提供していて、冒頭のモデルもこれを利用して作っているそうです。

github.com

  • PruneMeのRepoには「Mistral-7B-instruct-v0.2」を例に実際にレイヤー間の類似性を検証したグラフが掲載されています。
  • 元論文で主張されているとおり、前半の浅いレイヤーは類似性が低く、後半の深いレイヤーで類似性が高くなる様子がプロットされています。

PruneMeによるモデル刈り込みを試す

  • ではGoogle Colabで簡単に試してみます。Repoをクローンし、依存関係をインストールします（TransformersはColabにプリインストール済み）。
# レポジトリのクローン
!git clone https://github.com/arcee-ai/PruneMe

# 依存関係
!pip install bitsandbytes datasets accelerate
  • ここではMistral 7Bベースの日本語チャットモデルを使って、レイヤー類似性を測定します。
  • とりあえずサンプルと同じ英語データセットを使って測定しています（レイヤーの類似性が判定できればよく日本語を使う必要はない気がするので）。
  • 「--layers_to_skip」で、削除したいレイヤーの数を指定します。ここでは32レイヤーのうち1/4の8レイヤーを削除する想定です。
  • なお計算時間の問題で、データセットサイズはサンプルより大幅に減らしています。
!python ./compute_block_similarity/layer_similarity.py --model_path "TFMC/Japanese-Starling-ChatV-7B" \
                      --dataset "arcee-ai/sec-data-mini" \
                      --dataset_column "text" \
                      --batch_size 8 \
                      --max_length 1024 \
                      --layers_to_skip 8 \
                      --dataset_size 200 \
                      --dataset_subset "train"
  • 計算が終わると、CSVファイルで結果が出力されます（上記設定で、標準GPUで20分ほど）。測定結果は以下の通りでした。

  • サンプルとよく似た結果が得られました。21-29のブロックが最も冗長なようです。同じMistral 7Bベースのモデルなので当然といえば当然かもしれません。
    • いずれにしてもレイヤーごとの冗長性は追加学習では特に変わらない（アーキテクチャが同じなら同じ）可能性がありそうです。
  • あとは、この特定したレイヤーブロックを削除するため、MergeKitを導入します。
# MergeKitの導入
!git clone https://github.com/cg123/mergekit.git
%cd mergekit
!pip install -e .
  • 次に「./PruneMe/slice_with_mergekit/slice.yaml」を適宜編集します。MergeKitのpassthorough手法（いわゆるフランケンマージを作る手法）を流用しているようです。
slice.yaml
--------------------------------------
slices:
  - sources:
      - model: TFMC/Japanese-Starling-ChatV-7B
        layer_range: [0, 21]
  - sources:
      - model: TFMC/Japanese-Starling-ChatV-7B
        layer_range: [29,32]
            
merge_method: passthrough
dtype: bfloat16
  • 最後にスライスを実行すると、刈り込みモデルが出力されます。
# 刈り込みの実行
!python /content/PruneMe/slice_with_mergekit/merge_me.py

刈り込みモデルのテスト

  • ということで、モデルの刈り込みができたので実際に出力してみます。以下が出力の一例です。
<s> [INST] <<SYS>>\nあなたは役立つアシスタントです。\n<</SYS>>\n\n空が青いのはどうして？ [/INST] 

空が青いのは、一般には空の中での色の相物が効嵑たぼいるにお効がて効ぼとしいともい職て効㩉効い〾くでいてのいい。も樾詖倴棧倴壯〾棧倴壬費塏壯〳ぞ度、度桴壯〷度眸竇疎ゾ度、度諷耡紐ゾお填ゾとしい度熱ぼこいでいてまい臦諓ぞ度、...
  • 残念ながら、モデルの出力は途中から破綻してしまいます。
  • よく考えてみると、冒頭のモデルも元論文で扱っている例もいずれもBaseモデルだけで、Instructモデルには触れていませんでした...
  • 「刈り込みしても性能が落ちにくい」というのは、あくまで選択式のMMLUスコアで測定されるようなBaseモデルの単純性能の話で、Instructモデルはもっと繊細なようです。
    • あるいは、指示追従にとって重要なパラメータが実は冗長な深いレイヤーに存在していたりもするのでしょうか？
  • ということで、基本的には「InstructモデルではなくBaseモデルに刈り込みして、あらためて指示追従の追加学習をやる」という手順が必要になりそうで、お手軽感はありません。
  • 既成のモデルを単にサイズ圧縮するのが目的なら、量子化の質にこだわったほうが期待できそうです。
  • とはいえ、例えばMoEモデルならもっと刈り込める余地がありそう？とか、色々と面白そうな手法ではあるので、有効な活用方法が見いだされてほしい気もします。