PaperlayerView exact revision on arXiv

Source-linked AI summary

Continuous Adversarial Flow Models

Shanchuan Lin, Ceyuan Yang, Zhijie Lin, Hao Chen, Haoqi Fan

arXiv:2604.11521v1cs.LGcs.CV

TL;DR

Flow matching can generate samples that diverge from the underlying data distribution because its fixed Euclidean objective does not capture manifold structure. The paper introduces continuous adversarial flow models, which use a learned discriminator for continuous-time training and are primarily applied to post-train existing flow-matching models. Across ImageNet and text-to-image experiments, CAFM post-training improves reported generation metrics, including guidance-free FID and text-to-image benchmark scores.

  • Problem

    Flow matching may produce out-of-distribution samples, while guidance improves quality but alters the sampling distribution; more faithful data-distribution modeling remains open.

  • Method

    CAFMs are continuous-time flow models that replace flow matching’s fixed Euclidean criterion with a learned discriminator and can post-train existing flow-matching models or train from scratch.

  • Results

    CAFMs improve ImageNet generation and text-conditional image generation, including guidance-free and guided FID as well as GenEval and DPG scores.

  • Takeaways & Limitations

    Adversarial training can be integrated into continuous-time flow modeling to improve sample fidelity and distribution matching, including through efficient post-training.

  • Takeaways & Limitations

    CAFMs have no guarantee of generalizing to the true underlying distribution, especially in low-density regions containing outliers.

Abstract

from arXiv · show

We propose continuous adversarial flow models, a type of continuous-time flow model trained with an adversarial objective. Unlike flow matching, which uses a fixed mean-squared-error criterion, our approach introduces a learned discriminator to guide training. This change in objective induces a different generalized distribution, which empirically produces samples that are better aligned with the target data distribution. Our method is primarily proposed for post-training existing flow-matching models, although it can also train models from scratch. On the ImageNet 256px generation task, our post-training substantially improves the guidance-free FID of latent-space SiT from 8.26 to 3.63 and of pixel-space JiT from 7.17 to 3.57. It also improves guided generation, reducing FID from 2.06 to 1.53 for SiT and from 1.86 to 1.80 for JiT. We further evaluate our approach on text-to-image generation, where it achieves improved results on both the GenEval and DPG benchmarks.

1 Introduction

Flow matching can generate out-of-distribution samples because its Euclidean training criterion does not faithfully reflect the data manifold. CAFMs replace that fixed criterion with adversarial training in continuous time and substantially improve reported image-generation results.

  • Flow matching often produces out-of-distribution image and video samples unless guidance is applied, but guidance changes the sampling distribution.
  • Squared L2 training measures Euclidean velocity errors rather than manifold-aware distances, which can induce incorrect generalization relative to the data distribution.
  • Prior remedies alter the operating space, require a manually specified manifold, or rely on perceptual distances from frozen feature networks.
  • CAFMs extend adversarial flow models to continuous time while retaining probability-flow velocity learning through a simulation-free objective.
  • CAFMs use a learned discriminator instead of a fixed Euclidean criterion, producing more in-distribution samples perceptually and according to multiple metrics.
  • 8.26 to 3.63 guidance-free FID for latent-space SiT and 7.17 to 3.57 for pixel-space JiT after CAFM post-training on ImageNet 256px.The post-training uses only 10 epochs; guided FID also improves from 2.06 to 1.53 for SiT and from 1.86 to 1.80 for JiT, while text-to-image GenEval and DPG also increase.

2 Background

