Denoising Score Matching

The marginal score $\nabla \log p_t$ is unknown, but it can be learned from samples without knowing it.

Denoising score matching (Vincent 2011). For any noise kernel $q(x_t\mid x_0)$,

$${\mathbb{E}}_{x_t}\Vert s_{\theta }(x_t)-\nabla \log p_t(x_t){\Vert }^2={\mathbb{E}}_{x_0,x_t}\Vert s_{\theta }(x_t)-{\nabla }_{x_t}\log q(x_t\mid x_0){\Vert }^2+\text{const}.$$

Regressing onto the conditional score, which is known in closed form, gives the marginal score at the optimum.

$\epsilon$-prediction loss (DDPM). Inserting the Gaussian kernel of the DDPM Forward Process and parametrising $s_{\theta }=-{\epsilon }_{\theta }/\sqrt{1-{\bar{\alpha }}_t}$ gives the simple training objective

$$\mathcal{L}(\theta )={\mathbb{E}}_{t,x_0,\epsilon }[w(t)\Vert {\epsilon }_{\theta }(\sqrt{{\bar{\alpha }}_t}{x}_0+\sqrt{1-{\bar{\alpha }}_t}\epsilon ,t)-\epsilon {\Vert }^2],$$

with $t\sim \mathcal{U}[0,T]$, $x_0\sim p_{\text{data}}$, $\epsilon \sim \mathcal{N}(0,I)$ and weight $w(t)=1$ in DDPM. The network sees the noisy image and the time (via a Sinusoidal Time Embedding) and predicts the noise. Its optimum is ${\epsilon }_{\theta }(x_t,t)=\mathbb{E}[\epsilon \mid x_t]$.

Training loop. Sample a batch of images, times and noise; form $x_t$; take a gradient step on the squared error. No simulation of the SDE and no ODE solver are needed, which makes training cheap compared with Neural ODEs. Other parametrisations predict $x_0$ or $v=\sqrt{\bar{\alpha }}\epsilon -\sqrt{1-\bar{\alpha }}{x}_0$; they differ only in the implied weighting $w(t)$.

References

1. P. Vincent (2011). A Connection Between Score Matching and Denoising Autoencoders. Neural Computation 23(7), 1661–1674. doi:10.1162/NECO_a_00142
2. J. Ho, A. Jain, P. Abbeel (2020). Denoising Diffusion Probabilistic Models. NeurIPS 33. arXiv:2006.11239
3. T. Salimans, J. Ho (2022). Progressive Distillation for Fast Sampling of Diffusion Models. ICLR 2022. arXiv:2202.00512