TL;DR

This gonna be a short blog to share some of my thoughts on the distillation in on-policy distillation (during my discussion with Yuchen Zhu), which is obvious but anyway I want to write it down as a blog post.

In one sentence, I found that the famous diffusion distillation framework (DMD, Diff-Instruct, etc.) and the recently proposed forward-process based diffusion RL fine-tuning methods are actually doing the same thing.

The Distillation in On-Policy Distillation

The Illustration of the On-Policy Generator and the Objective

As shown in the figure below, the on-policy generator is actually a model $\theta$ that take inputs noise $x_0$ and context $c$ to generate outputs $x_1$. For diffusion model, this generation can take multiple $T$ steps, using corresponding samplers such as DDIM, DPM solver or even CM sampler (for the generator trained with consistency model objective or DMD objective).

Let’s denote the generator’s output distribution as $p_\theta(x_1\mid c)$, then we can write the objective of the on-policy distillation as minimizing some divergence between the generator’s output distribution and a target distribution $p_{\mathrm{target}}(x_1\mid c, c_\mathrm{ext})$, where $c_\mathrm{ext}$ is the external signals (e.g., reward model, privileged information, expert demonstrations, etc.).

The objective can be written as:

where $D$ is some divergence measure, e.g., reverse KL divergence, JS divergence, etc.

Note that during implementation, we sample $x_1$ from the generator and compute the loss at the sample level. We can think of this objective as training the generator to produce outputs that are close to the target distribution, thus we omit the input of the generator in the following figures for simplicity.

Diffusion (Step) Distillation

For diffusion distillation, the target distribution is usually the teacher model distribution. According to my previous blog, the distillation can be seen as minimizing the KL divergence between the student marginal distribution $q_{t}(x_t)$ and the teacher marginal distribution $p_t(x_t)$ at each time step $t$, which is illustrated in the figure below.

As shown in the above figure, the gradient of the divergence can be estimated in two ways: (a) the score-based approach (e.g., DMD) and (b) the discriminator-based approach (e.g., Diffusion GAN).

Diffusion Reinforcement Learning

For diffusion RL fine-tuning, the target distribution is usually defined with a reference model and a reward model, which can be seen as a Boltzmann distribution $p_{\mathrm{target}}(x_1\mid c, c_\mathrm{ext}) \propto p_{\mathrm{ref}}(x_1\mid c) \exp(r(x_1))$, where $r(x_1)$ is the reward function defined by the reward model.

In this case, the objective can be written as:

Expanding using reverse KL divergence and splitting the log:

Splitting the expectation and rearranging terms, we have two terms: the reward term $\mathbb{E} _ {x_1 \sim p_\theta}[-r(x_1)]$ and the KL regularization term $\mathbb{E} _ {x_1 \sim p_\theta} \left[ \log \frac{p_\theta(x_1 \mid c)}{p_{\mathrm{ref}}(x_1 \mid c)} \right]$. With importance sampling and the REINFORCE log-derivative trick $\nabla_\theta \mathbb{E} _ {p_\theta}[f] = \mathbb{E} _ {p_\theta}[f \cdot \nabla_\theta \log p_\theta] $, we can rewrite the reward term as $\mathbb{E} _ {x_1 \sim p_{\mathrm{ref}}} \left[ - \frac{p_\theta(x_1 \mid c)}{p_{\mathrm{ref}}(x_1 \mid c)} r(x_1) \right]$.

As a result, the objective can be rewritten as a standard RL objective with KL regularization:

For diffusion model, the density ratio $p_\theta / p_{\mathrm{ref}}$ can be estimated by the diffusion ELBO estimator, which is the same as the weighted denoising score matching (DSM) loss. To be specific, in the ELBO estimator, we calculate the density ratio at a randomly sampled time step $t$ as $r_t(x_t) = \frac{p_{\theta,t}(x_t|c)}{p_{\mathrm{ref},t}(x_t|c)}$, where $q_{\theta,t}$ is the generator marginal and $p_{\mathrm{ref},t}$ is the reference model marginal at time step $t$. The diffusion RL training can be summarized in the below figure.

Summary

In this blog, we show that diffusion distillation and diffusion RL fine-tuning are both instances of on-policy distillation, minimizing a divergence $D(p_{\mathrm{target}} | p_\theta)$, differing mainly in the choice of target distribution and the approach to estimate the divergence.

Acknowledgements

The author thanks Yuchen Zhu for his insightful discussion.

A Short Discussion on RL and ICL.

Recently I had a discussion with my friend Jiwen Yu, and I was trying to explain the benefits of RL, which in my opinion, is simply tilting the output distribution of the model towards high-reward regions $p^* = p_{\mathrm{ref}}\exp(\beta r)$.

“So basically, RL can increase the pass@1, hence the user will get a better answer with higher probability”

“Yeah, but in practice I found that for my specific question, the LLM will give me much better answers with more details I provided in the prompt.”

Self-Distillation Fine-Tuning (SDFT) and Self-Distillation Policy Optimization (SDPO)

Later I read the paper “Self-Distillation Enables Continual Learning” and I realized that the SDFT is trying to distilling the expert demonstrations in context into the model itself.

Concurrent work, “Reinforcement Learning via Self-Distillation”, also proposed a similar idea of using self-distillation to improve the model’s performance. To be specific, this time it does not rely on the expert demonstrations but the expert feedback on initial model outputs, then the expert feedback is fed back to the model in the form of context to generate better outputs for self-distillation. This work also demonstrated an elegant approach to utilize text rewards for RL fine-tuning.

Everything is about Distillation

I was supposed to make this a title for the whole blog, but perhaps it is too clickbaity. Nevertheless, here I would like to emphasize that the core idea is, for every method that is scalable at test time (with the help of external reward models or expert signals $c_\mathrm{ext}$), we can always find a way to distill the knowledge back into the model itself, so that at test time, the model can perform better without relying on external signals.

From another perspective from the first principal by Rishabh Agarwal on twitter, “Teacher doesn’t have to be a bigger neural net, just something better than the student”.

