Understanding Drifting Models Through Distribution Matching
Paper: arXiv · Code (my PyTorch implementation): Drifting-Model-Pytorch
Kaiming’s Generative Modeling via Drifting [1] addresses a straightforward yet long-standing challenge: how to turn inference from a multistep iterative process into a single step (1-NFE), while maintaining high generation quality.
Unlike diffusion and flow frameworks, which evolve the distribution during inference, the central idea of Drifting Models is to move distribution evolution into training: during training, the generated distribution \(q_\theta\) gradually “drifts” toward the data distribution \(p_r\) as the parameters are updated. At inference time, only a single forward pass is needed: \(x=f_\theta(\epsilon)\).
This article takes a distribution matching perspective, connecting the drifting field \(\mathbf V\) in Drifting Models to the difference between the scores of the real and generated distributions, to understand the motivation and mathematical form of the method.
1. One-Step Generators and Their Distributions
- Data space: \(\mathbf x \in \mathbb R^d\), with the real data distribution denoted by \(p_r(\mathbf x)\).
- Noise space: \(\boldsymbol\epsilon \in \mathbb R^c\), with the prior noise distribution denoted by \(p_{\text{noise}}(\boldsymbol\epsilon)\).
- One-step generator: \(f_\theta:\mathbb R^c\to \mathbb R^d\), producing samples as follows:
- The pushforward distribution induced by the generator (often written as \(q_\theta\) or \(p_f\) in the paper’s notation):
The goal is to make \(q_\theta \approx p_r\).
2. Distribution Matching Based on Reverse KL
A classical distribution matching objective is reverse KL:
\[D_{\mathrm{KL}}(q_\theta \,\|\, p_r) = \int q_\theta(\mathbf x)\log\frac{q_\theta(\mathbf x)}{p_r(\mathbf x)}\,d\mathbf x = \mathbb E_{\mathbf x\sim q_\theta}\big[\log q_\theta(\mathbf x)-\log p_r(\mathbf x)\big]\]Substituting \(\mathbf x=f_\theta(\boldsymbol\epsilon)\) gives:
\[D_{\mathrm{KL}}(q_\theta \,\|\, p_r) = \mathbb E_{\boldsymbol\epsilon\sim p_{\text{noise}}} \big[\log q_\theta(f_\theta(\boldsymbol\epsilon))-\log p_r(f_\theta(\boldsymbol\epsilon))\big]\]3. Score Difference
First, define the scores:
\[\mathbf s_r(\mathbf x)=\nabla_{\mathbf x}\log p_r(\mathbf x), \qquad \mathbf s_q(\mathbf x)=\nabla_{\mathbf x}\log q_\theta(\mathbf x)\]Taking the gradient with respect to \(\mathbf x\), treating \(\mathbf x\) as a “particle position” variable, gives:
\[\nabla_{\mathbf x}\big(\log q_\theta(\mathbf x)-\log p_r(\mathbf x)\big) = \mathbf s_q(\mathbf x)-\mathbf s_r(\mathbf x)\]Thus, from the particle-position perspective, if we can obtain this direction at each \(\mathbf x_i\), we can perform gradient descent:
\[\mathbf x_{i+1} = \mathbf x_i - \eta\big(\mathbf s_q(\mathbf x_i)-\mathbf s_r(\mathbf x_i)\big) = \mathbf x_i + \eta\big(\mathbf s_r(\mathbf x_i)-\mathbf s_q(\mathbf x_i)\big)\]This equation is crucial: moving samples in the direction of the real distribution’s score while canceling out the generated distribution’s own score pushes the distribution gradually from \(q\) toward \(p_r\).
4. The VSD/DMD Approach: Using Diffusion to Learn \(\mathbf s_r\) and \(\mathbf s_q\)
Methods such as VSD (Variational Score Distillation) [2] and DMD (Distribution Matching Distillation) [3] use the score-estimation capabilities of diffusion models to approximately obtain:
- \(\mathbf s_r(\mathbf x)\): the score of the target (real) distribution, learned using a teacher diffusion model.
- \(\mathbf s_q(\mathbf x)\): the score of the generated distribution, learned online using a student diffusion model.
5. The Drifting Model Approach: Constructing a Drifting Field \(\mathbf V_{p,q}\) with Kernels Instead of Diffusion Scores
Drifting Models take a different approach: using kernels to directly estimate an attraction–repulsion drifting field around sampled points, driving the sample distribution to evolve during training and naturally allowing one-step inference. The core form of the drifting field in the paper is as follows (for ease of comparison, we use \(p\) for the real distribution \(p_r\) and \(q\) for the current generated distribution \(q_\theta\)):
\[\mathbf V_{p,q}(\mathbf x) = \underbrace{ \frac{\mathbb E_{\mathbf y^+\sim p}\big[k(\mathbf x,\mathbf y^+)(\mathbf y^+-\mathbf x)\big]} {\mathbb E_{\mathbf y^+\sim p}\big[k(\mathbf x,\mathbf y^+)\big]} }_{\text{attraction toward real samples}} - \underbrace{ \frac{\mathbb E_{\mathbf y^-\sim q}\big[k(\mathbf x,\mathbf y^-)(\mathbf y^--\mathbf x)\big]} {\mathbb E_{\mathbf y^-\sim q}\big[k(\mathbf x,\mathbf y^-)\big]} }_{\text{repulsion from generated samples}}\]Here, \(k(\mathbf x,\mathbf y)\) is typically an RBF/Gaussian kernel. One common form used in the paper is parameterized by a temperature \(\tau\):
\[k(\mathbf x,\mathbf y)=\exp\Big(-\tau^{-1}\|\mathbf x-\mathbf y\|_2^2\Big)\]Normalization is implemented using softmax, which essentially corresponds to the denominators above.
6. The Connection Between Score Difference and the Drifting Field
Why does the \(\mathbf V\) above behave like \(\mathbf s_r-\mathbf s_q\)? The key is that the score under kernel density estimation (KDE) has a closed-form expression. Let the kernel density estimate of a distribution \(p\) be:
\[\hat p(\mathbf x)=\frac1{n}\sum_{j=1}^n \exp\left(-\frac{\|\mathbf x-\mathbf y_j\|_2^2}{2\sigma^2}\right)\]Its score is then:
\[\nabla_{\mathbf x}\log \hat p(\mathbf x) = \frac{\sum_{j} \exp\left(-\frac{\|\mathbf x-\mathbf y_j\|_2^2}{2\sigma^2}\right)\cdot \frac{(\mathbf y_j-\mathbf x)}{\sigma^2}} {\sum_{j} \exp\left(-\frac{\|\mathbf x-\mathbf y_j\|_2^2}{2\sigma^2}\right)} = \frac1{\sigma^2}\cdot \frac{\sum_{j} k_\sigma(\mathbf x,\mathbf y_j)(\mathbf y_j-\mathbf x)} {\sum_{j} k_\sigma(\mathbf x,\mathbf y_j)}\]Comparing the expressions, we see that:
\[\underbrace{ \frac{\sum_{j} k_\sigma(\mathbf x,\mathbf y_j)(\mathbf y_j-\mathbf x)} {\sum_{j} k_\sigma(\mathbf x,\mathbf y_j)} }_{\text{mean-shift vector}} \approx \sigma^2 \nabla_{\mathbf x}\log \hat p(\mathbf x)\]Therefore, if we perform KDE separately using real samples \(\mathbf y^+\sim p_r\) and generated samples \(\mathbf y^-\sim q_\theta\), we obtain:
\[\mathbf V_{p,q}(\mathbf x) \approx \sigma^2\Big(\nabla_{\mathbf x}\log \hat p_r(\mathbf x)-\nabla_{\mathbf x}\log \hat q_\theta(\mathbf x)\Big) \approx \sigma^2\big(\mathbf s_r(\mathbf x)-\mathbf s_q(\mathbf x)\big)\]This aligns the core design of Drifting Models with the “score-difference direction” of reverse KL in distribution matching:
- The particle descent direction for reverse KL: \(\mathbf s_r-\mathbf s_q\).
- The drifting field in Drifting Models: \(\mathbf V_{p,q}\), constructed using kernels/mean-shift, approximates \(\mathbf s_r-\mathbf s_q\).
Could Other Forms of Drifting Field Work?
The original Drifting Models paper includes an ablation study of the antisymmetry of the drifting field, demonstrating its importance. Intuitively, beyond score differences, it seems possible to construct many other distribution matching update rules (drifting fields). The DMVAE paper [4] discusses different forms of distribution matching loss—devised heuristically by its authors—and their results. Interested readers can refer to that paper; I will not go into further detail here.
7. Training Drifting Models
Drifting Models do not perform explicit iterations over \(\mathbf x\) directly. Instead, they formulate a fixed-point update during training:
\[\mathbf x_{t+1}=\mathbf x_t+\mathbf V_{p,q_t}(\mathbf x_t)\]Here, \(\mathbf x_t=f_{\theta_t}(\boldsymbol\epsilon)\). To teach the network to produce the “drifted samples,” the paper uses a stop-gradient regression formulation (fixed-point bootstrapping):
\[\mathcal L(\theta) = \mathbb E_{\boldsymbol\epsilon} \left\| f_\theta(\boldsymbol\epsilon) - \operatorname{sg}\big( f_\theta(\boldsymbol\epsilon)+\mathbf V_{p,q_\theta}(f_\theta(\boldsymbol\epsilon)) \big) \right\|_2^2\]Intuitively:
- Use the current \(f_\theta\) to generate \(\mathbf x\).
- Compute the drift \(\mathbf V(\mathbf x)\) using a real batch and a generated batch.
- Set the target to \(\mathbf x+\mathbf V(\mathbf x)\), applying stop-gradient so that it acts as a fixed target.
- Train \(f_\theta\) to regress to this target.
As training proceeds, the generated distribution \(q_\theta\) is continually pushed toward \(p_r\) until equilibrium is reached: if \(q_\theta=p_r\), attraction and repulsion cancel out, giving \(\mathbf V\approx 0\).
Summary
To summarize the perspective developed in this article:
- Reverse-KL updates tell us that the ideal direction for matching \(q\) to \(p\) is \(s_p-s_q\).
- DMD uses diffusion to learn this score difference.
- Drifting Models construct an antisymmetric drifting field \(V_{p,q}\) using kernels/mean-shift. Through fixed-point iterations, SGD drives distribution evolution during training, naturally allowing one-step inference.
Viewed through distribution matching, Drifting Models therefore have three key aspects:
- A substitute for the score difference: \(V_{p,q}\approx s_p-s_q\).
- Antisymmetry: ensuring that \(p=q\Rightarrow V=0\) (equilibrium).
- Distribution evolution during training: the multistep pushing process is folded into the training iterations, leaving only one forward pass at inference time.
References
[1] Generative Modeling via Drifting.
[2] ProlificDreamer: High-Fidelity and Diverse Text-to-3D Generation with Variational Score Distillation.
[3] One-step Diffusion with Distribution Matching Distillation.