はじめに
初期の拡散モデルのサンプラーにおいて線形多段法が存在し、その一種であるAdams-Bashforth法が現れたりする。線形多段法は過去の勾配の推移から高次の現在の正確な勾配を算出する手法である。
しかしながら、StableDiffusion3や最近の動画生成AI(HunyuanVideo、Wan2.1やWan2.2)においてはTimestepが等間隔ではなく、shift関数を用いて最初は緩やかに、後半は大きくTimestepを動かしている。
以下のスケジューラーにおいてtimestepにshift関数を用いてるのにAdams-Bashforth法の係数が見えた。
if len(self.velocity_buffer) >= 4 and adaptive_order >= 4:
velocity = (
(55/24) * self.velocity_buffer[-1] -
(59/24) * self.velocity_buffer[-2] +
(37/24) * self.velocity_buffer[-3] -
(9/24) * self.velocity_buffer[-4]
)
elif len(self.velocity_buffer) >= 3 and adaptive_order >= 3:
velocity = (
(23/12) * self.velocity_buffer[-1] -
(16/12) * self.velocity_buffer[-2] +
(5/12) * self.velocity_buffer[-3]
)
elif len(self.velocity_buffer) >= 2 and adaptive_order >= 2:
velocity = 1.5 * self.velocity_buffer[-1] - 0.5 * self.velocity_buffer[-2]
else:
velocity = model_output
このように非等間隔のTimestepにおいてこのような旧来の線形多段法(Adams-Bashforth法)は成立するのかについて考えてみたい。
Adams-Bashforth法
StableDiffusionにおけるLMS(Linear MultiStep Method)
StableDiffusionにおいて初期にはLMSとかPLMSとか呼ばれる線形多段法が存在した。現在はこのサンプリング法はほとんど用いられない。
このサンプリング手法の係数は固定値であり、これはAdams-Bashforth法の係数である。
e_t = get_model_output(x, t)
if len(old_eps) == 0:
# Pseudo Improved Euler (2nd order)
x_prev, pred_x0 = get_x_prev_and_pred_x0(e_t, index)
e_t_next = get_model_output(x_prev, t_next)
e_t_prime = (e_t + e_t_next) / 2
elif len(old_eps) == 1:
# 2nd order Pseudo Linear Multistep (Adams-Bashforth)
e_t_prime = (3 * e_t - old_eps[-1]) / 2
elif len(old_eps) == 2:
# 3nd order Pseudo Linear Multistep (Adams-Bashforth)
e_t_prime = (23 * e_t - 16 * old_eps[-1] + 5 * old_eps[-2]) / 12
elif len(old_eps) >= 3:
# 4nd order Pseudo Linear Multistep (Adams-Bashforth)
e_t_prime = (55 * e_t - 59 * old_eps[-1] + 37 * old_eps[-2] - 9 * old_eps[-3]) / 24
等加速度運動
まず簡単な例を考えてみたい。
等加速度運動する物体があり、この物体の移動距離を計測する事は出来ず、1時間ごとに現在の物体の速度を測定する。過去の速度の測定値は80km/h→90km/h→100km/hと推移しており、おそらく加速度は10km/h^2であり、一時間後の速度は110km/hに到達するであろう。この時、t=0からt=1の進んだ距離は何km(厳密には積分を使う)で、これを一時間で進むにはどのくらいの速度であるか？
等加速度運動で知られるのは以下の式であるから$v_0=100km/h,a=10km/h^2$とすればt=0とt=1の距離$y$は105kmであり、これを一時間で進むには平均時速105kmが必要である。これは中学生でも計算できる。
y=v_0t+\frac{1}{2}at^2
a=v_t-v_{t-1}
さてこのt=0～1を進む推定速度$v$を現在(t=0)の速度$v_t=v_0$と一個前(t=-1)の速度$v_{t-1}=v_0-a$によって示せば
v=v_t+\frac{1}{2}(v_t-v_{t-1})=\frac{3}{2}(v_0)-\frac{1}{2}(v_{0}-a)
これが２次精度のAdams-Bashforth法である。
現在の速度100km/hと一時間前の速度90km/hから一時間後の移動距離105kmを一時間で移動する速度105km/hを計算するには以下の通りである。加速度aを求めずとも計算出来る。
$100*\frac{3}{2}-90*\frac{1}{2}=150-45=105$
等加加速度運動
次に加速度が線形で増加していく運動を考えたい。等加速度運動では$a=10km/h^2$で速度の増加量は等しかったが、この加速度も徐々に大きくなっていくとする。
さてここで速度の定義を以下のようにおいてみる。
v_t=v_0+at+b\frac{t(t+1)}{2}
$t=-2,-1,0,1$を代入したとき
$v_{t=-2}＝v_0-2a+b$,
$v_{t=-1}＝v_0-a$,
$v_{t=0}＝v_0$,
$v_{t=1}＝v_0+a+b$,
から上記の仮定に従えば$v_0=100,a=10,b=2$となる。
次にこの速度のt=0～1での積分を求めると
$\int_0^1v_tdt=[v_0t+\frac{1}{12}(6a+b(2t+3))t^2]^1_0=v_0+\frac{1}{12}(6a+5b)$となる。
t=0~1の距離は105.8333…kmとなる。これは中学生では計算できない。
一方で適当に係数を決定し、$\frac{23}{12}(v_0)-\frac{16}{12}(v_0-a)+\frac{5}{12}(v_0-2a+b)=v_0+\frac{1}{12}(6a+5b)$であるから３次精度のAdams-Bashforth法の係数が上述の等加加速度運動の係数が一致するように定義されているのが確認できる。
現在の速度100km/hと一時間前の速度90km/h、二時間前の速度82km/hから一時間後の移動距離yを一時間で移動する速度を計算するには以下の通りである。加速度aやbを求めずとも過去の速度の線形結合で定積分を計算出来る。
$100*\frac{23}{12}-90*\frac{16}{12}+82*\frac{5}{12}=191.6666..-120+34.166666..=105.83333..$
４次精度のAdams-Bashforth法
適当に以下の三次加速度$c$を定義し、この定積分を計算する。
v_t=v_0+at+b\frac{t(t+1)}{2}+c\frac{t(t+1)(t+2)}{6}
$\int_0^1v_tdt=[v_0t+\frac{1}{24}(12a+2b(2t+3)+c(t+2)^2)t^2]^1_0=v_0+\frac{1}{24}(12a+10b+9c)$となる。
$t=-3,-2,-1,0,1$を代入したとき
$v_{t=-3}＝v_0-3a+3b-c$,
$v_{t=-2}＝v_0-2a+b$,
$v_{t=-1}＝v_0-a$,
$v_{t=0}＝v_0$,
$v_{t=1}＝v_0+a+b+c$,
一方で適当に係数を決定し、$\frac{55}{24}(v_0)-\frac{59}{24}(v_0-a)+\frac{37}{24}(v_0-2a+b)-\frac{9}{24}(v_0-3a+3b-c)$を計算すれば$=v_0+\frac{1}{24}(12a+10b+9c)$であるから４次精度のAdams-Bashforth法の係数が上述の速度の定積分と一致するように定義されているのが確認できる。
係数算出
一方でAdams-Bashforth法の係数$z_0,z_1,z_2,z_3$とすればこの行列は以下より
\begin{pmatrix}
v_0 & a & b & c
\end{pmatrix}
\begin{pmatrix}
1 & 1 & 1 & 1 \\
0 & -1 &-2 & -3 \\
0 & 0 & 1 & 3\\
0 & 0 & 0 & -1
\end{pmatrix}
\begin{pmatrix}
z_0 \\ z_1 \\ z_2 \\ z_3
\end{pmatrix}=
\begin{pmatrix}
v_0 & a & b & c
\end{pmatrix}
\begin{pmatrix}
1 \\ \frac{12}{24} \\ \frac{10}{24} \\ \frac{9}{24}
\end{pmatrix}
この行列を解けば$z_0=2.29166667...=\frac{55}{24}$である４次精度のAdams-Bashforth法の係数が算出されるのが確認できる。
import numpy as np
a = np.array([[1,1,1,1],[0,-1,-2,-3],[0,0,1,3],[0,0,0,-1]])
b = np.array([1,12/24,10/24,9/24])
print(np.linalg.inv(a) @ b)
-------------------------------------
[ 2.29166667 -2.45833333 1.54166667 -0.375 ]
tが非等間隔の場合
ここからは、時間間隔が等しくない場合に同様の導出が可能かを検討する。
前述の等加速度運動の図からt=-1の点を削除するとどうなるだろうか？
この場合、二時間で速度が20km変わったという考慮は必要だがこれ自体は前述の加速度aの導出式を点の間隔である2で割ればいいだけで、以降は$v_t=v_0+at$のt=0~1の定積分という式は変わらない。
$v_t=v_0+at$
$\int_0^1v_tdt=[v_0t+\frac{1}{2}at^2]^1_0=v_0+\frac{1}{2}a$となる係数は
$t=-2,0,1$を代入したとき
$v_{t=-2}＝v_0-2a$,
$v_{t=0}＝v_0$,
$v_{t=1}＝v_0+a$,
$\frac{5}{4}(v_0)-\frac{1}{4}(v_0-2a)=v_0+\frac{1}{2}a$
次にt=3まで進むときt=0～3の積分を計算してみると
$\int_0^3v_tdt=[v_0t+\frac{1}{2}at^2]^3_0=3v_0+\frac{9}{2}a$となる。
この速度換算を二個の速度で求めるための係数は$\frac{21}{4}(v_0)-\frac{9}{4}(v_0-2a)=3v_0+\frac{9}{2}a$となる。この係数はAdams-Bashforth法の係数とはだいぶ違う。
一般化して(2次)
$t=-\sigma_1,0,\sigma$を取るとする。
v_t=v_0+at
$v_{t=-\sigma_1}＝v_0-\sigma_1a$,
$v_{t=0}＝v_0$,
$\int_0^\sigma v_tdt=[v_0t+\frac{1}{2}at^2]^\sigma_0=\sigma v_0+\frac{\sigma^2}{2}a$となる。
\begin{pmatrix}
v_0 & a
\end{pmatrix}
\begin{pmatrix}
1 & 1 \\
0 & -\sigma_1
\end{pmatrix}
\begin{pmatrix}
z_0 \\ z_1
\end{pmatrix}=
\begin{pmatrix}
v_0 & a
\end{pmatrix}
\begin{pmatrix}
\sigma \\ \frac{\sigma^2}{2}
\end{pmatrix}
を解けば以下の通り。$\sigma=3,\sigma_1=2$とすれば前述した$\frac{21}{4}$と$-\frac{9}{4}$が行列計算から求まる。
非等間隔の点ではAdams-Bashforth法の係数は有用ではない。
import numpy as np
s1= 2
s = 3
a = np.array([[1,1],[0,-s1]])
b = np.array([s,1/2*s*s])
print(np.linalg.inv(a) @ b)
-------------------------------------
[ 5.25 -2.25]
一般化して(4次)
$t=-\sigma_3, -\sigma_2,-\sigma_1,0,\sigma$を取るとする。
v_t=v_0+at+b\frac{t(t+\sigma_1)}{\sigma_2(\sigma_2-\sigma_1)}+c\frac{t(t+\sigma_1)(t+\sigma_2)}{\sigma_3(\sigma_3-\sigma_1)(\sigma_3-\sigma_2)}
$v_{t=-\sigma_{3}}＝v_0-\sigma_3 a+ \frac{\sigma_3(\sigma_3-\sigma_1)}{\sigma_2(\sigma_2-\sigma_1)}b-c$,
$v_{t=-\sigma_{2}}＝v_0-\sigma_2 a+ b$,
$v_{t=-\sigma_{1}}＝v_0-\sigma_1 a$,
$v_{t=0}＝v_0$,
ここで$\int_0^\sigma v_tdt$を求めたい。wolfram alphaによれば
従って任意の$\sigma_3, \sigma_2,\sigma_1,\sigma$が定義される場合、
import numpy as np
s = 0.755
s1 = 0.7
s2 = 1.5
s3 = 2.4
a = np.array([[1,1,1,1],
[0,-s1,-s2,-s3],
[0,0,1,s3*(s3-s1)/(s2*(s2-s1))],
[0,0,0,-1]])
b = np.array([s,s*s*6/12, s*s*(4*s+6*s1)/(s2*(s2-s1))/12, s*s/12*(3*s*s+4*s*(s1+s2)+6*s1*s2)/(s3*(s3-s1)*(s3-s2))])
print(np.linalg.inv(a) @ b)
s = 1
s1 = 1
s2 = 2
s3 = 3
a = np.array([[1,1,1,1],
[0,-s1,-s2,-s3],
[0,0,1,s3*(s3-s1)/(s2*(s2-s1))],
[0,0,0,-1]])
b = np.array([s,s*s*6/12, s*s*(4*s+6*s1)/(s2*(s2-s1))/12, s*s/12*(3*s*s+4*s*(s1+s2)+6*s1*s2)/(s3*(s3-s1)*(s3-s2))])
print(np.linalg.inv(a) @ b)
print(np.array([55/24, -59/24, 37/24, -9/24]))
------------------------------------------------------------
[ 1.76502389 -1.75079481 0.9303404 -0.18956947]
[ 2.29166667 -2.45833333 1.54166667 -0.375 ]
[ 2.29166667 -2.45833333 1.54166667 -0.375 ]
このようにして非等間隔における求めた係数はAdams-Bashforth法の4次係数$([ 2.29166667 -2.45833333 1.54166667 -0.375 ])$と異なるのが確認できる。
よって、timestepが非等間隔である場合、固定係数のAdams-Bashforth法は厳密には成立せず、この係数は非等間隔timestep間長さに依存して再計算される必要があると考えられる。
UniPC
Wanだと高速化loraによって推論stepは4～8stepにされるため推論stepが少なすぎて多段法が有用かどうかは逆に疑問である。
dpm++やunipcが非線形のTimestepにおける多段法の行列を解いて係数を求めているのかはよく分からない。
一応、逆行列を求める関数linalg.invは見当たらなかった。しかし、逆行列はガウスジョルダン法で計算できるので逆行列を求める関数が見当たらないのは理由にはならない。
以下の解説によれば行列を解いて係数を求めているように見える。
ただし、この式は$v_t$を以下のようにおいて$t=k_3h,k_2h,k_1h,-h$を代入したときの行列のように思う。これにtの値を代入すればVandermonde行列が出てくる。
v_t=v_0+\frac{at}{h}+\frac{bt^2}{h^2}+\frac{ct^3}{h^3}
ところで、数式に階乗が出てくるのは４次精度のAdams-Bashforth法の分母$(6=3!)$であったが、これは各点が1,2,3と等間隔だから出たのではないのか？
v_t=v_0+at+b\frac{t(t+1)}{2}+c\frac{t(t+1)(t+2)}{6}
非等間隔では
v_t=v_0+at+b\frac{t(t+\sigma_1)}{\sigma_2(\sigma_2-\sigma_1)}+c\frac{t(t+\sigma_1)(t+\sigma_2)}{\sigma_3(\sigma_3-\sigma_1)(\sigma_3-\sigma_2)}
となり階乗は出ない。
非線形多段法？
あくまで計算したのはtimestepが非等間隔の線形多段法であり、これを非線形多段法とは多分呼ばないと思う。
何故ならおそらく非線形とかいうのは$v_t=v_0+e^{-at}$のように係数aが線形結合で表せないのを示すからではないだろうか。
まとめ
Adams-Bashforth法の係数は等間隔のtimestepでは有用だが非等間隔のtimestepにおいては成り立たないのを示した。
参考：