If we formulate the generation process as $p_\theta(x\mid c)$, where $c$ is the context, then we can always find a better teacher distribution of the same model $p(x\mid c, c_\mathrm{ext})$ that leverages the external signals $c_\mathrm{ext}$ to provide better outputs. Then we can distill the knowledge from the teacher back to the student model by minimizing some divergence between the two distributions, e.g. KL divergence.

Some examples include:

• RL: distill the reward model and the policy into the model itself.
• SDFT: distill the expert demonstrations in context into the model itself.
• SDPO: distill the expert feedback as text reward in context into the model itself.

To end, everything is about distillation!

Acknowledgements

The author thanks Jiwen Yu and Yao Teng for their valuable feedbacks (for me to distill my thoughts).

Offline DiffusionNFT[1] as KL-Regularized Reward Tilting

In my personal view, DiffusionNFT is a significant work because it is the first native algorithm for diffusion RL without relying on differentiable reward models.

The paper derives an optimal velocity field update and interprets it as a positive/negative split, but the corresponding density-level optimum can be written explicitly as a reward-tilted distribution. This section provides a detailed derivation of this closed-form solution.

Setup (From DiffusionNFT)

Introduce a binary optimality variable $o\in{0,1}$ and a prompt/context $c$. DiffusionNFT defines the reward as an optimality probability

Given a reference (old) model $\pi_{\mathrm{old}}(x_0\mid c)$, define the positive (optimality-conditioned) distribution

and the corresponding forward-noised marginals at diffusion time $t$ under a fixed kernel $\pi_{t\mid 0}(x_t\mid x_0)$:

DiffusionNFT’s training objective¹ yields an optimal velocity field of the form

with guidance direction $\Delta(x_t,c,t)=\alpha(x_t,c)\bigl(v^{+}(x_t,c,t)-v^{\mathrm{old}}(x_t,c,t)\bigr)$.

And the mixture coefficient is defined as

Derivation of $\alpha(x_t,c)=p(o=1\mid x_t,c)$

Apply the forward diffusion to the positive distribution:

where we use the identity

we obtain

Note that $\alpha(x_t,c)=\mathbb{E}_{x_0\sim \pi _{\mathrm{old}}(\cdot \mid x_t,c)}[r(x_0,c)]$. This shows that the mixture coefficient equals the optimality posterior distribution at the noisy state $x_t$ under the old model.

Remark

This equality is a direct consequence of the definition of $\pi_{+}$ as an optimality-conditioned distribution and Bayes’ rule under the fixed forward noising kernel. It may be omitted or only implicit in the original DiffusionNFT paper. With this finding, the Lemma A.2 (Posterior Split) in the paper becomes obvious.

DiffusionNFT optimal distribution at each step

In order to compute the optimal distribution, we need to simplify the residual term $\Delta(x_t,c,t)=\alpha(x_t,c)\bigl(v^{+}(x_t,c,t)-v^{\mathrm{old}}(x_t,c,t)\bigr)$.

Using the Bayes relation for the positive marginal

and i) utilizing the relation between velocity fields and score functions under fixed Gaussian noising², and ii) applying log and gradient, one has

Thus we can rewrite the guidance direction as:

Closed-form optimal distribution (density-level)

Since $\alpha(x_t,c)=p(o=1\mid x_t,c)$ and $\alpha(x_t,c)\nabla\log\alpha(x_t,c)=\nabla\alpha(x_t,c)$, we can rewrite eq(\ref{optimal_v}) in the form of score and have:

This implies the explicit density:

At $t=0$ this reduces to $p^{*}(x_0\mid c)\propto p^{\mathrm{old}}(x_0\mid c)\exp(\frac{2}{\beta}r(x_0,c))$.

Thus, DiffusionNFT with fixed $p^{\mathrm{old}}=p^{\mathrm{ref}}$ can be viewed as learning a KL-regularized exponential tilt of the reference model, where the “reward” is the optimality posterior at the noisy state.

Online DiffusionNFT Leads to Reward Hacking via Accumulated Tilting

In practice, DiffusionNFT can be run online: after fitting $v_\theta$ for one epoch, the old model is updated (e.g., by copying weights (hard) or by EMA (soft)) and the next epoch is trained against this updated old model. This induces a recursion over the initial reference distributions.

Idealized recursion with exact per-epoch optima

Let $p^{(k)}_t(x_t\mid c)$ denote the old marginal at epoch $k$ (the distribution induced by the current old velocity field). Treat the reward/optimality posterior $\alpha_t(x_t,c)=p(o=1\mid x_t,c)$ as fixed across epoches.

From the closed-form optimum, the population-optimal update at epoch $k$ is

Unrolling the recursion yields

As $K\to\infty$, the previous expression concentrates mass on the set of maximizers of the reward.

Remarks on EMA references

If the reference is updated via EMA in parameter space, the induced distribution recursion is not exactly the simple multiplicative update above. Nevertheless, EMA typically acts as a trust-region mechanism that interpolates between keeping $p^{(k)}_t$ fixed and fully replacing it by the latest student, effectively reducing the rate at which the tilt coefficient grows with $k$. The qualitative conclusion remains: in the absence of an anchoring term to the initial model, repeated online improvement tends to accumulate reward tilt and can become increasingly peaked around high-reward regions and thus prone to reward hacking.

Online DiffusionNFT with an Additional KL to the Initial Reference

To prevent unbounded drift from the original model, one can augment the per-epoch objective with an additional regularizer that penalizes deviation from the initial reference distribution $p^{(0)}_t(\cdot\mid c)=p^{\mathrm{ref}}_t(\cdot\mid c)$.

At the distribution level, consider the KL-regularized problem at epoch $k$ [2,3]:

Here $\eta_1$ controls the trust region to the current reference and $\eta_0$ controls anchoring to the initial reference.

The unique optimizer has the closed form

with weights and effective tilt coefficient

Thus, the per-epoch solution is an exponential tilt of a geometric mixture of the current and initial references.

It is straightforward to see that the iterates converge to a limiting distribution:

Notably, the limiting distribution depends only on the anchoring strength $\eta_0$; the trust-region parameter $\eta_1$ affects only the dynamics (convergence rate and stability) through $w$.

