Clarification on the training objective of denoising diffusion models

Question

I'm reading the Denoising Diffusion Probabilistic Models paper (Ho et al. 2020). And I am puzzled about the training objective. I understood (I think) the trick regarding the reparametrization of the variance in terms of the noise:

$$\mu_\theta(x_t, t) = \frac{1}{\sqrt{\alpha_t}}\left(x_t - \frac{\beta_t}{\sqrt{1 - \bar{\alpha_t}}}\epsilon_\theta(x_t, t)\right)$$

But what I do not understand fully is the training objective:

$$\nabla_\theta\lVert \epsilon - \epsilon_\theta\left(\sqrt{\bar{\alpha}}x_0 + \sqrt{1- \bar{\alpha}_t}\epsilon, t \right) \rVert$$

It looks to me like $\epsilon$ led from $x_0 \rightarrow x_t$, so it is unclear how learning to predict this noise actually forces the model to learn to undo the forward process from $x_{t-1} \rightarrow x_t$ only. The sampling instead is very clear.

---

The same question arose while reading another paper that uses a different formulation to build a diffusion generative model over graphs. In the paper "DiGress" (Vignac et al 2023), they define a discrete version of the forward process for application to graphs $G=(X, E), \space X \in \mathbb{R}^{n \times a}, \space E \in \mathbb{R}^{n \times n \times b} $. Sampling is done by a categorical distribution defined by matrices: $$[Q_X^t]_{ij} = q(x^t = i | x ^ {t - 1} = i) \\ [Q_E^t]_{ij} = q(e^t = i | e ^ {t - 1} = i) \\ \bar{Q}^t_X = Q^1_X...Q^t_X\\ \bar{Q}^t_E = Q^1_E...Q^t_E$$

and $$q(G^t| G^{t-1})=(X^{t-1}Q^t_X, E^{t-1}Q^t_E) \space\text{and}\space q(G^t| G)=(X\bar{Q}^t_X, E\bar{Q}^t_E)$$

The main difference here, besides the forward distributions to be categorical, is the fact that the model is required to compute the probabilities of the categorical distributions directly: $p_\theta^G=(p^X_\theta, p^E_\theta)$. And the loss is computed by comparing the predicted probabilities of the model against the original graph.

$$ \nabla_\theta\left( \sum_{l\leq i\leq n} \text{cross-entropy}(x_i,p_{i \space \theta}^{X}) + \sum_{l\leq i,j\leq n} \text{cross-entropy}(e_{ij},p_{ij \space \theta}^{E})\right) $$

My doubt is, here as well, how can it be correct that the model predicts directly the distribution of the original graph and still somehow this process should be equivalent to predicting only the backward to the timestep $t-1$, as in all diffusion models.

Answer 1

The key intuition I think is that high dimensional standard normals have almost constant length $\sqrt{n}$ due to Law of Large Numbers. And also, random high dimensional vectors are perpendicular to each other (unlike 2D where two random unit vectors are 30 degrees apart). Thus, repeated sampling gives us points that are efficiently spread out on a hollow ball of constant radius. I think the denoising NN really uses these regularities to find structure in the image set.

What I’m am imagining is this:

1. Every image in the library is a vector in a high dimensional space. The tips are leaves of an unknown tree. We’ve to identify the tree. (The tree is imaginary and helps me think. I will get rid of it later.)

2. By repeatedly adding small orthogonal constant norm noise to the tips, we are adding little balls to the tips.

3. Imagine you cut off all the portions of the tree inside these balls. Now you’ve fewer branch tips cuz the finer branchings are all gone. The NN basically identifies these fewer tips. Ie, fed a point on the ball around a final leaf, the NN will give the tip or fewer new point(s) that are common to all the nearby leaves.

4. Now take the truncated tree and attach relatively bigger balls at the new tips. Cut the tree again within these balls. You’ve still fewer tips which the NN learns cuz it knows this is a different size ball. Note that I can get similar results by successively adding balls to balls to the original leaves and cutting the tree all at once.

5. You do this with bigger and bigger balls. Suppose the final step causes the NN to learn, say 3 tips that are the ends the now very diminished tree.

6. To generate, start with a point on a unit ball. Knowing the scale, the NN will try to push it towards on of the three stumps and then up the tree towards the finer stumps and so on.

7. Basically this way of thinking takes the library image set as points in the n dimensional space. The NN reduces dimensions by finding fewer points that are close to the original points within successive balls of various radii. The orthogonality and constant norm of high dimensional random vectors makes it super easy to generate these balls. And the NN structure is highly suited to find common points within these balls.

8. Ball with larger norms means more nonzero pixels ie broad brush features. Balls with smaller norms means few nonzero pixels ie minute detail. So training happens in any order by successively adding multiple balls in one step, but reconstruction is sequential from large ball to small.