Adversarial Flow Models

TL;DR

Adversarial Flow Models (AFM) is a class of generative models that unifies GANs with flow matching: the generator is trained adversarially but constrained to learn the same deterministic noise-to-data transport plan that flow matching’s linear interpolation + squared-distance loss implicitly defines. The constraint is enforced by adding an optimal-transport loss L_OT = E[‖x_θ(ε) - ε‖²] (with a time-decayed weight λ) on top of a standard relativistic GAN objective. This breaks the symmetry in the adversarial loss landscape, prevents generator drift, and stabilizes training of standard DiT-style transformers — something prior GAN scaling work needed non-standard architectures or external feature priors for. On ImageNet-256 class-conditional, AFM-XL/2 reaches FID 2.38 at 1 NFE (best single-step result against consistency-based baselines under the same VAE + DiT setup), and depth-repeated 56/112-layer 1-NFE models hit FID 2.08 / 1.94, surpassing their multi-step counterparts. Without guidance, AFM beats flow matching trained for 250+ NFE, because the learned discriminator provides a better distributional criterion than Euclidean MSE.

Key claims

The adversarial objective alone has no unique optimization target: it only requires distribution matching, not a specific transport plan, so the generator drifts during training; this is identified as a primary cause of GAN training instability with transformer architectures [§1, §3.2].
Adding an OT loss L_OT = E_ε[‖x_θ(ε) - ε‖²] makes the generator’s transport plan unique — the one that minimizes squared-Wasserstein-2 between prior and data, which is exactly the flow-matching transport with linear interpolation [§3.2, Eq. 12].
Without L_OT, training of standard DiT diverges regardless of gradient-penalty strength; with L_OT, training converges across model sizes (B/2 → XL/2) [Table 1, §4.2].
The OT-loss weight λ must be decayed over training (cosine 0.2 → 0.01 over 100 epochs): too large forces identity output (hurts distribution matching), too small fails to regularize; constant λ=0.2 gives FID 29.4 vs decayed 8.5 [Table 2, §4.2, Fig. 2].
The method generalizes to multi-step generation by giving the generator source and target timesteps (s, t) and the discriminator the target timestep t only; conditioning D on the source x_s is incorrect because it tells D that the same x_s should match all x_t, which a deterministic generator cannot satisfy [§3.3, §3.4].
The discriminator must not be conditioned on the source sample — formulating D(x_s, x_t, t) causes training to oscillate or diverge because x_s and x_0 are randomly paired during training [§3.4].
A gradient-normalization technique on the generator’s adversarial gradient (EMA-tracked norm, rescaled to data dimension) is required to make λ transfer across model sizes; without it the ratio between adversarial and OT losses varies with architecture / init / GP strength [§3.4, Eq. 23].
1-step training without other-timestep supervision conserves model capacity: B/2 AFM approaches the performance of consistency-based XL/2 models because consistency methods waste capacity learning intermediate-timestep predictions for propagation [§4.3, “Single-step with guidance”].
XL/2 AFM hits FID 2.38 at 1 NFE with guidance on ImageNet-256 latent space (SD-VAE), a new best result against sCM, iMM (the “IMM” Jiaming references), Shortcut, MeanFlow, R3GAN, GAT, and concurrent works [Footnote 5, §4.3].
Depth-repetition gives 1-NFE models that beat their multi-step counterparts: 56-layer 1-NFE FID 2.08 (vs the same params at 2 NFE) and 112-layer 1-NFE FID 1.94 (vs the same params at 4 NFE), end-to-end trained without any intermediate supervision [§3.8, §4.3].
Without guidance, AFM-XL/2 outperforms flow matching at 250+ NFE, because the learned discriminator captures semantic/perceptual distance on the data manifold while flow matching’s MSE measures Euclidean distance — flow matching’s “blended-looking” guidance-free samples are a generalization failure of the loss, not a convergence failure [§3.5, Footnote 5 no-guidance setting].
Flow-based classifier guidance (CG on noised samples at multiple timesteps) reproduces CFG-like effects without requiring CFG training; gives small additional FID lift (2.40 → 2.36) at the optimal range [§3.6, Table 3].
For consistency-based / moment-matching methods (iMM, sCM, Shortcut, MeanFlow): the paper claims their pointwise / moment-matching losses cause blurriness because finite-capacity 1-step networks lack the capacity to reproduce many-step targets, motivating distributional matching (adversarial or score-distillation) as a final refinement step [§1, §2].
Practical training techniques that matter independently of the AFM formulation: EMA-of-generator weight averaging (replaces online weights when EMA plateaus), discriminator augmentation for the guidance setting, and checkpoint reloading when training stalls in the no-guidance setting [§3.9].