In practice, the finetuned model in DiffusionNFT is initialized from a pre-trained diffusion model without CFG. Moreover, the initial reference model is also the pre-trained diffusion model without CFG. Thus, with large inital KL strength, the online DiffusionNFT procedure effectively regularizes the finetuned model toward the pre-trained diffusion model without CFG and the learned model can only generate blurry samples; with small initial KL strength used in the experiments in the paper, the finetuned model generates samples with high reward but low diversity and pure color background, which suggests severe reward hacking. Based on the analysis above, we suggest adding an medium-strength initial KL regularization and inference the pre-trained model with CFG as the initial reference to mitigate reward hacking in online DiffusionNFT.

Acknowledgements

The author thanks Huayu Chen for his insightful work, helpful discussions, and valuable feedback. The author also thanks Ruiqing Wang for proofreading.

References

It feels like yesterday when I wrote my New Year Resolution for 2025. While many (mostly good) things have happened in 2025, at this moment, I have the resolutions firmly set for 2026 and wish to avoid the mistakes made in the previous year.

Time, is the most fair resource people have. The feeling of time is related to the infomation input we receive, thus we need to experience more. Due to the finite memorey capacity of human brain, it is more important to collect precious memories, better with treasureable people.

In this blog, I want to show a simple but intuitive connection between Variational Score Distillation (VSD/DMD) [1,2,3] and Diffusion-GAN [4] for diffusion distillation. This connection is noticed when I was trying to extend our recent workshop paper Di-Bregman [5].

The final conclusion is not bonded to Di-Bregman, I just want to use Di-Bregman as an example to illustrate the connection.

Preliminaries

In Di-Bregman, we derived a new distillation loss, whose loss gradient extend the DMD loss with an extra coefficient $h^{\prime\prime}(r_t(x_t))\,r_t(x_t)$:

where

• $r_t(x_t) = \frac{q_{\theta,t}(x_t)}{p_t(x_t)}$ is the density ratio between the student marginal $q_{\theta,t}$ and the teacher marginal $p_t$ at time $t$;
• $G_\theta(\epsilon)$ is the student generative model that maps noise $\epsilon$ to clean sample;
• $h(r)$ is a convex function defining the Bregman divergence $\mathbb{D}_h$;
• $w(t)$ is a time-dependent weight function;
• $\epsilon \sim \mathcal{N}(0,I)$, $t \sim \mathcal{U}(0,1)$, and $x_t = \alpha_t G_\theta(\epsilon) + \sigma_t z_t$ with $z_t \sim \mathcal{N}(0,I)$.

From Scores to a Discriminator

One difference between Di-Bregman and DMD is that we need to estimate the density ratio $r_t(x_t)$ in the extra coefficiency $h^{\prime\prime}(r_t(x_t))\,r_t(x_t)$, usually achieved by training a discriminator $D(x_t,t)$ to distinguish samples from $p_t$ and $q_{\theta,t}$.

With the usual logistic convention (where $D_t(x_t)$ outputs the probability that $x_t$ comes from $p_t$), we have the following optimal discriminator $D_t^*(x_t)=\frac{p_t(x_t)}{p_t(x_t)+q_{\theta,t}(x_t)}$.

A natural question is: can we rewrite the Di-Bregman loss gradient entirely in terms of the discriminator $D_t$, eliminating explicit score functions?

Step 1: Substitute the density-ratio gradient identity

We have the identity

Substituting into the Di-Bregman gradient gives

Step 2: Express $\nabla_{x_t} r_t$ and $r_{t}$ in terms of $D_t$

Since $r_t = \frac{q_{\theta,t}}{p_t} = \frac{1-D_t}{D_t}$, we have $\frac{dr_t}{dD_t} = -\frac{1}{D_t^2}$, hence

Substitute this back:

Step 3: Convert $\nabla_{x_t} D_t$ to $\nabla_\theta D_t$

Using the chain rule, we have:

The conversion from $\nabla_{x_t}D_t$ to $\nabla_\theta D_t$ introduces the factor $\alpha_t$, which is absorted into $w’(t)$.

The boxed equation shows that the Di-Bregman loss gradient can be computed by backpropagating through the discriminator $D_t$ only, without explicitly estimating the score functions. This suggests that Di-Bregman / VSD and Diffusion-GAN are closely connected.

Verify with Special Case: reverse KL Divergence (DMD)

When $h(r)=r\log r$, we have $h^{\prime\prime}(r)=\frac{1}{r}$, hence the coefficiency in the boxed equation becomes $\frac{h^{\prime\prime}(\frac{1-D_t}{D_t})}{D_t^2} = \frac{1}{D_t(1-D_t)}$, and the loss gradient becomes

This resonates with the GAN generate loss for KL divergence $\mathrm{KL}(q_\theta||p)$:

Given that both DMD and Di-Bregman are equivalent to diffusion-GAN training, and with the belief that scalable pre-training requires algorithms that directly optimize the data likelihood or an associated ELBO, I am skeptical that these approaches can serve as general-purpose pre-training methods. For example, it seems unlikely that we could pre-train a next-frame prediction video generative model using methods such as self-forcing.

References

Second-order MeanFlow: Combining Forward and Backward Distillation

In this blog, We show how the MeanFlow (backward distillation) identity and the tri-consistency constraint can be combined to produce a second-order identity for the student velocity in diffusion/flow distillation. Similar to MeanFlow, this 2-order MeanFlow contains remarkably simple and elegant target: the average of two end point velocities and a second-order correction term.

1. Setup: Notation and Goals

• $x_t$ denotes the state at time $t$ on the diffusion/flow trajectory.
• $v_\phi(x_t,t)$ is the pre-trained teacher velocity field (teacher’s instantaneous velocity at $(x_t,t)$).
• $v_\theta(x_t,t,s)$ is the student average velocity intended to take a sample at time $t$ to the time $s$ in a single step (a short-cut/one-step generator). We treat $s > t$.
• We will often use small increments $ds$ with $s_2 = s_1 + ds$.

Two consistency principles underly the derivation:

