複数サイズのネットワークを混ぜて拡散モデルの生成を速くする「Multilevel Euler–Maruyama」法
この論文は、拡散モデルの生成に使う確率微分方程式(SDE)や常微分方程式(ODE)を、計算コストの異なる複数の近似器を組み合わせて解く新しい方法を示します。研究者たちは「Multilevel Euler–Maruyama(ML-EM)」と名付けた手法で、計算の重い大きなネットワーク(UNet)をごく少数回だけ評価し、計算の軽い小さなネットワークを多数回使うことで全体の計算量を減らします。実験では顔画像データセット(CelebA)を64×64に縮小した場合で、最大で約4倍の速度向上を確認しました。論文は標準の方法と比べて、ある条件下で「多項式的な」速さの改善が可能だと主張します。
研究者たちが行ったことを簡単に言うと、微分方程式の中で使う「ドリフト項」(方程式を進める力に相当)を、精度とコストが段階的に上がる複数の近似器 f^1,...,f^k で置き換えました。各時間刻みでどの近似器を使うかは確率的に選びます(ベルヌーイ試行で選ぶ仕組み)。拡散モデルの文脈では、これらの近似器はサイズの違うUNetで、ML-EMは大きなUNetを少数回、 小さなUNetを多数回使うことを許します。基準手法はEuler–Maruyama(オイラー・マルヤマ)法で、これを基に計算量と誤差の関係を比較しています。
仕組みを高いレベルで説明すると、論文は「Harder than Monte Carlo(HTMC)」と呼ぶ状況を仮定します。これは、ドリフトを誤差εまで近似するのに必要な計算量がε^{−γ}で表され、ここでγ>2である場合を指します。この場合、従来のEuler–Maruyamaでは SDE の解を誤差 ε で得るのに計算量が約 ε^{−(γ+1)} になりますが、ML-EM は ε^{−γ} にまで改善できると理論的に示します。言い換えれば、誤差目標を達成するのに要する計算量が、大きなネットワークを一回評価するのと同じオーダーに落ちる可能性がある、ということです。手法の考え方は多重レベルのモンテカルロ(Multilevel Monte Carlo)法に似ており、各レベルの差分を確率的にサンプリングして合成します。
なぜこれが重要かというと、拡散モデルの計算コストは大きな深層ニューラルネットワーク(DNN)の評価回数で支配されます。評価回数を減らせば、生成にかかる時間とエネルギーが減ります。特にSDEでは高次の高速な刻み方(Runge–Kuttaに相当する手法)が使いにくく、刻み数を減らす工夫が重要です。論文の数値実験では、DDPM(Denoising Diffusion Probabilistic Models)にML-EMを適用して、同じ平均二乗誤差(MSE)を出すのに従来法より最大4倍高速であったと報告しています。DDIM(Denoising Diffusion Implicit Models)に対しても改善はあるが効果は小さめでした。実験では“1000ステップの真のサンプル”を参照としてMSEを測っています。
ただし重要な注意点もあります。まず、理論はHTMC(γ>2)という仮定に依存します。著者は画像タスクなどで経験的にスケーリング則が観察されると述べていますが、この係数γの値は問題やモデルにより変わります。実験で測った値は γ≈2.5 で、これは条件ぎりぎりを満たす範囲です。さらに、論文の保証は期待値での計算量に関するもので、実際の誤差にはベルヌーイ試行のサンプリングに起因するばらつきが見られます。実装面では著者が「best-of-15」(15回試して一番良いものを選ぶ)といった手続きを用いており、これが現実運用でどの程度必要かは追加の検討が必要です。最後に、実験は64×64の縮小画像での結果であり、著者は「より大きなネットワークではさらに大きな利得が期待できる」と述べていますが、それは理論的予測と経験的観察に基づく期待であり、ここで示された範囲を超える実際の効果はさらに検証が必要です。