Method

The generator x_θ(ε) is a standard DiT (B/2, M/2, L/2, XL/2 from MeanFlow’s sizing) operating on SD-VAE latents at ImageNet-256. The discriminator D_φ is nearly identical to the generator — same DiT — except for a learnable [CLS] token prepended to the input and a final LayerNorm + linear projection that produces the relativistic logit. Both networks use AdaLN-style conditioning on the class label.

The objective stacks three pieces. (1) A relativistic GAN loss with R1+R2 gradient penalties on D (computed via finite-difference approximation, applied to 25% of samples per batch) and a logit-centering penalty to prevent unbounded drift. (2) The optimal-transport loss L_OT = E_ε‖x_θ(ε) - ε‖² with weight λ that decays from 0.2 → 0.01 cosine over 100 epochs. (3) A gradient-normalization wrapper around the adversarial loss term that EMA-tracks the gradient norm and rescales to a unit-norm-times-√d signal, so that the relative weight of the OT loss is stable across architectures.

For multi-step generation, the generator becomes x_θ(x_s, s, t) and D becomes D_φ(x_t, t). During training, s, t ~ U[0,1] and x_s = (1-s)x_0 + s·ε, and the OT loss generalizes to weighted squared distance between source and target. The model is typically trained on a discrete set of timesteps used at inference rather than any-step continuously, to conserve capacity; any-step training works but converges more slowly.

For extra-deep models, the depth is increased only via block repetition: the same transformer block stack is run multiple times (with a timestep-like embedding only to differentiate repetition iterations), end-to-end trained for 1-NFE generation without intermediate supervision. This matches multi-step parameter counts but avoids the projection-and-reinterpretation overhead between data-space evaluations and the manual timestep discretization choice.

For guidance, AFM uses a flow-based classifier guidance — a time-conditioned classifier trained on noised samples evaluated at multiple timesteps along the probability flow, whose accumulated gradient is added to D’s logits. This recovers CFG-like effects from a 1-step model. CFG itself is also incorporable in the AFM-as-post-training mode by deriving an implicit classifier from the existing flow-matching model.

Additional stabilization: keep an EMA of the generator; once EMA-FID plateaus toward end-of-training, replace online weights with EMA and repeat each epoch; decrease LR during this phase. When training stalls without guidance, reload from an earlier checkpoint (avoids the implicit biases of discriminator augmentation in this setting). With guidance, use DA freely.

Results

ImageNet-256 class-conditional, latent space (SD-VAE), DiT-style transformers, 1-NFE with guidance [Footnote 5 / Table comparison rows]:

AFM-B/2: FID approaches or exceeds many XL/2 consistency-based models (B/2 is capacity-tight, AFM doesn’t waste it on intermediate timesteps).
AFM-XL/2: FID 2.38 at 1 NFE — new best under same VAE + DiT setup, beating sCM, iMM, Shortcut, MeanFlow, R3GAN, GAT, and concurrent works.
AFM-56-layer (1-NFE): FID 2.08, surpassing the 2-NFE multi-step counterpart.
AFM-112-layer (1-NFE): FID 1.94, surpassing the 4-NFE multi-step counterpart [§4.3].

Few-step setting (designated discrete timesteps to conserve capacity): AFM also achieves better FIDs than baselines at matched NFE [Footnote 5]. Any-step training underperforms designated-step training because of batch-size / capacity dilution.