1. MeanFlow (backward) identity (backward distillation): for $s > t$,

   $$ v(x_t,t,s) = v_\phi(x_t,t) + (s-t)\frac{d}{dt}\bigl[v(x_t,t,s)\bigr], $$

   which expresses the student average velocity from timestep $t$ to $s$ in terms of the (teacher’s) instantaneous velocity at $t$ and the time derivative of the student velocity with respect to the local timestep $t$.

2. Tri-consistency (additivity of short segments): for $s_2 > s_1 > t$,

   $$ (s_1-t)\,v(x_t,t,s_1) + (s_2-s_1)\,v(x_{s_1},s_1,s_2) \;=\; (s_2-t)\,v(x_t,t,s_2). $$

   This says that two short forwards ($t$ to $s_1$ to $s_2$) should equal the single forward (from $t$ to $s_2$).

Goal: derive a practical target for $v_\theta(x_t,t,s)$ that consider both the forward and backward constraints.

2. A Intuitive Derivation

Start from tri-consistency and replace the second segment ($s_1$ to $s_2$) with teacher PF-ODE simulation:

For small step $ds = s_2 - s_1$ the one-step Euler approximation to the teacher ODE gives

Now set $s_1 = s$ and $s_2 = s + ds$. Substitute the MeanFlow identity (backward) for the two student velocities $v(x_t,t,s)$ and $v(x_t,t,s_2)$. After algebra and cancellation of terms proportional to $ds$, we arrive at the following mixed relation:

We can substitute the MeanFlow identity to eliminate the first time derivative and rearrange to isolate $v(x_t,t,s)$, yields the second-order MeanFlow identity:

Remarks

• Equation (★) is the second-order MeanFlow expression: the student velocity equals the average of the two teacher velocities plus a second-order correction involving the mixed partial $\dfrac{d^2}{dt\,ds}v$.
• (Quick validation) Similar to MeanFlow, when $t=s$, Equation (★) degrades to flow matching loss.
• Unlike forward and backward distillation loss, Equation (★) can not be simply interpreted as a special case of tri-consistency.
• In practive, we expect to use this loss along with the first-order MeanFlow loss to improve the few-step performance (requires more computation).

2. Generalized 2-order Loss

Given the backward and forward distillation losses:

we can combine them with our second-order MeanFlow loss using weights $\alpha$ and $\beta$ to get a generalized loss:

For $\beta = 0$, $\alpha = 1$ or $0$, we recover backward or forward distillation loss; for $\beta = 1$, $\alpha = \frac{1}{2}$, we recover the second-order MeanFlow loss (★).

References

It’s a time to reflect on the past year and set new goals for the next year. While I tried to summarize my research journey in 2024, I finally commented out the whole content because it was a bit lengthy and not well-organized, which might not be interesting to read.

But anyway, here comes my new year resolution for 2025.

1. Stay close to my family and friends.
2. Keep healthy and exercise regularly.
3. Do some interesting research (overcome necessary engineering difficulties).
4. Have more discussions with my colleagues and other researchers, inspire more sparks.

Introduction

In the previous post, I presented some of the most classical MCMC method, Metropolis-Hastings (MH) and its variants. In this post I’d like to focus on another important MCMC method: Unadjusted Langevin Algorithm (ULA) in Bayesian statistics or Langevin Monte Carlo (LMC) in machine learning. LMC is a MCMC method which utilize Langevin dynamics for obtaining random samples from a probability distribution $\pi$ for which direct sampling is difficult (high-dimension and large number of data).

Langevin dynamics provides an MCMC procedure to sample from a distribution $p(x)$ using only its score function $\nabla_x \log p(x)$. Specifically, it initializes the chain from an arbitrary prior distribution $x_0 \sim \pi(x)$, and then iterates the following:

$$ \begin{align*} x_{i+1} \leftarrow x_i + \epsilon \nabla_x \log p(x) + \sqrt{2 \epsilon} z_i \label{Des_LD}\tag{1} \end{align*}$$

—— Generative Modeling by Estimating Gradients of the Data Distribution by Yang Song. Personal Blog, 2021.

Langevin Dynamics Sampling

Langevin Diffusion

The overdamped¹ Langevin Itô diffusion can be written as the following Stochastic Differential Equation (SDE):

where $U(X)$ is the (time-dependent) potential energy function on $\mathbb{R}^d$ and $B_t$ is a d-dimensional standard Brownian Motion (or called the Wiener process). We assume that $U(X)$ is $L$-smooth (or $\nabla U(X)$ is $L$-Lipschitz): i.e. continuously differentiable and $||\nabla U(x)-\nabla U(y)|| \leq L||x-y||, \exists L > 0$.

In order to sample the diffusion paths, we can discrete eq(\ref{LD}) using the Euler-Maruyama (EM) scheme as following² (similar to eq(\ref{Des_LD})):

where $z_i$ is i.i.d $\mathcal{N}(0,I_d)$ and $\gamma_{i}$ is the stepsize, either constant or decreasing to 0.

The subject of mixing time in MCMC algorithms is quite complex and necessitates substantial analysis skills. Hence, we won’t delve into it in this blog post.

Fokker-Plank Equation

The Fokker-Planck (FP) equation, a form of partial differential equation (PDE), outlines how a probability distribution changes over time due to the influences of deterministic drift forces and stochastic fluctuations.

Let’s denote the law of $X_t$ as $p_t$ and the FP equation for $p_t$ in eq(\ref{LD}) can be written as^3,4:

To verify that $\pi$ is the unique stationary distribution for this FP equation (\ref{FP}) and hence the corresponding Langevin diffusion — which motivate us to use Langevin dynamics for sampling from $\pi$ — we can do any of the following⁵:

• substitute $p_t$ with $\pi$ to verify that it’s a stationary distribution ($\partial_t \pi = 0$)
• assume we already have a stationary distribution $p_\infty(x), s.t. \partial_t p_\infty=0$ and verify that $p_\infty(x) \propto \exp(-U(x))$.

Unadjusted Langevin Algorithm

When we using eq(\ref{Des_LD2}) for sampling we are implementing the ULA. Even though the stationary distribution of the Langevin diffusion is $\pi$, the stationary distribution of ULA is not, due to the discretization⁶! Note that when $\gamma_i$ is constant, the value of it controls the trade-off between the convergence speed and the accuracy⁷. ULA can be biased easily with larger step-size $\gamma_i$ and one could use a decreasing step-size or introduce the HM scheme to eliminate the bias.

