StableDiffusionのLoRAでLoRA+を試す

はじめに

LoRA+と呼ばれる新しい手法が提案されました。
arxivの論文によると、最大で二倍の収束の高速化と1～2%の計算速度の改善が期待できるようです。

LoRAとLoRA+の簡単な解説

LoRAはウェイトをAとBの二つの行列に分解してから学習することでメモリ消費量を激減させています。
しかし、AとBで同じLRを使用しており、学習の効率が悪いとのことです。そこで、AとBで異なるLRを適用することで、効率改善を図ります。
これにより、最大で二倍の学習高速化と1～2%の計算速度の改善が期待できるようです。
すでにsd-scirptsにdevブランチで実装されています。近いうちにmainブランチにも来るでしょう。 sd-scriptsのSD3対応作業でしばらく来なさそう。

LoRA vs LoRA+

使用するには、sd-scriptsをdevブランチに切り替えてから、--network_argsに以下の引数を追加する。
・TEとUnetの両方指定: "loraplus_lr_ratio=倍率(整数)"
・Unetの個別指定: "loraplus_unet_lr_ratio=倍率(整数)"
・TEの個別指定: "loraplus_text_encoder_lr_ratio=倍率(整数)"

学習コマンド

accelerate launch --num_cpu_threads_per_process 1 sdxl_train_network.py --pretrained_model_name_or_path "Genimagine-XL-3.1.safetensors" --train_data_dir "DATASET DIR" --output_dir "OUTPUT DIR" --network_module "networks.lora" --xformers --gradient_checkpointing --persistent_data_loader_workers --cache_latents --cache_latents_to_disk --max_data_loader_n_workers 1 --enable_bucket --save_model_as "safetensors" --lr_scheduler_num_cycles 4 --mixed_precision "fp16" --learning_rate 0.0001 --resolution 1024 --train_batch_size 2 --max_train_epochs 14 --network_dim 4 --network_alpha 1 --shuffle_caption --keep_tokens 1 --optimizer_type "Lion" --lr_warmup_steps 100 --output_name "utaha_xl_loraplus" --save_precision "fp16" --lr_scheduler "cosine_with_restarts" --min_bucket_reso 512 --max_bucket_reso 2048 --caption_extension ".txt" --seed 42 --logging_dir "LOGGING DIR" --fp8_base --loraplus_lr_ratio 16 --save_every_n_epochs 2

LR Ratioはとりあえず推奨されていた16に設定。

lossの変化

青線がLoRA、赤線LoRA+です。

データセットが大きい分変化が小さいが、LoRA+のほうがわずかにlossが小さい。

出力の比較

minusが従来のLoRAでplusがLoRA+です。

Seed:10000

フルサイズの画像はこちら

Seed:20000

フルサイズの画像はこちら

わーお、LoRA+ってすっごく速いね！448ステップで再現できてるじゃん！
1344ステップがいい感じかな？先生。

普通のLoRAは2688ステップ以上、LoRA+は1344ステップ以上で個人的に満足のいく結果です。これはProdigy(大体2000ステップが最適)より速いです。

フルサイズの画像はこちら

フルサイズの画像はこちら

LoRAは1344ステップから再現し始めるのに対してLoRA+はわずか448ステップから再現し始めています。超高速。
LoRAは1792ステップ以上、LoRA+は896ステップ以上が良い感じです。

Ratioの違いを比較

minusが通常のLoRA、plus~~がLoRA+、wdがweight_decay有効

ratioが上がるにつれて学習が速くなるが品質も低下する？

ratio8以下の比較

Lion LR7e-05、データセットの数は256、consine_with_restart(num_scheduler_cycles=4)、warmup_steps=100

横軸はステップ数、縦軸はratioの値

ratio=4までは大きな変化無し。安定性を考慮して4が良さそう。

ratio=16でLRの違いを比較

画像最下段は参考としてTransformerの多い層のみで学習したもの、1epoch=300steps

LR0.0001で常用できる上限に見える。0.00015は高すぎる。
7e-05前後が最適?

まとめ

論文の主張通り、2倍以上の高速化が確認できました。しかも負荷が増加しないので非常に有用な技術です。
何度か試した感じratioは4が最適なようです。
ただ、学習が速い分ピーキーな印象を受けます。データセットによっては品質が微妙だったり以下の画像のように見切れやすくなったりします。結果が微妙なときは普通のLoRAでやった方がいいです。

試行錯誤をしていたら普通のLoRAの倍の時間がかかってしまった。そしてこのざま

LoRA+の注意事項

D-Adaptation系(Prodigy含む)Optimizerはパラメータグループごとに異なるLRを指定できない仕様によりLoRA+と併用できません。
でもLoRA+はDAdaptationより計算と収束が速いから問題なし！