No-guidance setting on ImageNet-256: AFM 1-NFE and 2-NFE outperform flow matching at 250+ NFE in the same SD-VAE + DiT setup [Footnote 5 no-guidance rows]. The team interprets this as evidence that the learned discriminator’s semantic / perceptual criterion is a better distributional matching signal than flow matching’s Euclidean MSE — flow matching produces “out-of-distribution-looking” guidance-free samples because of the loss function, not the architecture.

Ablations [Tables 1–3, §4.2]:

No L_OT: training diverges across all γ ∈ {0.1, 0.25, 0.5} settings tested.
Constant λ_OT = 0.2: FID 29.4 vs decayed 0.2 → 0.01 cosine: FID 8.5 (B/2 100-epoch setting).
Optimal initial λ_OT = 0.25 in the ablation grid; too small fails to break loss symmetry, too large hurts distribution matching.
Flow-based CG gives minor improvement over standard CG (2.40 → 2.36 at optimal range).

Why it’s interesting

This is the direct precursor to Continuous Adversarial Flow Models (CAFM, April 2026, same author team) — the same five authors (Lin / Yang / Lin / Chen / Fan) extend AFM to continuous time and reframe it as a post-training objective applied to pretrained flow-matching models, where the OT loss is set to 0 and a JVP-based discriminator operates in the velocity-derivative space. AFM is the discrete-time, from-scratch, OT-regularized version; CAFM is the continuous-time, post-training, no-OT version. CAFM’s no-CFG FID drops on JiT-H/16 (7.17 → 3.57) and Z-Image (GenEval 0.60 → 0.71) are the production-scale evidence that AFM’s core insight — that the adversarial-loss landscape needs to be anchored to a deterministic transport plan, and that the learned discriminator is a manifold-aware criterion — generalizes beyond ImageNet.

Jiaming’s “IMM but with GANs?” framing is sharp: the paper explicitly positions itself against iMM (Inductive Moment Matching, cited as [107] in §2) and other consistency-based methods (Mean Flows for One-step Generative Modeling MeanFlow, Shortcut, sCM). The thesis is that pointwise or moment-matching losses on a 1-step network can’t reproduce many-step targets at finite capacity, so they get blurry; the production fix is to combine them with adversarial training as a refinement (cited prior work: SeedVR2, ArcFlow). AFM’s contribution is to argue you don’t need the moment-matching loss at all — you just need the adversarial objective plus the OT regularizer to make the transport plan unique. So yes, it’s “iMM but with GANs” in the sense that it replaces moment-matching distributional loss with an adversarial one, but more precisely: it argues moment matching was never the right primitive, and the right primitive is “adversarial + flow transport plan.”

AFM sits next to two adjacent recipes in the wiki’s Diffusion Distillation and one-step image cluster. Scalable GANs with Transformers (GAT, Sep 2025) is the immediately prior “GANs scale to transformers” attempt: it solves the architecture-scaling problem with multi-level noise-perturbed guidance + width-aware LR + REPA-on-discriminator, but still uses a non-standard transformer architecture and an external DINOv2 prior. AFM uses standard DiT with no external feature network and locates the fix in the loss landscape — the OT loss is what GAT’s MNG and REPA were implicitly substituting for. Mean Flows for One-step Generative Modeling (MeanFlow) is the consistency-based 1-step frontier MeanFlow paper compares against; AFM-XL/2 at 2.38 FID beats MeanFlow-XL/2’s 3.43 at 1 NFE under matched backbone, validating the “consistency models waste capacity on intermediate timesteps” claim.

For Classifier-Free Guidance variants, AFM gives a clean reframing parallel to Random noise augmentation during diffusion training eliminates the need for CFG (random noise augmentation removes the need for CFG) and CAFM: the paper argues CFG is partly compensating for the Euclidean-MSE generalization failure of flow matching, and that with a learned discriminator the guidance scale can be smaller (or, with CG on the flow, optimal at much smaller scales than typical flow matching).