Flow matching transports noise to data by learning a probability-flow velocity field, typically with squared L2 loss. Adversarial flow models replace this fixed criterion with discriminator-guided training, but discrete-time adversarial objectives become unstable as timestep intervals vanish.

  • Flow matching transports samples from a prior noise distribution to the data distribution along a predefined probability flow.
  • With linear interpolation, the conditional velocity is −x + z, and the generator is trained to match this velocity using a criterion d.
  • Criteria satisfying the stated expectation condition recover the marginal probability-flow velocity, while the common squared L2 criterion is an MSE variant.
  • Finite-capacity models generalize beyond memorized training samples, and the loss objective changes that generalized distribution.
  • Flow matching’s isotropic Euclidean loss lacks data-manifold awareness, leading to incorrect generalization and out-of-distribution generation.
  • Adversarial flow models train a generator and discriminator in a minimax game so generated and real intermediate states approach the same distribution.
  • AFMs add an optimal-transport objective that favors targets near the source and supports a unique transport for stable training.
  • AFM training uses discriminator regularizers and remains unstable as discrete timestep intervals approach zero, with gradient vanishing requiring mitigation.

3 Method

CAFMs extend adversarial training to continuous-time flow models by discriminating velocity directions through discriminator directional derivatives. Their objectives combine adversarial and centering terms, with optional optimal-transport regularization and practical training strategies for efficient implementation.

  • Continuous Adversarial Flow Models: CAFMs use a generator G(x_t, t) to predict velocity and a discriminator D(x_t, t) whose directional derivative discriminates real and generated flows.The generator has the same form as in flow matching, while discrimination occurs in derivative space to reflect velocity as a position derivative.
  • Continuous Adversarial Flow Models: The discriminator JVP evaluates D using (x_t, t) as primal variables and (v̄_t, T) or G(x_t,t) as tangent variables, producing a scalar discrimination logit.The JVP combines derivatives with respect to x_t and t along the supplied tangent direction.
  • Continuous Adversarial Flow Models: CAFMs train D and G with bounded contrastive adversarial objectives that compare real-flow and generated-flow JVPs, while gradients are backpropagated through the JVP.The contrastive function is f(a,b) = (a - 1)^2 + (b + 1)^2.
  • Continuous Adversarial Flow Models: A centering penalty keeps the discriminator’s absolute value near zero because the adversarial objectives penalize derivatives but leave D’s offset unconstrained.The penalty is E[D(x_t,t)^2].
  • Continuous Adversarial Flow Models: For high-dimensional flows, an optional minimum-norm regularizer addresses discriminator-induced ambiguity by encouraging G to avoid exploiting the null space.This regularizer is related to optimal-transport regularization.
  • Training strategy and implementation: CAFMs remove gradient penalties from the final objectives, can mitigate gradient vanishing, and benefit from updating D for N steps per generator update.The final objectives combine adversarial and centering terms for D with adversarial and optional optimal-transport terms for G.
  • Training strategy and implementation: Forward-mode automatic differentiation computes JVPs efficiently, while vmap evaluates multiple tangents and remains compatible with DDP, FSDP, and gradient checkpointing.G has no architectural restrictions, whereas RMSNorm improves discriminator training stability in the reported experiments.
  • Training strategy and implementation: Because CAFMs require an extra discriminator, JVP computation, and multiple discriminator updates, the authors primarily recommend post-training flow-matching models.CAFMs can also be trained from scratch, but this is less efficient than flow matching.

4 Experiment

Experiments evaluate CAFM post-training on latent- and pixel-space ImageNet models and text-to-image generation, with additional studies of training from scratch and hyperparameters. CAFM improves guidance-free and guided generation, while training from scratch converges more slowly than FM.

  • ImageNet Generation Post-training: CAFM post-training improves both guidance-free and guided ImageNet generation for latent-space SiT and pixel-space JiT, measured by FID and IS.
  • ImageNet Generation Post-training: CAFM achieves the best compared performance in SD-VAE latent space under both guided and guidance-free settings without DINOv2.
  • ImageNet Generation Post-training: In pixel space, CAFM achieves competitive performance, but varying architectures and settings make method-specific attribution harder.
  • Text-to-Image Generation Post-training: CAFM post-training significantly improves guidance-free and guided text-to-image generation on both GenEval and DPG-Bench.
  • ImageNet Generation Trained from Scratch: Hyperparameter ablations find that λot should decrease and N should increase during training for better performance.
  • ImageNet Generation Trained from Scratch: CAFM can train from scratch, but it converges more slowly than FM under the same number of epochs.