Metropolis-Adjusted Langevin Algorithm

The well-known MALA can be considered as MH with improved proposals or in this post ULA with an extra rejection step:

where the additional term $\tau \nabla \tilde{\pi}(x)$ drives the samples toward areas with high probability density, which makes the proposed samples more likely to be accepted. One significant advantage of MALA over Random Walk Metropolis (RWM) lies in the optimal step-size selection. In MALA, the asymptotic value for this parameter is larger compared to the RWM equivalent, leading to a decrease in the dependence observed between subsequent data points.

Stochastic Gradient Langevin Dynamics

Stochastic Gradient Langevin Dynamics (SGLD) is a method proposed by Max Welling and Yee Whye Teh where the traget density $\pi$ is the density of the posterior distribution $p(\theta|x)\propto {p(\theta)}\prod_{i=1}^{N}{p(x_i|\theta)}$. The authors initially observe the parallels between stochastic gradient algorithms (that utilize a batch of samples in each iteration) and Langevin dynamics (that employ all available samples in each iteration; one can consider ULA as gradient descent with some Gaussian noise added to the gradient at each iteration), and then propose a novel update rule that merges these two concepts with unbiased estimates of the gradient:

where the step-sizes $\epsilon_{t}$ decrease towards zero, n is the batch size, and $\eta_{t} \sim \mathcal{N}(0,\epsilon_{t})$.

Yes, SGLD is proposed to train a model given the dataset⁸, and the idea is quite simple: We train the model using regular SGD, but add some Gaussian noise to each step. There are two sources of randomness: estimates of the gradient and Gaussian added noise to sample.

From Langevin Dynamics to Diffusion Models

Pitfalls with Unadjusted Langevin Dynamics

Assume now we have access to oracle $s_\theta \approx \nabla \log p_\infty = -\nabla U$ through score matching⁹:

are we safe to use Langevin Algorithm (\ref{Des_LD2}) for image generation in high dimension? No, there are pitfalls¹⁰!

• The manifold hypothesis: the score function $s_\theta$ is ill-defined when $x$ is confined to a low dimensional manifold in a high-dimensional space; and the score matching objective in eq(\ref{sm}) will provides a inconsistent score estimator when the data reside on a low-dimensional manifold.
• The scarcity of data in low density regions can cause difficulties for both score estimation with score matching and MCMC sampling with Langevin dynamics (suffers from slow mixing time due to separate regions of data manifold).

Annealed Langevin Dynamics

To overcome these issues, Yang Song et al. proposed Noise Conditional Score Network (NCSN) $s_\theta(x,\sigma)$, which learns the Gaussian-perturbed data distribution $p_{\sigma_i}(x) = \mathbb{E}_{x’ \sim p(x’)}\left[p_{\sigma_i}(x|x’)\right] = \int p(x’) \mathcal{N}(x; x’, \sigma_i^2 I) \mathrm{d} x’ $ with various levels of noise $\sigma_i$. $\{\sigma_i\}^T_{i=1}$ is a positive geometric sequence that satisfies $\frac{\sigma_1}{\sigma_2}=\frac{\sigma_i}{\sigma_{i+1}}=\frac{\sigma_{T-1}}{\sigma_T}>1$. Since the distributions $\{p_{\sigma_i}\}^T_{i=1}$ are all perturbed by Gaussian noise, their supports span the whole space and their scores are well-defined, avoiding difficulties from the manifold hypothesis. To generate samples with annealed Langevin dynamics, we start with the largest noise level $\sigma_0$ and anneal the noise level to $\sigma_T\approx 0$:

where $\alpha_{i}=\frac{\gamma}{2} \frac{\sigma_i^2}{\sigma_{i+1}^2}$ is the step-size. The only difference between the above eq(\ref{ALD}) and eq(\ref{Des_LD2}) is that the drift term $s_\theta(x_i,\sigma_{i+1}) \approx -\nabla U_i(x_i)$ is no longer fixed and will change with time.

The sampling algorithm of NCSN, which includes a double loop, can be located in ¹⁰. The outer loop is responsible for determining the noise levels, while the inner loop takes $T$ steps to guarantee that the samples are from the distribution $p_{\sigma_i}$.

To learn the score function of the perturbed data distribution, we can use the objective of denoising score matching⁹:

From score matching to denoising score matching, the difference is a constant term independent of $\theta$¹¹:

By matching the score, we are actually minimize the KL divergence between the parameterized and the target (data) probability distribution: $D_\mathrm{KL}(p_{data}(x)||p_{\theta}(x))$.

With the learnt scores $s_\theta(x,\sigma_i) \approx \nabla_x \log p_{\sigma_i}(x)$, we can generate samples according to eq(\ref{ALD}). When $\sigma_i$ is large, modes in $p(x)$ are smoothed out by the Gaussian kernel and $s_\theta(x,\sigma_i)$ points to the mean of the modes; as $\sigma_i$ annealed down, the dynamic will be attracted to the actual modes of the target distribution¹².

In the next section, we will explore another family of sampling algorithm: reverse SDEs.

Score-based Stochastic Differential Equations

The continuous SDE form of eq(\ref{ALD}) can be written as:

where $U_t\propto -\log p_t$ now also depends on time $t$ comparing to $U\propto -\log p_\infty$ in eq(\ref{LD})¹³.

In another wonderful work by Yang Song et al.¹⁴, the authors construct forward diffusion process, which maps the target distribution to a (usually) simple distribution, in a more general form of SDEs:

Each pair of $f_t$ and $g_t$ define the unique forward process and the corresponding marginal $p_t$ given the initial distribution.

From now on, we will swap the notation of time $t$, such that for $t=0$ we have target distribution $p_0 = \pi$ and for $t=T$ we have simple distribution $p_T$.

Furthermore, we can initiate with samples from the simple distribution $p_T$, and generate samples from target distribution $p_0$ following the dynamics of the corresponding reverse SDE:

