拡散モデルの先読み勾配ガイダンス
NeurIPS 2024 に採択された論文 “Gradient Guidance for Diffusion Models: An Optimization Perspective” (Guo et al. 2024) を読んでみたので紹介します。 拡散モデルにおいて Classifier Guidance / Classifier-free Guidance などを用いたガイダンス手法が広く用いられているなか、この論文はユーザーが定義した任意の目的関数 $f$ を最大化するようにガイダンスを行うための手法について検討しています。 この論文では単純な勾配 $\nabla f$ を用いた勾配法による最適化では生成物が崩壊してしまうことを示し、“Look-ahead Loss” という損失を定義し、生成されたデータによるスコア関数のファインチューニングを繰り返しながら最小化するようなガイダンス手法を提案しています。 1. 背景と課題 拡散モデルにおいて、拡散過程を $\mathrm{d}\mathbf{x} = f( t)\mathbf{x}\mathrm{d}t + g(t)\mathrm{d}\mathbf{w}$ と表す場合の逆拡散過程は次のような確率微分方程式(SDE)で表されます。 ($\mathrm{d}t$: 無限小ステップ, $\bar{\mathbf{w}}$: 時刻 $T \rightarrow 0$ まで逆向きにたどった際の標準ウィーナー過程) $$\mathrm{d}\mathbf{x} = [f(t)\mathbf{x}-g(t)^2\nabla_\mathbf{x}\log p_t(\mathbf{x})]\mathrm{d}t + g(t) \mathrm{d}\bar{\mathbf{w}}$$論文では拡散過程・逆拡散過程を以下のように表記しています。 $$\mathrm{d}{X}_t = -\frac12 q(t){X}_t\mathrm{d}t + \sqrt{q(t)}\mathrm{d}{W}_t\tag{1}$$ $$\mathrm{d}X_t^{\leftarrow} = \left[ \frac{1}{2}X_t^{\leftarrow} + \nabla \log p_{T-t}(X_t^{\leftarrow}) \right] \mathrm{d}t + \mathrm{d}\overline{W}_t \tag{2}$$ここで、ベイズの定理より条件付きスコアが以下のように展開できることを考えます。 ...
離散拡散モデルの分類器フリーガイダンス
Simple Guidance Mechanisms for Discrete Diffusion Models (Schriff et al., 2025) という ICLR2025 に採択された論文で離散拡散モデルに分類器ガイダンス(CG)および分類器フリーガイダンス(CFG)を適用する手法が提案されていたので読んでみました。 連続拡散モデルでのCG/CFG 連続拡散モデルでは $x_{t-1}$ の条件付き確率分布が以下のように表されます。 $$p^\gamma(x_{t-1} \mid y, x_t) \propto p(y \mid x_{t-1})^\gamma \, p_\theta(x_{t-1} \mid x_t)$$ 対数勾配をとると以下のようになります。(分類器ガイダンス) $$\nabla_{x_{t-1}} \log p^\gamma(x_{t-1} \mid y, x_t) = \gamma \nabla_{x_{t-1}}\log p(y \mid x_{t-1}) + \nabla_{x_{t-1}}\log p_\theta(x_{t-1} \mid x_t)$$ ここでベイズの定理より $$\log p_\theta(x_{t-1} \mid y, x_t) = \log p(y \mid x_{t-1}) + \log p_\theta(x_{t−1}\mid x_t)−\log p(y\mid x_t)$$ が成り立ちます。 両辺の勾配をとって $$\nabla_{x_{t-1}}\log p_\theta(x_{t-1}\mid y,x_t) = \nabla_{x_{t-1}}\log p(y\mid x_{t-1}) + \nabla_{x_{t-1}}\log p_\theta(x_{t-1}\mid x_t)$$ これを先ほどの分類器ガイダンスの式に代入して ...
離散DDPMの実装
離散拡散モデルについての理解を深めるために既存の実装を参考にしながらタンパク質配列生成を試してみました。 参考にした論文 Structured Denoising Diffusion Models in Discrete State-Spaces (NeurIPS 2021) 1. 離散拡散モデルについて 本論文(D3PMと呼びます)で紹介されている離散拡散モデルの仕組みについて紹介します。 順過程 離散拡散モデルの拡散過程では以下のようにノイズを載せて拡散していきます。 $$ q(x_t \mid x_{t-1}) = \text{Cat}(\boldsymbol{x_t}; \boldsymbol{p}=\boldsymbol{x_{t-1}}\boldsymbol{Q_t}) $$ここで $x_t$ は時刻 $t$ におけるタンパク質配列のワンホットベクトル、$Q_t$ は時刻 $t$ における遷移行列を表します。遷移行列の作成方法は様々ですが、今回は論文中で紹介されている Absorbing-state (吸収状態) と呼ばれるものを使います。吸収状態とはアミノ酸が [MASK] トークンで置換された状態を指します。$Q_t$ は語彙サイズを $K$ として $(K, K)$ の行列であり各要素は以下のようにして決定されます。$m$ は既に [MASK] トークンに置換されている状態を指します。 $$ [\boldsymbol{Q}_t]_{ij} =\begin{cases}1 & \text{if } i = j = m \quad \\1 - \beta_t & \text{if } i = j \neq m \quad \\\beta_t & \text{if } j = m, i \neq m \quad \end{cases} $$$\beta_t$ の決め方はノイズスケジューリングと呼ばれ様々な方法がありますが、D3PMでは ...