5 Related Work

Related work positions CAFM at the intersection of adversarial training, continuous-time flow modeling, altered generalization, guidance, and divergence objectives. Unlike manifold-based approaches, CAFM retains Euclidean trajectories while changing the training objective.

  • Unifying Adversarial and Flow Modeling: CAFM extends adversarial flow models from discrete time to continuous time.
  • Adversarial Post-Training: Unlike distillation-focused adversarial post-training, CAFM uses adversarial post-training to induce different model generalization in continuous-time flow models.
  • Generalization Behavior: CAFM retains Euclidean-space trajectories and the same ground-truth flow as standard flow matching while altering generalization through the loss objective.
  • Guidance: Guidance steers sampling toward a modified distribution and has effects similar to low-temperature sampling.
  • Divergence Measures: Flow matching minimizes forward KL, whereas GAN objectives can minimize different divergences that influence generalization.

6 Conclusion

The paper introduces CAFMs as continuous-time flow models trained with an adversarial objective and demonstrates their use for post-training flow-matching models. The method improves ImageNet and text-conditional image generation performance.

  • CAFM is a continuous-time flow model trained with an adversarial objective.
  • The objective can be efficiently used to post-train flow-matching models and improves performance on ImageNet and text-conditional image generation.

A Additional Results on ImageNet Post-training

Additional ImageNet results provide full metrics, ablations, qualitative comparisons, and implementation details for CAFM post-training on SiT and JiT. They also identify stable hyperparameter choices and show substantial guidance-free and guided FID improvements.

  • Ablations: Post-training ablations select N = 16 and λot = 0, while higher learning rates cause gradient spikes and divergence.
  • Qualitative Comparisons: Qualitative comparisons cover guidance-free and guided outputs for both SiT and JiT.
  • SiT-XL/2 Results: SiT-XL/2 guidance-free generation improves from FID 8.26 with FM to 3.63 with CAFM.
  • SiT-XL/2 Results: SiT-XL/2 guided generation improves from FID 2.06 with FM to 1.53 with CAFM.
  • JiT-H/16 Results: JiT-H/16 guidance-free generation improves from FID 7.17 with FM to 3.57 with CAFM.
  • JiT-H/16 Results: JiT-H/16 guided generation improves from FID 1.86 with FM to 1.80 with CAFM.

B Additional Results on Text-to-Image Post-training

The text-to-image experiments apply CAFM post-training to ZImage and compare it with flow-matching baselines on GenEval and DPG-Bench. CAFM improves GenEval against the FM-finetuned baseline, while DPG-Bench comparisons require care because some models use different datasets.

  • Architecture: CAFM post-training uses the unchanged ZImage generator architecture, while the discriminator adds a cross-attention layer to project its logit.The experiments use a pre-distillation, 6B-parameter ZImage checkpoint for continuous-flow training.
  • Results: CAFM-finetuned ZImage beats both the original and FM-finetuned models on GenEval.The original ZImage model is included for reference, but the FM-finetuned model is the fair comparison baseline.
  • Results: CAFM-finetuned ZImage performs worse than the original ZImage model on DPG-Bench, which the authors attribute to different datasets.The main-text comparison removes the original model because it is not a directly comparable baseline.
  • Qualitative analysis: Figures 10 and 11 provide qualitative DPG-Bench comparisons and examples of failures in guidance-free text-to-image generation.Figure 11 specifically documents failure cases rather than only successful generations.
  • Flow-matching objective: For flow matching, any strictly positive definite quadratic criterion converges theoretically to the same marginal velocity but can change finite-capacity generalization.The squared L2 criterion is the special case M = I; different criteria can therefore induce different practical distributions.

D On Discriminator JVP Designs