This SDE is only meant for time flows backwards from $T$ to 0, and $dt$ is an infinitesimal negative timestep. Now we can sample from the target distribution use this reverse SDE as long as we have access to $s_\theta \approx \nabla \log p_t$.

Note that eq(\ref{ALD2}) and eq(\ref{reverseSDE}) are two different sampling strategies, and we can apply both for sampling¹⁴ (aka. corrector and predictor): we use the corrector to ensure $x_t \sim p_t$ and use the predictor to jump to $p_{t-1}$.

Still, for $f_t=\nabla_{x} \log p_t({x})$ and $g_t=\sqrt{2}$, we get eq(\ref{ALD2}) as a special case of eq(\ref{reverseSDE}) representing backward sampling (note that eq(\ref{reverseSDE}) is reversed in time). But now the forward diffusion does not lead to the simple Gaussian distribution that we can sample from easily. In other words, we can’t enjoy the fast sampling trajectory from Gaussian to target distribution anymore in this scenario.

For score-based models (or diffusion models) with forward SDE in the form of eq(\ref{SDE}), we can write the corresponding FP equation in the form of the continuity equation¹⁵:

where we define the vector field as $w_t = f_t - \frac{g_t^2}{2} \nabla\log p_t$, $g_t$ is somehow related to the temperature of the stationary distribution, and $p_\infty$ is the initial distribution (not necessarily Gaussian).

The corresponding Ordinary Differential Equations (ODE) of a particle moving along this vector field is the so-called probability flow ODE (PFODE)¹⁴:

And we can use this probability flow ODE to travel either forward or backward in time.

Given $f_t$ and $g_t$ which define the forward diffusion process, we can travel backward in time with learned score function $s_\theta \approx \nabla \log p_t$ or the learned vector field $v_\theta \approx f_t - \frac{g_t^2}{2} \nabla\log p_t$.

It’s worthy noting that, in order to sample from a target distribution according to the reverse SDEs (\ref{reverseSDE}), the score functions with different $t$ are explicitly required (no matter the form of $f_t$ and $g_t$) rather than only the score function of the target distribution. Moreover, my understanding to the success of diffusion models is that, diffusion models explicitly define paths in both the data space and the probability measure space: given the initial state $x_0$, for each time-step $t$ on the paths, we know where we are ($p_t$) and where we are aiming for ($v_t$).

Introduction

In the previous post, I showed that we can use MCMC methods to draw samples from a target distribution $\pi$. In this post, I will present the most classical MCMC method, Metropolis-Hastings (MH) and its variants. Prior to MCMC methods, I want to first get the readers familiar with another Monte Carlo approach named acceptance-rejection method, which is also able to generate samples from a (unnormalized) distribution $\tilde\pi$ and whose shortcomings motivate us to use Markov chain.

Acceptance-Rejection Method

To sample directly from the target distribution $\pi(x)$ is unfeasible. The technique named acceptance-rejection sampling, also known as rejection sampling, provides a solution to this issue. The fundamental concept here is to propose a different probability distribution, $G$, with a density function $g(x)$, which not only have an efficient sampling algorithm¹ readily available, but should also span (at least) the same domain as $\tilde\pi(x)$ (Otherwise, there would be parts of the curved area we want to sample from $\pi(x)$ that could never be reached.). Consequently, it’s seem natural now to construct a decision rule² to determine whether to accept samples from the proposed distribution $g$, of course based on our knowledge of $\tilde\pi(x)$. This decision rule should ensure that the long-term fraction of time spent in each state is proportional to $\tilde\pi(x)$.

Let first have a look at the algorithm, which is quite simple by repeating the following two steps until enough samples are accepted:

1. sample $x$ from proposal $g$.
2. accept with probability $p_{\mathrm{accept}}(x) = \frac{\tilde\pi(x)}{Mg(x)}$.

where $M$ is a sufficient large constant to ensure $\left[\frac{\tilde\pi(x)}{Mg(x)}\right]$ can be interpreted as probability, that is $0 < p_{\mathrm{accept}} < 1$ for all values of $x$. This also require $g(x)>0$ whenever $\tilde\pi(x)>0$ as mentioned above.

Let’s now establish that the accepted samples indeed come from the distribution $\pi$ (asymptotically correct). Suppose that the act of accepting a sample is denoted as $A$. Therefore, for any given sample $s$, we can state that:

The denominator $p(A)$ is the unconditional acceptance probability and is equal $\frac{Z}{M}$ (For a more thorough explanation, please refer to the detailed derivation available at this wiki):

The unconditional acceptance probability $p(A)$ tell us, we need on average $\frac{M}{Z}$ draws to get an accepted sample, which can be very inefficient when $M$ is very large. And especially when we have random variable $x$ in high dimension, rejection sampling suffers from the “curse of dimensionality”³.

Moreover, in acceptance-rejection method, each draw is independent and information from the last draw is discarded completely (i.e, we may want to explore its neighbor if the sample from last draw get accepted). It’s natural to consider adding some correlation between the closest samples, and this is when Markov chain (or the Markov transition kernel $p(x’|x)$) is introduced.

Recap of Markov Chain

A (discrete-time) Markov chain is a sequence of random variables $X_1, X_2, X_3, …$ with the Markov property:

and a Markov process is uniquely defined by its transition probabilities $p(x’| x)$.

In our case, we always want the Markov chain to have a unique stationary distribution that is indeed $\pi$.

Markov Chain Properties

• Prior $p(X_1)$ and transition probabilities $p(X_{t+1} | X_{t} )=p(x’| x)$ independent of $t$.
• Ergodic Markov Chains: there exists a finite $t$ such that every state can be reached from every state in exactly $t$ steps.
• Stationary Distributions: $\lim\limits_{T \rightarrow \infty} p\left(X_{T}=x\right)=\pi(x)$, which is independent of $p(X_1)$.

Detailed Balance Equation

A chain satisfies detailed balance⁴ if we have:

