Interpreting the Diffusion Policy Loss Function

Diffusion Policy: Visuomotor Policy Learning via Action Diffusion (“Original Diffusion Policy Paper”)

“As shown in Ho et al. (2020), minimizing the loss function in Eq 3 also minimizes the variational lower bound of the KL-divergence between the data distribution and the distribution of samples drawn from the DDPM using Eq 1.”

What does this mean…? I wanted to understand intuitively, so I derived this claim…

[The actual diffusion policy paper might be helpful background]

• Formulation: (Eq 1)

  • represents timestep
  • is the noise prediction network, learned params
    • is observation at timestep
  • is the DDPM noise added at each sampling step
  • empirically set slightly smaller than 1, improving stability
  • set by noise schedule
  • Note: = denoising step

• Training

  • trained with loss: (Eq 3)
    • Mathematical Interpretation:
      • • i.e. the denoising process is a Markov Chain; next state; is a prior; each then is independent of all other given and .
        • is the joint probability density of predicting the whole chain of latent semi-denoised “images”
      • Take the marginal to get the probability density of predicting the unnoisy image:
        • • This is an integral over all possible denoising paths
      • We want to minimize the KL divergence between the data distribution and DDPM distribution of sampled clean images (in other words, maximize log-likelihood; this is ML estimation):
      • Minimizing the KL Divergence directly is intractable —
        • Even computing the log likehood for a given is intractable because computing the integral in is intractable; firstly, we are integrating over a huge neural net, not some clean math function; there’s no analytical soln. But, even if you discretized and approximated the integral as a sum, you’d have to sum over all possible (discretized) denoising paths , of which there are still tons.
      • Therefore, we minimize the (negative) variational lower bound (VLB) on the true log-likelihood; a computational proxy for the actual true log-likelihood:
      • Note that we introduce as the forward-noising Markov Chain we used to create the training data
      • At this step, we sub in (using Bayes Rule):
        • Bayes Rule derives the analytically ideal reverse step from the known forward-noising process. We can then reformulate the VLB in a much more interpretable way, as matching with .
      • The last equality is non-obvious, but if we expand the case of the sum, we can see this:
      • The last term of VRB, is also trivial and so it’s dropped
        • Both are just a standard Gaussian prior; there are no parameters to learn here.
      • This yields the final objective (there are a few more steps to simplify this to the MSE loss described in the paper):
      • The next key realization is that, during training: “we randomly select denoising iteration and then sample a random noise with appropriate variance for iteration ”
        • Basically, we never compute the loss for the whole denoising sequence ; we monte-carlo sample a single denoising step and compute the loss just for that.
        • (With enough monte-carlo samples, optimizing using the sampled per-denoising step loss should be equivalent to optimizing using the sum of losses from all denoising steps).
        • Thus, we basically optimize, for a single per step of training:
      • The last key realization is that both distributions: and are Gaussian. Note that the mean here represents the semi-noised image and the variance is fixed by the noise scheduler.
        • is by design of the diffusion noising process as adding Gaussian noise.
        • also is by design of the denoising process.
      • The KL-Divergence between 2 Gaussians (w/ fixed, diagonal covariance) has closed-form:
      • Typically, we simply drop the terms and minimize . This works in practice, though we are no longer minimizing exactly the VLB.
        • The multiplier is effectively a weighting depending on the current denoising step; Ho et al. (original DDPM paper) find empirically that dropping this weighting is actually better.
      • However, the model doesn’t predict the “image” (i.e. ) directly — it predicts the noise vector between and . Thus, we must take one final step to convert the MSE error of to an MSE error of the noise vectors predicted by the network.
        • In particular, Diffusion Policy uses a “Variance-Preserving” Noise Schedule that is applied like so:
          • This implies:
          • Then, is equivalent to:
          • This is exactly the MSE loss proposed in the paper.
  • Notes:
    • This is a skewed way of looking at it, but: diffusion-policy training is kind like the EM algorithm with a trivial E-step. We already know the ground-truth used to generate the training data; so there is no need to optimize as we do in EM. Rather, we just perform the M-step of optimizing .
    • The VLB becomes tighter as the model better learns to approximate using ; this is not obvious from the math shown above; but see my 6.7810 notes on EM; is exactly the gap between VBL and the true log likelihood.
      • This is not true of EM in general; only true here because the ground-truth is known. Also, in practice, in EM, we often restrict to something simple like a Gaussian, even though the ground-truth may not be a Gaussian. Thus the VLB gap is not closable.
    • The underlying graphical model is a directed Markov Chain:
    • and the forward noising process can also be seen as a Markov Chain (though nothing to learn in this case):
    • The parameters live in the neural network, but represent (using a neural net) the probability distributions