The paper compares a naive discriminator over velocity predictions with a discriminator defined through JVPs in state-time space. JVP-based discrimination avoids inaccessible marginal targets and provides trajectory-consistent optimization behavior.

  • Naive formulation: The naive discriminator formulation has two problems: the marginal velocity is inaccessible, and conditional targets do not generally preserve the marginal target through a nonlinear discriminator.Replacing the marginal velocity with conditional velocities can repeatedly pull the deterministic generator toward incompatible targets.
  • JVP design: CAFM instead defines D(x_t, t) and performs discrimination in JVP space using adversarial objectives for conditional target and generator velocities.The JVP discriminator evaluates directional derivatives along the flow rather than directly comparing velocity vectors through a nonlinear discriminator.
  • JVP design: JVP equilibrium is reached only when the generator velocity matches the marginal velocity at every x_t.The derivation uses linearity of the JVP and shows that the expected discriminator response depends on v_t - G(x_t,t).
  • Trajectory consistency: The JVP discriminator remains globally consistent along trajectories, and the authors find that these constraints improve optimization.A separate-network parameterization does not yield good results on high-dimensional data.

E On the Vanishing-Gradient Problem

The paper analyzes vanishing gradients when real and generated distributions have disjoint support and argues that JVP linearization can preserve generator gradients. CAFMs extend adversarial flow modeling to continuous time, but the underlying support problem motivates careful interpretation.

  • Problem: When real and generator supports do not overlap, an optimal discriminator can separate them perfectly and provide little useful gradient near generated samples.This issue is especially relevant in high-dimensional spaces where data concentrate near low-dimensional manifolds.
  • Existing mitigations: Instance noise gives the convolved distribution support everywhere, but larger noise can make the model learn a noisy rather than intended data distribution.The useful noise scale depends on the distance between manifolds and discriminator capacity.
  • CAFMs: CAFMs extend adversarial flow models from piecewise learning to continuous time as the discretization interval approaches zero.The paper presents this limit as a way to further mitigate the vanishing-gradient problem.
  • CAFMs: Under an optimal discriminator, JVP linearization gives the generator a nonzero gradient whenever the discriminator loss gradient and input Jacobian are nonzero.The paper connects a nonzero Jacobian to discrimination whenever G(x_t,t) differs from the marginal velocity.
  • Gradient analysis: The least-squares contrastive function has a nonzero derivative with respect to its first input whenever that input differs from 1.For an optimal discriminator, the generator’s velocity prediction receives the corresponding nonzero contrastive gradient.

F On Implementation of JVP

The implementation wraps forward-mode JVP and vectorization around the model in ways compatible with distributed training and checkpointing. The placement differs across distributed strategies.

  • Core implementation: PyTorch implementations use torch.func.jvp and torch.func.vmap for forward-mode JVP and vectorization.The functions are compatible with DDP, FSDP, and gradient checkpointing with specialized arrangement.
  • DDP: For DDP, JVP should wrap the network before DDP rather than wrapping the DDP module itself.The recommended order is ddp(jvp(D)) because applying JVP to DDP conflicts with gradient synchronization logic.
  • FSDP and checkpointing: For FSDP and gradient checkpointing, wrapping only top-level submodules is sufficient for sharding and checkpointing.The paper contrasts this with wrapping every neural-network module, which it considers unnecessary and excessive.
  • Attention: PyTorch’s math fused scaled-dot-product attention kernel supports JVP and Vmap natively for image-generation training.This avoids requiring a separate attention implementation for the JVP-enabled training path.

G On LayerNorm and RMSNorm

RMSNorm improves discriminator training stability by avoiding the large gradient-norm spikes seen with LayerNorm. The method also uses repeated discriminator updates, with post-training incurring additional computational cost.

  • Normalization and stability: RMSNorm avoids the large discriminator gradient-norm spikes observed with LayerNorm under equal settings.This comparison is shown for pre-training SiT-B/2 on ImageNet 256px.
  • Discriminator update schedule: The training schedule performs N discriminator updates for each generator update to keep the discriminator near its local optimum.
  • Computational cost: 4.8× more wall-clock time per epoch is required by CAFM than FM for ImageNet SiT-XL/2 post-training.The overhead comes from the additional discriminator and its forward and backward JVP computations.
Loading 2604.11521v1…