从 Distribution Matching 看 Drifting Model

论文: arXiv · 代码（我的 PyTorch 复现）: Drifting-Model-Pytorch

Kaiming 的 Generative Modeling via Drifting [1] 想解决一个非常直接但长期棘手的问题：如何在保持高生成质量的同时，把推理（inference）从多步迭代变成 one-step（1-NFE）。

与 diffusion / flow 体系“推理时演化分布”不同，Drifting Model 的核心是把分布演化搬到训练时：训练过程中生成分布 \(q_\theta\) 随着参数更新逐步“漂移（drift）”到数据分布 \(p_r\)，而推理时只需一次前向 \(x=f_\theta(\epsilon)\)。

本文尝试从 Distribution Matching 视角，把 Drifting Model 的漂移场 \(\mathbf V\) 与“真实/生成 score 的差”联系起来，从而理解其设计动机与数学形式。

1. one-step generator 与分布

数据空间：\(\mathbf x \in \mathbb R^d\)，真实数据分布记为 \(p_r(\mathbf x)\)
噪声空间：\(\boldsymbol\epsilon \in \mathbb R^c\)，先验噪声分布记为 \(p_{\text{noise}}(\boldsymbol\epsilon)\)
one-step 生成器：\(f_\theta:\mathbb R^c\to \mathbb R^d\)，生成样本：

\[\mathbf x = f_\theta(\boldsymbol\epsilon),\quad \boldsymbol\epsilon\sim p_{\text{noise}}\]

生成器诱导的 pushforward 分布（论文记法常写 \(q_\theta\) 或 \(p_f\)）：

\[q_\theta = (f_\theta)_\# p_{\text{noise}}\]

目标是让 \(q_\theta \approx p_r\)。

2. 基于 reverse KL 的 Distribution Matching

一个经典的分布匹配目标是 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]\]

把 \(\mathbf x=f_\theta(\boldsymbol\epsilon)\) 代入，可写为：

\[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

先定义 score：

\[\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)\]

对在 \(\mathbf x\) 处求梯度（把 \(\mathbf x\) 看成“粒子位置”变量），有：

\[\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)\]

因此，从“粒子位置”角度，如果我们能在每个 \(\mathbf x_i\) 上获得这个方向，就可以做梯度下降：

\[\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)\]

这条式子非常关键：让样本朝着真实分布的 score 方向走，同时抵消生成分布自身的 score，就会推动分布逐步从 \(q\) 靠近 \(p_r\)。

4. VSD/DMD 的思路：用 diffusion 去学 \(\mathbf s_r\) 与 \(\mathbf s_q\)

VSD (Variational Score Distillation) [2] 和 DMD（Distribution Matching Distillation）[3] 等方法的思路是利用 diffusion model 提供的 score 估计能力，近似获得：

\(\mathbf s_r(\mathbf x)\)：目标（真实）分布的 score，使用一个 teacher diffusion model 学习。
\(\mathbf s_q(\mathbf x)\)：生成分布的 score，使用一个 student diffusion model 在线学习。

5. Drifting Model 的思路：不用 diffusion score，而用 kernel 构造一个“漂移场” \(\mathbf V_{p,q}\)

Drifting Models 选择了一条不同路线：直接用 kernel 在采样点附近估计一种“吸引-排斥”的漂移场（drifting field），在训练时驱动样本分布演化，推理时自然 one-step。论文中漂移场的核心形式（为便于对照，这里用 \(p\) 表示真实 \(p_r\)，用 \(q\) 表示当前生成分布 \(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）}} - \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）}}\]

其中 \(k(\mathbf x,\mathbf y)\) 通常取 RBF / Gaussian kernel。论文使用的一个常见写法是（以温度 \(\tau\) 参数化）：

\[k(\mathbf x,\mathbf y)=\exp\Big(-\tau^{-1}\|\mathbf x-\mathbf y\|_2^2\Big)\]

并用 softmax 实现归一化（本质对应上式分母）。

6. Score Difference 与 Drift Field 的联系

为什么上面的 \(\mathbf V\) 会像 \(\mathbf s_r-\mathbf s_q\)？核心在于：KDE 下的 score 有闭式形式。令对分布 \(p\) 的核密度估计为：

\[\hat p(\mathbf x)=\frac1{n}\sum_{j=1}^n \exp\left(-\frac{\|\mathbf x-\mathbf y_j\|_2^2}{2\sigma^2}\right)\]

则其 score 为：

\[\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)}\]

对照可见：

\[\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 向量}} \approx \sigma^2 \nabla_{\mathbf x}\log \hat p(\mathbf x)\]

因此，如果我们分别用真实样本 \(\mathbf y^+\sim p_r\) 与生成样本 \(\mathbf y^-\sim q_\theta\) 做 KDE，那么：

\[\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)\]

这就把 Drifting Model 的核心设计，与 Distribution Matching 中 reverse KL 的“score 差方向”对齐起来：

reverse KL 的粒子下降方向：\(\mathbf s_r-\mathbf s_q\)
Drifting 的漂移场：用 kernel/mean-shift 构造的 \(\mathbf V_{p,q}\)，近似 \(\mathbf s_r-\mathbf s_q\)

其他形式的 Drift Field 行不行？

在 Drifting Model 的原始论文中，对 Drifting Field 的反对称性进行了 ablation，说明了其重要性。直观地想，除了 score difference，我们似乎还可以构造出许多其他形式的 Distribution Matching update rule（Drift Field），论文 DMVAE [4] 中讨论了不同形式（作者拍脑门想出来的）Distribution Matching Loss 形式和其结果，感兴趣的朋友可以去看原论文，在这里不多赘述。

7. Drifting Model 的训练

Drifting Models 不直接做“对 \(\mathbf x\) 的显式迭代”，而是把它写成一个训练期固定点更新：

\[\mathbf x_{t+1}=\mathbf x_t+\mathbf V_{p,q_t}(\mathbf x_t)\]

其中 \(\mathbf x_t=f_{\theta_t}(\boldsymbol\epsilon)\)。为了让网络学会产生“漂移后的样本”，论文采用 stop-gradient 的回归形式（固定点自举）：

\[\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\]

直观理解：

先用当前 \(f_\theta\) 生成 \(\mathbf x\)。
用真实 batch 与生成 batch 计算漂移 \(\mathbf V(\mathbf x)\)。
把目标设为 \(\mathbf x+\mathbf V(\mathbf x)\)（但 stopgrad，使其作为“固定 target”）。
训练 \(f_\theta\) 去回归这个 target。

这样随着训练迭代推进，生成分布 \(q_\theta\) 被持续“推向” \(p_r\)，直到达到平衡：若 \(q_\theta=p_r\)，则吸引与排斥抵消，\(\mathbf V\approx 0\)。

总结

总结本文视角：

reverse KL 更新告诉我们：匹配 \(q\) 到 \(p\) 的理想方向是 \(s_p-s_q\)。
DMD 用 diffusion 来学这个 score difference。
Drifting Model 则用 kernel/mean-shift 构造一个反对称的 drift field \(V_{p,q}\)，并通过固定点迭代让 SGD 在训练时推动分布演化，从而推理时天然 one-step。

因此，从 distribution matching 的角度看 Drifting Model，有三个关键点：

score difference 的替代物：\(V_{p,q}\approx s_p-s_q\)
反对称性：保证 \(p=q\Rightarrow V=0\)（equilibrium）
训练时分布演化：把多步推送“藏进”训练迭代里，推理阶段只需 1 次 forward。

Reference

[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.

[4] Distribution Matching Variational AutoEncoder.