This means in the in-flow to state $x’$ from $x$ (the probability of being in state $x$ times the transitional probability from $x$ to $x’$) is equal to the out-flow from state $x’$ back to $x$, and vice versa (Or each elementary process is in equilibrium with its reverse process at equilibrium / stationary). If a chain with transitional probability $p\left(\mathbf{x}^{\prime} | \mathbf{x}\right)$ satisfies the above detailed balance, then $\pi(\mathbf{x})$ is its stationary distribution. And it worth to mention that detailed balance is in general sufficient but not necessary condition for stationarity.

Metropolis-Hastings Methods

In the previous section, we learnt that the key of rejection sampling is proposal and decision rule. These principles remain applicable in MH techniques. In MH methods, instead generate samples independently from proposal $g$, the proposal (with Markov property) $g(x_{t+1}|x_t)$ generate new candidate based on the current sample value (The samples are correlated.). And we still need the decision rule to ensure $\pi$ is the stationary distribution.

Again, let’s have a look at the algorithm first, which is similar to rejection sampling. For each current state $x_t$:

1. sample $x’$ from proposal $g(x’|x=x_t)$.
2. accept $x_{t+1}=x’$ with probability $p_{\mathrm{accept}}(x’) := p_{\mathrm{accept}}(x’|x)$, set $x_{t+1}=x_{t}$ with probability $1-p_{\mathrm{accept}}(x’)$.

where the probability $p_{\mathrm{accept}}(x’)$ is defined as:

The transition probability of the Markov chain defined in the MH algorithm can be written as:

We need to ascertain that our proposal along with the decision rule satisfies the detailed balance equation (\ref{Detailed_Balance}). To accomplish this, we can verify (left as exercise) that the left and right sides of the equation are equal using the target distribution $\pi$ and the transition probability $p(x’|x)$, which is defined in eq(\ref{transition_kernel}).

Rejection Sampling as Special Case

When the proposal $g(x’|x)$ is independent of previous state $x_t$ — in other words, $g(x’|x) = g(x’)$ is an independence sampler — we end up with the method rejection sampling introduced in the beginning.

Random Walk Metropolis-Hastings⁵

When the proposal $g(x’|x)$ is Gaussian, we get random walk Metropolis-Hastings (RWM), where the new state is a random walk based on the previous state. This proposal can be written as following:

where $\tau$ is a scale factor chosen to facilitate rapid mixing, often referred to as the random walk step size. This step size plays a crucial role in balancing the trade-off between exploration, which is essential for covering the distribution, and maintaining an acceptable acceptance rate.

RR01b prove that, if the (target) posterior is Gaussian, the asymptotically optimal value is to use $\tau^2 = 2.382/D$, where $D$ is the dimensionality of $x$; this results in an acceptance rate of 0.234, which (in this case) is the optimal tradeoff …

—— Probabilistic Machine Learning: Advanced Topics by Kevin Patrick Murphy. MIT Press, 2023.

Since the proposal $g(x’|x)$ is now symmetric with $g(x’|x)=g(x|x’)$, we have simplified $p_{\mathrm{accept}}(x’) = \min \left(1, \frac{\tilde\pi\left(x^{\prime}\right)}{\tilde\pi(x)}\right)$. This simplified acceptance probability is more easily understood: if $x’$ is more probable than $x$, we unquestionably transition there, otherwise, we may still move there anyway by change, depending on the relative probabilities.

However, in high-dimensional space⁶, this classical MH algorithm is not applicable as it will end up without moving at all for a long time^7,8. This is because due to the concentration of measure, most of the probability mass are concentrating around the typical set which is a narrow crust shape area away from the (high density) mode; almost all proposed jumps will be outside — outside the crust and away from the mode because the ratio between volume of n-ball and n-cube tends to 0⁶ (thus there is almost 0 probability of falling inside the ball from the sphere) — the typical set and thus will be rejected for mode-seeking algorithms such as MH.

Metropolis-Adjusted Langevin Algorithm (MALA)

The well-known Metropolis-Adjusted Langevin Algorithm, which we will revisit in the next post, can be considered as MH with improved proposals:

where the additional term $\tau \nabla \tilde{\pi}(x)$ drives the samples toward areas with high probability density, which makes the proposed samples more likely to be accepted.

Remarks

In this post we mainly discussed the basic idea of MCMC method, especially the propose-accept scheme in Metropolis-Hastings algorithm. In the next posts, we will explore a more intuitive family of methods, the dynamic-based MCMC.

The last post was written around one year ago, when I decided to switch my semester project topic from style transfer with normalizing flow to image restoration with diffusion models.

In my current engineering-oriented master’s thesis, I find myself longing for the elegance of the theory of diffusion models. As a solution, I have made the decision to dedicate my spare time to learning Bayesian sampling (sampling method for bayesian inference).

Hence I will write a series of posts to record this learning process and to improve my understanding of this topic by reorganizing my knowledge. In this post, I will start with some basics and a brief introduction to this topic.

Introduction

Bayesian Inference Problem & Challenging

Figure 1: Visual Representation of Bayes' Theorem.

In the probabilistic approach to machine learning, all unknown quantities — be they predictions about the future, hidden states of a system, or parameters of a model — are treated as random variables, and endowed with probability distributions. The process of inference corresponds to computing the posterior distribution over these quantities, conditioning on whatever data is available.

—— Probabilistic Machine Learning: Advanced Topics by Kevin Patrick Murphy. MIT Press, 2023.

To be specific, we assume the dependence between unknown random latent variable $z$ and the available data $x$ is probabilistic and what we want to do is to estimate $z$ given $x$ with $p(z| x)$, which we call posterior. Most of the time we only have some prior knowledge about $z$ as $p(z)$ and the likelihood model $p(x|z)$, or equivalently the deep latent variable model $p_\theta(x, z)=p_\theta(z)p_\theta(x|z)$ through Bayesian Modeling, where $\theta$ can be estimated with maximizing likelihood (Maximizing the expected log-likelihood is equivalent to minimizing the Kullback-Leibler (KL) divergence between $p_{\mathrm{data}}(\mathbf{x})$ and $p_{\theta}(\mathbf{x})$):

where the constant term is the entropy of the data distribution. While the entropy of the empirical data distribution goes to infinity, it is independent of $\theta$ and can be ignored in optimization.

We can compute the posterior $p_{\theta}(z| x)$ using Bayes’s rule (see Figure 1 for visual illustration):

