Epsilon-prediction diffusion objective estimator

Fix finite integers $d,N,T\ge 1$, a data set $D=(x_0^{(1)},\dots ,x_0^{(N)})\in ({\mathbb{R}}^d{)}^N$, forward variances ${\beta }_1,\dots ,{\beta }_T\in (0,1)$, reverse variances ${\sigma }_1^2,\dots ,{\sigma }_T^2\in (0,\mathrm{\infty })$, a Borel map ${ϵ}_{\theta }:{\mathbb{R}}^d\times \{1,\dots ,T\}\to {\mathbb{R}}^d$, and a mode from $\{\mathtt{\text{weighted-elbo}},\mathtt{\text{simple}}\}$. Set ${\alpha }_t=1-{\beta }_t$ and ${\bar{\alpha }}_t={\prod }_{s=1}^t{\alpha }_s$. Use the noise-prediction means ${\mu }_{\theta ,t}(y)=\frac{1}{\sqrt{{\alpha }_t}}(y-\frac{{\beta }_t}{\sqrt{1-{\bar{\alpha }}_t}}{ϵ}_{\theta }(y,t)),1\le t\le T,$ and the reverse Gaussian model with these means, variances ${\sigma }_t^2$, and prior $N_d(0,I_d)$ to define $p_{\theta }$. In weighted-elbo mode, require for each $x_0^{(i)}$ that $0<p_{\theta }(x_0^{(i)})<\mathrm{\infty }$, the forward path law be absolutely continuous with respect to the normalized reverse posterior path law, and ${\mathbb{E}}_{q(\cdot \mid x_0^{(i)})}\text{\hspace{0.17em}⁣}[\mathrm{\mid }\log q(x_{1:T}\mid x_0^{(i)})\mathrm{\mid }+\mathrm{\mid }\log p_{\theta }(x_{0:T})\mathrm{\mid }]<\mathrm{\infty }\mathrm{.}$ In simple mode, require for every $i$ and $t$ that ${\mathbb{E}}_{ϵ\sim N_d(0,I_d)}\text{\hspace{0.17em}⁣}\left[{\parallel ϵ-{ϵ}_{\theta }\text{\hspace{0.17em}⁣}(\sqrt{{\bar{\alpha }}_t}{x}_0^{(i)}+\sqrt{1-{\bar{\alpha }}_t}ϵ,t)\parallel }_2^2\right]<\mathrm{\infty }\mathrm{.}$ This procedure returns one stochastic loss estimate; it does not prescribe a parameter optimizer or claim optimizer convergence.

1. Choose $I$ uniformly from $\{1,\dots ,N\}$, choose $J$ uniformly from $\{1,\dots ,T\}$ independently of $I$, and draw $ϵ\sim N_d(0,I_d)$ independently of both.

2. Set $x_J=\sqrt{{\bar{\alpha }}_J}{x}_0^{(I)}+\sqrt{1-{\bar{\alpha }}_J}ϵ\mathrm{.}$

3. If the mode is $\mathtt{\text{simple}}$, return ${\parallel ϵ-{ϵ}_{\theta }(x_J,J)\parallel }_2^2\mathrm{.}$

4. If the mode is $\mathtt{\text{weighted-elbo}}$, first compute the exact terminal-prior term $K_T(x_0^{(I)})=\frac{1}{2}[{\bar{\alpha }}_T\mathrm{\parallel }{x}_0^{(I)}{\mathrm{\parallel }}_2^2+d(-{\bar{\alpha }}_T-\log (1-{\bar{\alpha }}_T))]\mathrm{.}$ If $J=1$, return $K_T(x_0^{(I)})-T\log p_{\theta }(x_0^{(I)}\mid x_J)\mathrm{.}$ If $J\ge 2$, return $K_T(x_0^{(I)})+T[C_J+w_J{\parallel ϵ-{ϵ}_{\theta }(x_J,J)\parallel }_2^2],$ where $w_J$ and $C_J$ are the exact per-step coefficient and variance constant in the DDPM objective theorem.

Each branch returns one sampled objective contribution. The $J=1$ weighted branch uses the positive-variance learned reconstruction likelihood; it never evaluates a density for the degenerate forward conditional $q(x_0\mid x_1,x_0)$. This algorithm specifies the sampling procedure only; it makes no optimizer-convergence claim.

Finite Gaussian chain and reverse likelihood: Finite Gaussian forward chain and Gaussian reverse model.

Terminal KL and weighted per-step terms: DDPM variational objective and weighted noise-prediction terms.

Gaussian KL formula used for the terminal term: Reverse KL between isotropic Gaussian laws.

Stochastic epsilon-prediction objective and simple-loss convention: Ho, Jain, and Abbeel, Algorithm 1 and Eqs. (12)–(14).

Pedagogical treatment of weighted noise-prediction terms and the simplified training objective; its ELBO discussion approximates the terminal-prior term by zero: Prince, Chapter 18, “The diffusion ELBO” and “Simplifying it”.

Finite expectation and the exact objective estimated in each mode are proved in lem-diffusion-objective-estimator-expectation.

Expected output of the diffusion objective estimator ↗