Generative Modeling via Drifting

URL

• paper: https://arxiv.org/pdf/2602.04770

TL;DR

• 本文提出一种单步图像生成的算法 drifting，可以做到一步从噪声得到图片，无需像 diffusion / flow matching 一样多轮迭代得到图片

Algorithm

损失函数设计

$$\mathcal{L} = \mathbb{E}_{\epsilon} \left[ \left\| \underbrace{ f_\theta(\epsilon) }_{\text{prediction}} - \underbrace{ \operatorname{stopgrad}\!\left( f_\theta(\epsilon) + V_{p,q_\theta}\!\left(f_\theta(\epsilon)\right) \right) }_{\text{frozen target}} \right\|^2 \right]$$

• $\epsilon$ 是随机噪声
• $f_\theta()$ 是生成模型
• $V_{p,q_\theta}$ 是漂移场

漂移场的计算

$$V_{p,q}(x) = \underbrace{V_p^+(x)}_{正样本吸引} \underbrace{-V_q^-(x)}_{负样本排斥}$$

• $p$ 表示真实数据分布
• $q$ 表示模型生成的数据分布

正样本吸引（来自真实数据）

$$V_p^{+}(x) = \frac{1}{Z_p} \mathbb{E}_{y^{+} \sim p} \left[ \underbrace{k(x, y^{+})}_{标量权重} \, \underbrace{(y^{+} - x)}_{向量方向} \right]$$

负样本排斥（来自生成分布）

$$V_q^{-}(x) = \frac{1}{Z_q} \mathbb{E}_{y^{-} \sim q} \left[ \underbrace{k(x, y^{-})}_{标量权重} \, \underbrace{(y^{-} - x)}_{向量方向} \right]$$

• $y^+$ 真实样本
• $y^-$ 生成样本
• $k(x,y)$ 相似度核函数，输出标量

相似度核函数

$$k(x,y)=exp(-\frac{\|x-y\|}{\tau})$$

• $\tau$ 表示温度

用一个例子说明损失函数的作用方式

• 当前生成样本（要算 drift 的点）

$$x=(0,0)$$

• 正样本（来自真实分布 $p$）两个点

$$y_1^+=(2,0),\ y_2^+=(2,2)$$

• 负样本（来自生成分布 $q$）两个点（通常就是 batch 里的其他生成样本）

$$y_1^-=(-1,0),\ y_2^-=(0,-1)$$

Step 1: 算 kernel 权重（谁更“像” x）

• 正样本的权重
  • 距离：
    • $\|x-y_1^+\|=2$
    • $\|x-y_2^+\|=\sqrt 8\approx 2.83$
  • kernel:
    • $k(x,y_1^+)=e^{-2}\approx0.135$
    • $k(x,y_2^+)=e^{-2.83}\approx0.059$
• 负样本的权重：
  • $k(x,y_1^-)=k(x,y_2^-)=e^{-1}\approx0.368$

Step 2: 算方向向量（往哪里走）

• 正样本方向：

$$y_1^+-x=(2, 0)\\ y_2^+-x=(2, 2)$$

• 负样本方向：

$$y_1^--x=(-1, 0)\\ y_2^--x=(0, -1)$$

Step 3: 算正向漂移 $V_p^+$

• 加权求和

$$\sum k(x,y^+)(y^+-x)=0.135(2, 0)+0.059(2, 2)=(0.388, 0.118)$$

• 归一化

$$Z_p=0.135+0.059=0.194\\ V_p^+(x)=\frac{1}{0.194}(0.388, 0.118)\approx (2.0, 0.61)$$

Step 4: 算负向漂移 $V_q^-$

• 和 $V_p^+$ 同理

$$V_q^-(x) = (-0.5,-0.5)$$

Step 5: 合并成为最终漂移 $V_{p,q}$

$$V_{p,q}(x)=V_p^+-V_q^-=(2.0, 0.61)-(-0.5,-0.5)=(2.5, 1.11)$$

每个模块的组成

• Generator: 采用 DiT (Diffusion Transformer) 架构。输入是噪声 + 类别条件。
• Tokenizer: 使用 Stable Diffusion 的 VAE 将图像压缩到 Latent Space ($32 \times 32 \times 4$)。
• Conditioning: 使用 AdaLN-Zero (Adaptive Layer Norm) 注入类别信息。这与标准 DiT 一致。
• Feature Extractor: 用于计算损失函数。论文指出损失是在特征空间中计算的，而非直接在像素空间。用的模型是 MAE 自监督训练得到的。

Thoughts

• 有很多 MoCo 的影子，比如 stop_grade 作为监督，对比学习等
• 还有 MAE，kaiming 大佬把自己的作品都串起来了，tql…