连续对抗流模型(Continuous Adversarial Flow Models) Continuous Adversarial Flow Models
在导数空间用对抗目标训练连续流模型,仅10个epoch后训练即大幅降低FID
前置知识
流匹配 (Flow Matching)
流匹配是一类连续时间生成模型,沿连接数据分布与噪声分布的线性插值路径 $x_t=(1-t)x+tz$ 训练神经网络预测速度场 $v_t=E[\bar v_t\mid x_t]$。训练目标是最小化与条件速度 $\bar v_t = -x+z$ 的均方误差,采样时通过反向积分ODE把噪声运送到数据。
本文正是对流匹配模型做后训练,读者必须先理解流匹配的速度场预测、ODE采样与MSE损失的具体含义,才能体会作者想替换的'欧氏判据'为何是错误的推广来源。
雅可比向量积 (Jacobian-Vector Product, JVP)
JVP是前向模式自动微分的运算: 给定函数 $D:\mathbb R^n\to\mathbb R$ 在 $x$ 处的雅可比 $J_D(x)$ 与切向量 $v$, 返回 $J_D(x)\cdot v$。本论文用 torch.func.jvp 计算 $D$ 沿速度方向的方向导数, 用作判别 logit。
CAFM的核心创新就是把判别器 $D(x_t,t)$ 的输出当作'标量势能',再用 JVP 取其在真实/生成速度方向上的方向导数作为判别 logit。读不懂JVP就无法理解导数空间对抗目标如何实现。
生成对抗网络 (GAN) 与对抗判别
GAN通过判别器 $D$ 与生成器 $G$ 的极小极大博弈训练: $D$ 区分真假样本, $G$ 试图生成让 $D$ 失灵的样本。常见对比损失 $f(a,b)=-\log\sigma(a-b)$,为稳定训练常需加入梯度惩罚 $R_1/R_2$ 与centering penalty等正则项。
CAFM本质上是把 GAN 的对抗机制搬到了连续流匹配的框架里。读者需要熟悉判别器损失与生成器损失的形式,以及为什么需要 centering penalty、为什么本文可以省去 gradient penalty。
连续归一化流 (Continuous Normalizing Flow, CNF) 与概率流ODE
CNF把生成过程建模为常微分方程 $\frac{dx_t}{dt}=G(x_t,t)$ 的反向积分,通过神经网络参数化速度场即可描述任意连续变换。流匹配是CNF的一种训练友好的特例,采用'条件流匹配'目标规避直接拟合边缘速度。
CAFM 的生成器 $G(x_t,t)$ 就是标准的 CNF 速度场网络。理解'CAFMs 与 FMs 学到相同 ground-truth flow,只因泛化方式不同而表现不同'这一论点,必须先在 CNF 框架下理解 marginal velocity 与 conditional velocity 的区别。
Fréchet Inception Distance (FID) 与 Inception Score (IS)
FID把生成图像与真实图像的 Inception-v3 池化特征视为高斯分布,计算两者的 Fréchet 距离,数值越低代表分布越接近。IS基于 Inception 输出的类别条件熵与边缘熵,数值越高代表样本既清晰又多类。两者是文生图与类条件生成的标配指标。
本文所有表格(Tab. 1-4)都以 FID 为主指标,读者需理解为什么 FID 大幅下降(如8.26→3.63)是生成质量上台阶的明确证据,以及 guidance-free 与 guided 两套指标各自的衡量重点。
研究动机
流匹配 (FM) 在图像、视频等视觉模态上取得了显著成功,但基于 MSE 的 L2 损失隐含了'欧氏距离'判据,与数据真实所在的高维流形结构并不一致。理论上,只要优化 $\arg\min_a \mathbb E_b[d(a,b)] = \mathbb E[b]$ 形式的判据(如平方 L2),网络都能收敛到 ground-truth marginal velocity field $v_t=E[\bar v_t\mid x_t]$;但有限容量网络必须'泛化',不同判据诱导出不同泛化方向。在视觉生成中,FM 常出现明显的 out-of-distribution 现象: 纹理模糊、边缘发虚,只有引入 classifier-free guidance (CFG) 强行压低采样温度才能让样本'看起来像数据',但代价是采样分布被人为扭曲。代价可量化: SiT-XL/2 在 ImageNet 256px 上 guidance-free 的 FID 高达 8.26,只有 CFG=1.5 才能压到 2.06,意味着模型并未真正学到数据流形,只是靠 guidance 临时矫正。
本文的目标是本文提出连续对抗流模型 (Continuous Adversarial Flow Models, CAFMs),目标是把对抗训练机制扩展到连续时间流模型中,让生成器在训练阶段就以'流形感知'的方式泛化,从而显著降低 guidance-free 的 FID。具体地,作者希望对已有的 FM 预训练模型(SiT-XL/2、JiT-H/16、Z-Image 文生图模型)做轻量后训练,只用 10 个 epoch 就把 SiT 的 guidance-free FID 从 8.26 压到 3.63、JiT 的从 7.17 压到 3.57,并把 guided 模式下的最佳 FID 也分别从 2.06/1.86 进一步压到 1.53/1.80,在文生图基准 GenEval (0.81→0.85) 与 DPG-Bench (83.7→85.2) 上同样取得稳定提升。
与已有工作不同的是，已有对抗流工作 (Adversarial Flow Models, AFMs) 只在离散时间下设计对抗目标,需要切分固定步长 $|t-s|$,在 $|t-s|\to 0$ 时训练会变得不稳定、目标函数也会失效;另一条路线 (Representational Autoencoders / RFM) 通过更换隐空间间接缓解问题,但需要在 VAE 空间上重新训练。本文选择了一条此前没人走通的路:把判别器 $D(x_t,t)$ 视作一个标量势函数,改用其在真实/生成速度方向上的方向导数 (JVP) 进行对抗判别,让对抗机制自然嵌入到连续 ODE 框架里。这套'导数空间对抗'既保持了和 FM 完全相同的 ground-truth flow,又仅通过替换训练目标让模型用一种更接近流形的方式泛化,与 RFM 改空间、AFM 离散化的路线都不同。
核心方法
CAFM 的整体思路是'同一份 ground-truth flow, 不同的泛化判据'。生成器 $G(x_t,t)$ 与 FM 一样预测连续速度场 $v_t$,训练时不再用 MSE,而是引入判别器 $D(x_t,t):\mathbb R^n\times[0,1]\to\mathbb R$ 作为'标量势能'。判别器不直接判 $x_t$ 真假,而用 Jacobian-Vector Product 取 $D$ 在 $(x_t,t)$ 处沿真实速度 $\bar v_t$ 与生成速度 $G(x_t,t)$ 的方向导数 $D_{\rm jvp}$ 作为判别 logit。直觉上 $D$ 学到势能曲面,真实流方向势能增加更快; $G$ 优化方向就是让势能增加最快,当 $G$ 收敛到 ground-truth flow 时, $D$ 势能曲面变平坦, $D_{\rm jvp}\to 0$。实现上 JVP 通过 torch.func.jvp 一次前向,配合 vmap 批量求值,用 $f(a,b) = (a-1)^2+(b+1)^2$ 计算对抗损失,配合 centering penalty 与 OT 正则。
核心创新点是把对抗判别从'输出空间'(判 $x_t$ 真假)搬到'导数空间'(判 $D$ 在速度方向上的方向导数大小),由此解决了三个本质问题:第一,沿用了 FM 的连续时间形式,不需要像 AFM 那样切分步长 $|t-s|$,因此在 $|t-s|\to 0$ 时目标函数仍然有效;第二,判别器只对'速度方向上的势能变化率'敏感,因此天然抓住了 $x_t$ 沿流场移动时的'瞬时一致性',这与视觉流形上'小步移动仍属于同一样本'的直觉一致;第三,通过 JVP 一次性前向算出 $D$ 与 $D_{\rm jvp}$,可以完全省去 AFM 中必需的 $R_1/R_2$ 梯度惩罚,作者在实验中明确报告 CAFM 训练不再依赖 gradient penalty,这在经验上大幅简化了调参。
方法步骤详情
CAFM 后训练流程: 输入 $x$、$z$、$t\in[0,1]$、$c$。第 1 步构造 $x_t=(1-t)x+tz$ 与 $\bar v_t=-x+z$。第 2 步 $G(x_t,t)$ 前向得到 $u_t$。第 3 步(更新 $D$)用 vmap+jvp 一次性得到 $D(x_t,t)$ 与 $D_{\rm jvp}$ 在两切向的值。第 4 步 $D$ 损失 $\mathcal L^D=\mathbb E[(d\bar v-1)^2+(du+1)^2]+\lambda_{cp}\,\mathbb E[o^2]$。第 5 步(更新 $G$) $\mathcal L^G=\mathbb E[(du-1)^2]+\lambda_{ot}\,\mathbb E[\\|u_t\\|^2/n]$,后训练 $\lambda_{ot}=0$。第 6 步 $G$ 用 EMA (0.99) 维护平均; $D$ 与 $G$ 按 $N=16:1$ 交替更新,前 2/4 epoch 冻结 $G$ 预热 $D$。$G$ 沿用 FM 预训练, $D$ 同结构但 LayerNorm 换 RMSNorm。
技术新颖性
技术新颖性可以归纳为三点。第一,把判别从'点比较'升级为'方向导数比较',这在 GAN/AFM 文献里均无先例,作者也强调这是首次把对抗训练系统地引入连续时间流模型。第二,通过 JVP 的前向模式自动微分,模型可以一次前向同时得到 $D$ 与 $D_{\rm jvp}$,再配合 vmap 对多切向批量求值,使得 CAFM 在实现上只需 30 行 PyTorch 代码 (Algorithm 1),并与 DDP/FSDP/gradient checkpointing 兼容,这一点对大规模训练非常关键。第三,把'最优传输正则'连续化为 $\frac{1}{n}\\|G(x_t,t)\\|_2^2$ 并论证其在 $|t-s|\to 0$ 极限下与 AFM 离散正则一致,为'为什么在导数空间下欧氏范数仍能起到 OT 作用'提供了形式化衔接。
实验结果
三组任务验证 CAFM 有效性。第一组 SiT-XL/2 后训练: guidance-free 时 SiT 原始 FID 8.26、IS 131.65, 经 10 epoch CAFM 后训练后 FID 降至 3.63、IS 升至 178.08 (提升 56%),带 CFG 时最佳 FID 从 2.06 压到 1.53; 对照 SiT+FM FID 漂到 8.64, 收益完全来自 CAFM。第二组 JiT-H/16 像素空间: guidance-free FID 从 7.17 压到 3.57, guided 最佳 FID 从 1.86 压到 1.80。第三组文生图: GenEval 整体分从 0.81 升到 0.85, DPG-Bench 从 83.7 升到 85.2。横向比较: SiT+CAFM 以 675M 把隐空间无 CFG FID 做到 3.63, 优于 DiT (9.62)、DDT-XL (6.27);JiT+CAFM 以 956M 把无 CFG FID 压到 3.57, 仅落后 2B SiD (2.77)。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| ImageNet 256px 类条件生成 (无guidance, SiT-XL/2 latent) | FID↓ / IS↑ | 3.63 / 178.08 (SiT+CAFM, 1400+10 epoch) | 8.26 / 131.65 (SiT-XL/2 原始) | FID 降低 56.1% (8.26→3.63), IS 提升 35.3% |
| ImageNet 256px 类条件生成 (有guidance, SiT-XL/2) | FID↓ (最佳CFG档位) | 1.53 (SiT+CAFM, CFG=1.3) | 2.06 (SiT-XL/2, CFG=1.5) | FID 降低 25.7%, 最佳CFG 档位从 1.5 下移到 1.3 |
| ImageNet 256px 类条件生成 (无guidance, JiT-H/16 pixel) | FID↓ / IS↑ | 3.57 / 198.08 (JiT+CAFM, 600+10 epoch) | 7.17 / 151.54 (JiT-H/16 原始) | FID 降低 50.2% (7.17→3.57), IS 提升 30.7% |
| ImageNet 256px 类条件生成 (有guidance, JiT-H/16) | FID↓ (最佳CFG档位) | 1.80 (JiT+CAFM, CFG=1.8) | 1.86 (JiT-H/16, CFG=2.2) | FID 降低 3.2%, 最佳CFG 档位从 2.2 下移到 1.8 |
| GenEval 512px 文生图 (整体分, with PE & CFG) | Overall ↑ | 0.85 (CAFM) | 0.81 (FM) | 整体分提升 +0.04, 单对象 0.99→0.99, 双对象 0.89→0.92, 颜色 0.69→0.71, 计数 0.77→0.81, 颜色绑定 0.89→0.94 |
| DPG-Bench 512px 文生图 (整体分, with CFG) | Overall ↑ | 85.21 (CAFM) | 83.67 (FM) | 整体分提升 +1.54, Global 90.34→89.55, Entity 90.56→89.83, Attribute 88.98→89.99, Relation 88.17→91.20, Other 90.71→91.88 |
局限与改进
作者在文中坦率指出两点局限。第一,虽然 CAFM 显著改善了 guidance-free 生成,但没有理论保证收敛到的就是真实数据分布,特别是在训练数据低密度、含异常点的区域,生成质量仍可能下降;这与作者强调'ground-truth flow 仍相同,只是泛化方式不同'的论点一致,但也意味着在分布外生成上 CAFM 并不天然优于 guidance。第二,后训练阶段额外引入了一个与生成器规模相当的判别器、一次前向 JVP、以及 $N=16$ 次判别器更新,导致训练成本至少是单纯 FM 目标的 10 倍以上(Fig. 4 显示从零训练时 CAFM 慢于 FM),因此'CAFM 仅在已有 FM 预训练模型上做后训练'是出于实用性考虑,不是理论最优;从零预训练在 1000 epoch 内 FID 仍在 20-30 量级,远未收敛。我自己的观察是:由于判别器复用生成器权重初始化,文章没有报告'判别器架构选择是否关键'的系统 ablation,例如换成更小的 D、或用预训练视觉编码器冻结做 D 是否仍能保持收益,这块仍是开放问题。
独立分析的弱点
独立分析本文有三点可改进之处。第一, 判别器更新频率 $N=16$ 是从 AFM 经验中沿用的'魔数',Fig. 5 (a)(d) 显示 $N$ 在 4→8 之间变化对收敛影响明显, 但没有给出 $N$ 与 batch size、模型规模的标度关系, 实际部署到更大模型时如何调参仍是经验性的; 改进方向是引入自适应 $N$,根据 D/G loss 差值动态调整。第二, 文章所有结果都建立在'FM 预训练 + CAFM 后训练'两阶段流程上,当预训练数据本身有偏, CAFM 只能改'泛化方式'而不能改'见过的内容', 在罕见概念、长尾类别上仍受限于预训练覆盖; 改进方向是把判别器与 CLIP/DINOv2 等外部语义对齐信号耦合, 在低密度区域显式施加结构约束。第三, 文章没有给出 CAFM 失败案例的细致分析, 改进方向是提供 per-class FID 与 diversity-FID 的完整消融,以便更精确识别后训练何时引入回归。
未来方向
作者明确给出了三个方向:一是探索更优的 divergence 度量(Tab. 12 比较了不同对抗损失对应的训练行为),把对比函数 $f$ 从二次换成 hinge、softmax 等形式,看是否能进一步提升分布匹配;二是把 CAFM 推广到预训练阶段,即让 FM 训练一开始就采用 JVP-based 判别,从而省去'后训练'的额外开销,但需要解决 Fig. 4 中展示的'CAFM 从零训练慢'的问题,可能的方向是 curriculum: 先用 FM 训 200 epoch, 再切到 CAFM 训 800 epoch, 模拟后训练流程但合并到一次训练里;三是把'导数空间对抗'这一思想推广到更广义的概率流场景,例如 stochastic differential equation (SDE)、Riemannian 流形上的流匹配,以及视频/3D 等更高维模态。
复现评估
复现评估整体偏中等。算法上 Algorithm 1 已给出 30 行 PyTorch 实现, 核心代码几乎可原样搬用;关键超参都已明确: 学习率 $1\times 10^{-5}$、Adam $\beta=(0,0.95)$、batch size 256 (SiT)/1024 (JiT)、EMA 0.99、$\lambda_{cp}=0.001$、$\lambda_{ot}=0$ (后训练)、$N=16$、$D$ 预热 epoch 2/4。数据与预训练模型: SiT-XL/2 与 JiT-H/16 均为官方公开 checkpoint, ImageNet 256px 可直接下载。算力门槛: SiT 后训练 10 epoch 在 8×H100 上约需数天 (batch 256), JiT 后训练 batch 1024 需更大显存。最大复现障碍: 字节跳动未承诺开源 Z-Image 模型本身与本文使用的判别器 checkpoint, 文生图实验(Tab. 5/6) 的验证需要先有等量级 Z-Image 复现; SiT/JiT 部分代码与数据完全可复现, 建议把主要工作放在这两组实验上。
论文图表
上方是标准 Flow Matching 在 ImageNet 256px 上无 guidance 生成的一组样本,下方是 CAFM (本文方法) 在完全相同设置下的样本;可以直观看到 CAFM 生成的动物、物体具有更清晰的轮廓、更逼真的纹理和更少的人造瑕疵,在分布匹配上明显更接近真实数据。
这张图是整篇论文最重要的'视觉证据',直接展示了 CAFM 的核心承诺——'替换欧氏判据后,生成器的泛化方向更靠近数据流形',即便不使用 CFG 也能产出高质量样本。