Continuous Adversarial Flow Models

TL;DR

Continuous Adversarial Flow Models (CAFMs) extend the discrete-time Adversarial Flow Models (AFMs) recipe into continuous time, replacing flow matching’s fixed MSE criterion with a learned discriminator that operates in the velocity-derivative space. The discriminator is implemented via a Jacobian-Vector Product (JVP) that computes both the generator’s velocity and the discriminator’s directional derivative in a single forward pass, with forward-mode autodiff. Intended primarily as a post-training objective for existing flow-matching models (10 epochs of finetuning), CAFM more than halves guidance-free FID on ImageNet-256 — SiT-XL/2 drops from 8.26 → 3.63, JiT-H/16 from 7.17 → 3.57 — and improves guided generation as well (SiT best CFG-FID 2.06 → 1.53; text-to-image Z-Image GenEval 0.81 → 0.85). The framing: flow matching converges to the right ODE in theory but generalizes badly with finite capacity under Euclidean loss; a learned discriminator approximates the data manifold better, so the same probability flow generates more in-distribution samples.

Key claims

Flow matching’s MSE criterion is a generalization problem, not a convergence problem: under infinite capacity the objective fits the data perfectly, but with finite capacity the choice of criterion (Euclidean vs manifold-aware) controls how the model generalizes off the training set [§1, §2.1 “The limitation of flow matching”].
A learned discriminator operating in continuous time can act as the manifold-aware criterion that flow matching lacks; this is the first continuous-time application of adversarial training to flow modeling [§1, §3.1].
Adversarial training is done in the derivative space of the discriminator: the JVP $\dot D(x_t, t)[v(x_t,t), 1]$ is used as the discrimination logit, explicitly reflecting that velocity is a derivative of position [§3.1, Eq. 17–19].
CAFM avoids AFMs’ vanishing-gradient problem and removes the need for the gradient penalties $R_1$ / $R_2$ that AFMs required [§2.2 “The limitation of AFMs”, §3.1, Appendix 0.E].
Optimal-transport regularization on the generator (continuous limit of AFMs’ OT term) prevents the discriminator from exploiting the null space of the high-dimensional → scalar projection, accelerating convergence; for post-training $\lambda_{OT}$ is set to 0 to eliminate Euclidean bias [§3.1, Eq. 24].
Post-training SiT-XL/2 on ImageNet-256 for 10 epochs: guidance-free FID 8.26 → 3.63, guided best FID 2.06 → 1.53 (at CFG 1.3 instead of 1.5); a controlled FM-objective trial at matched compute does not move FID, isolating the gain to the CAFM objective [Table 2, §4.1].
Post-training JiT-H/16 (pixel-space) on ImageNet-256: guidance-free FID 7.17 → 3.57, guided best FID 1.86 → 1.80; outperforms SiD2 / PixelFlow-XL / PixNerd-XL at equal or smaller parameter count in the no-CFG setting [Table 4, §4.1].
Post-training Z-Image (text-to-image): GenEval overall 0.81 → 0.85 (with PE+CFG); without CFG, GenEval 0.60 → 0.71 and DPG 72.25 → 77.21 — the guidance-free gains are substantially larger than the guided gains, consistent with the “fix the generalization, then you need less guidance” framing [Tables 5–6, §4.2].
Switching LayerNorm → RMSNorm in the discriminator is necessary for stable JVP-based training; without it the JVP-involving forward pass diverges [§3.2].
CAFM can train from scratch but converges more slowly than FM under the same epoch budget; the recipe is therefore explicitly proposed as a post-training method for efficiency [§3.3, §4.3, Fig. 5].
Hyperparameter schedule for from-scratch training: $\lambda_{OT}$ should decrease over training, $k$ (discriminator updates per generator update) should increase over training — confirming the AFM finding that the OT term acts as an early-training stabilizer [§4.3, Fig. 5].

Method

CAFM keeps the generator $v_\theta(x_t, t)$ in the exact form used by flow matching (velocity prediction along a linear interpolation between noise and data). The new piece is a discriminator $D_\phi$ that consumes $(x_t, t)$ and outputs a scalar potential, with the discrimination logit defined as the directional derivative of $D_\phi$ along the predicted velocity:

$\text{logit} = \dot D_\phi(x_t, t)[v_\theta(x_t,t), 1] = \nabla_{x_t} D_\phi \cdot v_\theta + \partial_t D_\phi$

Intuitively, $D_\phi$ learns a scalar field whose gradient flow points toward higher-realism directions; $v_\theta$ is trained to maximize alignment with that gradient, $D_\phi$ is trained to discriminate real-data flow directions from generator flow directions. The bounded contrastive function (least-squares GAN style, Eq. 22) is used so the objective remains well-posed.

