为漂移模型学习判别性几何 Learning Discriminative Geometry for Drifting Models
持续表示学习让像素空间漂移模型 FID 降低约 82–95%,无需预训练编码器
前置知识
Drifting Models(漂移模型)
一类单步生成模型,把扩散模型中推理阶段的迭代分布细化搬到训练阶段:训练时在当前 batch 的真实样本与生成样本之间构造核加权「漂移场」(对真实样本吸引、对生成样本排斥),让生成器向移动后的 stop-gradient 目标做回归;训练完成后推理只需一步前向即可出图。
本文的全部改进都围绕漂移模型展开,必须先理解其「由表示几何决定核权重、由核权重决定漂移」的训练机制,才能读懂为什么表示空间如此关键。
核密度估计(KDE)
非参数密度估计:用高斯核 $k_\tau(x,y) = (\pi\tau)^{-d/2}e^{-\|x-y\|_2^2/\tau}$ 在每个样本周围「涂抹」概率质量再求和。文中用 batch 内 KDE 分别平滑真实分布和生成分布,两者的对数密度比 $\ell(u)=\log\hat{p}_\tau-\log\hat{q}_\tau$ 同时充当真假分类打分与漂移方向的来源。
论文的核心洞察是漂移场本质上是 KDE 密度比的梯度($\hat{V}=\frac{\tau}{2}\nabla_u\ell$),而表示空间的距离决定核权重,理解这一点才能读懂「判别性几何」的含义。
Wasserstein 梯度流与密度比
Wasserstein 梯度流描述概率分布在最优传输几何下最速下降某个能量泛函(如 KL 散度)的演化方程,其速度场为 $\nabla\log p-\nabla\log q$。密度比 $p/q_\theta$ 刻画真实分布相对生成分布的过剩概率,是发散度 GAN 中最优判别器所编码的对象。
漂移场是 KL 梯度流的平滑近似,速度裁剪的动机(势函数局部正则性上界控制速度)直接建立在这个理论框架上。
GAN 与 Lipschitz 约束
生成对抗网络训练判别器区分真假样本,最优判别器编码密度比;谱归一化、梯度惩罚等手段约束判别器的 Lipschitz 常数以控制其输入梯度幅度。f-GAN 与 MonoFlow 表明判别器输出经函数变换后对应某种散度的变分形式,其输入梯度为生成器提供分布匹配方向。
本文的 KDE 分类损失是判别目标的无参数版本,Proposition 1 把它与非参数漂移打通;速度裁剪与谱归一化的对照实验正是对「该约束什么」这一 GAN 老问题的回应。
FID 与 IS
FID(Fréchet Inception Distance)用 Inception 网络特征上真实与生成分布的高斯拟合距离衡量生成质量,越低越好;IS(Inception Score)同时衡量样本类别清晰度与多样性,越高越好。两者都是图像生成的标准评估指标,论文在 5 万张生成图上计算。
所有核心结论都通过这两个指标表达,例如 CIFAR-10 FID 从 124.94 降到 12.81,读不懂 FID 就无法评估论文贡献的大小。
残差参数化与 DRN
编码器采用 $E_\phi(x)=x+R_\phi(x)$ 的残差形式,$R_\phi$ 为膨胀残差网络(膨胀率 1,2,4,1、128 隐藏通道、SiLU 激活、无下采样、无归一化层),最后一层投影零初始化使 $E_{\phi_0}$ 严格等于恒等映射。
恒等初始化保证训练从标准像素漂移平滑出发、表示几何随训练渐进演化,是从零学习表示而不破坏早期训练稳定性的实现关键。
研究动机
Drifting Models 通过把迭代分布细化搬进训练阶段实现一步生成,但在复杂数据集上的成败几乎完全取决于漂移场的构造空间:直接在像素空间构造时训练基本失效——CIFAR-10 上 5 万步 FID 高达 124.94、IS 仅 3.18,SVHN 上 FID 61.56;而把漂移场放进预训练特征空间(如 MoCo-v2)则质量大幅提升。这一对比背后的机制此前不清楚:为什么预训练表示能诱导更有效的漂移?固定像素几何下 KDE 密度比几乎无法区分真假样本(分类准确率接近 0.5 的随机水平),导致生成器收到的监督信号近乎无效;而且这并非核带宽调参能解决的问题——把温度 $\tau$ 从 0.002 扫到 2,像素空间 FID 仍徘徊在 134–376 之间。
本文的目标是本文要回答两个层层递进的问题。机制层面:表示空间究竟如何塑造漂移场的估计质量,进而影响 KDE 判别力与生成器更新——为什么预训练特征有效而像素空间失败。算法层面:能否不依赖任何固定外部特征,让合适的表示几何随生成器一起从漂移目标本身涌现,即从像素出发端到端训练出高质量的单步生成器。同时,论文希望补上密度比 GAN 优化与非参数漂移之间缺失的理论联系,并对表示持续演化引入的训练不稳定给出直接的控制手段。
与已有工作不同的是,已有漂移研究(Generative Modeling via Drifting、DriftXpress、W-Flow、基于预训练特征的单步蒸馏)全部使用固定表示,表示如何随训练自适应是公认的开放问题;GAN 文献虽已证明判别所用的表示至关重要(KDD-GAN 用学习表示、Projected GAN 用固定预训练表示),但漂移与密度比 GAN 之间的算法联系尚未建立。本文的独特切入是把漂移场重写为 KDE 密度比势函数的梯度,并证明匹配条件下两种损失的生成器梯度精确相等,从而把「学表示」(等价于学判别器)与「控速度」(等价于 Lipschitz 控制)统一进漂移训练循环。
核心方法
直觉上,漂移场的质量取决于「谁离谁近」:表示空间中的距离决定每个真实/生成样本获得的核权重,进而决定局部密度估计与漂移方向。本文让这个几何变得可学习并随生成器共同演化:引入恒等初始化的轻量残差编码器 $E_\phi$,每个生成器更新前先用 KDE 二分类 BCE 损失更新编码器 $K$ 步,再冻结编码器、在特征空间构造漂移场并回归,如此交替贯穿训练。理论上,高斯核下 KDE 比值损失与漂移回归损失在匹配条件下梯度完全相等:$\nabla_\theta\hat{L}^{KDE}=\nabla_\theta\hat{L}^{Drift}=-\frac{2}{B}\sum_i J_G^\top J_E^\top\hat{V}_\phi(u_i)$,把判别几何与生成器监督直接挂钩。又因 $\hat{V}_\phi(u)=\frac{\tau}{2}\nabla_u\ell_\phi(u)$,势函数局部正则性控制漂移速度,据此引入径向投影式速度裁剪稳定训练。
核心创新是持续表示学习:编码器参数跨 batch 保留、逐批更新,使判别几何随生成器共同演化,而 KDE 与漂移场每步从当前 batch 重建,历史信息保留在共享参数中而无需存储旧样本。这与已有工作的本质区别在于:Drifting Models 及其后续(包括特征空间蒸馏版本)把表示当作固定的外部资源,本文则证明表示本身可以被漂移目标持续改进——这解释了预训练特征有效的真正原因是其判别几何更好,而非「语义」本身有魔力。消融还揭示一个反直觉结论:持续性必须放在表示上,仅持续标量判别器反而严重恶化(Figure 5 中 FID 冲破 500),与 GAN 中持续训练判别器的常规直觉相反,说明表示级持续性与标量势持续性有本质差异。
方法步骤详情
完整训练流程(Algorithm 1,每个迭代):第一步,循环 $K$ 次编码器更新——采样 $z_i\sim\mu$ 与真实样本 $\tilde{x}_j$,生成样本经 stop-gradient 得 $\hat{x}_i=sg(G_\theta(z_i))$,映射为特征 $u_i=E_\phi(\hat{x}_i)$、$v_j=E_\phi(\tilde{x}_j)$;第二步,以等类先验的 KDE 后验 $P(y=real|u)=\sigma(s^{KDE}_\phi(x))$ 为分类 logit,计算二分类 BCE 损失(式 7),梯度穿过 KDE 更新 $\phi$(主实验 $K=1$、lr $10^{-4}$);第三步,重新采样新鲜 batch,在特征空间构造漂移场 $\hat{V}_\phi(u)$(式 5);第四步,速度裁剪 $\hat{V}_{\phi,clip}=\hat{V}_\phi/\max(1,\|\hat{V}_\phi\|_2/V_{max})$(式 8,$V_{max}=0.005$);第五步,以 $\hat{u}_i=sg(u_i+\hat{V}_{\phi,clip}(u_i))$ 为回归目标,用式 6 更新生成器。生成器为 U-Net,每 GPU batch 1500,4×GH200 训练 5 万步。
技术新颖性
三点新颖性。其一,首次在漂移框架内让表示随生成器联合学习并保持跨批持续性:恒等初始化、从零训练的几何即可把像素漂移的 FID 降低 82–95%,闭合了 Zhang et al. 留下的「无预训练特征从头训练」开放问题。其二,Proposition 1 给出匹配条件下(相同锚点、带宽与自交互处理,仅经生成查询反传)KDE 比值损失与漂移回归损失的当前步梯度精确等式,首次把 f-GAN/MonoFlow 一脉的密度比生成器优化与非参数漂移在算法层面打通,并精确指出编码器通过核权重与 $J_E^\top$ 拉回两个机制改变生成器监督。其三,把 Lipschitz 控制的位置从判别器权重(谱归一化)转移到显式漂移速度:实验显示谱归一化有害(FID 31.82 对无约束 27.41)而速度裁剪有益(21.57),为对抗式训练「约束什么」提供了新证据。
实验结果
主基准(Table 1,5 万步、4×GH200):SVHN 上本文 FID 2.88/IS 3.24 对标准漂移 61.56/2.91;CIFAR-10 上 12.81/8.69 对 124.94/3.18;CIFAR-100 上 20.61/8.52 对 113.14/3.46,全线超过 SNGAN(28.59/7.22 等),FID 降幅约 82–95%。机制分析(Figure 1、2):8 模环加 30 维干扰的玩具实验中,学习表示的漂移场与语义参考场余弦对齐显著更高;CIFAR-10 上像素 KDE 分类准确率贴近 0.5 随机线,冻结 MoCo-v2 同时改善判别(损失降至约 1.25)与生成质量,自适应其末残差块进一步提升。消融(Figure 5、Table 2、3,1 万步):无持续性 FID 150.49,仅速度裁剪 126.33,仅持续表示 27.41/IS 7.50,叠加 $V_{max}=0.005$ 达 21.57/7.80;谱归一化恶化至 31.82;持续标量判别器劣于不用。Figure 9 证明像素失效与带宽无关(FID 134–376);CKA 分析(Figure 7)显示表示从恒等渐进演化(50k 步 CKA 0.46,晚期检查点间大于 0.95);等时对比(Figure 10)5 小时内 21.57 对 124.94。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 单步图像生成(SVHN,5 万步) | FID ↓ | 2.88 | 标准漂移 61.56;SNGAN 13.16 | 较标准漂移约 -95%,同时优于 SNGAN |
| 单步图像生成(CIFAR-10,5 万步) | FID ↓ | 12.81 | 标准漂移 124.94;SNGAN 28.59 | 较标准漂移约 -90%,优于 SNGAN |
| 单步图像生成(CIFAR-100,5 万步) | FID ↓ | 20.61 | 标准漂移 113.14;SNGAN 30.33 | 较标准漂移约 -82%,优于 SNGAN |
| 单步图像生成(CIFAR-10,5 万步) | IS ↑ | 8.69 | 标准漂移 3.18;SNGAN 7.22 | 较标准漂移提升约 2.7 倍 |
| 消融:持续表示 + 速度裁剪(CIFAR-10,1 万步) | FID ↓ / IS ↑ | 21.57 / 7.80 | 无持续性 150.49 / 2.94;仅速度裁剪 126.33 / 3.44;仅持续表示 27.41 / 7.50 | 表示持续性贡献主要提升,速度裁剪再降约 21% FID |
| 等时训练 5 小时(CIFAR-10) | FID ↓ | 21.57 | 标准漂移 124.94 | 约 -83%,编码器额外开销被质量增益远超 |
局限与改进
作者承认两点局限:其一,从头学习的表示仍落后于预训练编码器,训练动态尚未恢复外部预训练所提供的全部有用结构;其二,受算力限制,实验只覆盖 SVHN、CIFAR-10、CIFAR-100 等 32×32 小规模数据集,方法在更大规模上的有效性尚未确立。我的补充观察:$V_{max}$ 是新增的敏感超参数(最优 0.005 比次优 0.05 的 FID 还低约 1.9),带来额外调参成本;持续标量判别器为何有害只有现象没有机制解释;KDE 依赖超大 batch(每 GPU 1500、共约 6000 样本),对显存受限场景不友好;每步增加编码器前后向开销,单位步成本更高,只是等时比较占优;此外预训练特征自适应实验只微调末残差块、学习率低至 $10^{-6}$,说明表示适应仍然相当脆弱。
独立分析的弱点
弱点一:学到的几何仍逊于预训练特征——CIFAR-10 从零训练 FID 12.81,而自适应 MoCo-v2 在更早训练阶段就给出更低损失与更好 FID,说明从零几何缺少外部知识,可把对比学习或掩码建模等自监督目标作为辅助损失注入编码器更新。弱点二:$V_{max}$ 是手工预算且敏感(五个预算中 0.005 最优),可改为按训练进度或当前漂移范数分布(Figure 8 显示中位数约 0.70)自动校准的分位数裁剪。弱点三:持续标量判别器灾难性失败(FID 冲破 500)缺乏解释,可能因标量势的输入梯度无界或与 KDE 打分几何冲突,值得做梯度范数追踪与理论刻画。弱点四:KDE 需要超大 batch(4×1500),限制可扩展性,可探索跨 batch 锚点缓存或低秩/随机特征近似来降低每步样本需求。弱点五:只在 32×32 分辨率验证,而 KDE 核带宽与维度相关,高分辨率下恒等残差编码器是否仍够用存疑。
未来方向
作者未单列未来工作段落,但从成果可自然延伸:其一,扩展到 ImageNet 级别高分辨率数据与更大生成器,检验持续表示学习与 KDE 的可扩展性,必要时结合论文已支持的多分辨率特征组(从输入、stem、各残差块输出提取);其二,把持续表示与一致性模型、渐进蒸馏等其他单步生成范式结合,作为其表示前端;其三,从理论上解释「持续标量判别器有害而持续表示有益」,可能需要分析两类势函数输入梯度的正则性差异;其四,设计自适应速度控制,如按样本或按训练阶段动态调整 $V_{max}$,用漂移统计反馈闭环;其五,探索更丰富的预训练适配策略(当前仅解冻末残差块、学习率 $10^{-6}$),例如分层解冻或 LoRA 式低秩适配。
复现评估
复现条件较好:论文给出开源代码链接(github.com/aalto-icl/DiscriminativeGeometryforDrifting),附录 A 的 Table 4 完整列出主基准与消融配置(U-Net 128 基础通道、AdamW、生成器学习率 $2\times10^{-4}$、编码器学习率 $10^{-4}$、每 GPU batch 1500、KDE 温度 0.05、$V_{max}=0.005$、训练 5 万步),附录 A.1 给出玩具实验(8 模环、32 维观测、学习率 $10^{-3}$、batch 256、温度 0.5)、持续判别器(SNGAN 结构、梯度范数裁剪 1)与 MoCo-v2 自适应(800 epoch 官方权重、仅末残差块、学习率 $10^{-6}$)的全部超参,附录 D/E 提供证明与生成器更新伪代码,透明度高。主要门槛是算力:主结果需 4×GH200 训练 5 万步,且未报告多种子方差。总体评估:有 A100/H100 集群即可复现主表,消融趋势可在单卡缩减 batch 下验证。
论文图表
玩具实验(8 个高斯模态的二维语义环加 30 维干扰维度):(a) 对同一批生成样本,原始像素空间的漂移场被干扰坐标搅得杂乱无章,而持续表示空间的漂移场与环形目标结构对齐、方向一致;(b) 训练过程中,表示空间漂移场与仅用语义坐标构造的参考场的余弦对齐度远高于原始空间并持续领先。
这是论文核心论点最直观的可视化:表示几何决定漂移场质量。它把抽象的「判别性几何」变成可看的向量场图,是理解动机(为什么像素漂移失败)的第一张图。
完整训练伪代码:外层循环内先用 $K$ 步新鲜 batch 以 KDE 分类损失更新编码器 $\phi$,再重新采样、以特征空间裁剪后漂移场构造 stop-gradient 目标 $\hat{u}_i=sg(u_i+\hat{V}_{\phi,clip}(u_i))$,用特征空间回归损失更新生成器 $\theta$。
方法的完整算法描述,把交替更新、stop-gradient、速度裁剪三个实现细节放在一处,是复现和真正理解训练循环的权威参考。