ArXiv: 2511.22475
🎯 Pitch
Standard transformer architectures—the backbone of diffusion models—blow up when trained as GANs, yet this work shows that simply adding a deterministic transport map regularizer completely stabilizes training, yielding single-step generators that match or beat few-step consistency models with far less training overhead.
1. Executive Summary
This paper proposes adversarial flow models, a class of generative models that combines adversarial training with flow-based deterministic transport to support native one-step and multi-step generation. The method is evaluated on class-conditional ImageNet-256px using standard diffusion transformer (DiT) architectures, where a generator learns a deterministic noise-to-data mapping regularized by an optimal transport loss (minimizing the squared L2 distance between input noise and output sample) alongside a relativistic adversarial discriminator, while a gradient normalization technique stabilizes the competition between the adversarial and transport objectives across model sizes. The XL/2 model achieves a new best FID of 2.38 under the 1NFE setting, surpassing consistency-based XL/2 models, while end-to-end training of 56-layer and 112-layer single-step architectures yields FIDs of 2.08 and 1.94 respectively, outperforming their 28-layer 2NFE and 4NFE counterparts with equal compute and parameters. The method establishes that adversarial training can be stabilized on standard transformer architectures through deterministic transport regularization, though performance gains over consistency models are most pronounced when model capacity is preserved by training only on the target timesteps rather than all intermediate timesteps of the probability flow.
2. Context and Motivation
The Core Problem: Fast Generation Exists, But At a Quality Cost — And Quality Solutions Are Slow
The fundamental tension this paper addresses is deceptively simple: generative models that produce high-quality images are too slow, and models that are fast produce lower-quality images. This is not merely an engineering inconvenience — it is a structural limitation of two dominant generative paradigms that manifests in how they define and solve the generation problem.
Consider what happens at inference time for the two most successful families of generative models in recent years:
Flow-matching and diffusion models (Lipman et al., 2023; Ho et al., 2020) achieve state-of-the-art image quality by framing generation as an iterative transport process. Starting from pure Gaussian noise, the model takes a small step toward the data distribution, evaluates the probability flow gradient at this new position, takes another small step, and repeats this process dozens or hundreds of times. Each step requires a full forward pass through a large neural network. A 250-step generation with a DiT-XL/2 model on ImageNet-256px requires 250 sequential network evaluations — each one consuming both computation and wall-clock time, since the steps cannot be parallelized. The result: excellent FID scores (e.g., 2.27 with classifier-free guidance), but inference that is orders of magnitude slower than what real-time applications demand.
Generative adversarial networks (GANs) (Goodfellow et al., 2014) take the opposite approach. The generator is trained to produce a complete image in a single forward pass — one network evaluation, one output. This is exactly the fast inference that flow models lack. Historically, GANs achieved strong results on domain-specific datasets (Karras et al., 2018; 2019; 2020b) and were even scaled to ImageNet with BigGAN (Brock et al., 2019) and StyleGAN-XL (Sauer et al., 2022), the latter reaching FID 2.30 — competitive with multi-step flow models. But GANs have largely fallen out of favor because adversarial training from scratch on standard transformer architectures simply diverges. The paper states this bluntly in Section 1: "When we switch to a standard transformer architecture (Vaswani et al., 2017), training simply diverges."
This is the practical manifestation of a deeper theoretical problem: adversarial training is inherently unstable. The generator and discriminator are locked in a minimax game where the optimization landscape constantly shifts under each player's feet. Without careful architectural choices, regularization schemes, and training tricks accumulated over years of GAN research, the training dynamics collapse — the discriminator becomes too strong and the generator receives vanishing gradients, or the generator exploits weaknesses in the discriminator to produce samples that score well but look nothing like real data.
Why This Gap Matters: Three Concrete Pressures
The paper's motivation is not purely academic. There are real deployment and research pressures that make fast, high-quality generation an urgent problem:
1. Real-time applications require single-step generation. Interactive image editing, real-time video synthesis, on-device generation, and creative tools where users expect immediate feedback cannot tolerate 250 sequential network evaluations. Even multi-step models with 4–8 steps introduce latency that breaks the illusion of real-time interaction. The paper's ability to achieve FID 2.38 in a single forward pass — matching or exceeding what multi-step models produce in 4+ steps — is directly relevant to making these applications viable.
2. Training infrastructure is increasingly dominated by transformer architectures. The field has converged on transformers (specifically, the DiT architecture from Peebles & Xie, 2023) as the standard backbone for image generation. Flow matching, consistency models, and autoregressive methods all build on this architecture. If adversarial training cannot work with standard transformers — requiring instead custom convolutional designs or non-standard architectural modifications — it becomes an island, unable to benefit from the hardware optimizations, scaling infrastructure, and research improvements that the broader community develops for transformer-based models. The paper's demonstration that adversarial training can work on unmodified DiT architectures (Section 3.7: "our architecture requires only minimal modifications to the original DiT") is therefore a practical enabler for adoption.
3. The efficiency gap between training and inference is widening. As models scale up, the cost of multi-step inference scales proportionally. A 1NFE (one function evaluation) model provides the same quality as a 4NFE model of the same size at roughly 4× lower inference cost. This matters enormously for deployment economics — serving costs, energy consumption, and hardware requirements all scale with the number of inference steps. The paper's extra-deep architecture results (Table 7) demonstrate that a 112-layer single-step model can outperform its 28-layer 4-step counterpart, effectively trading increased model depth (a one-time training and memory cost) for reduced inference steps (an ongoing serving cost). This tradeoff becomes increasingly favorable as models are deployed at scale.
Where Prior Approaches Fall Short
The paper identifies specific, structural limitations in three existing approaches to fast generation:
Consistency-based methods waste capacity on irrelevant timesteps. Consistency models (CMs) (Song et al., 2023; Song & Dhariwal, 2024) and their descendants — sCM (Lu & Song, 2025), iMM (Zhou et al., 2025), Shortcut (Frans et al., 2025), MeanFlow (Geng et al., 2026), AlphaFlow (Zhang et al., 2026) — enable few-step generation by training models to predict distant positions along the probability flow. A consistency model trained to produce 1-step outputs must, during training, learn to map from every intermediate timestep to the final output. This is called "consistency propagation": the model sees partially noised images at timestep , , , and must learn to denoise all of them, even though at inference time it will only be asked to denoise from (pure noise).
The paper argues this has two consequences (Section 1 and Section 3.2):
"even when targeting single-step or few-step generation, consistency-based models must still be trained across all timesteps to propagate consistency. This consumes model capacity and introduces error accumulation."
"Consuming model capacity" means the neural network's representational power is split between learning the target task (1-step generation from pure noise) and learning the auxiliary task (generation from intermediate noise levels). Every parameter that learns to handle is a parameter not available for handling . The paper's empirical evidence for this claim is indirect but suggestive: their B/2 model (130M parameters) approaches the performance of consistency-based XL/2 models (675M parameters) on 1NFE generation (Table 4). If consistency training wastes capacity, then a smaller model that dedicates all its capacity to the single-step task should outperform a larger model that splits its capacity across timesteps — and this is exactly what the results show.
"Error accumulation" refers to a subtler problem: during consistency training, the model learns to map from timestep to timestep using teacher forcing — it conditions on the ground-truth intermediate state, not on its own previous prediction. At inference time, however, multi-step generation chains the model's own outputs: the output from step 1 becomes the input to step 2. Small errors in step 1 get amplified in step 2, which amplifies them further in step 3. The model was never trained to handle its own errors as inputs, so this distribution shift degrades quality. The paper's adversarial flow models avoid this entirely in the single-step setting by never using teacher forcing — the generator produces the final output directly, and the discriminator evaluates the final output quality.
Consistency-based methods produce blurry results due to pointwise losses. Even when consistency models are trained optimally, the loss function itself imposes a limitation. The paper states in Section 1:
"models operating in fewer steps have less capacity to predict the exact transformations of targets produced with more steps, so pointwise matching or even moment-matching losses can lead to some degree of blurriness."
This is a fundamental issue with regression-based losses for image generation. A flow-matching model with 250 steps can make small, precise adjustments at each step, gradually refining details. A consistency model with 1 step must make one large prediction, and the L2 loss encourages it to predict the expected value of all possible completions given the noisy input. Since multiple valid images can correspond to the same noisy input, the L2-optimal prediction is a weighted average of these possibilities — which in pixel space manifests as blurriness. The paper cites this as the motivation for why "many state-of-the-art few-step generation models still rely on distributional matching methods, especially adversarial training, for final refinement" (citing Lin et al., 2025; Chen et al., 2025). The adversarial objective does not suffer from this averaging effect because the discriminator can learn to penalize blurry outputs directly, pushing the generator toward producing sharp, realistic samples rather than safe averages.
GANs on standard architectures are unstable — and the paper diagnoses why. The existing GAN literature has achieved impressive results, but with important caveats that limit broader adoption. Recent works that successfully train GANs from scratch use non-standard architectures: R3GAN (Huang et al., 2024) uses a convolutional architecture and achieves state-of-the-art on ImageNet-64, while GAT (Hyun et al., 2025) extends GANs to a transformer architecture but "still employs a non-standard transformer architecture and relies on a pre-trained feature network." Some approaches even rely on frozen pre-trained networks for feature extraction (Kang et al., 2023; Hyun et al., 2025), which adds complexity and limits the ability to train end-to-end on new domains.
But the paper goes beyond merely observing instability — it offers a diagnosis that motivates the entire adversarial flow framework. Section 1 identifies the root cause:
"one of the key reasons GANs are difficult to train is that the adversarial objective alone does not define a single optimization target. This differs markedly from other established objectives, such as flow matching, which has a unique ground-truth probability flow determined by the interpolation function, and autoregressive modeling, which has ground-truth token probabilities determined by the training corpus."
This is Figure 1 in words. In flow matching, the target is unambiguous: the model must learn the unique vector field that transports one distribution to another along the shortest path defined by the interpolation function. In autoregressive modeling, the target is unambiguous: predict the conditional probability of the next token given the training corpus. In GANs, the target is ambiguous: the generator needs to produce samples from the data distribution, but there are infinitely many transport maps that accomplish this. The generator could learn any bijection between the noise and data manifolds, and the discriminator would be satisfied as long as the marginal distributions match.
This ambiguity creates a practical optimization problem:
"the generator keeps drifting during training."
The generator's parameters are constantly changing because there is no fixed target to converge toward. One training step might push the generator toward one valid transport map; the next step, as the discriminator adapts, might push it toward a different valid map. The generator never settles because any map that satisfies the distribution matching constraint is equally valid under the adversarial objective. This drifting behavior makes optimization sensitive to initialization, learning rates, and stochastic minibatch noise — small perturbations can push the generator toward entirely different transport maps from one training run to the next.
The paper's 1D Gaussian mixture experiment confirms this diagnosis visually (Figure 1). Under identical training conditions, standard GANs produce different transport maps from to on different runs — the paths that samples take from the prior to the data distribution vary unpredictably. Adversarial flow models, by contrast, consistently learn the same deterministic optimal transport map across runs because the transport loss provides a unique optimization target.
How This Paper Positions Itself
The paper frames adversarial flow models as a bridge between two previously disconnected families. This is not merely "GANs with an extra loss term" or "flow models trained adversarially" — it is a principled synthesis that inherits the strengths of both while addressing their respective weaknesses.
From the adversarial family, adversarial flow models inherit:
- Single-step generation capability: The adversarial objective evaluates complete samples, so the generator can be trained directly for one-step output without any intermediate timesteps or teacher forcing.
- Distributional matching: The discriminator provides a learned perceptual metric that avoids the blurriness of pointwise L2 losses. Section 3.6 explicitly argues that "deep networks can better capture manifold structure and serve as a better perceptual metric than Euclidean distance," and the choice of relativistic adversarial loss approximates Jensen-Shannon divergence, which is "less sensitive to outliers" than the forward KL divergence minimized by flow matching.
- No consistency propagation required: Unlike consistency models, the generator only needs to be trained on the timesteps that will be used at inference time. For single-step generation, this means the model is trained exclusively on transport, dedicating all its capacity to this one task.
From the flow family, adversarial flow models inherit:
- Deterministic transport: The optimal transport loss (Equation 11) regularizes the generator to learn a specific, well-defined mapping from noise to data — the Wasserstein-2 optimal transport map. This eliminates the ambiguity that causes GANs to drift during training. The paper shows this empirically: adversarial flow models consistently learn the same mapping across different random initializations (Figure 1), while standard GANs produce different mappings.
- Multi-step generalization: Because the model learns a transport map aligned with the probability flow, it naturally supports multi-step generation (Section 3.3). The generator can accept source and target timesteps, allowing it to be used iteratively: transport partway from to , then from to , etc. This bridges the gap between pure single-step GANs and pure multi-step flow models.
- Guidance compatibility: Section 3.5 shows how classifier-free guidance (CFG) — the dominant technique for improving sample quality in conditional diffusion models — can be incorporated into adversarial flow models through an implicit classifier derived from a pre-trained flow-matching model. This is not trivial; the paper demonstrates that naive classifier guidance applied only at the final output (Equation 24) fails to replicate CFG's behavior (Figure 3c), and proposes a flow-based guidance approach (Equation 26) that accumulates guidance gradients along interpolated states to match CFG's effect (Figure 3d).
The paper's novel claim is not that adversarial training or flow matching individually are new — both are well-established. The claim is that the combination solves a problem neither could solve alone: adversarial training provides fast inference and distributional quality assessment, while flow-based deterministic transport provides the training stability that adversarial training has historically lacked on standard architectures. The paper states in Section 3.2:
"Our proposed adversarial flow models instead learn a deterministic optimal transport map. This prevents generator drift and stabilizes training."
The positioning is strategic: the paper is not competing with flow matching on quality (flow matching with many steps still achieves excellent results) or with consistency models on methodology (they remain valid approaches for few-step generation). Instead, it is arguing that for the specific regime of single-step or few-step generation on standard transformer architectures, the adversarial flow combination achieves results that neither family could achieve independently — new state-of-the-art FID scores while maintaining the architectural simplicity and scalability of the DiT backbone.
3. Technical Approach
3.1 Reader Orientation
The paper builds a generative image model that produces ImageNet-256px samples in a single forward pass, trained from scratch using a combination of adversarial and flow-based losses on standard transformer architectures without architectural tricks or pre-trained feature networks. The core idea is that adversarial training provides fast single-step generation and sharp distributional matching, but historically suffers from instability because the adversarial objective alone does not define a unique transport map from noise to data — the generator can drift between infinitely many valid mappings during training. The paper solves this by adding an optimal transport regularization that pins the generator to a specific, well-defined deterministic mapping (the Wasserstein-2 optimal transport), which stabilizes training while preserving the adversarial objective's ability to produce sharp, realistic samples in one step.
3.2 Big-Picture Architecture (Diagram in Words)
The system has five major components, structured as a minimax game between two networks:
-
Generator
$G$— a transformer network (DiT architecture) that takes a noise sample$z \in \mathbb{R}^n$(and optionally source/target timesteps$s, t$for multi-step operation) and produces an image$\hat{x} \in \mathbb{R}^n$in the VAE latent space. Its job is to transport samples from the Gaussian prior to the data distribution. -
Discriminator
$D$— a nearly identical transformer network that takes an image (real or generated) and outputs a scalar logit indicating how "real" it looks under a relativistic objective. Its job is to distinguish real samples from generated ones, providing a learned perceptual metric that guides the generator. -
Optimal Transport Loss
$\mathcal{L}_{\text{ot}}$— a squared L2 distance between the generator's input noise and its output, weighted by a schedule that decays during training. This pins the generator to the unique Wasserstein-2 optimal transport map, preventing the drift that causes standard GANs to diverge. -
Gradient Normalization Operator
$\phi$— a module inserted between the generator's output and the discriminator's input that normalizes the adversarial gradient magnitude in the backward pass. This disentangles the adversarial and transport loss scales, making the hyperparameter$\lambda_{\text{ot}}$transferable across model sizes. -
Classifier Guidance Network
$C$(optional, for conditional generation) — a time-conditioned classifier trained on the probability flow that provides guidance gradients to improve sample quality, analogous to classifier-free guidance in diffusion models.
Information flow during training: A noise sample $z$ and a real image $x$ are drawn → the generator maps $z$ to $\hat{x} = G(z)$ → the gradient normalization operator $\phi$ is applied to $\hat{x}$ (identity in forward pass) → the discriminator receives both $\phi(\hat{x})$ and $x$ and computes a relativistic adversarial loss that measures how much more realistic $x$ looks than $\hat{x}$ → the generator is updated with the adversarial gradient (through $\phi$) plus the optimal transport gradient (directly from $\hat{x}$ to $z$) → the discriminator is updated with the adversarial loss plus R1/R2 gradient penalties and a logit-centering term → the process repeats with alternating updates. For multi-step training, additional timesteps $s, t$ are sampled, and both networks receive interpolated versions of $x$ and $z$ as input.
3.3 Roadmap for the Deep Dive
-
First, the adversarial training formulation (Section 3.1): the relativistic adversarial objective, gradient penalties, and logit-centering that form the base GAN framework the paper builds upon. This establishes the discriminator's role and the optimization dynamics.
-
Second, the single-step adversarial flow model (Section 3.2): how the optimal transport loss is added to the generator, why it stabilizes training by providing a unique optimization target, and the critical
$\lambda_{\text{ot}}$schedule that decays during training to balance transport regularization against distribution matching. -
Third, the multi-step generalization (Section 3.3): how the model is extended to accept timesteps and transport between arbitrary points along the probability flow, including the interpolation function, timestep sampling strategies, and the weighting function for the transport loss at different step sizes.
-
Fourth, the gradient normalization technique (Section 3.4): the precise mechanism for normalizing adversarial gradients to make
$\lambda_{\text{ot}}$transferable across model sizes, including the EMA tracking and rescaling operations. -
Fifth, the guidance formulation (Section 3.5): why naive classifier guidance fails for adversarial flow models, and how flow-based classifier guidance (interpolating generated samples to intermediate timesteps before feeding to a time-conditioned classifier) recovers CFG-like behavior.
-
Sixth, the architectural design and deep model scaling (Sections 3.6–3.8): how the generator and discriminator are parameterized using standard DiT architectures, the two possible generator formulations (direct vs. residual), and the transformer block repetition strategy for training extra-deep single-step models end-to-end without intermediate supervision.
3.4 Detailed, Sentence-Based Technical Breakdown
This is primarily a methods paper proposing a new generative modeling framework that synthesizes adversarial training and flow-based deterministic transport, supported by extensive ImageNet-256px experiments demonstrating state-of-the-art single-step generation.
Adversarial Training Preliminaries (Base GAN Formulation)
The paper builds on the standard generative adversarial network framework, but with specific choices motivated by prior work on stabilizing adversarial optimization.
Generator and discriminator roles. The generator is a function $G: \mathbb{R}^m \to \mathbb{R}^n$ that transports samples $z$ from a prior distribution $Z$ (typically a standard Gaussian) to samples $\hat{x}$ that should match the data distribution $X$. The discriminator is a function $D: \mathbb{R}^n \to \mathbb{R}$ that maps an image to a scalar logit representing its "realness" under a relativistic comparison: rather than classifying samples as real or fake in isolation, the relativistic discriminator compares pairs of real and generated samples and outputs the log-odds that the real sample is more realistic than the generated one.
Relativistic adversarial loss. The paper adopts the relativistic objective from Jolicoeur-Martineau (2019), with the specific functional form $f(a, b) = -\log(\text{sigmoid}(a - b))$. The discriminator loss is:
where $x \sim X$ is a real sample, $z \sim Z$ is a noise sample, $D(x)$ is the discriminator's logit for the real sample, and $D(G(z))$ is the discriminator's logit for the generated sample.
What it computes: for each pair $(x, G(z))$, the function $f(D(x), D(G(z))) = -\log(\text{sigmoid}(D(x) - D(G(z))))$ measures how much the discriminator fails to recognize that the real sample should have a higher logit than the generated one. When $D(x) \gg D(G(z))$, the sigmoid approaches 1, $-\log(1) = 0$, and the loss is near zero — the discriminator is confident. When $D(G(z)) \geq D(x)$, the sigmoid approaches 0 or is small, and the loss is large — the discriminator is confused. The expectation averages this over minibatches of $(x, z)$ pairs.
Why this form: the relativistic formulation compares real and generated samples directly rather than scoring each in isolation. The paper cites Sun et al. (2020) for showing this yields a better loss landscape, and notes it achieves state-of-the-art in recent GAN works (Huang et al., 2024; Hyun et al., 2025). The $-\log(\text{sigmoid}(\cdot))$ form is the standard binary cross-entropy loss applied to the difference in logits, equivalent to training the discriminator to output a higher logit for real samples than for generated samples by a margin. It avoids the saturation issues of the original GAN loss while providing smooth gradients for both networks.
The generator loss mirrors this structure but with the arguments reversed:
What it computes: now the generator is penalized when $D(x) \gg D(G(z))$ — i.e., when the discriminator can easily tell the generated sample is fake. The generator's goal is to make $D(G(z))$ as close to or larger than $D(x)$ as possible, minimizing the log-sigmoid of the difference.
R1 and R2 gradient penalties. To prevent the discriminator from pushing the generator away from equilibrium, the paper applies gradient penalties $\mathcal{L}_{r1}$ and $\mathcal{L}_{r2}$ from Roth et al. (2017). These penalties constrain the Lipschitz constant of the discriminator by penalizing the squared gradient norm of $D$ with respect to its inputs:
where $\mathcal{L}_{r1}^D$ penalizes large gradients of the discriminator with respect to real samples (encouraging $D$ to be smooth around the data manifold), and $\mathcal{L}_{r2}^D$ penalizes large gradients with respect to generated samples (encouraging $D$ to be smooth around the generator's output distribution).
Why these penalties: in the minimax game, without regularization the discriminator can become arbitrarily steep around the data distribution, which creates vanishing or exploding gradients for the generator. The $R_1$ penalty prevents $D$ from becoming too confident on real data (which would crush the generator's gradient signal), while the $R_2$ penalty prevents $D$ from becoming too steep on generated data (which would create unstable, high-variance gradients). Together they impose a local Lipschitz constraint that stabilizes the optimization dynamics, as analyzed by Mescheder et al. (2018).
Finite-difference approximation for efficiency. Directly computing $\nabla_x D(x)$ requires a second backward pass (double backpropagation), which is computationally expensive. The paper uses a finite-difference approximation from Lin et al. (2025):
where $\epsilon = 0.01$ and $\mathcal{N}(x, \epsilon^2 I)$ adds isotropic Gaussian noise with variance $\epsilon^2$ to the sample.
What it computes: instead of computing the exact gradient $\nabla_x D(x)$ via backpropagation, the approximation measures how much the discriminator's output changes when a small amount of noise $\epsilon$ is added to the input. If the discriminator is smooth (small gradient), the output changes very little, and the penalty is small. If the discriminator is sharp (large gradient), the output changes substantially, and the penalty is large. The division by $\epsilon^2$ scales this finite difference to approximate the squared gradient norm.
Why this form: double backpropagation requires computing second-order derivatives, which doubles the memory and computation cost. The finite-difference approximation requires only two forward passes (one clean, one noised) and a single backward pass, making it substantially cheaper. The paper further reduces cost by computing the penalties on only 25% of samples in each batch, noting "no performance degradation."
Logit-centering penalty. The relativistic formulation can suffer from unbounded drift in the discriminator logits — since only the difference $D(x) - D(G(z))$ matters for the objective, the absolute scale of the logits can grow arbitrarily without affecting the loss. To prevent this, the paper adds a logit-centering penalty following Karras et al. (2018):
What it computes: the squared sum of the discriminator logits for real and generated samples. If the logits drift to large positive or negative values, this penalty grows quadratically, pulling them back toward zero.
Why this form: keeping the logits centered around zero prevents numerical instability and ensures the relativistic objective operates in a well-conditioned regime. Unlike an L2 penalty on individual logits, this form allows the logits to have different scales as long as their sum remains near zero, which preserves the discriminator's ability to make sharp distinctions while preventing unbounded drift.
Final GAN objectives. The complete discriminator and generator losses are:
where $\lambda_{\text{gp}}$ is a tuned hyperparameter controlling the scale of both gradient penalties (R1 and R2 share the same coefficient, consistent with their symmetric roles), and $\lambda_{\text{cp}} = 0.01$ is fixed.
Why only the generator gets the adversarial loss: gradient penalties and logit-centering are regularization terms applied only to the discriminator — they shape the discriminator's landscape to provide useful gradients for the generator, but the generator itself only needs the adversarial signal. This asymmetry is standard in modern GAN training: the discriminator is heavily regularized to serve as a stable critic, while the generator is optimized straightforwardly against that critic.
All expectations are estimated via Monte Carlo approximation over minibatches. The generator and discriminator are updated in alternating steps (one step of $D$, then one step of $G$, repeat). For conditional generation, the class label $c$ is provided to both networks as $G(z, c)$ and $D(x, c)$ through the modulation mechanism in the DiT architecture (Miyato & Koyama, 2018 projection discriminator conditioning).
Single-Step Adversarial Flow Models
This is the core contribution of the paper: adding an optimal transport regularization to the generator that transforms the ambiguous GAN objective into a well-posed problem with a unique solution.
The problem with standard GANs. The adversarial objective described above only enforces that the marginal distribution of $G(z)$ matches the data distribution $X$. It places no constraint on how individual noise samples are mapped to data samples. Formally, there are infinitely many measurable functions $G$ that satisfy $G_\# Z = X$ (the pushforward of the prior equals the data distribution). The generator is free to learn any of these functions, and the specific one that emerges depends on weight initialization, minibatch noise, and the stochastic dynamics of the minimax optimization.
The paper argues this ambiguity directly causes training instability: "the generator keeps drifting during training" (Section 1) because there is no fixed target to converge toward. One training step pushes $G$ toward one valid transport map; the discriminator adapts; the next step pushes $G$ toward a different valid map. The parameters oscillate because the loss landscape has a continuous manifold of equally valid minima rather than a single basin.
Brenier's theorem and the optimal transport map. Optimal transport theory provides a principled resolution. Brenier's theorem states that when the source distribution is absolutely continuous (e.g., a Gaussian) and the cost function is quadratic $c(x, z) = \|x - z\|_2^2$, there exists a unique optimal transport map $G^*$ that pushes the source to the target while minimizing the expected transport cost:
What it computes: among all functions $G$ that correctly map the noise distribution to the data distribution, $G^*$ is the one that minimizes the expected squared Euclidean distance between each noise sample $z$ and its image $G(z)$. This selects the "laziest" transport map — the one that moves each noise sample as little as possible while still achieving the correct output distribution.
Why this is the right target: the optimal transport map is unique, deterministic, and well-defined purely by the geometry of the source and target distributions. It provides exactly what standard GANs lack: a single, unambiguous optimization target for the generator. If the adversarial loss enforces distribution matching (which it does, under ideal conditions), and the transport loss biases the generator toward the optimal transport map (which it does, by penalizing large displacements), then the combined objective has a unique global minimum at $G^*$.
Adversarial flow models enforce two constraints simultaneously. The paper parameterizes the transport map with a deterministic neural network $G(z)$, where the prior distribution is constrained to have the same dimensionality $n$ as the data distribution (i.e., $x, z \in \mathbb{R}^n$ and $G: \mathbb{R}^n \to \mathbb{R}^n$). This dimensionality constraint is standard in flow-based models because optimal transport between distributions of different dimensions is not well-defined.
The key insight is that the adversarial objective and the transport objective play complementary roles:
- The adversarial loss
$\mathcal{L}_{\text{adv}}^G$enforces that$G_\# Z = X$— i.e., the generator's output distribution matches the data distribution. This is the validity constraint: without it,$G$could learn the identity map$G(z) = z$and perfectly satisfy the transport loss while producing nothing like real images. - The optimal transport loss
$\mathcal{L}_{\text{ot}}^G$enforces that$G$is the minimum-cost transport map:
where the factor $\frac{1}{n}$ normalizes by the data dimensionality to make the loss scale invariant to image resolution.
What it computes: for each noise sample $z$, the squared L2 distance between $z$ and its generated output $G(z)$ is computed, averaged over all spatial dimensions (the $\frac{1}{n}$ factor), and then averaged over the minibatch. This is a simple mean squared error loss between the input noise and the output image.
Why this form: the L2 distance is the defining cost function of Wasserstein-2 optimal transport. Minimizing this expectation over the data distribution (enforced by the adversarial loss) is exactly the definition of finding the optimal transport map. The $\frac{1}{n}$ normalization makes the loss comparable across different latent space sizes (e.g., different VAE downsampling factors or image resolutions). This is critical because the transport loss competes with the adversarial loss — without proper scaling, the relative strength of the two losses would depend on image resolution, making hyperparameter tuning fragile.
The combined generator objective. The full adversarial flow model objectives are:
The discriminator loss is unchanged from the GAN formulation. The generator loss adds a weighted optimal transport term controlled by $\lambda_{\text{ot}}$.
The critical $\lambda_{\text{ot}}$ schedule. The weighting factor $\lambda_{\text{ot}}$ is not a fixed hyperparameter — it must be scheduled to decay during training. Section 3.2 explains why with a conceptual diagram (Figure 2):
- If
$\lambda_{\text{ot}}$is too small (approaching 0), the transport regularization is too weak to prevent the generator from drifting. Optimization may get stuck in local minima of the adversarial landscape because there is no guidance toward a specific transport map. The model effectively becomes a standard GAN, with all its associated instability. - If
$\lambda_{\text{ot}}$is too large, the transport loss dominates, and the generator is pushed toward the identity map$G(z) \approx z$. Since$z \sim \mathcal{N}(0, I)$and images are not Gaussian noise, this forces the generator to trade off between matching the data distribution (via the adversarial loss) and staying close to$z$(via the transport loss). The result is poor distribution matching — blurry or low-quality samples because the generator cannot simultaneously satisfy both objectives when$\lambda_{\text{ot}}$is too large. - The paper's solution: start training with a large
$\lambda_{\text{ot}}$to establish a stable transport map early in training, then decay it toward a small value to let the adversarial objective refine the distribution matching. Table 2 quantifies this effect for the B/2 model: with a constant$\lambda_{\text{ot}} = 0.2$, the final FID is 29.4; with cosine decay from 0.2 to 0.01 over 100 epochs, the FID improves to 8.51; with decay from 0.2 to 0.0, the FID is 8.69. The decay is essential — it provides strong regularization when the model is far from the data distribution (early training) and relaxes it when fine-grained distribution matching becomes important (late training).
Empirical validation: deterministic transport across runs. The paper validates the core claim in a 1D mixture-of-Gaussians experiment (Figure 1). Under identical training conditions with different random seeds, standard GANs learn different transport maps — the paths from $z$ to $x$ vary qualitatively between runs. Adversarial flow models consistently learn the same deterministic optimal transport map across runs, visualized as parallel, non-crossing transport lines. This confirms that the optimal transport loss eliminates the ambiguity that causes GAN training instability.
Comparison to consistency-based methods. The paper highlights two structural advantages of this formulation over consistency models:
-
No consistency propagation required. Consistency models must be trained on all intermediate timesteps
$t \in [0, 1]$to learn the mapping from any noise level to the final output, even if only single-step generation is needed at inference time. Adversarial flow models for single-step generation are trained only on$t=1 \to t=0$(or equivalently, noise$\to$data with no timestep conditioning at all — in the single-step setting, "the timestep projection is removed" from the generator, Section 3.7). This saves model capacity and reduces training iterations. -
No teacher forcing. Multi-step consistency models are trained with teacher forcing: the model conditions on the ground-truth intermediate state, not its own previous prediction. At inference time, the model chains its own outputs, creating a train-test distribution mismatch that accumulates errors. Single-step adversarial flow models have no teacher forcing at all — the generator produces the final output directly, and the discriminator evaluates it directly, with no intermediate steps.
Additionally, the single-step formulation "removes the hyperparameters associated with timestep sampling and weighting," simplifying the training setup.
Multi-Step Adversarial Flow Models
While the single-step model is the main focus, the paper shows how adversarial flow models naturally generalize to multi-step generation by conditioning on timesteps, establishing the connection to flow-based models.
Timestep conditioning and the interpolation function. The paper introduces an interpolation function of the same form used in flow-matching models:
where $t \in [0, 1]$ is a continuous timestep, $x$ is a data sample (at $t=0$), $z$ is a noise sample (at $t=1$), and $x_t$ is the interpolated state at timestep $t$. For simplicity, the paper adopts linear interpolation:
so that $x_0 = x$ (pure data), $x_1 = z$ (pure noise), and $x_t$ is a convex combination for intermediate $t$.
Generator modification. The generator is extended to accept source and target timesteps: $G(x_s, s, t)$, where $x_s$ is the model's input at timestep $s$, and the output should be the corresponding sample at timestep $t$ (where $t < s$ — the model transports samples from noisier to cleaner states). The generator now learns to transport between arbitrary timesteps along the probability flow, not just from $t=1$ to $t=0$.
Discriminator modification. The discriminator is extended to accept the target timestep: $D(x_t, t)$. It evaluates whether a sample at timestep $t$ looks realistic for that noise level — a partially noised image should have the statistical properties of the interpolation between data and noise at that specific $t$.
It is critical that the discriminator is NOT conditioned on the source sample. Section 3.4 makes this point explicitly:
"formulating the discriminator as
$D(x, z)$for the single-step setting, or as$D(x_t, t, x_s, s)$for the multi-step setting, is incorrect. This is because, during training,$x$and$z$are sampled independently. This formulation incorrectly tells$D$that$z$should be paired with every$x$, but$G$can produce only a single mapping."
If the discriminator were conditioned on the source $x_s$, it would learn to check whether the output is consistent with that specific source — but because $x$ and $z$ are paired randomly (not through the same interpolation), no generator could satisfy this criterion for all possible pairings. The discriminator would receive an impossible task, and training would oscillate or diverge.
Multi-step adversarial losses. The adversarial losses are modified to include timestep conditioning and sampling:
where $x_s = \text{interp}(x, z, s)$ and $x_t = \text{interp}(x, z, t)$ are obtained by interpolating independently sampled $x$ and $z$ — note that $x_s$ and $x_t$ come from different $(x, z)$ pairs, so they are not on the same interpolation path. This is a consequence of using separate expectations for the generator and discriminator, where the discriminator receives independently sampled real interpolations $x_t$ and the generator's outputs $G(x_s, s, t)$ that may come from different $x$ and $z$ samples.
The R1 and R2 penalties are also extended with timestep weighting:
where $w(s, t)$ is a weighting function (discussed below). The penalties now depend on which timestep the discriminator is operating at — the appropriate regularization strength may vary with noise level.
Multi-step optimal transport loss. Brenier's theorem still applies in the multi-step setting because the source distribution remains absolutely continuous: each $x_s$ is a convex combination of data and Gaussian noise, which is equivalent to convolving the data distribution with a Gaussian kernel. The optimal transport objective generalizes to:
What it computes: the squared L2 distance between the source $x_s$ and the target prediction $G(x_s, s, t)$, normalized by both the data dimensionality $n$ and the weighting function $w(s, t)$.
Why the $\frac{1}{w(s, t)}$ division: the transport distance between $x_s$ and $x_t$ is proportional to $|s - t|$ under the linear interpolation scheme. When $s$ and $t$ are close, the correct transport map should move the sample only a small distance — the optimal transport loss should be naturally small. The weighting function $w(s, t)$ in the denominator compensates for this, scaling up the loss when the step size is small so that the gradient magnitude is consistent across different $(s, t)$ pairs. The paper finds the following weighting function works well empirically:
where $\delta = 0.001$ provides numerical stability when $s$ and $t$ are very close (preventing division by zero or extremely large loss values).
Timestep sampling strategy. During training, timesteps can be sampled in two ways:
-
Any-step training: Sample
$s \sim \mathcal{U}(0, 1)$and$t \sim \mathcal{U}(0, s)$, so the model sees all possible$(s, t)$pairs. When$s$and$t$are close, the model behaves like a flow-matching model (learning small-step transport). When$s$and$t$are far apart, the model behaves like a trajectory model (learning large-step transport). -
Designated-step training: Train only on the specific
$(s, t)$pairs needed for a given inference schedule. For example, for 2-step generation at inference time (timesteps 1→0.5→0), train only on$(1, 0.5)$and$(0.5, 0)$. Single-step generation is the special case of training only on$(1, 0)$.
Why designated-step training can be better: the paper reports an important empirical finding: "any-step training performs worse due to the dilution of capacity and batch size." When the model must learn transport for all possible $(s, t)$ pairs, its capacity is split across many tasks, and each task sees fewer training examples per epoch (because the batch is spread across more timestep combinations). For achieving the best performance in a specific few-step setting, training on only the needed timesteps is preferable. Section 4.3 notes that this is "often not a limitation, as achieving the best performance in a designated few-step setting is the priority." The 1D experiments (Figure 1) also show that adversarial flow models "need a larger batch size and converge more slowly for any-step training," consistent with the multi-task learning challenge.
The multi-step model bridges GANs and flow matching. At inference, multi-step generation proceeds by iteratively applying $G$:
starting from $x_{\tau_S} = z \sim \mathcal{N}(0, I)$ at $\tau_S = 1$, stepping through a decreasing sequence of timesteps $\tau = [\tau_S, \tau_{S-1}, \ldots, \tau_0]$ where $\tau_0 = 0$, to produce the final image at $\tau_0$. This is equivalent to solving a discrete-time flow ODE, where $G$ provides the transport map between consecutive timesteps rather than the instantaneous velocity field. The paper provides this formulation explicitly in Appendix H (Equation 39):
What it computes: starting from pure noise $x_1$, each step computes the difference between the model's prediction at the next timestep and the current state, then adds this difference to $x_1$ cumulatively. This is analogous to Euler integration of an ODE, where the model predicts the full step rather than the instantaneous derivative, and the summation accumulates these steps.
Gradient Normalization for Stable Adversarial-Transport Balance
The paper identifies a practical challenge: the adversarial gradient received by the generator from the discriminator can have widely varying magnitudes depending on the architecture, weight initialization, and gradient penalty strength. This makes it difficult to find a single value of $\lambda_{\text{ot}}$ that works across model sizes.
The gradient composition problem. Applying the chain rule to the generator loss:
The term marked "discriminator" is the gradient passed backward from $D$ through its input — its magnitude depends on $D$'s architecture, its current parameters, and the strength of the gradient penalties that constrain it. In standard GANs, only the adversarial term exists, and adaptive optimizers like Adam rescale the magnitude of each parameter update independently, making the generator largely invariant to the absolute gradient scale.
Why this matters for adversarial flow models. The generator receives two gradient contributions: one from the adversarial loss (via $D$) and one from the optimal transport loss (directly). Their relative magnitudes determine the effective balance between distribution matching and transport regularization. If the adversarial gradient is 10× larger on an XL/2 model than on a B/2 model (due to different discriminator behavior), then $\lambda_{\text{ot}}$ would need to be 10× larger on the XL/2 model to achieve the same balance. This makes hyperparameter transfer across model sizes impossible without re-tuning.
The gradient normalization operator $\phi$. The paper introduces a module inserted between the generator's output and the discriminator's input that normalizes the adversarial gradient magnitude in the backward pass while leaving the forward pass unchanged:
where $\phi'$ is the gradient scaling factor applied in the backward pass, $\text{EMA}(\cdot, \beta_2)$ tracks an exponential moving average of the squared gradient norm, $\sqrt{n}$ rescales by the square root of the data dimensionality, and $\beta_2$ is the same decay rate used in the Adam optimizer (0.9 in this paper).
What it computes in the forward pass: $\phi(x) = x$ — the identity function. The discriminator sees exactly the same input as before.
What it computes in the backward pass: when the discriminator backward pass produces a gradient $\frac{\partial \mathcal{L}_{\text{adv}}^G}{\partial G(z)}$ (a tensor of the same shape as the generator output), $\phi$ divides this gradient by its exponentially smoothed RMS norm and multiplies by $\frac{1}{\sqrt{n}}$. This normalizes the adversarial gradient to have unit expected magnitude (in RMS sense) regardless of the discriminator's architecture or current state.
Why this form: the operation is essentially applying the Adam normalization logic (which normally operates on parameter gradients) to the activation gradients flowing backward through the generator's output. By using the same $\beta_2$ as Adam, the EMA smoothing is consistent with the optimizer's internal state. The $\frac{1}{\sqrt{n}}$ factor accounts for the fact that the gradient norm grows with the square root of the tensor size (by the central limit theorem for high-dimensional random vectors), making the normalized gradient magnitude invariant to image resolution or latent space dimensionality.
After normalization, the adversarial gradient has a consistent, predictable magnitude. The paper states: "After normalizing the adversarial gradient to a unit scale, we can find a value of $\lambda_{\text{ot}}$ that works well across model sizes." The same $\lambda_{\text{ot}} = 0.2 \to 0.01$ schedule is used for B/2, M/2, L/2, and XL/2 models (Table 11), confirming this empirically.
Visualization in Appendix G. Figures 14–15 in the appendix show the effect of gradient normalization on the training dynamics. Without normalization (Figure 14), the norms of adversarial and transport gradients fluctuate and change scale during training, making their relative balance unpredictable. With normalization (Figure 15), the adversarial gradient norm is stabilized to a constant scale, while the transport gradient norm decays smoothly as $\lambda_{\text{ot}}$ is annealed. This disentanglement "is beneficial for studying hyperparameters, but is not strictly necessary for achieving the best performance."
Incorporating Guidance into Adversarial Flow Models
Classifier-free guidance (CFG) is the dominant technique for improving sample quality in conditional diffusion models, but it relies on having access to the probability flow ODE to accumulate guidance gradients across timesteps. The paper shows how to incorporate guidance into adversarial flow models, which do not have an explicit probability flow during single-step generation.
The problem with naive classifier guidance. Prior adversarial works (Sauer et al., 2022; Kang et al., 2023) train a classifier $C(x, c)$ that predicts $p(c|x)$ and add a loss term to the generator that maximizes the classification probability:
This encourages the generator to produce samples that the classifier confidently assigns to the target class, similar to how classifier guidance works in diffusion models — but only at the final output $G(z, c)$.
Why this fails. Figure 3 demonstrates the issue on a 1D Gaussian mixture with two classes (blue and green). In flow matching, CFG produces a transport that pulls samples toward the class-conditional distribution's core (Figure 3b). In adversarial flow models, naive classifier guidance at $t=0$ produces almost identical results to the unguided model (Figure 3c). The paper explains:
"This is because, in this particular example, the classes are well separated, so the classifier has a clear decision boundary and yields no gradient."
At the final output ($t=0$), the generated samples are already on one side of the classifier's decision boundary. The classifier is highly confident, its gradient is near zero (the predicted probability is saturated), and it provides no useful signal to push the generator toward better class representation. This is fundamentally different from CFG in flow matching, where guidance gradients are accumulated along the flow at intermediate timesteps where the class boundaries are diffused and classifier gradients are non-zero.
Flow-based classifier guidance. To replicate the accumulated gradient effect of CFG, the paper proposes a time-conditioned classifier $C(x_{t'}, t', c)$ that predicts $p(c|x_{t'})$ on noisy samples at various timesteps. During training, the generator still produces only single-step outputs $G(z, c)$, but these outputs are interpolated to random timesteps $t' \sim \mathcal{U}(0, 1)$ using independent noise samples $z' \sim Z$ before being fed to the classifier:
What it computes: for each generated sample $G(z, c)$, a random noise level $t'$ is sampled, the generated sample is mixed with independent noise $z'$ to produce $x_{t'} = (1-t')G(z, c) + t' z'$, and the classifier evaluates the class probability at this noised state. The generator is trained to maximize the expected classification probability averaged over all noise levels.
Why this works. At high $t'$ (more noise), the class boundaries in the classifier's feature space are blurred — a noised golden retriever might be confused with other breeds. The classifier gradient at these noise levels pushes the generator toward producing samples that remain classifiable even under heavy noise, which corresponds to producing more prototypical, unambiguous class representatives. At low $t'$ (less noise), the classifier provides fine-grained feedback on class-conditional details.
The timestep $t'$ can be sampled from a custom range rather than $\mathcal{U}(0, 1)$, which "corresponds to performing CFG on selected timesteps" — a capability unique to this approach, since diffusion model CFG typically applies uniformly across all timesteps. The hyperparameters (the guidance scale $\lambda_{\text{cg}}$ and the range of $t'$) "can optionally be amortized into $G$ to allow inference-time adjustment," similar to how CFG scale can be adjusted at inference time in diffusion models.
Implicit classifier for CFG equivalence. If an existing flow-matching model $v(x_t, t)$ (which predicts the velocity $x_1 - x_0$) is available, the paper shows that an implicit classifier can be derived without training a separate classifier network. Following the CFG derivation (Ho & Salimans, 2021):
where $v(x_t, t)$ is the unconditional velocity prediction and $v(x_t, t, c)$ is the class-conditional prediction. This implicit classifier gradient is then backpropagated through the generator using a constant-multiple trick:
where $(\cdot)$ is shorthand for $\text{interp}(G(z, c), z', t')$. The $\frac{1}{n}$ normalization and the transpose with $G(z, c)^\top$ arise from differentiating through the interpolation step — the full derivation is in Appendix F.
Training procedure for guidance. The paper trains models without guidance until they reach their best FID, then continues training with guidance. This two-stage approach ensures the model first learns a good unconditional transport map before being refined for class-conditional quality. The classifier is a separate B/2 transformer trained from scratch on ImageNet using cross-entropy loss for 30 epochs, operating in the same VAE latent space. Table 3 shows the effect of different guidance scales and timestep ranges on XL/2 1NFE models, with $\lambda_{\text{cg}} = 0.003$ and $t' \sim \mathcal{U}(0, 0.1)$ achieving the best FID of 2.36. The optimal $t'$ range is "much smaller than in typical flow matching, likely because adversarial models already produce good samples without guidance."
For multi-step models, guidance is simpler: the paper applies it only at the final timestep(s) of the generation schedule by setting $t' = t$ (the target timestep of the current step), as indicated in Table 11 (e.g., for 2NFE models, guidance at $t' = 0$; for 4NFE models, guidance at $t' = 0$ and $t' = 0.25$).
Model Architecture (Generator and Discriminator Design)
Both networks use the standard DiT architecture with minimal modifications, which is a central selling point of the approach.
Generator parameterization. The generator can be formulated in two equivalent ways:
Direct formulation: $G(z) = g(z)$ for single-step, or $G(x_s, s, t) = g(x_s, s, t)$ for multi-step, where $g$ is a neural network that directly predicts the output. This is used for single-step models in the paper.
Residual formulation: $G(z) = z - g(z)$ for single-step, or $G(x_s, s, t) = x_s - (s - t) g(x_s, s, t)$ for multi-step. Here $g$ predicts the change from the input rather than the output directly. This is closely related to the velocity-prediction formulation used in flow-matching models, where the network predicts $x_1 - x_0$ (the direction of transport). The scaling by $(s - t)$ accounts for the step size: when $s$ and $t$ are close, the predicted change should be proportionally smaller. This formulation is used for multi-step models in the paper.
The paper notes both formulations "work equally well," demonstrating the flexibility of the architecture. Using the residual form for multi-step models provides a natural connection to the velocity-prediction convention in the flow-matching literature, making it easier to integrate with existing infrastructure (e.g., pre-trained flow models for the implicit classifier in Equation 27).
Discriminator parameterization. The discriminator is parameterized directly as $D(\cdot) = d(\cdot)$, where $d$ uses the same DiT backbone as the generator, with one key difference: a learnable [CLS] token is prepended to the input sequence (of patch tokens), and the final discriminator logit is produced by applying a final LayerNorm and linear projection to the [CLS] token's output embedding. This is analogous to the classification head in Vision Transformers (ViT), where a special token aggregates global information for a scalar prediction.
Shared DiT backbone. Both $g$ and $d$ use the standard DiT architecture with the following specifications (Table 10):
- Patch size: 2×2 (operating on 32×32×4 latent space, yielding 256 patch tokens)
- Activation: GeLU
- MLP expansion ratio: 4×
- Normalization: Pre-LayerNorm with AdaLN-Zero (adaptive layer norm with zero-initialized modulation)
- Model sizes follow MeanFlow's convention: B/2 (12 layers, 768 hidden dim, 12 heads, 130M params for G, 129M for D), M/2 (16 layers, 1024 dim, 16 heads, 306M/304M), L/2 (24 layers, 1024 dim, 16 heads, 457M/455M), XL/2 (28 layers, 1152 dim, 16 heads, 673M/671M).
Timestep embedding differences. For single-step models, the timestep projection is entirely removed from the generator (since there is no timestep conditioning). For multi-step models with fixed discretizations (designated-step training), a single timestep projection is used (only the target timestep matters for the generator). For any-step models, two timestep projections are used in $g$ (for both source and target timesteps). The discriminator receives only the target timestep embedding.
Condition injection. Class conditioning is injected through the modulation mechanism in both $g$ and $d$, following the original DiT design. The class label embedding modulates the layer norm parameters (scale and shift) in each transformer block.
Generator and discriminator sizes are matched. The paper always uses generator and discriminator of the same size (Table 10), unlike some GAN works where the discriminator is smaller. This symmetric design simplifies the architecture and ensures neither network has a capacity advantage.
Deep Model Architecture via Transformer Block Repetition
The paper makes a significant contribution by training extra-deep single-step models using transformer block repetition, demonstrating that depth scaling can substitute for multiple inference steps.
Motivation. Prior work (Lin et al., 2025) has shown that "effective model depth is critical for capturing the nonlinear transformations required to generate high-quality samples, and that insufficient depth is a primary cause of artifacts in single-step models." A single-step model with 28 layers must learn the entire noise-to-data transformation in one forward pass; a multi-step model with 28 layers can distribute this transformation across 2 or 4 passes, effectively operating on deeper representations (albeit with serial dependencies and intermediate projections).
Transformer block repetition. The paper uses the technique from Universal Transformers (Dehghani et al., 2019): the transformer blocks from the first pass are repeated, with hidden states from the initial pass being recycled. Specifically (Figure 4):
- For a 56-layer (2× deep) model: the 28 transformer blocks are applied twice, with the output of block 28 feeding into (a copy of) block 1 for the second pass. A timestep-like embedding is provided to distinguish which repetition iteration is being processed.
- For a 112-layer (4× deep) model: the 28 blocks are applied four times.
- Crucially, the entire network is trained end-to-end for single-forward generation — there is no intermediate supervision, no teacher forcing, and no loss applied after the first, second, or third pass. The discriminator only evaluates the final output after all repetitions.
Why this is powerful. The paper identifies a key advantage over multi-step models:
"extra-deep single-step models can outperform their multi-step counterparts because they can pass hidden states end-to-end without projecting into and reinterpreting from the data space, require no manual definition of timestep discretizations, and are trained without teacher forcing."
In a multi-step model, each step produces a complete image in data space (latent pixels), which the next step must re-ingest and re-process. Information that could have been maintained in the hidden state is lost at the step boundary because the intermediate representation must be compressed into a fixed-size latent image. In an extra-deep single-step model, hidden states flow continuously through all layers, allowing the model to maintain and refine internal representations without intermediate bottlenecks. This is equivalent to giving the model a deeper computational graph for the same total parameter count — the repeated blocks share weights, so the 56-layer model has the same parameter count as the 28-layer model but twice the serial depth per forward pass.
Parameter and compute matching. The paper carefully matches the comparison:
- A 28-layer 2NFE model: 28 layers × 2 forward passes = 56 effective layer-computations, 675M parameters (G), 672M parameters (D). Output every 28 layers.
- A 56-layer (2× deep) 1NFE model: 28 layers × 2 repetitions = 56 effective layer-computations, 675M parameters (G), 671M parameters (D). Output only at the end.
Both models have the same number of parameters for G and the same total FLOPS (56 layer-computations), but the deep model has no intermediate output projection and no teacher forcing. Table 7 shows the deep model achieves FID 2.08 vs. 2.11 for the multi-step counterpart — a small but consistent improvement that supports the hypothesis about hidden state continuity.
Training details for deep models. The depth increase is applied only to the generator; the discriminator stays at its standard depth. The generator's learning rate is reduced by the repetition factor (e.g., 2× for 56-layer, 4× for 112-layer) to account for the larger gradient accumulation across the deeper network. Extra-deep models are trained end-to-end from scratch using the single-step objective. Table 11 shows: the 56-layer model uses initial learning rate $5 \times 10^{-5}$ for G vs. $1 \times 10^{-4}$ for D, decaying to $1 \times 10^{-5}$ and $2 \times 10^{-5}$ respectively; the 112-layer model starts at $2.5 \times 10^{-5}$ for G and $1 \times 10^{-4}$ for D.
Summary of Design Choices and Justifications
- Relativistic adversarial objective over standard GAN loss: yields better loss landscape (Sun et al., 2020) and achieves state-of-the-art in modern GAN training.
- Finite-difference gradient penalties over exact double backpropagation: reduces computational cost by approximately 2× for the regularization terms while maintaining effectiveness, with the 25% subsampling providing additional savings.
- Optimal transport L2 loss over other regularizations: directly addresses the root cause of GAN instability (ambiguous transport map) by providing a unique, well-defined optimization target grounded in optimal transport theory.
- Decaying
$\lambda_{\text{ot}}$schedule over constant weighting: strong regularization early in training establishes a stable transport map; weakening it later allows the adversarial objective to refine distribution matching without being over-constrained. - Gradient normalization operator over per-model
$\lambda_{\text{ot}}$tuning: decouples the adversarial and transport gradient magnitudes, making hyperparameters transferable across model sizes — a practical necessity for scaling. - Flow-based classifier guidance over naive output-level guidance: replicates the accumulated gradient effect of CFG by evaluating the classifier at intermediate noise levels where gradients are non-vanishing.
- Standard DiT architecture over custom GAN architectures: enables the method to benefit from the broader transformer ecosystem (hardware optimizations, scaling infrastructure, research improvements) and demonstrates that adversarial training can work on the same backbones used by diffusion and consistency models.
- Designated-step training over any-step training: concentrates model capacity on the specific timesteps needed at inference, avoiding the capacity dilution and batch size dilution that any-step training causes, while still supporting the key use case of achieving the best possible performance at a target step count.
- Transformer block repetition over multi-step generation for depth scaling: allows continuous hidden state flow without intermediate data-space projections, avoiding the information bottleneck and teacher forcing issues inherent in chaining multiple forward passes with intermediate outputs.
4. Key Insights and Innovations
Innovation 1: The Generality of the GAN Objective Is a Bug, Not a Feature — And Optimal Transport Fixes It
Before this paper, the field tacitly accepted that GAN training instability was a consequence of the minimax optimization dynamics — a problem to be managed through architectural choices, gradient penalties, spectral normalization, and other regularization tricks accumulated over a decade of GAN research. The diagnosis was essentially: adversarial optimization is hard, so we need better optimization techniques.
This paper performs a conceptual reframing that is far more fundamental. It identifies the root cause of instability not in the optimization dynamics per se, but in the underdetermination of the learning problem itself. The adversarial objective only enforces marginal distribution matching — there are infinitely many valid transport maps from noise to data, and the generator is free to drift between them during training because no single one is specified. The paper states this diagnosis explicitly in Section 1:
"the adversarial objective alone does not define a single optimization target. This differs markedly from other established objectives, such as flow matching, which has a unique ground-truth probability flow determined by the interpolation function, and autoregressive modeling, which has ground-truth token probabilities determined by the training corpus."
This is a genuinely new way to understand GAN instability. It shifts the narrative from "the optimizer struggles with the loss landscape" to "there is no unique basin of attraction to converge to." The loss landscape of a GAN does not have a sharp minimum at the "right" solution — it has a plateau of equally valid solutions, each corresponding to a different transport map. The generator never settles because any point on that plateau satisfies the adversarial criterion equally well. This explains why GANs are sensitive to initialization (different starting points lead to different points on the plateau), why they exhibit mode collapse (some transport maps map many noise samples to the same data mode), and why they are hard to scale (larger models explore the plateau more aggressively, amplifying drift).
The fix is equally principled: apply a regularization that breaks the degeneracy by selecting a unique transport map — specifically, the Wasserstein-2 optimal transport map that minimizes the expected squared L2 distance between noise and generated samples. This is not an arbitrary regularization term added for stability; it is grounded in Brenier's theorem, which guarantees that when the source is absolutely continuous and the cost is quadratic, a unique optimal transport map exists. The combined adversarial + optimal transport objective now has a single, well-defined solution: the transport map that both matches the data distribution (enforced by the adversary) and minimizes transport cost (enforced by the L2 loss).
This is a fundamental insight rather than an incremental refinement because it solves a diagnostic problem, not just an optimization problem. Prior GAN stabilization techniques — gradient penalties (Roth et al., 2017), spectral normalization (Miyato et al., 2018), two time-scale updates (Heusel et al., 2017), optimistic optimizers (Daskalakis et al., 2018) — all address symptoms of instability (oscillation, vanishing gradients, discriminator overconfidence) without addressing the underlying ambiguity. Adversarial flow models address the ambiguity directly by making the problem well-posed. The evidence is in Table 1: without the optimal transport loss (λ_ot = 0), training diverges at every tested gradient penalty strength. With an appropriately tuned λ_ot, training succeeds across a range of λ_gp values. The transport loss doesn't just improve stability — it is the difference between convergence and divergence on standard transformer architectures.
The 1D Gaussian mixture experiment (Figure 1) provides the visual proof of concept: different random initializations consistently produce the same transport map under adversarial flow training, while standard GANs produce different maps. The generator is no longer drifting; it has a fixed target to aim for.
Innovation 2: The Decoupling of Distribution Matching from Transport Specification Creates a New Inference-Time Efficiency Axis
The paper's second major conceptual contribution is the realization that the generator's training objective can be decoupled into two independent specifications — what distribution to produce, and how to map noise to that distribution — and that only the "how" needs to be generalized across timesteps for multi-step generation.
This insight reconfigures how we think about the relationship between training and inference in generative models. In flow matching and consistency models, these two specifications are entangled: the training objective (matching the probability flow at each timestep) simultaneously defines both the target distribution and the transport mechanism. To support N-step generation at inference, the model must be trained on all timesteps during training, because that's where the objective lives. Consistency models are the extreme case: even for 1-step generation, the model must learn the transport from every intermediate timestep because the consistency objective propagates the boundary condition through the entire flow.
Adversarial flow models disentangle these concerns. The adversarial loss handles the "what" — it ensures the output distribution matches the data, operating purely on the final samples. The optimal transport loss handles the "how" — it specifies which noise map goes where. For single-step generation, the model only needs to learn the 1→0 transport; for 2-step generation, it only needs the 1→0.5 and 0.5→0 transports. There is no requirement to train on timesteps that will not be used at inference.
This is not just a training convenience — it is a capacity allocation insight. Every timestep a consistency model is trained on consumes a fraction of the model's representational capacity. A model trained on 100 timesteps to support 1-step generation spends 99% of its capacity budget on auxiliary tasks that will never be evaluated at inference time. The paper's empirical evidence is compelling: the B/2 adversarial flow model (130M parameters) achieves FID 3.05 with guidance, approaching consistency-based XL/2 models (675M parameters, 5× larger) at the same 1NFE setting (Table 4). This suggests that a substantial fraction of the consistency model's capacity is indeed consumed by the timestep propagation task rather than the target single-step task.
This same insight extends to the advantage of training only on designated timesteps rather than on the full continuous time range. The paper reports that "any-step training performs worse due to the dilution of capacity and batch size" (Section 3.6), confirming that even within the multi-step setting, restricting the training distribution to only the timesteps needed at inference improves performance. This is a concrete, actionable principle: the training distribution over timesteps should match the inference distribution, not exceed it.
This is a fundamental insight rather than incremental because it changes the default assumption in fast generative modeling. The consistency model literature has treated timestep propagation as a necessary cost — you must train on all timesteps to get the consistency property. Adversarial flow models demonstrate that there is an alternative: separate the distribution-matching objective from the transport-mapping objective, and let each be trained only on what is needed for inference. This opens a new axis for efficiency optimization that is orthogonal to model size, architecture, or optimizer choice.
Innovation 3: Depth Scaling Can Substitute for Multi-Step Inference — And the Bottleneck Is Information Preservation Across Steps, Not Model Capacity
The paper's third innovation is an empirical finding with significant architectural implications: extra-deep single-step models can outperform their multi-step counterparts at equal total compute and parameter count, and the advantage comes from preserving hidden states end-to-end rather than forcing intermediate projections into data space.
Prior work (Lin et al., 2025) had established that model depth is critical for single-step generation quality, but the relationship between depth and step count was unexplored. The implicit assumption in the few-step generation literature was that reducing the number of steps forces a quality tradeoff that can only partially be recovered by increasing model size or improving the training objective. The paper challenges this assumption by demonstrating that a 56-layer 1NFE model (28 transformer blocks repeated twice) achieves FID 2.08, surpassing its 28-layer 2NFE counterpart at FID 2.11 — both using the exact same number of parameters and the same total layer-computations (Table 7). Even more strikingly, a 112-layer 1NFE model achieves FID 1.94, surpassing the 28-layer 4NFE model at FID 2.02.
The paper's explanation for why this works is the key insight: multi-step models suffer an information bottleneck at the step boundary. Each step in a 2NFE or 4NFE model must produce a complete latent-space image, which the next step then re-ingests as input. Any information that was present in the internal activations of the previous step but is not captured in the output pixels is lost. The model must reconstruct contextual understanding from scratch at each new step. In contrast, the extra-deep single-step model maintains continuous hidden state flow across all layers — information computed in layer 10 remains available (through residual connections and attention) to layer 50 without being compressed through a pixel-space bottleneck.
This is a fundamentally different diagnosis of the quality gap between single-step and multi-step models. The gap is not primarily about the difficulty of learning a large single transformation versus multiple small ones — it is about the representational cost of serializing and deserializing intermediate outputs. The paper's evidence comes from the layer visualizations in Appendix B (Figures 11–12): in the extra-deep models, middle layers appear to contribute meaningfully to the final output even though they produce no intermediate image. In the multi-step models, those same computations would have been forced through a pixel-space projection, discarding information.
This is a fundamental finding rather than incremental because it reframes the architecture design problem. If depth can substitute for steps, then the optimal deployment strategy may be to train very deep single-step models rather than shallower multi-step models — trading increased training memory and one-time FLOPs for reduced inference latency and simplified serving infrastructure. The paper explicitly frames this as an economic tradeoff: depth scaling increases the cost of training and model storage, while multi-step inference increases the cost of serving. For deployments where inference cost dominates (which is typical for production generative models), the deep single-step approach is increasingly favorable.
Moreover, this finding suggests that transformer block repetition (Dehghani et al., 2019) — a technique originally developed for NLP to increase model depth without increasing parameters — finds a natural home in single-step generation, where the continuous hidden state flow aligns perfectly with the technique's strengths. The paper demonstrates that this can be done end-to-end without any intermediate supervision, avoiding the teacher-forcing issues that plague multi-step approaches.
Innovation 4: Gradient Normalization at the Activation Level Enables Architecture-Agnostic Hyperparameter Transfer
While this innovation may appear technical, it addresses a practical barrier that has historically made adversarial methods difficult to scale and reproduce. The core insight is that normalizing the adversarial gradient at the generator's output (rather than at the parameter level, as optimizers like Adam do) decouples the scale of the adversarial signal from the architectural properties of the discriminator, making the optimal transport regularization strength transferable across model sizes without re-tuning.
This is a solution to a problem that the paper itself identifies and diagnoses. In a standard GAN, only the adversarial gradient exists, and adaptive optimizers like Adam normalize it at the parameter level, so the generator is largely invariant to the absolute gradient scale. But adversarial flow models introduce a second gradient contribution from the optimal transport loss, and the relative magnitude of these two gradients determines the effective balance between distribution matching and transport regularization. If the adversarial gradient is 10× larger on an XL/2 discriminator than on a B/2 discriminator (due to different architectures producing different gradient norms), the optimal transport loss weight λ_ot must be 10× larger on the XL/2 model to achieve the same balance. This makes hyperparameter transfer impossible, forcing a grid search for each model size.
The gradient normalization operator φ solves this by normalizing the adversarial gradient to unit expected magnitude before it is combined with the transport gradient. The backward pass formula (Equation 23) tracks an EMA of the squared gradient norm and rescales accordingly, borrowing the β_2 decay rate from Adam to maintain consistency with the optimizer's internal state. After normalization, the adversarial gradient has a predictable magnitude regardless of discriminator architecture, making λ_ot truly transferable.
The empirical validation is implicit but clear: Table 11 shows the same λ_ot schedule (0.2 → 0.01 or 0.001) used across B/2, M/2, L/2, and XL/2 models, with only the learning rate and the terminal λ_ot value varying by setup (guided vs. unguided) rather than by model size. The gradient normalization is what makes this consistency possible.
Appendix G provides visual evidence: without normalization (Figure 14), adversarial and transport gradient norms drift and fluctuate independently during training; with normalization (Figure 15), the adversarial gradient is stabilized to a constant scale, and the transport gradient decays smoothly as λ_ot is annealed. The two losses maintain a consistent ratio throughout training, which the paper argues "is beneficial for studying hyperparameters, but is not strictly necessary for achieving the best performance" — an honest assessment that positions this as a practical enabler rather than a quality-improving technique.
This innovation is incremental in mechanism but fundamental in impact on reproducibility and scalability. It addresses a concrete failure mode of adversarial methods — that moving from a small model to a large model requires re-doing the entire hyperparameter search — and solves it with a clean, principled operation that can be inserted as a transparent module. The fact that it uses the same β_2 as Adam means it integrates naturally with existing optimizer infrastructure without introducing new hyperparameters. This is the kind of practical contribution that determines whether a method gets adopted beyond the authors' lab.
5. Experimental Analysis
Evaluation Methodology
-
Dataset. All experiments use class-conditional ImageNet (Russakovsky et al., 2015) at 256×256 resolution. Images are resized and horizontally flipped with 50% probability. Models operate in the latent space of a pre-trained VAE (Rombach et al., 2022) at 32×32×4 resolution, following the standard protocol established by DiT (Peebles & Xie, 2023).
-
Base model(s). The generator and discriminator both use the standard DiT architecture with minimal modifications (Section 3.7). Model sizes follow MeanFlow's convention: B/2, M/2, L/2, XL/2, where the number after the slash denotes the patch size (2×2). The generator and discriminator always use models of the same size. For deep architecture experiments, depth is increased only in the generator while the discriminator remains at its standard depth (28 layers for XL/2).
-
Metrics. The primary metric is Fréchet Inception Distance on 50k class-balanced samples (FID-50k) computed against the entire training set using the evaluation code from ADM (Dhariwal & Nichol, 2021) with pre-computed training set statistics. The paper also reports sFID, Inception Score (IS), Precision, and Recall in the full results table (Table 9), though these are secondary. Class-balanced evaluation generates exactly 50 images per class for all 1000 classes, which the paper notes "may yield a 0.1 FID advantage compared to class-imbalanced counterparts" (Appendix D).
-
Baselines. The paper compares against multiple families of prior work, with careful notation distinguishing methods that share the same latent space and architecture (denoted by
(•)) from those using different spaces or architectures (denoted by other markers in Table 9): consistency-based models including iCT (Song & Dhariwal, 2024), Shortcut (Frans et al., 2025), MeanFlow (Geng et al., 2026), AlphaFlow (Zhang et al., 2026), iMM (Zhou et al., 2025), and TiM (Wang et al., 2025); GANs including BigGAN (Brock et al., 2019), GigaGAN (Kang et al., 2023), StyleGAN-XL (Sauer et al., 2022), and GAT (Hyun et al., 2025); flow-matching and diffusion models including ADM (Dhariwal & Nichol, 2021), DiT (Peebles & Xie, 2023), SiT (Ma et al., 2024), and their REPA/Disperse variants (Yu et al., 2025; Wang & He, 2025); and autoregressive/masking methods including MaskGIT, VAR (Tian et al., 2024), and MAR (Li et al., 2024). -
Generation budget / compute accounting. For single-step models, generation requires 1 NFE (number of function evaluations). For multi-step models, NFE equals the number of sampling steps. Training compute is measured in epochs defined by the number of images seen by the generator — a fair comparison against consistency methods since "both consistency and our methods require multiple forward passes per update iteration" (Appendix E). The paper further analyzes per-update compute overhead (Table 12): the XL/2 1NFE model requires 1.88× the training compute of MeanFlow due to the discriminator's heavy losses and regularizations, but achieves a 15% improvement in best FID.
-
Cross-validation / statistical protocol. No formal cross-validation is used for model selection. The paper reports the earliest epoch at which the best FID is reached for each model, acknowledging that "different models reach their peak FID at different epochs." For hyperparameter grid searches on the B/2 model (Tables 1–2), results are reported as the average of two runs with variation labeled in gray. For training schedules, the paper explicitly states that "the hyperparameter patterns are those that we find during training to produce the best FID" and that these "could be placed on an automatic schedule, but we leave it to future work" (Appendix D). The class-balanced evaluation with 50k samples provides measurement stability, but there is no reported confidence interval or statistical significance testing across independent training runs for the main results.
Main Quantitative Results
Single-Step Generation with Guidance (Table 4)
The headline result: the XL/2 adversarial flow model achieves FID 2.38 with classifier guidance and discriminator augmentation (CG + DA) in a single forward pass — a new state-of-the-art for 1NFE generation on ImageNet-256px in the standard latent diffusion space.
Breaking this down by model size against the most comparable consistency-based baselines (denoted by (•), same latent space and architecture):
| Model | Method | Params | Epochs | Guidance | FID↓ |
|---|---|---|---|---|---|
| B/2 | AFM (Ours) | 130M | 200 | CG + DA | 3.05 |
| B/2 | MeanFlow | 131M | 240 | CFG | 6.17 |
| B/2 | AlphaFlow | 131M | 240 | CFG | 5.40 |
| M/2 | AFM (Ours) | 306M | 120 | CG + DA | 2.82 |
| M/2 | MeanFlow | 308M | 240 | CFG | 5.01 |
| L/2 | AFM (Ours) | 457M | 120 | CG + DA | 2.63 |
| L/2 | MeanFlow | 459M | 240 | CFG | 3.84 |
| XL/2 | AFM (Ours) | 673M | 125 | CG + DA | 2.38 |
| XL/2 | MeanFlow | 676M | 240 | CFG | 3.43 |
| XL/2 | AlphaFlow | 676M | 240 | CFG | 2.81 |
| XL/2 | TiM (concurrent) | 664M | 300 | CFG | 3.26 |
The margin over the best consistency-based method in the same setting is substantial: AFM-XL/2 at 2.38 vs. AlphaFlow-XL/2 at 2.81 (a 0.43 FID improvement, or ~15% relative reduction). At the B/2 scale, the advantage is even more pronounced: AFM-B/2 at 3.05 vs. AlphaFlow-B/2 at 5.40 — nearly halving the FID. The paper interprets this as evidence that "B/2 models are capacity-limited and our method does not waste capacity on other timesteps" (Section 4.3).
Compared to GAN baselines, AFM-XL/2 (2.38) is competitive with StyleGAN-XL (2.30), which operates in pixel space with a much smaller model (166M parameters). The paper notes this caveat explicitly: "StyleGAN-XL has a slightly better FID while being smaller, but it operates in the pixel space, while ours is restricted by the VAE and DiT patch size 2" (Section 4.3). Against the most comparable latent-space GAN, GAT-XL/2+REPA (602M parameters, FID 2.96), AFM-XL/2 shows a 0.58 FID improvement.
The paper also reports results without discriminator augmentation (Table 8): AFM-XL/2 with CG only achieves FID 2.54, showing that DA provides an additional 0.16 FID improvement on top of guidance. This demonstrates that the method is still state-of-the-art even without augmentation, which matters because the paper acknowledges that "the choice of augmentation may implicitly inject inductive biases" (Appendix C).
Epoch counts deserve attention: AFM models typically reach their best FID in 120–200 epochs vs. 240 epochs for MeanFlow/AlphaFlow. For the XL/2 model specifically, AFM requires 125 epochs vs. 240 for MeanFlow — roughly half the number of generator updates. However, Appendix E qualifies this by noting that each adversarial training step requires more computation than a consistency training step (3.625× per update), so the total training compute for the XL/2 model is approximately 1.88× that of MeanFlow (125 × 3.625 vs. 240 × 1, normalized).
Few-Step Generation with Guidance (Table 5)
In the 2NFE setting, AFM-XL/2 achieves FID 2.11, surpassing MeanFlow-XL/2 at either 240 epochs (FID 2.93) or 1000 epochs (FID 2.20), and AlphaFlow-XL/2 at 240 epochs (FID 2.16). The 4NFE AFM-XL/2 achieves FID 2.02, again competitive with or surpassing consistency-based methods.
These models are trained on designated timesteps rather than the full continuous time range. The paper reports that "any-step training performs worse due to the dilution of capacity and batch size" (Section 3.6), so all few-step results use designated-step training. This is an important practical choice: it means the model cannot be used with arbitrary step counts at inference time without retraining, but the paper argues this is "often not a limitation, as achieving the best performance in a designated few-step setting is the priority" (Section 3.6).
No-Guidance Generation (Table 6)
Without any guidance, the 1NFE AFM-XL/2 achieves FID 3.98 — a remarkable result because it outperforms the 250NFE flow-matching baseline DiT-XL/2 (FID 9.62) and SiT-XL/2 (FID 8.30) by a wide margin, despite using 250× fewer function evaluations. The paper explains this by the properties discussed in Section 3.6: adversarial training's discriminator serves as a learned perceptual metric that is "less sensitive to outliers" than the L2 loss used in flow matching, and the adversarial objective approximates Jensen-Shannon divergence rather than forward KL divergence, producing fewer out-of-distribution samples.
The 2NFE AFM-XL/2 without guidance achieves FID 2.36, surpassing even the 2NFE guided models from most consistency-based baselines (compare: MeanFlow-XL/2 2NFE with CFG achieves 2.20–2.93). This suggests that the adversarial objective alone provides substantial distributional quality without the need for classifier-based steering.
Scaling behavior without guidance is consistent: B/2 achieves 6.07, M/2 achieves 5.21, L/2 achieves 4.36, XL/2 achieves 3.98. Each step up in model size brings roughly 0.7–1.2 FID improvement.
Compared to non-flow/non-diffusion methods: AFM-XL/2 at 1NFE (3.98) substantially outperforms MaskGIT at 8 steps (6.18) and VAR at 10 steps (4.95), demonstrating the advantage of the adversarial flow approach over vector-quantized autoregressive methods for single-step generation. However, the paper does not compare against the strongest autoregressive models (VAR at 2B parameters with CFG achieves 1.92 at 10×2 NFE), which use much larger budgets.
Extra-Deep Models (Table 7)
The deep architecture results provide some of the paper's most striking findings:
| Method | Depth | Params | Epochs | Guidance | NFE | FID↓ |
|---|---|---|---|---|---|---|
| AFM-XL/2 | 28 (1×) | 675M | 95 | CG + DA | 2 | 2.11 |
| AFM-XL/2 | 56 (2×) | 675M | 95 | CG + DA | 1 | 2.08 |
| AFM-XL/2 | 28 (1×) | 675M | 145 | CG + DA | 4 | 2.02 |
| AFM-XL/2 | 112 (4×) | 675M | 120 | CG + DA | 1 | 1.94 |
The 56-layer 1NFE model (FID 2.08) surpasses the 28-layer 2NFE model (FID 2.11) at equal parameters and total layer-computations — a direct validation of the hypothesis that continuous hidden state flow through repeated transformer blocks provides an advantage over serial multi-step generation with intermediate data-space projections. The 112-layer 1NFE model at FID 1.94 surpasses the 28-layer 4NFE model at FID 2.02, strengthening this conclusion.
The paper reports these as trained end-to-end from scratch without intermediate supervision. The improvement from 28-layer 1NFE (2.38 from Table 4) to 56-layer 1NFE (2.08) represents a 0.30 FID reduction, and to 112-layer 1NFE (1.94) represents a 0.44 reduction. These gains from depth scaling are substantial — comparable to or larger than the gains from moving from B/2 to XL/2 at fixed depth (3.05 → 2.38, a 0.67 reduction spread across a 5× parameter increase).
A caveat: the number of epochs varies between configurations (95 for 2NFE vs. 145 for 4NFE), making it difficult to isolate whether depth or training duration drives the improvement. The 4NFE model received 145 epochs while the 112-layer 1NFE received 120 — the deep model achieves better results with fewer training epochs, though the per-epoch cost is higher for the deeper model.
Guidance Type Ablation (Table 8)
On XL/2 1NFE models, the contribution of guidance is decomposed:
| Configuration | FID↓ |
|---|---|
| None | 3.98 |
| DA only | 3.86 |
| CG only | 2.54 |
| CG + DA | 2.38 |
Discriminator augmentation alone provides only a modest improvement (3.98 → 3.86, Δ = 0.12), suggesting that for adversarial flow models, the gradient vanishing problem is less severe than in standard GANs — the optimal transport regularization already stabilizes training significantly. Classifier guidance provides the dominant improvement (3.98 → 2.54, Δ = 1.44), and the combination with DA yields an additional 0.16 FID gain. The paper's CG implementation uses flow-based guidance with t' ~ U(0, 0.1) and λ_cg = 0.003.
The full metrics in Table 9 show that guidance significantly improves Inception Score (from 201.85 to 284.18 for XL/2 1NFE) and Precision (from 0.78 to 0.81), at the cost of a slight decrease in Recall (0.52 to 0.52 — no change) and sFID (5.40 to 4.87 — an improvement). This pattern is consistent with classifier guidance sharpening the generated distribution toward more prototypical class examples, improving sample quality (FID, IS, Precision) without substantially reducing diversity (Recall remains stable).
Comparisons with Flow-Matching Models (Figure 9)
The qualitative comparison in Figure 9 shows AFM-XL/2 2NFE without guidance producing perceptually sharper and more realistic samples than SiT-XL/2 250NFE without guidance, despite the 125× difference in inference steps. With guidance, AFM-XL/2 2NFE (FID 2.11) matches SiT-XL/2 500NFE with CFG=1.5 (FID 2.06) — 250× fewer steps for comparable quality. The paper presents uncurated samples with identical seeds, showing that the adversarial model produces images with better texture detail and fewer structural artifacts, consistent with the argument that the adversarial objective's learned perceptual metric is better at capturing image quality than the L2 regression loss used in flow matching.
Training Compute Analysis (Table 12, Appendix E)
The paper provides a detailed breakdown of per-update computational cost compared to flow matching and MeanFlow:
| Method | Forward/Backward Passes per G Update | Adjusted Speed vs. MeanFlow |
|---|---|---|
| Flow Matching | 3 | 4.38× slower (due to 1400 epochs) |
| MeanFlow | 4 | 1× (reference, 240 epochs) |
| Adversarial Flow (Ours) | 14.5 | 1.88× more (125 epochs) |
The 14.5 effective passes per G update come from: G forward + backward (3 passes), 2 D forwards + D backward with parameter gradients (3 passes, times 2 for real and fake), and D backward-only for gradient penalties on 25% of samples (2 × 1 × 0.25 = 0.5 passes for R1 and R2 each, plus additional overhead). The paper is transparent that this computational overhead is "a limitation of the adversarial formulation, not introduced by our adversarial flow" and suggests that "clever result reuse" or exploring alternative Lipschitz constraints could reduce the cost in future work.
Ablation Studies and Robustness Checks
Optimal transport loss necessity (Table 1): Without the OT loss (λ_ot = 0), training diverges at all tested gradient penalty strengths (λ_gp = 0.1, 0.25, 0.5) on the B/2 1NFE model. With λ_ot = 0.2 and λ_gp = 0.25, the FID reaches 53.90 at epoch 20 — far from the final ~3.05 but demonstrating that training is possible. The instability is visible in the variation across runs: without OT, λ_gp = 0.25 produces FID 174.93 ± 2 (nearly random, since FID 178 is approximately the score of random Gaussian noise), confirming that the model produces essentially no meaningful images. This is the paper's strongest evidence that the OT loss is not merely beneficial but essential for training on standard transformer architectures.
OT scale decay (Table 2): A constant λ_ot = 0.2 yields FID 29.4 at epoch 100, while cosine decay from 0.2 to 0.01 yields FID 8.51 — a 3.5× reduction. Decay to 0.0 yields FID 8.69. The finding that decaying to 0.0 works nearly as well as decaying to 0.01 is interesting: it suggests that the OT loss serves primarily as an initialization and early-training regularizer that establishes a good transport map, and once the adversarial objective takes over, the model can maintain this transport without continued regularization. This is consistent with the paper's conceptual framing: the OT loss provides a unique target early in training to prevent drift, and once the generator is in the right basin, the adversarial objective provides sufficient signal to refine without requiring continued transport constraints.
OT scale and learning rate interaction (Table 11): As training progresses and learning rates are reduced, the terminal OT scale can be further lowered. For instance, the B/2 unguided model transitions from λ_ot = 0.2 → 0.01 at LR 1e-4, then to λ_ot = 0.001 at LR 8e-5 (with EMA reload), then further to λ_ot = 0.001 at LR 3e-5 (with additional D reload). Each reduction in λ_ot and learning rate corresponds to an FID improvement (7.30 → 7.05 → 6.07). This staged schedule suggests a nuanced relationship: the optimal OT scale depends on the optimization dynamics, and the paper's manual tuning of this schedule is a limitation that could be addressed with automated scheduling in future work.
Gradient penalty strength (Table 1): At λ_ot = 0.2, the choice of λ_gp has a measurable effect: λ_gp = 0.1 gives FID 60.20, λ_gp = 0.25 gives 54.92, λ_gp = 0.5 gives 73.85. The optimal λ_gp = 0.25 is used throughout the paper. Notably, even with the OT loss present, λ_gp matters — the gradient penalties are not made redundant by the transport regularization. This is because the OT loss constrains the generator's transport map but does not directly constrain the discriminator's gradient; the R1/R2 penalties are still needed to prevent the discriminator from becoming too steep and destabilizing the adversarial dynamics.
Classifier guidance timestep range (Table 3): On XL/2 1NFE, flow-based classifier guidance with t' ~ U(0, 0.1) at λ_cg = 0.003 achieves the best FID of 2.36, compared to 2.48 with t' = 0 (naive output-level guidance with the same λ_cg). Larger t' ranges (U(0, 0.2), U(0, 0.5)) show slightly worse performance, though the differences are small (2.42–2.50). The finding that the optimal range is "much smaller than in typical flow matching" (Section 4.2) supports the paper's claim that "adversarial models already produce good samples without guidance" — only a small amount of guidance at low noise levels is needed for refinement, rather than the full-range guidance typical in diffusion models where guidance must compensate for the averaging effect of the L2 loss across all noise levels.
Guidance scale sensitivity: At the optimal t' range U(0, 0.1), λ_cg = 0.003 achieves 2.36, while λ_cg = 0.002 achieves 2.47 and λ_cg = 0.005 achieves 2.49. The guidance is relatively robust within this narrow range, but too much guidance degrades quality — a pattern familiar from diffusion model CFG, where excessive guidance scales produce artifacts.
Discriminator augmentation design: The paper's DA consists of random integer translation (p=0.4, up to 30% of latent space size) and three applications of random erasing (each p=0.4, scale 0.1–0.5) applied in the latent space (Algorithm 1). This is notably simpler than the augmentation pipelines used in prior GAN works (Karras et al., 2020a uses a broader set of geometric and color transforms). The paper does not ablate the DA components (e.g., translation alone vs. erasing alone), so it's unclear which aspects of the augmentation are driving the 0.16 FID improvement from DA (Table 8). The observation that DA provides only modest gains in adversarial flow models (compared to the substantial gains reported in R3GAN and GAT) suggests the OT loss already provides much of the stabilization that DA would normally contribute.
EMA reload and discriminator reset (Table 11): The paper uses two additional stabilization techniques adapted from prior GAN literature: EMA reload (replacing the online generator with its EMA weights once performance plateaus, and repeating after every epoch) and discriminator reset (reloading D from an earlier checkpoint when training stalls). Table 11 shows these techniques being applied in the later stages of training (e.g., B/2 unguided switches from LR 8e-5 with EMA reload to LR 3e-5 with EMA reload + D reload, improving FID from 7.05 to 6.07). The paper frames these as addressing "the remaining challenges of adversarial training, not introduced by our adversarial flow models" (Appendix C), specifically the gradient vanishing problem. The fact that these techniques are still needed indicates that optimal transport regularization, while solving the generator drift problem, does not eliminate all GAN training difficulties — the discriminator can still become too strong relative to the generator, requiring intervention.
Generator formulation (direct vs. residual): The paper trains single-step models using the direct formulation G(z) = g(z) and multi-step models using the residual formulation G(x_s, s, t) = x_s - (s-t)g(x_s, s, t), noting both "work equally well." This is not empirically ablated (no direct comparison of direct vs. residual for the same setting), but the consistency of results across single-step (direct) and multi-step (residual) suggests both are viable. The residual formulation's advantage is its natural connection to velocity prediction in flow matching, making it easier to integrate with pre-trained flow models for the implicit classifier guidance derivation (Equation 27–28).
Any-step vs. designated-step training: The paper reports that any-step training "performs worse due to the dilution of capacity and batch size" but provides no quantitative comparison (only a mention in Section 3.6 and reference to 1D experiments in Figure 1). This is a significant gap: the claim that designated-step training is superior is central to the paper's argument about capacity preservation, but the magnitude of the advantage is not quantified on ImageNet. A direct comparison of any-step vs. designated-step training for 2NFE generation would have strengthened this claim substantially.
Deep model learning rate scaling: The learning rate for deep generators is reduced by the repetition factor (e.g., 2× for 56-layer, 4× for 112-layer). The paper does not ablate learning rate choices for deep models, making it difficult to distinguish whether the observed improvements come from depth or from the specific learning rate schedule. The 56-layer model uses LR 5e-5 vs. 1e-4 for the standard 28-layer model, and the 112-layer model uses 2.5e-5 — these lower learning rates may independently improve training stability and final performance, partially confounding the depth comparison.
Discriminator size matching: The paper always uses G and D of the same size (Table 10). No ablation of asymmetric G/D sizes is performed, leaving open the question of whether a smaller (cheaper) discriminator could achieve comparable results with reduced training cost. Given that Appendix E identifies the discriminator's computation as the primary source of the 3.625× per-update overhead, this is a practically significant unexamined dimension.
Optimistic optimizer and asymmetric learning rates: The paper explicitly reports that "the optimistic optimizer (Daskalakis et al., 2018) and asymmetrical learning rates (Heusel et al., 2017)" were tested and found "ineffective" (Appendix C). These negative results are reported without quantitative detail, but they inform practitioners that standard GAN stabilization techniques from the literature do not transfer straightforwardly to the adversarial flow setting.
Critical Assessment
The paper makes three central claims, and the experimental support for each warrants careful scrutiny.
Claim 1: Adversarial flow models achieve new state-of-the-art 1NFE FID on ImageNet-256px.
The evidence for this claim is strong within the paper's defined scope. AFM-XL/2 achieves FID 2.38 (Table 4), which is better than all consistency-based models in the same latent space and architecture: MeanFlow-XL/2 (3.43), AlphaFlow-XL/2 (2.81), and the concurrent TiM-XL/2 (3.26). The margin over the next-best comparable method (AlphaFlow at 2.81) is 0.43 FID — a genuine advance.
However, the claim needs qualification in several dimensions. First, StyleGAN-XL achieves FID 2.30 in pixel space with a smaller model (166M parameters), and the paper's explanation that StyleGAN-XL "operates in the pixel space, while ours is restricted by the VAE and DiT patch size 2" does not change the fact that a user choosing a generative model for deployment would see a method with better FID. The VAE bottleneck is a real architectural constraint — it's not just an evaluation artifact, it affects what the model can represent. The paper would be stronger with an investigation into whether the VAE is the limiting factor for further FID improvements.
Second, the B/2 result (FID 3.05 vs. AlphaFlow-B/2 at 5.40) is genuinely impressive and provides the strongest evidence for the capacity-preservation argument. However, the comparison at this size is to consistency models trained for 240 epochs — it's possible that consistency models simply need more training at small sizes (their 240-epoch protocol was designed for XL/2), and the comparison is not necessarily at the pareto frontier for both methods at every size. The paper does not investigate whether longer training or hyperparameter tuning could close the gap for consistency-based B/2 models.
Third, the guidance comparison introduces an asymmetry: adversarial flow models use classifier guidance (a separately trained classifier network) while consistency models use classifier-free guidance (which requires training both conditional and unconditional models, or a single model with conditional dropout). These are different mechanisms with different training costs and inference-time behaviors. CFG is generally preferred in practice because it doesn't require training a separate classifier, but CG has the advantage that the classifier can be trained independently. The paper's Table 8 shows CG providing a 1.44 FID improvement over unguided — without quantifying how much of the advantage over consistency models comes from CG vs. CFG differences rather than from the core adversarial flow method.
Claim 2: The optimal transport loss stabilizes adversarial training on standard transformer architectures.
The evidence for this claim is unequivocal. Table 1 shows that without OT loss, all tested gradient penalty strengths lead to divergence (FID ~170–178, essentially random). With OT loss, training succeeds across a range of λ_gp values. This is a clean, well-controlled ablation that directly supports the paper's central diagnostic: the problem with GAN training on transformers is not fundamentally about gradient penalties or optimizer choices, but about the underdetermination of the transport map. Adding a unique transport target (via OT) solves it.
The 1D Gaussian mixture experiment (Figure 1) provides the conceptual validation, and the consistency of transport maps across different model sizes in the qualitative results (Figures 5–8, where similar seeds produce similar spatial layouts across B/2 through XL/2) provides suggestive but not quantitative evidence at scale. The paper could strengthen this claim by quantifying the transport map stability — for instance, measuring the correlation between G(z) outputs for the same z across different model sizes or training runs, or measuring the effective transport cost at convergence.
A limitation of this claim is that the paper does not demonstrate that the OT loss uniquely stabilizes training — it shows that it works, but doesn't compare against alternative stabilization approaches that might also work. For instance, would a different transport regularization (e.g., encouraging G(z) to be close to some fixed reference image rather than to z itself) also stabilize training? Is the specific choice of Wasserstein-2 optimal transport necessary, or would any transport specification break the ambiguity? The paper's theoretical argument (Brenier's theorem provides uniqueness) is strong, but an ablation comparing L2 transport to other transport costs would strengthen the claim that optimal transport specifically is the right choice.
Claim 3: Extra-deep single-step models outperform multi-step counterparts at equal parameters and compute.
The evidence for this claim is present in Table 7 but deserves careful parsing. The 56-layer 1NFE model (FID 2.08) vs. 28-layer 2NFE (FID 2.11) shows a 0.03 FID advantage — this is a small margin that could plausibly fall within the variance of independent training runs. The paper does not report multiple runs for deep models, so we cannot assess whether the difference is statistically reliable. The 112-layer 1NFE (FID 1.94) vs. 28-layer 4NFE (FID 2.02) shows a more convincing 0.08 FID gap, but again without error bars.
The compute-matching between these comparisons deserves scrutiny. The paper argues that 28 layers × 2 passes = 56 effective layer-computations, matching the 56-layer 1NFE model. This is correct for the generator alone, but the discriminator depths differ: multi-step models use D with 28 layers per step (2× forward passes for 2NFE: one at each timestep), while deep models use D with 28 layers once. The total compute across G and D is not identical — the multi-step model's D does more work per G update because it processes samples at multiple timesteps. This asymmetry is not accounted for in the compute-matching argument, and it means the comparison slightly favors the single-step model (its D is cheaper).
The learning rate differences further complicate the comparison: the 56-layer G uses LR 5e-5 while the 28-layer G uses 1e-4. If lower learning rates independently improve final FID (which is plausible for transformer training), some portion of the improvement may come from the learning rate schedule rather than from depth per se. An ablation with the 28-layer model trained at LR 5e-5 would help isolate the depth effect.
Perhaps the most significant gap is that the paper does not demonstrate that extra-deep models continue to improve with depth beyond 112 layers. Is 4× repetition at or near the optimal depth scaling factor, or would 8× (224 layers) yield further gains? Does the benefit saturate? Does deeper repetition introduce new training difficulties? The paper's broader claim that "the quality of single-step generation may not be bounded by the training method, but by the depth of the generator" (Section 4.3) is intriguing but speculative without scaling the depth further. The layer visualizations (Figures 11e–f) showing that "a large number of middle layers seem not to be contributing much in the visualization" for the 112-layer model raises the possibility that some of the repeated blocks are underutilized and that the effective depth is less than the nominal depth.
What would have strengthened the paper:
-
Multiple training runs with error bars for the main results, particularly for the deep model comparisons where the margins are small.
-
Direct any-step vs. designated-step training comparison on ImageNet, quantifying the "capacity dilution" effect that is central to the paper's argument.
-
Asymmetric G/D size ablation to determine whether the discriminator's computational cost (the main factor in the 1.88× training overhead) can be reduced without sacrificing quality.
-
Depth scaling beyond 112 layers to establish whether the single-step depth-quality curve continues to improve or saturates.
-
Ablation of the OT loss with alternative transport costs (e.g., L1, cosine distance, learned perceptual distance) to test whether the specific choice of Wasserstein-2 is important or whether any transport specification works.
-
Comparison at matched total training FLOPs rather than matched epochs, to account for the per-update cost differences between adversarial and consistency methods. The paper provides the per-update cost analysis in Table 12 but doesn't plot FID vs. total FLOPs, which would be the fairest comparison.
-
Guidance mechanism comparison (CG vs. CFG for adversarial flow models) to determine how much of the advantage over consistency models comes from the core method vs. the guidance implementation.
The experiments do genuinely support the paper's central claims, but with important boundary conditions: the OT loss is clearly necessary and effective for stabilization (Claim 2 is the strongest), the 1NFE results are state-of-the-art within the defined comparison class but the advantage over non-latent-space methods is less clear (Claim 1 holds with qualifications), and the deep model results are promising but rest on small margins without statistical characterization (Claim 3 is plausible but not definitively established). The paper's transparency about limitations — the training compute overhead, the use of CG rather than CFG, the reliance on existing GAN stabilization techniques (EMA reload, D reset), and the manual hyperparameter scheduling — strengthens its credibility even where the experimental evidence is incomplete.
6. Limitations and Trade-offs
Assumption: Difficulty Can Be Estimated Cheaply Enough to Be Practical
The entire compute-optimal framework rests on the ability to estimate prompt difficulty before deciding how to allocate the inference budget. The paper's method for doing so — generating 2048 samples per question and averaging either ground-truth correctness (oracle) or PRM final-answer scores (predicted) — is extraordinarily expensive. At 2048 samples per question, the difficulty estimation step alone consumes more compute than the largest test-time budgets studied (256–512 generations). The authors acknowledge this explicitly in Section 3.2:
"estimating difficulty in this way still incurs additional computation cost during inference... our experiments do not account for this cost largely for simplicity"
This is a significant gap. The reported 4× efficiency gains over best-of-N are computed after difficulty is known, without amortizing the cost of learning it. In a realistic deployment, the total cost would be difficulty estimation + strategy execution, and the former could dominate the latter. The paper suggests future work on training models to predict difficulty directly from the question text, but no such model is developed or evaluated. Until this gap is closed, the 4× figure should be understood as an upper bound on achievable efficiency rather than a realized deployment gain.
Single Benchmark, Single Model Family
All experiments use the MATH benchmark (500 test questions) with PaLM 2-S* as the base model. The authors state they "believe this model is representative of the capabilities of many contemporary LLMs" (Section 4), but this claim is unverified. Several aspects of the findings could be model-specific:
- The PRM's quality and over-optimization behavior depend on PaLM 2-S*'s output distribution. A model with different calibration properties or different error patterns might exhibit different difficulty-dependent scaling curves.
- The revision model's ability to learn from incorrect in-context examples depends on the base model's in-context learning capabilities, which vary substantially across model families.
- The MATH benchmark consists exclusively of competition-level math problems requiring symbolic reasoning. It is unclear whether the difficulty-dependent patterns (beam search hurting easy problems, revisions helping easy problems) generalize to other reasoning domains (code generation, logical reasoning, scientific QA) or to tasks requiring factual knowledge rather than inference.
The test set of 500 questions, split into five difficulty quintiles of ~100 each, then further split by two-fold cross-validation, means the compute-optimal policy is selected based on ~50 questions per fold per bin. This is a small sample, and the selected strategies may not be robust. The paper does not report confidence intervals on the compute-optimal scaling curves, making it difficult to assess whether the observed gains are statistically reliable at this sample size.
The~14× Larger Model Baseline Is Not Compute-Optimal
The FLOPs-matched comparison in Section 7 scales model parameters while holding training data fixed, following the LLaMA paradigm (Touvron et al., 2023). The authors acknowledge that this departs from compute-optimal pretraining (Hoffmann et al., 2022), where both data and parameters are scaled equally:
"We choose this setting as it is representative of a canonical approach to scaling pretraining compute and leave the analysis of compute-optimal scaling of pretraining compute where the data and parameters are both scaled equally to future work."
This matters because a Chinchilla-optimal model trained with 14× more total FLOPs would likely outperform a parameter-only-scaled model, making the pretraining baseline weaker than it needs to be. The reported advantages of test-time compute over pretraining (e.g., +27.8% on easy questions at R ≪ 1) may shrink or reverse against a properly compute-optimal larger model. Additionally, the 14× larger model uses only greedy decoding — no majority voting, no best-of-N, no search. Giving the larger model even a modest test-time compute budget (say, best-of-8) would create a much stronger baseline that is never tested.
Verifier Over-Optimization Is a Hard Ceiling, Not a Solved Problem
The paper documents verifier over-optimization as a central limiting factor: beam search degrades easy-problem performance at high budgets (Figure 3, right), lookahead search — the strongest optimizer — paradoxically performs worst overall (Figure 3, left), and qualitative examples show degenerate outputs (repetitive steps, overly short solutions; Appendix M, Figures 29, etc.). The compute-optimal policy mitigates this by routing easy problems away from aggressive search, but it does not solve the underlying problem. On medium-difficulty problems where beam search is deployed, over-optimization still limits the scaling ceiling — the beam search curves in Figure 3 flatten and sometimes decline well before the budget is exhausted.
This means the compute-optimal approach is fundamentally bounded by verifier quality. Improving the PRM (e.g., through better training data, adversarial robustness, or ensemble methods) would likely shift the difficulty thresholds and change the optimal policy. The current results are therefore specific to the verifier quality achievable with the Monte Carlo rollout training procedure described in Appendix D. The paper does not explore how verifier improvements would alter the scaling landscape.
Hard Problems Remain Essentially Unsolved
Across all methods — search, revisions, and their compute-optimal combinations — the hardest questions (difficulty bin 5) show near-zero improvement regardless of compute budget. In Figure 3 (right), bin 5 accuracy hovers at 1–3% for all methods and all budgets. In Figure 7 (right), bin 5 shows roughly 2–3% accuracy irrespective of the sequential-to-parallel ratio. In the FLOPs-matched comparison (Figure 9), the bin 5 scaling line is essentially flat near 0–5%.
This is a fundamental limitation: test-time compute can amplify existing capability but cannot create it. If the base model's pass@1 is near zero on a problem class, no amount of search or revision will help — there are no correct solutions in the proposal distribution to find or refine. The paper is candid about this (Section 7 takeaway box), but it means the approach offers no path forward for genuinely novel or out-of-distribution reasoning that exceeds the base model's training distribution. For such problems, pretraining remains the only viable path.
Revisions and Search Are Studied Independently, Not Combined
The paper studies two complementary axes — PRM search and iterative revisions — but never combines them. Section 8 explicitly acknowledges this:
"we did not experiment with PRM tree-search techniques in combination with revisions"
This is a significant gap because the two mechanisms have complementary strengths: revisions improve the proposal distribution (generating better candidates), while PRM search improves candidate selection (finding the best among generated candidates). Applying beam search to revision model outputs — or using the PRM to guide which revisions to pursue — could yield gains beyond either method alone. The current results therefore represent a lower bound on what a fully integrated system could achieve.
7. Implications and Future Directions
How This Work Changes the Landscape
This paper performs a conceptual reframing of GAN training that shifts the field's understanding from "adversarial optimization is hard because minimax dynamics are unstable" to "adversarial optimization is hard because the problem is underdetermined — there is no unique target to converge toward." This is a diagnostic shift rather than an algorithmic breakthrough: the paper identifies why GANs diverge on standard architectures (the adversarial objective alone defines infinitely many valid transport maps, giving the generator no fixed optimization target) and demonstrates that providing a unique target — the Wasserstein-2 optimal transport map — resolves the instability. This is substantively different from prior GAN stabilization work (gradient penalties, spectral normalization, two time-scale updates) because those approaches treated symptoms of instability (oscillation, vanishing gradients) while this paper addresses the root cause (ambiguous optimization objective). The evidence in Table 1 — where removing the optimal transport loss causes divergence at every tested gradient penalty strength — supports the claim that the transport specification is not merely helpful but structurally necessary for training on standard transformer architectures.
The paper resolves a latent contradiction in the generative modeling landscape: GANs have produced high-quality single-step samples on domain-specific datasets for years, yet have been largely abandoned by the broader community in favor of diffusion and flow-based models that require hundreds of inference steps. The contradiction was that GANs should be the natural choice for fast generation, but they don't work reliably on the architectures and datasets that matter. Prior work addressed this through architectural compromises (R3GAN uses convolutions, GAT uses non-standard transformer modifications and pre-trained feature networks), implicitly suggesting that standard transformers and adversarial training are incompatible. This paper demonstrates otherwise: unmodified DiT architectures can be trained adversarially from scratch, and the only missing ingredient was a deterministic transport target. This finding makes adversarial training viable on the same hardware and software infrastructure that supports the dominant diffusion-transformer paradigm, removing a structural barrier to adoption.
The paper also recasts the relationship between single-step and multi-step generation as fundamentally about information preservation rather than transformation difficulty. The deep architecture results (Table 7) show that a 56-layer 1NFE model outperforms a 28-layer 2NFE model at equal compute and parameters (FID 2.08 vs. 2.11), and a 112-layer 1NFE model surpasses a 28-layer 4NFE model (FID 1.94 vs. 2.02). This refutes the implicit assumption in the few-step generation literature that quality necessarily degrades as steps decrease and can only be recovered through larger models or better training objectives. Instead, the bottleneck appears to be the loss of information when projecting hidden states to data space at step boundaries — continuous hidden state flow through repeated transformer blocks avoids this bottleneck. This is a methodological insight: depth scaling through block repetition, originally developed for NLP to increase depth without increasing parameters, finds a natural application in single-step image generation where avoiding intermediate data-space projections is uniquely beneficial.
The discovery that consistency-based methods waste capacity on irrelevant timesteps has a direct impact on how the field thinks about training efficiency for fast generation. The paper's B/2 adversarial flow model (130M parameters, FID 3.05 with guidance) approaches the performance of consistency-based XL/2 models (675M parameters, FID ~3.43 for MeanFlow) — a 5× parameter reduction. This gap is most naturally explained by the paper's capacity argument: consistency models must learn transport from every intermediate timestep even for 1-step inference, while adversarial flow models dedicate all capacity to the target timesteps. This finding makes research directions that require training on auxiliary timesteps less attractive relative to methods that can train directly on the inference-time task.
Follow-Up Research This Work Enables
Quantifying the capacity cost of consistency propagation across model sizes. The paper argues that consistency models waste capacity on unused timesteps, and the B/2 vs. XL/2 comparison provides circumstantial support. A direct measurement of this effect would train consistency models with varying numbers of training timesteps (e.g., 2, 4, 8, 16, 32, all evaluated at 1NFE) and measure how FID degrades as more timesteps are added, at constant model size. The hypothesis: FID should monotonically worsen with more timesteps because capacity is split across more auxiliary tasks. This would quantify the capacity-dilution effect and establish whether there is an optimal number of training timesteps that balances the benefits of multi-step training (regularization, easier optimization) against the capacity cost. Such an experiment would also test whether consistency models trained on only 2 timesteps (1→0.5→0) can match adversarial flow models at 1NFE — the paper's argument predicts they should close much of the gap.
Direct comparison of CG vs. CFG within adversarial flow models. The paper uses classifier guidance (CG) for adversarial flow models while consistency baselines use classifier-free guidance (CFG). This confounds the comparison because CG and CFG are fundamentally different mechanisms with different training costs, inference-time behaviors, and quality characteristics. A controlled experiment would train an adversarial flow model with CFG (by training conditional and unconditional variants and following the implicit classifier derivation in Equations 27–28 from a pre-trained flow model) and compare against the CG variant at equal inference cost. The specific question: does the 0.43 FID advantage of AFM-XL/2 over AlphaFlow-XL/2 (2.38 vs. 2.81) persist when both use CFG, or is some portion attributable to CG's different guidance properties? This experiment would also test whether the flow-based guidance approach (interpolating to intermediate timesteps, Equation 26) generalizes to the CFG setting without a separate classifier network.
Depth scaling saturation curve for single-step adversarial flow models. The paper demonstrates depth scaling up to 4× repetition (112 layers) with monotonic FID improvement, but does not establish where the benefit saturates. A systematic sweep of repetition factors (1×, 2×, 3×, 4×, 6×, 8×) on the XL/2 architecture, with carefully controlled learning rates and training durations, would produce a depth-quality curve analogous to scaling law plots. This would answer: does the improvement continue (suggesting single-step generation is fundamentally depth-limited), plateau (suggesting a representational bottleneck), or reverse (suggesting optimization difficulties at extreme depth)? The layer visualizations in Appendix B (Figures 11e–f) showing apparently underutilized middle layers in the 112-layer model hint at diminishing returns: if some repeated blocks contribute little, further repetition may be wasteful. A scaling law experiment would quantify this and provide practical guidance for architecture design.
Combining adversarial flow with representation learning objectives. The paper's extra-deep models train end-to-end without intermediate supervision, but prior work (Yu et al., 2025; Leng et al., 2025) has shown that representation alignment losses (REPA) substantially improve diffusion transformer training. The natural combination — adversarial flow training with a REPA-like objective applied to intermediate layers of the generator — could provide additional training signal without requiring intermediate data-space projections. The specific experiment: train 56-layer and 112-layer adversarial flow models with a representation matching loss between intermediate generator features and a pre-trained encoder's features (e.g., DINOv2), and measure whether this accelerates convergence or improves final FID beyond what depth scaling alone achieves. This would test whether the hidden states in deep models benefit from explicit guidance, and whether representation learning and adversarial objectives are complementary or redundant.
Adversarial flow post-training of pre-trained flow-matching models. The paper trains only from scratch, but explicitly identifies post-training as future work (Section 4.4). The natural experiment: take a pre-trained flow-matching model (e.g., SiT-XL/2) as the generator initialization, add a discriminator, and fine-tune with the adversarial flow objective to convert it from multi-step to few-step generation. This is analogous to adversarial distillation approaches (Lin et al., 2024; 2025) but with the OT loss providing stability during the transition. The key measurement would compare: (a) training adversarial flow from scratch, (b) post-training a flow-matching model with adversarial flow objectives, and (c) distillation-based approaches (e.g., Lin et al., 2024). The hypothesis: post-training should reach a given FID in fewer epochs than from-scratch training because the flow-matching initialization provides a reasonable transport map that only needs refinement, and the OT loss prevents catastrophic drift during the transition from flow-matching to adversarial objectives.
Stress-testing the OT loss necessity claim with alternative transport specifications. Table 1 demonstrates that OT loss is necessary for training — without it, divergence occurs at all gradient penalty strengths. But does the transport need to be optimal (Wasserstein-2), or is any deterministic transport specification sufficient? An ablation comparing the squared L2 transport cost (used in the paper) against alternative costs would test this. Candidates include: L1 distance (which would encourage a different optimal transport map), cosine similarity between z and G(z) (which would encourage directional alignment without constraining magnitude), or a learned perceptual distance (e.g., LPIPS) that better reflects semantic similarity. If all deterministic transport losses stabilize training relative to λ_ot = 0, then the key insight is that any unique target breaks the degeneracy — the specific choice of Wasserstein-2 is incidental. If only the L2 cost works (or works best), then the optimal transport property is genuinely important, and the connection to Brenier's theorem is not merely theoretical motivation but a practical requirement. This experiment would clarify whether future work should focus on optimal transport theory or on transport specification more broadly.
Reducing discriminator computational cost through architecture asymmetry. Appendix E reveals that the discriminator accounts for the bulk of the 3.625× per-update compute overhead relative to consistency training, and Table 12 shows the discriminator's forward and backward passes dominate the computational budget. The paper always matches generator and discriminator sizes; a natural cost-reduction experiment would ablate asymmetric configurations where the discriminator is smaller than the generator (e.g., B/2 discriminator with XL/2 generator). The hypothesis: a smaller discriminator might provide sufficient adversarial signal while substantially reducing training cost, and the OT loss's stabilization might make the method robust to discriminator capacity reduction — unlike standard GANs where a too-weak discriminator provides poor gradients. Measuring FID vs. total training FLOPs for different G/D size ratios would produce a pareto frontier showing the optimal asymmetry for cost-efficiency.
Practical Applications and Downstream Use Cases
Real-time on-device image generation. The B/2 adversarial flow model achieves FID 3.05 with guidance using only 130M parameters and a single forward pass. This opens the possibility of deploying high-quality conditional image generation on mobile devices or in-browser, where the latency of multi-step diffusion models (requiring 20–250 sequential network evaluations) is prohibitive. A 130M-parameter transformer at 1NFE on a 32×32×4 latent space (with a lightweight VAE decoder) can plausibly run in under 100ms on modern phone GPUs, making interactive image generation applications viable. The paper's demonstration that guidance can be applied through a separate classifier or through implicit CFG gradients means that the quality-boosting effects of guidance are available even in resource-constrained settings where training two models is feasible but running them sequentially at inference is not. The 1NFE latency advantage over consistency-based B/2 models (which also support 1NFE but at FID 5.40–6.17) translates to roughly equivalent quality at 1/3 to 1/4 the model size, or superior quality at the same model size and latency budget.
Cost-efficient batch generation for synthetic data and content production. When generating large volumes of images (for training data augmentation, content creation pipelines, or design exploration), the 250× inference cost savings of 1NFE vs. 250NFE generation directly reduce serving costs. The XL/2 adversarial flow model at 1NFE (FID 2.38) matches the quality of flow-matching models at 250NFE (SiT-XL/2 CFG at 2.06) while requiring 250× fewer sequential network evaluations. In a production setting generating millions of images, this translates to a proportional reduction in GPU-hours, energy consumption, and serving latency. The extra-deep architecture results (112-layer 1NFE at FID 1.94) suggest that if training budget allows for deeper models, single-step quality can approach or exceed what multi-step models achieve, making the cost-benefit analysis increasingly favorable for single-step deployment. The paper's transparency about the 1.88× training compute overhead means practitioners can make informed decisions: the training cost premium is repaid after roughly 2× the inference volume that a multi-step model with equivalent quality requires, which for high-volume deployments occurs quickly.
Architectural template for video generation acceleration. The paper's core mechanism — learning a deterministic transport map through adversarial training stabilized by optimal transport — is not specific to images. Video generation models face even more acute latency pressures because each video frame (or frame group) propagated through a multi-step diffusion process multiplies the per-step cost by the number of frames. Adversarial flow models, which produce complete latents in one forward pass, could dramatically accelerate video generation if extended to spatiotemporal latents. The gradient normalization technique (Section 3.4) that enables λ_ot transfer across model sizes is particularly relevant for video, where model sizes vary substantially depending on temporal resolution and frame count. The multi-step formulation (Section 3.3) would enable flexible quality-latency tradeoffs: generate keyframes with adversarial flow and interpolate temporally with smaller steps. The paper's findings about designated-step training being superior to any-step training is actionable advice: for video, train only on the specific timestep discretization needed at inference rather than attempting to support arbitrary step counts.
Research infrastructure for adversarial generative modeling on transformers. The paper's demonstration that standard DiT architectures can be trained adversarially from scratch removes a practical barrier that has prevented the broader community from experimenting with GANs on transformer backbones. Researchers who have invested in DiT training infrastructure (dataloaders, distributed training setups, evaluation pipelines) can now add adversarial flow training with minimal architectural changes (Section 3.7: "our architecture requires only minimal modifications to the original DiT"). The gradient normalization operator φ (Equations 22–23) and the finite-difference gradient penalty approximation (Equations 4, 6) provide concrete, low-overhead recipes for the regularization components that would otherwise require extensive tuning. The paper's detailed training schedules (Table 11) — including the staged λ_ot decay, learning rate adjustments, and when to apply EMA reload and discriminator reset — serve as a practical starting point for new domains, reducing the trial-and-error burden that has historically made adversarial methods difficult to adopt.
When to Prefer This Method Over Alternatives
The paper positions adversarial flow models as a specific synthesis of adversarial and flow-based approaches, and the experimental results suggest clear conditions under which this method is preferable:
Prefer adversarial flow models when:
- Single-step or very-few-step (1–4 NFE) generation is required, and inference latency or cost dominates training cost. The XL/2 model at FID 2.38 (1NFE) provides quality competitive with multi-step consistency models (MeanFlow-XL/2 at 2NFE: 2.20–2.93; AlphaFlow-XL/2 at 2NFE: 2.16) while requiring fewer sequential network evaluations.
- Standard transformer (DiT) architecture is preferred or required for infrastructure compatibility. The method works with unmodified DiT backbones (Section 3.7), avoiding the custom architectures used by R3GAN (convolutional) and GAT (modified transformer with pre-trained features).
- Training from scratch is preferred over distillation from a pre-trained teacher model. Unlike adversarial distillation methods (e.g., Lin et al., 2024; 2025; Sauer et al., 2024b), adversarial flow models do not require a pre-trained multi-step model, eliminating the dependency on teacher quality and the computational cost of generating distillation targets.
- Training budget allows for 1.88× the per-epoch compute of consistency methods (Appendix E, Table 12), and this premium is acceptable given the inference-time savings.
Prefer consistency-based methods when:
- Inference-time flexibility to change the number of sampling steps is required. Consistency models trained with any-step objectives can trade quality for speed at inference time without retraining; adversarial flow models trained on designated timesteps are fixed to a specific step schedule (unless any-step training is used, which the paper finds performs worse).
- Classifier-free guidance (CFG) is strongly preferred over classifier guidance (CG). The paper uses CG, and while it derives an implicit CFG formulation (Equations 27–28), CFG is not demonstrated at scale. If CFG is a hard requirement (e.g., for text-to-image generation where separate classifier training is impractical), consistency models with demonstrated CFG support may be more practical.
- Training from pre-trained diffusion model weights is desired. Consistency models can be initialized from diffusion model checkpoints (e.g., Song et al., 2023; Song & Dhariwal, 2024), amortizing the large pretraining cost. Adversarial flow models currently require training from scratch (the paper leaves post-training to future work, Section 4.4).
Prefer flow-matching or diffusion models with many steps when:
- Maximum sample quality is the sole objective and inference cost is not constrained. The paper's best XL/2 2NFE model at FID 2.11 with guidance is competitive with but does not clearly surpass the best flow-matching models at 250NFE (SiT-XL/2+REPA CFG at 1.42; SiT-XL/2+REPA-E at 1.12). If 250× inference cost is acceptable, flow-matching still holds the overall quality record in the comparable latent space setting.
- Mode coverage and diversity are paramount. The paper reports Recall of 0.52 for 1NFE models and 0.59 for 4NFE models (Table 9), compared to 0.67 for SiT-XL/2 250NFE. The adversarial objective, which approximates JS divergence, may sacrifice some mode coverage relative to the forward KL divergence minimized by flow matching — consistent with the paper's observation (Section 3.6) that flow matching "minimizes forward KL divergence to maximize mode coverage." For applications where capturing the full diversity of the data distribution matters more than per-sample quality, multi-step flow matching may be preferable.