The JVP is computed with forward-mode autodiff (torch.func.jvp) — a single forward pass returns both $D_\phi(x_t, t)$ and the JVP scalar, after which both the centering penalty ( $D_\phi^2$ , keeps the absolute value bounded) and the adversarial loss can be derived. vmap is used to compute multiple tangents at the same primal when updating $D_\phi$ over multiple sampled velocities per state. The implementation is compatible with DDP, FSDP, and gradient checkpointing — a concise PyTorch implementation fits in Algorithm 1.

Training alternates updates: for each generator update, $D_\phi$ is updated $k$ times (with $k$ growing over training). For post-training, both networks initialize from the pretrained flow-matching weights (with the discriminator’s LayerNorm swapped to RMSNorm), and the discriminator is warmed up for the first 2–4 epochs with the generator frozen. Learning rates of 1e-5 with $\beta=(0.9, 0.999)$ , EMA decay of 0.99 on the generator, $\lambda_{OT}=0$ , and centering weight 0.1 are the working post-training hyperparameters.

Results

ImageNet-256 class-conditional, latent-space SiT-XL/2 post-trained for 10 epochs: guidance-free FID 8.26 → 3.63 (-56%), guidance-free IS 131.65 → 178.08 [Table 2]. Best guided FID drops 2.06 → 1.53 at CFG 1.3 (lower than the pretrained model’s optimal CFG of 1.5, indicating CAFM needs less guidance to reach the same quality). Compared to SOTA at 675M parameters under SD-VAE: beats DDT-XL (6.27), REPA (5.90), Disperse (7.43) on guidance-free FID; beats DiT-XL/2 (2.27), SiT-XL/2 (2.06), Disperse (1.97) on guided FID, only second to REPA (1.42) and DDT-XL (1.26) among methods using DINOv2 supervision [Table 4].

ImageNet-256, pixel-space JiT-H/16 post-trained 10 epochs: guidance-free FID 7.17 → 3.57 (-50%) at 956M params, vs SiD’s 2.77 at 2B params [Table 4]. Best guided FID 1.86 → 1.80, beating PixelFlow-XL/4 (1.98) and JiT-G/16 (2B, 1.82) at smaller scale.

Text-to-image post-training on Z-Image: GenEval overall lifts 0.81 → 0.85 with PE+CFG, but the most striking number is without CFG (with PE) the overall jumps 0.60 → 0.71 (Single Obj. 0.95 → 0.99, Two Obj. 0.66 → 0.83, Color Attr. 0.35 → 0.50, Position 0.40 → 0.52) [Table 5]. DPG-Bench overall 72.25 → 77.21 (no CFG) and 83.67 → 85.21 (with CFG) [Table 6]. The team’s interpretation: CAFM concentrates the model on the data manifold so guidance has less work to do.

From-scratch SiT-B/2 training: CAFM converges slower than FM under the same epoch budget on ImageNet — consistent with the team’s framing of CAFM as a post-training tool, not a from-scratch substitute [Fig. 5, §4.3].

Why it’s interesting

CAFM is a sharper version of the question that PixelGen: Improving Pixel Diffusion with Perceptual Supervision (PixelGen) opened: if the gap between models isn’t the flow but the generalization induced by the loss, what’s the right loss? PixelGen answers with fixed perceptual networks (LPIPS + frozen DINOv2 as P-DINO); CAFM answers with a learned discriminator that, by construction, adapts to the data manifold instead of borrowing one from ImageNet-pretrained features. The two recipes are not mutually exclusive — PixelGen’s perceptual losses and CAFM’s adversarial criterion are both ways of replacing isotropic Euclidean loss with a manifold-aware one — but they sit at opposite ends of the “static vs dynamic criterion” axis. Both target the same JiT pixel-space backbone and report comparable headline numbers (PixelGen FID 5.11 at 80 epochs from-scratch; CAFM FID 3.57 at JiT-H/16 + 10 epochs post-training), and CAFM further confirms the pixel-space-diffusion thesis at higher absolute quality.

This also reframes Classifier-Free Guidance variants: in Random noise augmentation during diffusion training eliminates the need for CFG Ostris argues CFG is patching a noise-distribution miscalibration; CAFM tells a related but distinct story — CFG is partly patching a generalization miscalibration that better criterion design fixes at the source. Z-Image’s GenEval/DPG numbers go up more without CFG than with it under CAFM post-training, which is the cleanest evidence yet that guidance and generalization are interchangeable knobs and that the field has been over-using one to compensate for the other.

Mechanically, CAFM also matters as the cleanest application yet of adversarial post-training that is not a distillation. Almost all prior adversarial post-training in the Diffusion Distillation cluster (DMD, ArcFlow, SDXL-Lightning, SeedVR2, etc.) uses the discriminator to enable step-reduction; CAFM keeps the full sampler (250-step Euler-Maruyama for SiT, 50-step Heun for JiT) and uses the discriminator purely to improve distribution matching. That separates two factors that have been entangled in the distillation literature — “adversarial signal helps few-step generation” vs “adversarial signal helps any generation” — and credits the latter explicitly.