where the evidence term (also called marginal likelihood) served as a normalization constant in the denominator can be formulated as:

Beyond point estimation (MLE, MAP), we can use the posterior distribution to get posterior expectations of any function $f(z)$, such as mean and marginals. For instance, predicting new output in Bayesian linear regression where $w$ represents the coefficient that we aim to estimate its posterior distribution given the data: $y^\star =\int p(y^\star|x^\star, w)p(w | X,y)dw$.

However, this integral is usually analytically intractable¹ to calculate or evaluate, which leads to intractable posterior, and most Bayesian inference requires numerical approximation of such intractable integrals.

Two Approaches: Variational Inference & (MCMC) Sampling

In this post, I go through the two primary methodologies utilized to address the problem of Bayesian inference: Variational Inference (VI) and Markov Chain Monte Carlo (MCMC). I will strive to cover the most significant concepts associated with VI, and also provide a brief introduction to MCMC.

Variational Inference

Variational inference (VI) is a method in machine learning that approximates complex probability distributions by finding the most similar, simpler and hence tractable distribution $q(z)$ from a specified family $\mathcal{Q}$, thereby enabling efficient computation and handling of uncertainty.

Most of the time when referring to Variational Inference (VI), we are discussing parametric VI. In parametric VI, we use a parameter $\phi$ to represent the variational distribution $q_\phi(z|{x})$. There is another type of VI called particle-based VI, which utilizes a set of particles $\{z^{(i)}\}_{i=1}^{N}$ to represent the variational distribution $q(z|{x})$.

Figure 2: Illustration of Variational Inference.

The main idea of variational methods is to cast inference as an optimization problem. The goal of VI is to approximate an intractable probability distribution, so as to find $q_{\phi} \in \mathcal{Q}$ that minimize some discrepancy $D$ (here we use the KL divergence) between $q_{\phi}({z}|{x})$ and $p_{\theta}({z}|{x})$:

The challenge here is that we still don’t know the true posterior $p_{\theta}({z}|{x})$, and the KL divergence is intractable to compute. Luckily, we can rewrite the KL divergence in a way that makes it easier to optimize:²:

where the log evidence $\log p_\theta({x})$ does not change with the choice of $\phi$ during variational inference.

For convention, we can rewrite this as:

Since the KL divergence between $q_{\phi}({z}|{x})$ and $p_{\theta}({z}|{x})$ is non-negative, the first term in the RHS of eq(\ref{evidence2}) is a lower bound of the log evidence term $\log p_\theta({x})$, which is named variational lower bound, also called the Evidence Lower BOund (ELBO):

Therefore, the common mission to find the optimal $q_{\phi}({z}|{x})$ that minimizes the KL divergence (approximates $p_{\theta}({z}|{x})$) is equivalent to maximize the ELBO (without worrying about the evidence term in $p_{\theta}({z}|{x})$), and we can optimize it w.r.t. both ${\phi}$ and ${\theta}$ (when $\theta$ is unknown) in algorithms such as variational EM.

As $\theta$ is the model parameters, optimizing $\theta$ is essential for learning the underlying model that generates the data. When $\theta$ is tunable, we can jointly optimize $\phi$ and $\theta$ to maximize the ELBO. This joint optimization can help in finding a more accurate posterior approximation ($\phi$) and model parameters that better explain the data ($\theta$ that gives tighter ELBO of likelihood).

We can rewrite the ELBO as follows³:

Additionally, the ELBO can be reorganized and interpreted as the following in Variational AutoEncoder (VAE)^4,5, where $\phi$ and $\theta$ represent the encoder and decoder, respectively:

Suppose $p_{\theta}({x}|{z})=\mathcal{N}(\mu_\theta(z),\sigma^2)$, and $q_{\phi}({z}|{x})$ is a deterministic mapping $\psi_{\phi}(x)$. The first term is a reconstruction error, proportional to $|| x−\mu_\theta(\psi_{\phi}(x))||^2$. And the training objective on given dataset is hence⁶:

The optimization of the above equation (\ref{objective}) usually involve taking gradient w.r.t. $\phi$, which is more difficult as we cannot swap the gradient and the expectation like when taking gradient w.r.t. $\theta$. To resolve this issue, we can use methods like score function estimator and the reparametrization trick.

Markov Chain Monte Carlo

Unlike VI which solves inference with optimization, MCMC tackles it via sampling techniques. More specifically, MCMC applies Monte Carlo methods to generate a sufficient number of samples for an accurate estimation of the posterior distribution. However, it is almost always impossible to directly do so. As a solution, we can use MCMC, which is aimed at simulating a Markov chain whose stationary distribution is $p_{\theta}({z}|{x})$ and hope a fast convergence. And guess what, we only need unnormalized probability density (e.g. $p(x,z)$) to simulate the chain!

A more general problem setting is: sampling (=generating new examples) from a target distribution $\pi$ over $\mathbb{R}^d$ whose density is known up to an intractable normalization constant $Z$⁷:

where $\tilde{\pi}$ is the known unnormalized distribution, $\beta$ is an arbitrary positive constant akin to an inverse temperature, and $U(\cdot)$ can be treated as energy function. To make the notation consistent, now we rewrite the problem (\ref{KL}) as:

where $D$ is a dissimilarity functional such as KL divergence, and $\mathcal{F}_{\pi}(\mu)$ is a shorthand of $D(\mu \lVert \pi)$. We can approximate integrals $\int f(\cdot) d\pi$ of any function $f(\cdot)$ with samples from the Markov chain as $\frac{1}{n}\Sigma_{i=b}^{b+n-1}f(x_i)$, where $b,n$ are sufficiently large integers, and $b$ is called the mixing time or burn-in time. Note that the initial samples from the chain should be discarded because they do not come from the stationary distribution; reducing this is one of the most important factors in securing a fast convergence.

The ultimate goal of this series of posts is exactly to learn various MCMC techniques to sample from $\pi^\star$!
It’s also highly recommended to try out this great website for fantastic MCMC animations first.

The future topics should include:

• MCMC basics
• Dynamics-based MCMCs
• Relation to Wasserstein gradient flows