Diffusion MNIST その２

はじめに

その１で試したDiffusion MNISTについて、 ノイズを乗せるステップの細かさを粗くするとどうなるのかを見てみます。

方法

https://github.com/MarceloGennari/diffusion_mnist をいくつか変更することで粗さを変えていきます。

スケジュール変更

DiffusionProcessの引数に渡すvariance_scheduleを変えていきます。 デフォルトでは、

variance_schedule = torch.linspace(1e-4, 0.01, steps=1000)

となっています。これをパターンAでは

variance_schedule = torch.linspace(1e-4, 0.1, steps=100)

と、パターンBでは

variance_schedule = torch.linspace(1e-4, 0.999, steps=10)

とします。

それぞれのスケジュールを使ったときのalphaは

[パターン デフォルト]
[0.99990, 0.99989, 0.99988, ... , 0.99001, 0.99000]

[パターン A]
[0.99990, 0.99889, 0.99788, ... , 0.90101, 0.90000]

[パターン B]
[0.99990, 0.88891, 0.77792, ... , 0.11199, 0.00100]

となります。

alpha_barは

[パターン デフォルト]
[0.9999, 0.9998, 0.9997, ... , 0.0064, 0.0063]

[パターン A]
[0.9999, 0.9988, 0.9967, ... , 0.0062, 0.0056]

[パターン B]
[9.9990e-01, 8.8882e-01, 6.9143e-01, ... , 9.5131e-04, 9.5130e-07]

となります。ここで重要なことは、最初の時刻(ノイズが乗っていない)をt=0、最後の時刻(完全にノイズ)をt=1とするとき、alpha_barはt=0では1に近く、t=1では0に近くなるようにvariance_scheduleを決める必要があるということです。 各時刻tにおけるノイズの強さがalpha_barで決まり、t=1のときに完全にノイズになっていないと拡散プロセスの前提が崩れてしまうためです。

実際、パターンAの

variance_schedule = torch.linspace(1e-4, 0.1, steps=100)

を

variance_schedule = torch.linspace(1e-4, 0.01, steps=100)

に変えると、alpha_bar の値は

0.99990, 0.99970, 0.99940, ... , 0.60857, 0.60248

となりますが、この場合、数字の画像をうまく生成できません。

学習時のtの値

デフォルトではmain.pyの

t = torch.randint(0, 1000, (image.shape[0],))

の1000のところを、パターンAでは100に、パターンBでは10にします。

生成時のtの値

デフォルトではinference_unet.pyの

for t in trange(999, -1, -1):

の999のところを、パターンAでは99に、パターンBでは9にします。刻む数が少なくなると（ステップの細かさを粗くすると）、その分だけ生成時間を短くできます。

結果

デフォルトの設定ではこのようになります（その１の再掲）

t=500

t=0

時刻を100個に刻んだパターンAでも特に変わりなく生成できています。

t=50

t=0

時刻を10個に刻んだパターンBだと、多少ノイズが残ってしまいますが、生成できないというほどではありません。

t=5

t=0

まとめ

デフォルトの1000ステップではなくても、MNIST程度なら生成できることが分かりました。