大規模言語モデル(LLM)のファインチューニング技術「LoRA」と生成型AI「Stable diffusion LoRA」
Hu et al., “LoRA: Low-Rank Adaptation of Large Language Models, ” In ICLR 2021.
Paper link: https://arxiv.org/abs/2106.09685
Github(Pytorch): https://github.com/microsoft/LoRA
この記事では、以下を紹介します：
- Microsoftが提供するLoRA技術により、大型言語モデルのファインチューニングのパラメータが大幅に削減できること。
- LoRA技術を使用する場合と使用しない場合のメモリ使用量の比較。
- LoRA技術の概要。
- LoRAをStable diffusionと組み合わせた場合のStable-diffusion-LoRA。
- Stable-diffusion-LoRAでは、LoRAのモデルを共有することで、自分のモデルをオンラインで共有し、パッチ/プラグインのアイデアを実現できます。
近年、大規模なデータセットや多数のパラメータを持つモデルを用いて、Large Language Models (LLMs)やFoundational Modelsとして知られる言語モデルが訓練されています。代表的な例としてGPT-3 (1750億パラメータ)があります。ChatGPTの登場も、一般的な問題においてLLMsが高い汎化能力を持つことを示しています。ただし、特定のドメインにおいては、コンテキストに沿った学習を数例 (few-shot)で達成できることがありますが、モデルをFine-tuningすることでより良い結果を得ることができます。
モデルがますます大きくなるにつれ、すべてのパラメータを直接Fine-tuningすることには重要なコストがかかります。そのため、最近では、研究者らはパラメータ効率の高いFine-tuning、つまりParameter-Efficient Fine-Tuning (PEFT)に注目しています。本記事では、Microsoftチームが提唱したLow-Rank Adaptation (LoRA)を紹介します。これは、事前学習されたモデル (例：GPT-3) のモデルを凍結し、小さなモデルでFine-tuningすることで、Adapterのコンセプトに類似した優れたFine-tuning結果を実現します。アイデアは、特定のレイヤーに挿入された小さなLoRAネットワークを使用して、モデルを異なるタスクに適応可能にすることです。
「AACL_2022_tutorial_PLMs」、以下の図に示すように、チュートリアルで詳細な説明が提供されています。
今後、大規模ニューラルネットワークモデルのパラメータを微調整する代わりに、より小さなモデルや重みを学習し、元のLLMの特定のレイヤーと組み合わせるアプローチが採用されるかもしれません。GPT-3モデルを微調整するよりも、この方法ではトレーニングパラメータが10,000倍少なく、GPU使用量も1/3になります。この技術はLLMにのみ適用されるだけでなく、Stable-Diffusion生成モデルなどの高解像度画像生成AIのトレーニングにも広く使用されています。
LoRA技術を使用する場合としない場合のメモリ使用量の比較は
以下の図に示されています。LoRAモデルを使用する場合、訓練パラメータの数（ Θ）はわずか0.5M〜11Mで、元のLLMモデル（ここではGPT-2 Mediumを使用し、パラメータ数は345M）よりもはるかに小さいことがわかります。さらに、LoRA技術を使用する場合、バッチサイズ=1の推論において、過去のAdapter技術よりも効率が良いことが示されています。
モデルのパラメータ/サイズがなぜ重要なのか？まず、モデルのトレーニング中にどの程度のGPUメモリが使用されるかを理解することが重要です。詳細については、Jacob Stern氏の「A comprehensive guide to memory usage in PyTorch」をご覧ください。
トレーニング時の最大モデル使用量(混合精度を考慮しない場合)：
- 計算式: モデルに必要なメモリ + Forward に必要なメモリ(柔軟性が高い) + Backwardに必要なメモリ(モデルに必要なトレーニングパラメーターのメモリ)+ Optimizer に必要な変数 * モデルに必要なトレーニングパラメーターのメモリ(通常は最大で、Adamはモデルに必要なメモリの2倍と見なすことができます)
- Forward 演算は、Batch size、入力コンテンツのサイズ、混合精度の有無によって決まります。また、PyTorchのCheckpoint機構を使用することで、この部分のメモリ消費を減らすことができます(調整可能です)。
- Optimizerに必要な変数は、異なるオプティマイザによって異なります(SGD：0、RMSProp：1、Adam：2)。一般的に使用されるAdamオプティマイザは、モデルの過去のEMAおよびモーメンタムのGradientを記録するため、Adamオプティマイザにとっては、2倍のモデルサイズのパラメーター数を格納する必要があります！
以下は、Adam Optimizer を使用し、混合精度を使用しない場合の推定値です(誤りがあれば指摘してください！)：
大型モデルのメモリ使用量が4グリッドを占めていると仮定すると、全体のモデルをトレーニングするために、Gradientもモデルサイズと同じ4グリッドのメモリを必要とします。そして、モデルをトレーニングするためには、Optimizer（Adam）が必要であり、Optimizer はまた、2つのモデルサイズのメモリが必要であるため、合計8グリッドを占有します。この場合、Forwardに必要なメモリは含まれていませんが、Forwardに必要なメモリもモデルサイズの4倍のメモリが必要になります。したがって、超大型モデルであるGPT-3の場合、175Bのモデルサイズを使用すると、必要なメモリ量は非常に大きくなります！！
実際には、ミキシングプレシジョン（混合精度）やいくつかのテクニックを組み合わせて使用することが多いですが、トレーニングに必要なメモリサイズはトレーニングするパラメーターの量に依存します。実際には、大型モデルを微調整する場合、通常はAdapterを使用して一部のパラメーターをトレーニングし、全体のモデルをトレーニングしない傾向があります。LLMのモデルウェイトをフリーズして、拡張モデルの一部だけをトレーニングすればよいため、この場合、私たちのオプティマイザーとBackwardに必要なストレージパラメーターの数は瞬時に大幅に減少します。以下の例を使用して説明します。LoRAモデルだけを使用する場合、通常、トレーニングするパラメーター（ Θ）の量はLLMの0.1％未満です。この場合、メモリ使用量は下の図のようになります。
実際には、トレーニング可能なパラメーターが非常に少ないため、さらに低いことになります。また、余分な演算コストも非常に低いです。
LoRAメソッド
LLMやFoundation Models（GPTシリーズなど）を様々なダウンストリームタスクに適用するために、モデル（Φ）をトレーニングする目的は、モデルが複数の異なるタスク（Z）を適切に処理できるようにすることでした。
以下の図は、GPT-1に使用されたダウンストリームタスクを示しており、分類、仮説検証、類似性比較、多肢選択問題などの一般的なNLPタスクが含まれています。モデルは、異なるプロンプトを入力として提供することでトレーニングされます。
過去、異なるタスクに対して行われたパラメーター効率的なファインチューニングの方法には2つのタイプがあります。
Adapter：モデルアーキテクチャの少量の追加とLLMモデルパラメータの凍結により、トレーニングが実行されます。
Prefixing：Prompt の前半部にトークンを追加することで、モデルが特定のタスクに対してより優れた結果を出すことができるようにします。
この記事で紹介されているLoRAはAdapterに属します。LoRAのコンセプトは、LLMが異なるタスクに適用されるため、モデルが異なるニューロン/フィーチャーを処理する必要があることを示しています。多数のフィーチャーの中から、特定のタスクに適したフィーチャーを見つけ、それらのフィーチャーを強化することができれば、特定のタスクに対してより良い結果を得ることができます。したがって、LLMモデル — Φを、もう一つのトレーニング可能なパラメータであるトレーナブルウェイト — Θ（ランク分解行列）と組み合わせることで、タスクの結果を最適化することができます。
右側のオレンジ色のモジュールは、トレーニングする必要がある LoRA モデルです，中間のランクrの制限を通じて、トレーニングするパラメータ数を大幅に減らし、特徴の次元数をr<<dにします。全体のパラメータ数は|Θ|=2×LoRA ×dmodel ×rとなります。ここで、LoRAはモデル全体で使用されるLoRAモジュールの数であり、論文ではTransformerアーキテクチャのAttentionにLoRAモジュールが挿入されています。また、rのサイズは異なるタスクに基づいて変化しますが、実験では通常、2〜4を使用すると良い結果が得られます。最終的に、タスクを最適化するためにLoRAモジュールを使用したモデルを作成したいと考えています。以下はその公式です。
実験
LoRAは、Fine-tuningの結果よりも多くの成果を上げ、トレーニングされるパラメータ数ははるかに少なくなっています。
他の効率的なFine-tuning方法と比較して、LoRAは最高の精度を獲得しています。
実験では、AttentionにLoRAモジュールを追加した場合の結果のみを評価し、パラメータ数を固定した状態で、Q、K、V、OブロックにLoRAを追加すると最良の効果が得られることを評価しました。
ランク数の選択。
LoRA コマンド
GitHub で確認することができます: https://github.com/microsoft/LoRA
現在、HuggingFace Parameter-Efficient Fine-Tuning（PEFT）に統合されています。
特定のレイヤーを LoRA に置き換える場合、モデルのアーキテクチャを調整する必要がありますが、呼び出しは非常に簡単です：
# ===== Before =====
# layer = nn.Linear(in_features, out_features)
# ===== After ======
import loralib as lora
# Add a pair of low-rank adaptation matrices with rank r=16
layer = lora.Linear(in_features, out_features, r=16)トレーニングの前に、元の LLM モデルを凍結し、LoRA のパラメーターのトレーニング可能であるように設定する必要があります。
import loralib as lora
model = BigModel()
# This sets requires_grad to False for all parameters without the string "lora_" in their names
lora.mark_only_lora_as_trainable(model)
# Training loop
for batch in dataloader:モデルを保存するときは、LoRA でトレーニングされたモデルのみを保存することができます。この機能により、自分のウェイト/パッチを共有するのが簡単になります。
# ===== Before =====
# torch.save(model.state_dict(), checkpoint_path)
# ===== After =====
torch.save(lora.lora_state_dict(model), checkpoint_path)LoRA または元の LLM のウェイトを読み込む場合、strict を False に設定する必要があります。
# Load the pretrained checkpoint first
model.load_state_dict(torch.load('ckpt_pretrained.pt'), strict=False)
# Then load the LoRA checkpoint
model.load_state_dict(torch.load('ckpt_lora.pt'), strict=False)LoRA 論文の概要：
LLMモデルを共有することは、将来のトレンドです。特定のタスクに適応する場合は、LoRAモジュールのみをトレーニングするだけで済み、交換性の利便性も提供されます。将来的には、LoRAのモデルを保存するだけで、タスクを共有したり切り替えたりできます。
訓練パラメータを大幅に減らすことにより、ハードウェアのトレーニングのハードルを大幅に下げ、完全なFine-tuningモデルと比較しても、推論速度の増加はほとんどありません。
欠点：LoRAは、モデルをトレーニングして、各タスクを最適化するためのものであり、1つのバッチに異なるタスクが含まれている場合、プログラム構造を調整して、異なる入力が対応するモジュールで推論されるようにする必要があります。
大型モデルをトレーニングするためのより少ない設備を使用するテクニックには、LoRA方法を含むParameter-Efficient Fine-Tuning（PEFT）の他に、MicrosoftのDeepSpeedにはZeRO-Offload：Democratizing Billion-Scale Model Trainingが統合され、単一のGPU上でGPT-3 175Bモデルを推論することができるFlexGenもあります。これらは、近年の技術であり、興味がある方は、詳細を確認することができます。
Stable-diffusion-LoRA(Low-rank Adaptation for Fast Text-to-Image Diffusion Fine-tuning)
近年、DALLE から Stable-diffusion まで、生成型 AI は高品質で高解像度の画像を生成することができることを示していますが、高解像度のモデルを訓練するには大量の計算リソースが必要であり、高解像度の拡散モデルを訓練するには相当量のメモリが必要です。Stable-diffusion は Pixel-level Diffusion を Latent Diffusion Model に変換することで訓練時のメモリ使用量を大幅に削減しましたが、まだ単一の 11 GB GPU 上で訓練することはできませんでした。しかし、今では LoRA 技術を Stable-diffusion に統合し、Stable Diffusion LoRA をリリースした人がいます！
Stable-diffusion に LoRA を統合することは、次の利点を直接もたらします：
- 訓練速度が大幅に向上します。
- 11GB VRAM 上で直接訓練できます。
- LoRA モデルの保存サイズが3MB〜200MBで、共有が簡単です。
HuggingFace は、LoRA チュートリアルの使用方法を共有しています。
この技術的な突破は、Stable Diffusion コミュニティに多数の生成モデルをもたらし、CivitAI ウェブサイトにモデルをアップロードすることができます。そこでは、LoRA を使用してトレーニングされた多数のモデルを見ることができます。
網路上には、Colabや個人PCを使用してモデルを生成/トレーニングするための多くのリソースがあります。最近、Stable diffusionのコミュニティは多くのプロジェクトをオープンソース化し、GUIインターフェイスを提供し、プログラムを理解する必要がなくなりました。
Stable-diffusionを試したい場合は、WebUIを使用することをお勧めします。公式のリリースモデルを使用できるだけでなく、CivitAIに直接リンクすることもでき、他の人の生成モデルを直接ダウンロードできます。
その他のリソース：
Stable-diffusion-webui Online Services
GitHub — AUTOMATIC1111/stable-diffusion-webui
GitHub — camenduru/stable-diffusion-webui-colab: stable diffusion webui colab
余談
最近は多くのAIツールが利用可能になっていますが、これらのツールを効果的に活用できるかどうかが、将来の生産性を決定することになります。ここでは、「Explainpaper」というウェブサイトをお勧めします。これは、調べたい論文を効率的に読むためのサイトで、理解できないテキストをハイライトするだけで、AIが記事全体に基づいて説明してくれます。私はたまにしか使わないし、説明がすべて良いわけではありませんが、トレンドです。論文はたくさんあり、AIに読ませるほうが速いです…
AIに取って代わられる前に、私の小さな知識の貢献が役立つことを願っています :)
参考文献：
Cheng-Han Chiang, Yung-Sung Chuang, Hung-yi Lee, “AACL_2022_tutorial_PLMs,” 2022.
同時に、Xiaoseanの個人ウェブサイトでも公開されました。
ご視聴いただきありがとうございました。
