A Blog Entry on Bayesian Computation by an Applied Mathematician
$$
$$
1 離散雑音除去拡散模型 (D3PM) (Austin et al., 2021)
1.1 はじめに
拡散モデルは画像 (Ho et al., 2020), (Y. Song and Ermon, 2020) と音声 (Kong et al., 2021), (Chen et al., 2021) の分野で大成功を収めたが,そのアプローチは連続空間上の拡散過程の離散近似を実装するというものに限られていた.
最初の拡散モデルの提案 (Sohl-Dickstein et al., 2015) では,拡散過程の他に,\(2=\{0,1\}\) 空間上での二項分布を用いたモデルも考慮されていた.
離散データ(言語と画像 segmentation のラベル)に対するフローベースのサンプリング法として,Argmax Flows と Multinomial Diffusion が (Hoogeboom et al., 2021) により提案されたが,state-of-the-art に迫るモデルが得られるわけではなかった.
画像や音声も本質的には離散的であるとはいえ,積極的にそのようにみなす利点があるとは思う者は少なかっただろう.
しかし離散空間上では遷移核 \(P_\theta,Q\) の分布形を Gauss 分布に限るといった制約は必要なくなり,より柔軟なノイズ過程の設計が可能になる(第 1.4 節).
つまり,(Hoogeboom et al., 2021) はノイズ過程を空間一様な Markov 連鎖に限っていた点で離散拡散過程のポテンシャルを見誤っていたと言える.
(Austin et al., 2021) は例えばトークン間の類似性を考慮に入れる(第 1.5.5 節)など,遷移核の柔軟な設計を可能にするという,離散拡散過程のポテンシャルを示した.
(Austin et al., 2021) の提案した D3PM は,特に [MASK] トークンを特別に扱うことも可能になり,その場合は離散拡散モデルは BERT (Lewis et al., 2020) などのマスク付き言語モデルと等価になることも指摘している.その後,近年は (Arriola et al., 2025) など,自己回帰型のモデルと組み合わせたブロック拡散モデルを用いた,巨大な言語モデルの作成も試みられている.
1.2 設定
1.2.1 離散状態空間 \([K]\)
離散状態空間を \([K]:=\{1,\cdots,K\}\) (Kategorie) とし,この one-hot encoding の空間 \[ \mathcal{K}:=\{x\in\Delta^K\mid\exists_{i\in[K]}\;x_j=\delta_{ij}\} \] \[ \mathcal{P}([K])\simeq\Delta^K:=\left\{x\in[0,1]^K\,\middle|\,\sum_{i=1}^Kx_i=1\right\} \] と適宜同一視する.
\([K]\) は特定の位置にあるトークンの値域の集合,または特定のピクセルにおける色の値域の集合などの抽象化である.1
\([K]\) 上の確率分布は全てある確率ベクトル \(p\in\Delta^K\) に関する(カテゴリカル)分布 \(\mathrm{Cat}(p)\in\mathcal{P}([K])\) として表せる.2
1.2.2 \([K]^L\) 上の Markov 連鎖
トークン列の長さ \(L\ge1\) (Length) を固定し, \[ x_t=\begin{pmatrix}x_t^1\\\vdots\\x_t^L\end{pmatrix}\in E:=[K]^L \] 上の Markov 過程 \(\{X_t\}_{t\in[0,1]}\) を考える.
この遷移核 \(Q^{s,t}:[K]\to\mathcal{P}([K])\) は,第 \(i\) 行を \(x_s=i\) だった場合の次の時刻の状態 \(x_t\) の確率分布とする行列 \(Q^{s,t}\in M_K(\mathbb{R})\) と同一視して, \[ \biggr(X_t|X_s=x_s\biggl)\sim\mathrm{Cat}(x_sQ^{s,t}) \] と理解できる.
1.2.3 約束
以降,基本的にローマ字 \(x\) やギリシャ文字の小文字 \(\rho,\mu\) はベクトル(成分となるスカラーは上付き添字 \(x^j\) で表す)を表し,大文字 \(P,Q\) は(確率)行列を表す.
\(E=[K]^L\) 上の座標や関数は縦ベクトルで,確率測度は横ベクトルで表す.
one-hot encoding には不思議な性質があり,\(\mathcal{V}\)-値確率変数 \(X_t\) の平均をとると,ちょうどその分布の質量関数になる: \[ \operatorname{E}_\mathbb{Q}[X_t]=(\rho Q^{0,t})^\top. \]
1.3 枠組み
1.3.1 動的定式化
データ分布を \(\rho\in\mathcal{P}([K])\) とする.これを \(\{Q^{s,t}\}\) が定める Markov 過程によって輸送して得る終端の分布を \[ \mu:=\rho Q^{0,1}\in\mathcal{P}([K]) \] と定める.
離散化点 \(0=t_0<t_1<\cdots<t_N=1\) を与えた下で,確率核の族 \((Q^{s,t})_{s<t\in[0,1]}\) は,見本道の空間 \(E^{N+1}\) 上に確率測度 \(\mathbb{Q}\in\mathcal{P}(E^{N+1})\) を定める: \[ \mathbb{Q}(dx_{0:1}):=\bigotimes_{i=0}^N\rho Q^{0,t_i}(dx_{t_i}) \]
この測度 \(\mathbb{Q}\) と同じ分布を持つが,\(X_1\sim\mu\) から始まる時間反転過程を与える確率核 \(P^{s,t}:[K]\to\mathcal{P}([K])\) を,ニューラルネットワークで係数付けられたモデル \(\{P_\theta\}\) の最尤推定により学習することを考える.
変分推論の枠組みの下では,次の目的関数(損失関数)を最小にすることで達成される (Section 2.4 Sohl-Dickstein et al., 2015): \[ \mathcal{L}:=\operatorname{E}_{\mathbb{Q}}\left[\sum_{i=1}^N\operatorname{KL}\biggr(\mathcal{L}[X_{t_{i-1}}|X_{t_i},X_0]\,\bigg|\,\mathcal{L}[P_\theta^{t_i,t_{i-1}}X_{t_i}]\biggl)\right]. \tag{1}\] ただし,上式は読みやすさのために特別な約束をしているものと考える.というのも,\(i=1\) の場合,KL 乖離度の項は \(X_0\) のモデリングの良さを表す重要な項となっているはずであるが,\(X_0\sim\rho,X_{t_1}\sim\rho Q^{0,t_1}\) が与えられている下では次の表示を持つことに注意: \[ \operatorname{KL}\biggr(\mathcal{L}[X_0|X_{t_1},X_0]\,\bigg|\,\mathcal{L}[P_\theta^{t_1,0}X_{t_1}]\biggl)=\operatorname{KL}\biggr(\delta_{X_0}\,\bigg|\,\mathcal{L}[P_\theta^{t_1,0}X_{t_1}]\biggl) \] \(\delta_{X_0}\ll \mathcal{L}[P_\theta^{t_1,0}X_{t_1}]\) は成り立つとは限らない.その場合この項は通常 \(\infty\) と定義されるが,それでは学習できないので,KL 乖離度と定数分の違いである 交差エントロピー の値であるとする: \[ \operatorname{E}_\mathbb{Q}\left[\operatorname{KL}\biggr(\delta_{X_0}\,\bigg|\,\mathcal{L}[P_\theta^{t_1,0}X_{t_1}]\biggl)\right]=\operatorname{E}_{\rho Q^{0,t_1}}\left[\mathcal{H}\biggr(\rho\,\bigg|\,\mathcal{L}[P_\theta^{t_1,0}X_{t_1}]\biggl)\right]=\operatorname{E}_\mathbb{Q}\left[\log p_\theta(X_0|X_{t_1})\right]. \]
ただし,\(p_\theta(x_0|x_{t_1})\) は \(P_\theta^{t_1,0}\) の質量関数であるとした.
1.3.2 静的定式化
以上の枠組みは全て,Markov 過程の見本道の空間 \(E^{N+1}\) 上で定式化された.実際,損失関数 (1) は道測度間の距離 \(\operatorname{KL}(\mathbb{Q},\mathbb{P}_\theta)\) を測っていると理解できる.
一方で,\(N=2\) として,両端の分布 \(\rho,\mu\) しか考えなかった場合,\(\mathbb{P}_\theta\) の特定は一般に最適輸送問題として理解される.
拡散模型や動的 Schrödinger 問題などは,ステップ数 \(N\) を増やすことによって逆に問題を解きやすくしているのである.
1.4 ノイズ過程の設計
上述の理論を実践に移すためには,次の点が重要である
であるとする.条件2と3により, \[
L_{t-1}(x_0):=\int_\mathcal{X}\operatorname{KL}\biggr(q(x_{t-1}|x_t,x_0),p_\theta(x_{t-1}|x_t)\biggl)\,q(x_t|x_0)\,dx_t
\]
の Monte Carlo 近似が可能になる.
連続空間の場合は,遷移核 \(P_\theta,Q\) を Gauss に取ることが慣例になっている.この選択により \(q(x_{t-1}|x_t,x_0)\) も Gauss になるため,KL 乖離度は既知である.ただしこの選択は,\(\mu:=\rho Q^{0,1}\) が近似的にしか Gauss にならない点が問題になる.
しかし離散空間の場合,\(\mu\) を一様分布などに自由に設定しても,2と3はいつでも(効率的とは限らないかもしれないが)計算できる.3 遷移行列 \(Q^{s,t}\) をどのように難しい分布にとっても所詮は行列である.
また条件付き分布は,ベイズの定理より \[ q(x_{t-1}|x_t,x_0)=\frac{q(x_t|x_{t-1},x_0)q(x_{t-1}|x_0)}{q(x_t|x_0)} \] であるから,次のような行列演算によって計算できる: \[ \mathcal{L}[X_{t_{i-1}}|X_{t_i},X_0]=\mathrm{Cat}\left(\frac{X_{t_i}Q^{t_i,t_{i-1}}\odot X_0Q^{0,t_{i-1}}}{X_0Q^{0,t_i}X_{t_i}^\top}\right). \tag{2}\] ただし \(\odot\) は成分ごとの積(Hadamard 積)とした.\(Q^{t_i,t_{i-1}}=(Q^{t_{i-1},t_i})^\top\) という関係に注意.
1.5 核 \(Q\) の設計
(Section 3.1 Austin et al., 2021) では次の4種類の枠組みに触れている.
1.5.1 (空間)一様核 (uniform noise)
(Sohl-Dickstein et al., 2015) での \(K=2\) の場合での定式化を一般化する形で,(Hoogeboom et al., 2021) は次のような遷移核を扱った: \[ Q^t:=(1-\beta_t)I_K+\frac{\beta_t}{K}\boldsymbol{1}_K\boldsymbol{1}_K^\top. \] \[ \therefore\qquad Q^t_{ij}=\begin{cases} 1-\frac{K-1}{K}\beta_t&i=j,\\ \beta_t/K&i\ne j. \end{cases} \]
これは \(\{\beta_t\}\) の存在により時間一様ではないが,各状態を等価に扱っており,空間一様であると言える.その意味で \([K]\) 上に一才の事前構造を仮定していないと言える.
\(Q^t\) は対称な確率行列であり,平衡分布は一様分布である.
1.5.2 吸着付きの核 (absorbing-state noise)
トークンの空間 \([J]\) のうち1つ \(m\in[J]\) (MASK) を吸着点とすることを考える.
このときの \(Q_t\) は次のような形にする: \[ Q_t:=(1-\beta_t)I_K+\beta_t\boldsymbol{1}_Ke_m^\top \] \[ (Q_t)_{ij}:=\begin{cases}1&i=m=j,\\ 0&i=m\ne j,\\ 1-\beta_t&m\ne i=j,\\ \beta_t&i\ne m=j, \end{cases} \]
\(Q_t\) は \(\{\beta_t\}\) のスケジュールで \(\rho\) を \(\delta_m\) に収束させる.この逆過程を学習することで,[MASK] トークンの infilling が学習できる.
1.5.3 正規分布様核
拡散過程を参考にし,状態間の離散的な距離を考慮に入れた核を構成することも当然可能である.
\([K]\) を順序数の集合と見た距離を用いるため,画像の RGB 値などに適する.
\[ Q^t_{ij}=\begin{cases} Z^{-1}\exp\left(-\frac{4\lvert i-j\rvert^2}{(K-1)^2\beta_t}\right)&i\ne j,\\ 1-\sum_{l\ne i}Q^t_{ij}&i=j. \end{cases}\qquad Z:=\sum_{n=-(K-1)}^{K-1}\exp\left(-\frac{4n^2}{(K-1)^2\beta_t}\right) \] という核を用いると,\(Z\) の存在により \(Q^t\) は二重確率行列になっており,平衡分布は一様分布である.
1.5.4 帯行列核
他に \([K]\) の順序数としての構造を導入した核として \(k>0\) を定めて \(k\)-近傍に一様に移動する核 \[ Q^t_{ij}=\begin{cases} \frac{\beta_t}{K}&0<\lvert i-j\rvert\le k,\\ 1-\sum_{l\ne i}Q_{il}^t&i=j, \end{cases} \] を考えることもできる.
1.5.5 距離学習核
\([K]\) 内の自然な距離行列が,上流タスク,例えば言語の埋め込みで考えた \(k\)-最近傍隣接行列 \[ G_{ij}=1_{\mathcal{N}_k(i)}(j),\qquad\mathcal{N}_k(i):=\left\{\text{トークン}\;i\;\text{の}\;k\text{-近傍トークン}\right\} \] によって \[ A:=\frac{G+G^\top}{2}\in S_K(\mathbb{R})^+ \] などと与えられていることがある.
このとき, \[ Q_t:=e^{\alpha_tR}\qquad R_{ij}:=\begin{cases} A_{ij}&i\ne j,\\ -\sum_{l\ne i}A_{il}&i=j, \end{cases} \] により遷移確率行列を定めることができる.
\(G\) を連結グラフに取れば,\(Q_t\) はやはり二重確率行列であり,平衡分布は一様分布である.
この設計により,\(K\)-トークンの空間上の \(k\)-近傍グラフ \(G\) 上の乱歩により,一様な平衡分布 \(\mu=\mathrm{U}([K])\) に至るという描像がある.
1.6 エルゴード性について
1.6.1 概観
確率行列 \(\{Q^{s,t}\}\) に関する制約として,有限時間内に既知の分布 \(\mu\) に収束することが望ましい.
正確には,最後のステップに確定的な遷移を加え,これの時間反転過程も含めて学習することにしても,撹拌性(初期値 \(X_0=x_0\) を忘れる性質)を持つことが望ましい.
1.6.2 平衡分布
\(\{Q^{t_{i-1},t_i}\}\subset M_K((0,1))\) をすべて正行列に取れば既約性と非周期性が保たれる.4 \(E\) は有限であるため,これで直ちに指数エルゴード性が確保される.
仮に \(\{t_i\}_{i=1}^N\) は等間隔で,\(\{Q^{t_{i-1},t_i}\}\) を一様に \(Q^{t_{i-1},t_i}=Q^{t_i-t_{i-1}}\) と取った場合,平衡分布は \(Q^\top\) の固有値 \(1\) に属するただ1つの全ての成分が正の長さ1固有ベクトルになる.
仮に時間一様でなくとも,正行列の族 \(Q^{t_{i-1},t_i}\) をさらに二重確率行列に取れば,平衡分布 \(\mu\) は一様分布であることが確保される.
1.6.3 ノイズスケジュール
収束を速めるための一般的な慣行は,分散 \(\beta_t\) を時間と共に大きくしていくことである.
(Nichol and Dhariwal, 2021) は cosine schedule を導入しており,(Hoogeboom et al., 2021) や (Austin et al., 2021) の(空間)一様核でも用いられている.
1.6.4 相互情報量スケジュール
(Austin et al., 2021) は他にも,\(X_0,X_t\) の相互情報量を線型近似した \[ \beta_t:=\biggr(1-\frac{t}{T}\biggl)H(x_0) \] というスケジュールを提案している.
これは (Sohl-Dickstein et al., 2015) が \(K=2\) の場合に提案したスケジュールである \(\beta_t=(T-t+1)^{-1}\) という設定に,吸着核を用いた場合一致する.
1.7 除去過程
1ステップの遷移確率 \(p_\theta(x_{t-1}|x_t)\) (のロジット)をモデリングするのではなく,時刻 \(0\) への遷移確率 \(\widetilde{p}_\theta(x_0|x_t)\) (のロジット)をモデリングし, \[ p_\theta(x_{t-1}|x_t)\,\propto\,\sum_{\widetilde{x}_0\in[K]}q(x_{t-1}|x_t,\widetilde{x}_0)\widetilde{p}_\theta(\widetilde{x}_0|x_t) \tag{3}\] は間接的にモデリングする.これは (Hoogeboom et al., 2021), (Austin et al., 2021) だけでなく,(Ho et al., 2020) から続くモデリングの慣行である.
これは (2) で与えられる条件付き分布 \(q(x_{t-1}|x_t,x_0)\) と \(p_\theta(x_{t-1}|x_t)\) との距離を最小化するという (1) の目的関数 \(\mathcal{L}\) の構造をよく反映する.
例えば \(Q\) としてスパースな遷移行列を取った場合,\(q(x_t|x_{t-1})\) が非零である場合に限って \(q(x_{t-1}|x_t,x_0)\) も非零になるが,上述のパラメトライゼーションは自然に \(p_\theta(x_{t-1}|x_t)\) に同様のスパース構造を導入することになる.
またステップ数を小さく取った場合でも,\(k\) ステップをまとめて \(p_\theta(x_{t-k}|x_t)\) をいきなりサンプリングするということも可能になるという美点がある,
1.8 損失関数について
これは随分なスキャンダルであると思うのだが,(1) の目的関数 \(\mathcal{L}\) を実際に訓練に用いている場合は少ない.
そもそも (Ho et al., 2020) は各項の重要度を変えることで大胆な簡略化を行なった \(\mathcal{L}_{\text{simple}}\) を用いており,(Nichol and Dhariwal, 2021) は \[ \mathcal{L}_{\text{hybrid}}=\mathcal{L}_{\text{simple}}+\lambda\mathcal{L} \] としている.
(Austin et al., 2021) でも,パラメトライぜーション (3) の学習を加速するために,\(\widetilde{p}_\theta\) の交差エントロピー項を追加した \[ \mathcal{L}_\lambda:=\mathcal{L}+\lambda\operatorname{E}_\mathbb{Q}\biggl[-\log\widetilde{p}_\theta(X_0|X_t)\biggr] \] というものを用いている.この追加された項も,\(\mathcal{L}\) 内の KL 項と同様,\(\widetilde{p}_\theta(X_0|X_t)=\delta_{X_0}\) を達成した場合に最小化される.
2 拡散言語モデル
2.1 はじめに
離散拡散モデルは言語や DNA (Avdeyev et al., 2023) などの生物情報に盛んに応用されている.本稿では言語モデリングへの離散拡散モデルの応用をまとめる.
2.2 BERT (Devlin et al., 2019) は one-step 離散拡散モデルである
\(N=1,t_1=1\) とした one-step diffusion model を考え,さらに \(Q_t\) としては一様核と脱落核を重ね合わせたものを取る: \[ Q=\alpha\boldsymbol{1}_Ke_m^\top+\frac{\beta}{K}\boldsymbol{1}_K\boldsymbol{1}_K^\top+(1-\alpha-\beta)I_K. \]
これにより,各トークンを各ステップで \(\alpha=10\%\) でマスクし,\(\beta=5\%\) で一様にリサンプリングし,これを元に戻す逆過程を学習することになる.
これが定める目的関数 \(\mathcal{L}\) (1) は BERT (Devlin et al., 2019) と全く同じ目的関数になる.
MaskGIT (Masked Generative Image Transformer) (Chang et al., 2022) も,画像をベクトル量子化した後に,全く同様の要領でマスク・リサンプリングをし,これを回復しようとする.これはトランスフォーマーなどの自己回帰的モデルを用いて逐次的に生成するより,サンプリングがはるかに速くなるという.
2.3 自己回帰モデルは離散拡散モデルである
\(N\) をトークン列の長さとし,\(Q^{t_{i-1},t_i}\) は各段階で \(i\) 番目のトークンを [MASK] にする決定論的な変換とする.
詳細な対応は (Austin et al., 2021, pp. 5–6) を参照.
続いて,Masked Language-Models も,[MASK] を次々に剥がしていく逆過程と見ることができ,目的関数 \(\mathcal{L}\) は MLM 目的関数の reweight と見ることができるという (A.3 Austin et al., 2021).
3 参考文献
(Ryu, 2024) に素晴らしい教育的リポジトリがある.D3PM の 425 行での PyTorch での実装を提供している.
(Campbell et al., 2024) は最新の論文の一つである.
References
Footnotes
Block diffusion (Arriola et al., 2025) では複数のトークンをまとめてブロックとして考える.↩︎
(Section 3 Austin et al., 2021) に簡単な文献紹介がある.(Sohl-Dickstein et al., 2015) は \(K=2\) の場合,(Hoogeboom et al., 2021) は一様な遷移核 \(Q_{s,t}=Q_{t-s}\) のみを考えた.(J. Song et al., 2021) も supplementary で同様の考察をしている.↩︎
(A.4 Austin et al., 2021) では \(n\) 個の行列積として定まる \(Q^{0,t_n}\) を効率的に計算する方法を議論している.↩︎
これは (Markov, 1910)-(Dobrushin, 1956) 条件という.↩︎