AI画像生成における学習 (その3)

AIの学習における問題点の理解と対策

はじめに

前回、ちょうど学習が開始するところまで終了しました。

今回はその学習結果をみながら、AIの学習における問題点と学習結果の評価について見ていきたいと思います。

1. 勾配爆発（発散）問題

勾配爆発（発散）問題は、発生すると作成したモデルがだめになってしまいます。

発散前のモデルを確認した結果、学習自体はまだする余地があるとした場合、これ以上学習がすすめられないことは問題になります。

前回で紹介した TensorBoard の出力を見れば、作成したモデルを確認しなくとも学習が発散しているのかどうかがわかるのでTensorBoardのグラフを確認しながら学習の状況を把握します。

機械学習では、目標とどれくらいの差があるのかを loss という形で評価します。どっちに向かって学習を進めていけば loss がより小さくなるかを傾きから得て、そちらに向かって1歩進んで、またその地点での傾きを見てどっちに進めば…と、繰り返して最終的に loss値が一番小さい地点を探します。

例えば上のようなパラメータの曲線があったとして、loss が最小の一番下の地点を目指すとします。

ランダムなスタート地点(今回は左上)からスタートして、loss 値とその地点の勾配(gradient) を計算して少なくなる方向に向かって( descent ) パラメータを修正して再度結果を見るということを繰り返していきます。学習率 (learning rate, lr) が 0.1 の場合は20 step くらいでだいたい一番下の地点で収束しました。

learning rate をもっと大きくするとどうでしょう。例えば 0.2 にすると

lr = 0.2 にすると 10 step 程度で一番下の地点に収束しました。

では、もっと調子に乗って lr = 1 ならどうでしょう。

一番下に収束するどころか、左と右を行ったり来たりするばかりで一向に収束しません。

これが更に大きくなると…

収束するどころか、どんどん値が大きくなってしまいました。これが「発散」している状態です。

学習率( lr ) が小さいと、発散はしないが収束するまでに時間がかかる。 lr が大きすぎると学習が進まなくなったり、発散したりするということがわかると思います。

2. スケジューラ

lr をいくつにすればよいかについては、元のモデルと学習する対象がどれだけパラメータが違っているかによるために、一概に「この値が best」といえる値がとれません。小さくしておけば発散することはないので安全ですが、学習に時間がかかってしまいます。

そこで、スケジューラの登場です。最初は比較的大きな値にしておき、学習が進むにつれてだんだん小さい値を lr に設定することで学習にかかる時間を短くしながら発散も防ぎます。

画像生成AIの学習で使われるスケジューラは以下のようなものがあります。

constant 常に一定の値を使います。あまり使われません。
linear 開始時 lr = lr で最終 epoch で lr = 0 になるように落ちる。
cosine 序盤ゆっくり減少、中盤で傾きがピーク、後半またゆっくり減少しまう。よく使われます。
cosine with restart cosine と同じように減少、0になったらまたスタート lr に戻る動作を複数回繰り返す　どのくらいの step / epoch が最適かわからない場合が発散を防ぐためによく使われます。