ArXiv: 2604.11521
🎯 Pitch
Swapping flow matching's MSE loss for a learned discriminator slashes ImageNet 256 FID from 8.26 to 3.63 without guidance—a 2× leap in raw generation quality from merely 10 epochs of post-training. The trick is discriminating in the Jacobian-vector product space, giving continuous normalizing flows a GAN-like sensitivity to perceptual detail.
1. Executive Summary
This paper introduces continuous adversarial flow models (CAFMs), a method that post-trains existing flow-matching generative models by replacing the fixed mean-squared-error objective with a learned discriminator that guides the model toward better generalization. Evaluated on ImageNet 256×256 generation, CAFM post-training substantially improves guidance-free FID for the latent-space SiT-XL/2 model from 8.26 to 3.63 and for the pixel-space JiT-H/16 model from 7.17 to 3.57—a more than 2× reduction in each case—while also improving the best guided FID from 2.06 to 1.53 for SiT and from 1.86 to 1.80 for JiT using only 10 epochs of finetuning. The method extends adversarial training into continuous-time flow modeling by discriminating in the Jacobian-vector product space of a learned discriminator (discriminating velocity predictions as directional derivatives of a scalar potential rather than discriminating positions directly), establishing that adversarial objectives can improve distribution matching for continuous normalizing flows without requiring discrete simulation steps, though the gains manifest primarily as improved generalization under finite model capacity rather than as convergence to a different target distribution.
2. Context and Motivation
The Core Problem: Flow Matching Generalizes Poorly Without Guidance
Flow matching has emerged as one of the dominant paradigms for training generative models, particularly in visual domains. The core idea is elegant: learn a velocity field that transports samples from a simple noise distribution to the complex data distribution by solving an ordinary differential equation (ODE). Training minimizes a simple mean-squared error (MSE) between the model's predicted velocity and the ground-truth conditional velocity along this probability flow. This formulation—popularized by Lipman et al. [46] and related work—has achieved significant success in image and video synthesis precisely because it is simulation-free (unlike earlier continuous normalizing flows that required expensive ODE solves during training) and conceptually straightforward.
Yet a critical problem persists, one that the paper identifies in its opening paragraphs and that practitioners in the field have likely observed: flow-matching models trained with the standard MSE objective produce out-of-distribution samples unless classifier-free guidance (CFG) is applied at inference time. This is not a subtle issue. As Figure 1 illustrates and as the ImageNet results quantify (FID 8.26 for SiT-XL/2 without guidance), unguided samples from flow-matching models exhibit visible artifacts, incorrect textures, and structural inconsistencies that make them perceptibly worse than guided samples from the same model. The paper states this directly in Section 1:
"models often produce out-of-distribution samples unless guidance is applied"
Guidance, while effective at improving sample quality, comes with a fundamental drawback: it alters the sampling distribution. Classifier-free guidance [21] steers generation toward a sharpened, lower-entropy version of the learned distribution—essentially trading diversity for fidelity. The resulting samples, while perceptually appealing, no longer faithfully represent the original data distribution that the model was trained to approximate. For applications where distribution fidelity matters (e.g., scientific simulation, data augmentation, or any setting where capturing the full diversity of the data is important), this is a significant limitation. The paper frames this tension clearly:
"While guidance improves sample quality, it alters the sampling distribution. How to more faithfully model the underlying distribution of the original data remains an open problem."
Why This Problem Matters
The gap between guided and unguided generation quality is not merely an academic concern. It has concrete practical implications:
Distribution fidelity for downstream tasks. When generative models are used to produce training data for other systems (a increasingly common practice in the era of synthetic data), samples that deviate from the true data distribution can introduce systematic biases. If the model consistently sharpens textures or simplifies structures, downstream models trained on these synthetic samples will inherit these biases.
Guidance as a crutch, not a solution. The widespread reliance on CFG reveals that the underlying models are not learning the data distribution as faithfully as the MSE objective suggests. The fact that guidance—which can be understood as low-temperature sampling that discards probability mass in low-density regions—improves FID so dramatically (from 8.26 to 2.06 for SiT) indicates that the unguided model's probability mass is poorly allocated. It places too much mass in regions that don't correspond to realistic images.
Theoretical implications for generative modeling. Flow matching's theoretical guarantee is that, with infinite capacity and infinite data, the MSE objective converges to the ground-truth marginal velocity field. The model would then reproduce the training data exactly—overfitting in the classical sense. In practice, finite-capacity neural networks must generalize, producing new samples rather than memorized training examples. The MSE objective, however, provides no guidance about how to generalize. It measures isotropic Euclidean distance in the velocity space without any awareness of the data manifold's structure. The paper articulates this insight in Section 2.1:
"models with finite capacity and generalizes [sic], and therefore results in the generation of new data samples. However, the squared L2 objective measures Euclidean distance rather than the manifold-aware distance, inducing incorrect generalization relative to the underlying data distribution."
This is the crux of the problem: the loss function that works perfectly in theory (infinite capacity) becomes a liability in practice (finite capacity) because it imposes the wrong inductive biases for generalization.
Prior Approaches and Their Shortcomings
The paper situates itself against several lines of prior work that have attempted to address the generalization problem in flow matching and related generative models. Each approach has identifiable limitations that motivate the CAFM formulation.
Representational Autoencoders and Latent Space Methods
Recent work has explored training flow models in learned latent spaces rather than directly in pixel space [83]. The intuition is straightforward: if the data manifold is complex and nonlinear in pixel space, perhaps an autoencoder can learn a representation where the manifold is more Euclidean and thus better suited to the MSE objective. Representational autoencoders have empirically reported improvements in generation quality, and the paper acknowledges this:
"Representational autoencoders convert the data space on which flow matching operates and have empirically reported improvements in generation quality"
However, this approach has a fundamental limitation that the paper identifies: it requires operating in a latent space rather than the original data space. This means the model never directly learns the pixel-space distribution; instead, it learns a latent distribution and relies on the decoder to map back to pixels. Any errors or artifacts introduced by the decoder become part of the generation pipeline, and the latent representation may not capture all aspects of the data distribution that matter for downstream use. Moreover, the choice of autoencoder architecture and training objective itself imposes generalization biases—the problem is shifted rather than solved.
Riemannian Flow Matching
Another line of work extends flow matching to non-Euclidean geometries [7, 52]. The idea is to define the probability flow on a manifold that better matches the data's intrinsic geometry, so that the Euclidean metric in that manifold corresponds more closely to semantically meaningful distances. The paper acknowledges this direction:
"Riemannian flow matching extends flow matching to non-Euclidean geometries"
However, this approach faces a critical practical barrier that the paper highlights:
"this requires manual definition of the data manifold, which is often unknown for general datasets"
The fundamental challenge is that the data manifold is not given a priori—it must be discovered from the data itself. For natural images, the manifold structure is complex, high-dimensional, and not amenable to explicit mathematical specification. Riemannian methods work well when the geometry is known (e.g., data on spheres, tori, or other analytically defined manifolds), but general image datasets resist such characterization. This makes Riemannian flow matching more of a specialized tool than a general solution to the generalization problem.
Perceptual Losses with Fixed Feature Networks
A more directly relevant line of work replaces the Euclidean loss with perceptual distances derived from frozen feature networks [44]. The motivation is compelling and empirically grounded: deep networks trained on natural images have been shown to serve as better perceptual metrics than pixel-wise Euclidean distance [81]. Their internal representations capture texture, structure, and semantic content in ways that simple L2 distances do not. Using such a network as a fixed criterion for flow matching training could, in principle, guide the model toward generalizations that better respect the data manifold.
The paper acknowledges this approach and its intuition:
"Other work replaces Euclidean loss with perceptual distances derived from frozen feature networks, motivated by the empirical finding that deep networks can serve as better perceptual metrics"
However, fixed perceptual losses introduce a well-known vulnerability that the paper identifies:
"a fixed criterion network can be exploited by the generator, leading to artifacts in the generated samples"
This vulnerability is a manifestation of the broader phenomenon of generator hacking (related to adversarial examples [18]): if the generator knows exactly what features the fixed network responds to, it can learn to produce outputs that score well on those features without actually being realistic. The fixed network was never trained to be robust against such exploitation—it was trained for a different task (e.g., image classification) and its feature responses may not perfectly align with perceptual quality across all regions of the image space. The generator can find "blind spots" in the fixed network's assessment, producing samples that satisfy the perceptual loss while containing subtle artifacts invisible to that specific network.
Generative Adversarial Networks (GANs)
GANs [17] represent a fundamentally different approach to generative modeling. Rather than learning to predict velocities or denoise samples, GANs train a generator and discriminator in a minimax game: the discriminator learns to distinguish real from fake samples, and the generator learns to fool the discriminator. This adversarial dynamic has proven remarkably effective for image synthesis, with modern GANs achieving strong performance on ImageNet benchmarks [26, 28, 43, 64].
The paper recognizes GANs' empirical strengths, particularly in a related context where they are already being used:
"They achieve strong performance on ImageNet benchmarks and are widely used in flow-model distillation for sharp image synthesis"
The paper hypothesizes that this advantage stems from the discriminator's sensitivity to perceptual details:
"We hypothesize that this advantage arises because the discriminator networks are more sensitive to perceptual details, e.g. texture, sharpness, contour, etc., than pointwise Euclidean losses, because they may have learned to better capture the manifold structure."
This hypothesis is central to the paper's motivation: if adversarial discriminators are better at capturing manifold structure than Euclidean losses, then integrating adversarial training into flow matching could address the generalization problem directly. The discriminator, being learned jointly with the generator, would adapt to the generator's current weaknesses and provide more targeted feedback than a fixed perceptual loss—mitigating the exploitation problem.
Critically, GANs have a well-known limitation: training instability. The minimax game can oscillate, the discriminator can become too strong (providing vanishing gradients) or too weak (providing uninformative feedback), and careful regularization is required. These challenges have historically limited GANs' applicability to high-resolution, diverse datasets where diffusion and flow models have since dominated.
Adversarial Flow Models (AFMs)
The most directly relevant prior work is adversarial flow models (AFMs) [43], which represent the first attempt to combine adversarial training with flow modeling. AFMs train a generator that transports samples between two points on the probability flow and a discriminator that differentiates real and generated samples at intermediate timesteps. The adversarial objective is a contrastive function (the paper uses ), and an optimal transport regularization term encourages the generator to predict targets closest to the source points.
AFMs demonstrated that adversarial objectives could improve flow model training, but they have a critical limitation that the paper identifies:
"AFMs are formulated in discrete time, leaving open the question of how to incorporate adversarial training into continuous-time flow modeling."
The discrete-time formulation creates several problems:
- Training instability at fine discretizations. As the timestep interval shrinks, training becomes increasingly unstable. The paper notes:
"Although the timestep interval can be made arbitrarily small, the training becomes increasingly unstable, and the objective breaks down when ."
This means AFMs cannot be extended to the continuous limit where flow matching operates, making them fundamentally a discrete approximation rather than a continuous-time method.
- Persistent vanishing gradient problem. Despite operating on the probability flow, AFMs still suffer from gradient vanishing, a classic GAN training pathology where the discriminator becomes too effective and provides near-zero gradient signal to the generator. The paper notes:
"AFMs still have the gradient-vanishing problem. They rely on gradient penalties, discriminator augmentation, and discriminator reset to mitigate the issue."
These workarounds add complexity and hyperparameter sensitivity without fully resolving the underlying issue.
- Discrete simulation at inference. AFMs generate samples by iteratively applying the generator over discrete timesteps (Equation 16), solving a difference equation rather than an ODE. This introduces discretization error and requires choosing a step count, losing the continuous-time formulation's flexibility to trade off computation for accuracy via adaptive ODE solvers.
How This Paper Positions Itself
The paper positions CAFM as filling a specific, well-defined gap: extending adversarial training to continuous-time flow modeling. This is not presented as an alternative to flow matching or to adversarial training, but as a synthesis that captures the strengths of both paradigms while addressing their individual weaknesses.
The key conceptual move is the recognition that flow matching and adversarial training can target the same ground-truth flow. The paper explicitly states this in Section 1:
"Although FMs and CAFMs target the same ground-truth flow, they differ in finite-capacity generalization because CAFMs use a learned discriminator rather than a fixed Euclidean criterion."
This is a crucial insight: the MSE objective and the adversarial objective both converge to the same velocity field in the infinite-capacity limit (assuming the Nash equilibrium of the adversarial game corresponds to the true flow). Where they differ is in how they generalize under finite capacity. The MSE objective penalizes all velocity errors equally in Euclidean space, regardless of whether those errors matter perceptually or semantically. The adversarial objective, through its learned discriminator, penalizes errors that are perceptually salient and ignores those that are not—inducing a generalization that better respects the data manifold.
The paper positions CAFM primarily as a post-training method for existing flow-matching models, not as a replacement for flow matching pre-training. This is both a practical choice (Section 3.3 acknowledges that CAFM training from scratch is "inherently less efficient than FMs due to the involvement of an extra discriminator network") and a conceptual one: it leverages the fact that the two objectives target the same flow, so the FM-pre-trained model already provides a good initialization for CAFM fine-tuning. The post-training approach is efficient—only 10 epochs in the ImageNet experiments—and can be applied to any existing flow-matching model without architectural changes to the generator.
The paper's positioning relative to prior work can be summarized along four axes:
Versus standard flow matching: CAFM addresses the generalization problem that causes flow matching to produce out-of-distribution samples without guidance, while remaining a continuous-time method that integrates ODEs rather than applying discrete steps.
Versus guidance: CAFM improves the base model rather than altering the sampling distribution at inference time. The paper explicitly notes that guidance can be applied orthogonally, and that improving the base model via CAFM also improves guided results—suggesting that the benefits are complementary rather than competing.
Versus representational autoencoders and Riemannian methods: CAFM operates in the same Euclidean space as standard flow matching, avoiding the need for learned latent spaces or manually specified manifolds. The discriminator learns the manifold structure implicitly from data.
Versus AFMs: CAFM extends adversarial training to continuous time, resolving the discrete-time limitations of AFMs. The continuous-time formulation eliminates the vanishing gradient problem (Section 3.1 and Appendix E), removes the need for gradient penalties, and enables standard ODE solvers for inference.
The Key Design Challenge: How to Discriminate in Continuous Time
The paper's central technical challenge is not whether to combine adversarial and flow training, but how to do so in continuous time. In discrete-time AFMs, the discriminator compares real and generated samples at specific timesteps—a straightforward discrimination task. In continuous time, the generator no longer predicts positions but velocities—a fundamentally different quantity. The discriminator must therefore operate on velocities rather than positions, and it must do so in a way that is consistent with the continuous-time flow.
The paper's solution—discriminating in the Jacobian-vector product (JVP) space of a learned scalar potential —is motivated by the physical interpretation of velocity as a derivative of position. Rather than discriminating velocities directly (which the paper shows in Appendix D has fundamental problems with marginal versus conditional targets), the discriminator learns a scalar function whose directional derivative along the velocity direction serves as the discrimination logit. This formulation ensures that discrimination is consistent with the ODE structure of the flow and that the equilibrium condition corresponds exactly to the generator matching the ground-truth marginal velocity field.
This design choice is the conceptual heart of the method, and understanding why direct velocity discrimination fails is essential to appreciating the JVP formulation. The naive approach—defining and training adversarially—encounters two problems (detailed in Appendix D): first, the ground-truth marginal velocity is inaccessible during training (only conditional velocities are available), and nonlinear discriminators do not satisfy , so matching conditional targets doesn't enforce matching the marginal target; second, even with access to the marginal velocity, the discriminator can separate real and generated velocities too easily when their supports don't overlap, leading to vanishing gradients.
The JVP formulation elegantly sidesteps both issues by exploiting the linearity of the derivative operator: , so the conditional targets are sufficient for learning the marginal velocity. This makes CAFM the first method to successfully apply adversarial training to continuous-time flow modeling.
3. Technical Approach
3.1 Reader Orientation
We are building a training procedure—a specific way to update the weights of a neural network—that replaces the standard mean-squared-error (MSE) loss in continuous-time flow models with an adversarial objective involving a jointly trained discriminator. The system is not a new model architecture or a new sampling method; it is a loss function (and the associated training loop) that can be applied to any existing flow-matching generator after pre-training.
The core problem this solves is that flow-matching models trained with MSE produce out-of-distribution samples—images with artifacts, incorrect textures, or structural inconsistencies—despite the MSE objective being theoretically correct in the infinite-capacity limit. The "shape" of the solution is a minimax game: a discriminator learns to tell apart real velocity samples from generated ones by evaluating directional derivatives of a learned scalar potential, and the generator learns to produce velocities that the discriminator cannot distinguish from real ones. Because both objectives converge to the same ground-truth flow in theory, the adversarial objective only changes how the model generalizes under finite capacity, not what it converges to in the limit.
3.2 Big-Picture Architecture (Diagram in Words)
The system has four major components that interact in a training loop (no architectural changes are needed at inference time):
-
Generator
$G(x_t, t)$— a neural network identical in form to a standard flow-matching model. It takes a noisy sample$x_t$at time$t$(where$x_t$is a blend of data$x$and noise$z$) and outputs a predicted velocity$u_t = G(x_t, t)$—a vector in the same space as$x$that estimates how the sample should move to progress toward the data distribution. This is the same input/output signature as in flow matching. -
Discriminator
$D(x_t, t)$— a neural network that takes the same noisy sample$x_t$and time$t$as input but outputs a single scalar value (a "potential"). Unlike standard GAN discriminators that classify samples as real or fake, this discriminator learns a scalar field over the flow space whose directional derivative along a velocity direction provides the discrimination signal. The network architecture mirrors the generator's (same transformer backbone), with modifications: LayerNorm replaced by RMSNorm for stability, a learnable[CLS]token prepended to the input, and projection layers that map to the scalar output. -
Jacobian-Vector Product (JVP) Evaluator — not a separate network but a computational operation applied to the discriminator. For any primal point
$(x_t, t)$and any tangent direction$(\dot{x}_t, \dot{t})$, the JVP computes the directional derivative of$D$along that tangent. This is implemented via forward-mode automatic differentiation (torch.func.jvp) and produces a scalar that serves as the discrimination logit—the value the adversarial loss acts on. The JVP is linear in the tangent by construction, which is the key property that makes the adversarial objective compatible with conditional velocity targets. -
Training Loop Controller — the orchestration logic that alternates between updating the discriminator and the generator. The discriminator is updated
$N = 16$times per generator update to keep it near its local optimum. Both networks see the same batch of$(x, z, t)$triplets. The loss functions involve centering penalties (to prevent discriminator logit drift) and optional optimal transport regularization (for training from scratch, but set to zero during post-training).
Information flow during a single training step: A batch of clean data $x$ and noise $z$ is sampled → a random timestep $t \sim \text{Uniform}(0,1)$ is drawn → the interpolated sample $x_t = (1-t)x + tz$ is computed → the conditional velocity $\bar{v}_t = -x + z$ is computed (the "real" velocity that the ground-truth flow assigns to this $(x, z, t)$ triplet) → the generator predicts $u_t = G(x_t, t)$ (the "fake" velocity) → for discriminator updates, the JVP of $D$ is evaluated twice (once with tangent $(\bar{v}_t, 1)$ as the real logit and once with tangent $(u_t, 1)$ as the fake logit), the adversarial loss compares them, and gradients flow back through the JVP into $D$'s parameters → for generator updates, only the JVP with $(u_t, 1)$ is evaluated, and the generator is optimized to make this logit match the real target of $+1$.
At inference time, nothing changes: The same ODE solver used for standard flow matching integrates $G(x_t, t)$ from noise to data. The discriminator is discarded after training.
3.3 Roadmap for the Deep Dive
-
First, why the standard MSE objective fails to induce correct generalization under finite capacity, and what property a "better" objective should have—this establishes the formal motivation for adversarial training as a replacement criterion.
-
Second, the continuous adversarial objective itself: how the JVP of a learned scalar potential serves as the discrimination logit, why this specific form is necessary (and what naive alternatives fail), and the equilibrium condition that ensures the generator recovers the ground-truth flow.
-
Third, the practical implementation of the training loop: the alternation scheme, the centering penalty, the optimal transport regularizer (and why it's set to zero during post-training), the choice of contrastive function, and the computational techniques (forward-mode AD, vmap) that make the JVP-based discrimination efficient.
-
Fourth, the post-training strategy that distinguishes CAFM from training-from-scratch: why post-training is both more efficient and more effective, how the FM-pre-trained generator is initialized, and how the discriminator is warmed up before joint training begins.
-
Fifth, the architectural choices for the discriminator network: why RMSNorm replaces LayerNorm, how the
[CLS]token and projection layers are added, and the constraints (or lack thereof) on the generator architecture.
3.4 Detailed, Sentence-Based Technical Breakdown
This is primarily a methodological paper whose core idea is that adversarial training can be extended to continuous-time flow models by discriminating in the derivative space of a learned scalar potential, and that this produces better generalization than the standard MSE objective when applied as a post-training step to existing flow-matching models.
Why the MSE Objective Generalizes Poorly Under Finite Capacity
The paper's starting point is a formal analysis of why flow matching's standard MSE loss, despite being theoretically correct, fails in practice. Understanding this analysis is essential because it motivates why an adversarial objective is the right intervention and what property the adversarial formulation must preserve.
Flow matching (Section 2.1) trains a generator $G(x_t, t)$ to predict the conditional velocity $\bar{v}_t$ at each point $x_t$ along the probability flow. The standard objective is:
where $x \sim X$ is a clean data sample, $z \sim \mathcal{Z}$ is a noise sample (typically $\mathcal{N}(0, I)$), $t \in [0, 1]$ is a randomly sampled timestep, $x_t = (1-t)x + tz$ is the interpolated sample at time $t$ along the linear probability path, and $\bar{v}_t = -x + z$ is the conditional velocity (the time derivative of $x_t$ for this specific $x$ and $z$ pair).
What it computes: for each $(x, z, t)$ triplet in the batch, the squared Euclidean distance between the generator's predicted velocity and the known conditional velocity, averaged over the batch and over the $n$ dimensions of the velocity vector.
Why this form: Under the linear interpolation $x_t = (1-t)x + tz$, the conditional velocity is the simple difference $\bar{v}_t = -x + z$. The squared L2 loss is the unique criterion (up to positive-definite linear transformations) satisfying the property:
where $d(a, b) = \|a - b\|^2_2$. This property guarantees that minimizing the conditional loss in expectation over $(x, z)$ pairs yields the marginal velocity $v_t = \mathbb{E}[\bar{v}_t | x_t]$, which is the correct target for the probability flow ODE. In other words, even though the model sees only conditional targets $\bar{v}_t$ (which depend on the specific $x$ and $z$), the minimizer of the expected loss is the marginal expectation over all possible $(x, z)$ that could have produced $x_t$.
The generalization problem. The paper's key insight, developed in Section 2.1 and Appendix C, is that any loss of the form:
where $M$ is a strictly positive-definite matrix (satisfying $x^\top M x > 0$ for all $x \neq 0$), also satisfies $\arg\min_a \mathbb{E}[d(a, b)] = \mathbb{E}[b]$. The standard MSE corresponds to the special case $M = I$ (the identity matrix), which applies equal weight to every dimension of the velocity vector and every spatial location. But in theory, any $M$ works—they all converge to the same optimal generator in the infinite-capacity, infinite-data limit.
In practice, with finite-capacity neural networks that must generalize, the choice of $M$ matters enormously. The squared L2 criterion $M = I$ measures isotropic Euclidean distance: an error of magnitude $\epsilon$ in one dimension of the velocity vector is penalized identically to an error of magnitude $\epsilon$ in any other dimension. But not all velocity errors are perceptually equal. An error that shifts a pixel by a tiny amount in a smooth region may be invisible; an error that creates a sharp edge where none should exist may be glaring. The Euclidean metric, having no awareness of the data manifold, treats these identically. The result is "incorrect generalization relative to the underlying data distribution"—the model allocates its finite capacity to minimize velocity errors uniformly rather than prioritizing errors that matter for perceptual quality.
The ideal $M$ would be a manifold-aware metric—one that penalizes velocity errors in directions orthogonal to the data manifold more heavily than errors tangent to it, and one that is spatially varying (errors matter more in textured regions than in smooth regions). But the data manifold is precisely what the model is trying to learn; it cannot be specified in advance. The paper therefore frames the problem as: can we learn the appropriate metric jointly with the generator, rather than fixing it?
Connection to adversarial training. In a GAN, the discriminator network $D$ is essentially a learned metric. It assigns high values to samples that are realistic and low values to samples that are not. The generator is optimized to maximize $D$'s assessment of its outputs—which is equivalent to minimizing a learned distance between the generated and real distributions. The paper's hypothesis, stated in Section 1, is that "discriminator networks are more sensitive to perceptual details, e.g. texture, sharpness, contour, etc., than pointwise Euclidean losses, because they may have learned to better capture the manifold structure." The goal of CAFM is to bring this learned-manifold sensitivity into continuous-time flow matching without sacrificing the theoretical guarantees of the flow formulation.
Why not simply replace the MSE with a fixed perceptual loss? The paper explicitly addresses this alternative in Section 1. Fixed perceptual losses from pre-trained networks (e.g., VGG features) have been explored in prior work and provide some improvement, but they suffer from generator hacking: the generator can learn to produce outputs that score well on the fixed network's features without actually being realistic, exploiting blind spots that the fixed network was never trained to defend against. A jointly trained discriminator adapts to the generator's current weaknesses, making exploitation harder—the discriminator is constantly being updated to detect the generator's latest failure modes.
The Continuous Adversarial Objective: Discriminating in JVP Space
The central technical challenge is: how do you perform adversarial training when the generator's output is a velocity field rather than a sample? In a standard GAN, the discriminator receives a complete sample (real or generated) and classifies it. In a continuous-time flow model, the generator at training time never produces a complete sample—it only predicts the velocity $G(x_t, t)$ at a single point $x_t$ along the flow. The discriminator must therefore evaluate whether a given velocity at a given position and time is consistent with the true data-generating flow.
The paper considers and rejects a naive formulation (detailed in Appendix D) where the discriminator would directly take $(x_t, t, v_t)$ as a three-part input:
This naive formulation has two fundamental problems:
-
Marginal vs. conditional target mismatch. During training, we only have access to conditional velocities
$\bar{v}_t = -x + z$—the velocity for a specific$(x, z)$pair. The true quantity we want the generator to learn is the marginal velocity$v_t = \mathbb{E}[\bar{v}_t | x_t]$—the expected velocity over all$(x, z)$pairs that could have produced this$x_t$. If the discriminator is a nonlinear function, then in general$\mathbb{E}[D(\bar{v}_t)] \neq D(\mathbb{E}[\bar{v}_t])$. This means that matching discriminator responses to conditional targets does not enforce matching to the marginal target. At equilibrium, the generator would need to represent the full distribution of conditional velocities at each$x_t$, but the generator outputs only a single deterministic velocity, so it gets pulled toward incompatible targets and training oscillates. -
Vanishing gradients from support mismatch. Even if we had access to the marginal velocity
$v_t$, the real velocity at any given$x_t$is a point mass (Dirac delta). If the generator's predicted velocity$u_t$is different from$v_t$, the supports of the real and fake velocity distributions don't overlap, and an optimal discriminator can separate them perfectly—saturating and providing zero gradient to the generator.
The JVP formulation solves both problems simultaneously by exploiting the linearity of the derivative operator.
The JVP discriminator. Define the discriminator as a scalar function $D(x_t, t) : \mathbb{R}^n \times [0, 1] \rightarrow \mathbb{R}$—a learned "potential field" over the flow space. The discrimination logit is not $D(x_t, t)$ itself but rather the Jacobian-Vector Product of $D$ along a velocity direction:
where $\frac{\partial D(x_t, t)}{\partial x_t} \in \mathbb{R}^{1 \times n}$ is the Jacobian (row vector) of $D$ with respect to the spatial input $x_t$, and $\frac{\partial D(x_t, t)}{\partial t} \in \mathbb{R}^{1 \times 1}$ is the partial derivative with respect to time. The tangent vector $(\dot{x}_t, \dot{t})$ specifies a direction in the $(x_t, t)$ space—in our case, this will be either the real velocity $(\bar{v}_t, 1)$ or the generated velocity $(G(x_t, t), 1)$, where the $\dot{t} = 1$ component ensures the time derivative is included.
What it computes: the instantaneous rate of change of the discriminator's scalar potential when moving from $(x_t, t)$ in the direction $(\dot{x}_t, \dot{t})$. Geometrically, if $D(x_t, t)$ is a surface over the flow space, $D_{\text{jvp}}$ is the slope of that surface in the specified direction—how fast the potential is increasing or decreasing along that velocity vector.
Why this form—linearity property: The critical property is that $D_{\text{jvp}}$ is linear in the tangent $(\dot{x}_t, \dot{t})$ while being nonlinear in the primal $(x_t, t)$ (through the Jacobians $\partial D / \partial x_t$ and $\partial D / \partial t$, which depend nonlinearly on $D$'s parameters). This means:
The expectation over conditional velocities commutes with the JVP because the JVP is linear in the tangent. This is the property that the naive nonlinear discriminator lacks. It means that training with conditional targets $\bar{v}_t$ is equivalent to training with the marginal target $v_t$ in expectation—exactly the property that the MSE loss has, but now with a learned metric (through the Jacobian $\partial D / \partial x_t$). The paper proves this formally in Appendix D (Equation 51).
Why this form—vanishing gradient mitigation: When the generator's predicted velocity $u_t$ differs from $v_t$, the JVP difference is:
This is a linear function of the velocity error. An optimal discriminator can still separate the conditional velocities, but the gradient with respect to $G(x_t, t)$ never vanishes as long as $\partial D / \partial x_t \neq 0$. Specifically (Appendix E, Equation 64):
where $g = \partial \mathcal{L} / \partial D_{\text{jvp}}$ is the gradient propagated from the loss. This gradient is nonzero whenever $g \neq 0$ and $\partial D / \partial x_t \neq 0$. The paper shows that whenever $G(x_t, t) \neq v_t$, an optimal discriminator will learn a nonzero Jacobian to enable discrimination, so gradients flow freely. Empirically, this means CAFM can train without gradient penalties (the $R_1$ and $R_2$ regularizations that AFMs require)—a significant simplification.
Why this form—global consistency: An additional benefit of the JVP formulation (Appendix D, Equation 58) is that it enforces consistency along trajectories:
The discriminator's potential difference between any two points equals the integrated directional derivative along any path connecting them—a consequence of the fundamental theorem of calculus for line integrals, enforced because $\partial D / \partial x_t$ and $\partial D / \partial t$ are the actual partial derivatives of a single scalar function. An alternative parameterization using separate networks $A(x_t, t)$ and $B(x_t, t)$ (Equation 53 in Appendix D) could achieve the same linearity but would not enforce this consistency constraint. The paper reports that "this separate formulation does not yield good results on high-dimensional data."
The adversarial objective. The adversarial losses for discriminator and generator are:
where $T = 1$ (a tensor of ones matching the batch dimension, representing $\dot{t} = 1$ for all samples), and $f(a, b)$ is a contrastive function.
What it computes: For the discriminator, the loss penalizes the difference between the JVP evaluated on the real velocity $\bar{v}_t$ (which should produce a high score) and the JVP evaluated on the generated velocity $G(x_t, t)$ (which should produce a low score). For the generator, the loss penalizes the opposite—encouraging the generated velocity's JVP to look like the real velocity's JVP. Both expectations are over the same distribution of $(x, z, t)$ triplets drawn from the training set, the noise distribution, and the uniform time distribution.
The contrastive function. The paper adopts a bounded least-squares contrastive function:
What it computes: the sum of two squared errors. The first term $(a - 1)^2$ penalizes deviations of the first argument from the target value $+1$; the second term $(b + 1)^2$ penalizes deviations of the second argument from the target value $-1$. For the discriminator loss, $a = D_{\text{jvp}}(x_t, t, \bar{v}_t, T)$ (real, should be $+1$) and $b = D_{\text{jvp}}(x_t, t, G(x_t, t), T)$ (fake, should be $-1$). For the generator loss, the arguments are swapped: $a = D_{\text{jvp}}(x_t, t, G(x_t, t), T)$ (fake, the generator wants this to be $+1$) and $b = D_{\text{jvp}}(x_t, t, \bar{v}_t, T)$ (real, the generator wants this to be $-1$).
Why this form: The paper compares three contrastive functions in an ablation study (Table 12): non-saturating (the standard GAN loss from Goodfellow et al., $f(a,b) = -\log(\sigma(a)) - \log(1-\sigma(b))$), hinge (from Geometric GAN, $f^D(a,b) = \max(0, 1-a) + \max(0, 1+b)$, $f^G(a,b) = -a + b$), and least squares (from LSGAN, shown above). The least squares loss produces the strongest results: FID 3.63 (guidance-free) and 1.53 (guided, CFG 1.3) for SiT-XL/2, compared to 3.54/1.58 for non-saturating and 4.00/1.57 for hinge. The authors adopt least squares as the primary loss. The bounded nature of least squares (unlike non-saturating, which can diverge) likely contributes to training stability, and the quadratic penalty provides stronger gradients when the discriminator is far from its target compared to hinge loss's linear penalties.
Equilibrium condition. At the Nash equilibrium of this minimax game, the discriminator cannot distinguish real from generated velocities, meaning:
for all $(x, z, t)$. Taking expectations over $x, z | x_t$ and using the linearity property:
Since $D$ is a learned function, the discriminator can in principle satisfy this equality even when $v_t \neq G(x_t, t)$ by setting the Jacobian $\partial D / \partial x_t$ to zero—producing flat potentials that are insensitive to the velocity direction. However, the discriminator is incentivized not to do this during training because making the Jacobian nonzero is what enables it to distinguish real from fake. The generator, in turn, is optimized to minimize the JVP difference, which drives $G(x_t, t)$ toward $v_t$. At the unique equilibrium where neither player can improve, $G(x_t, t) = v_t$ for all $x_t$ and $D$ becomes a flat potential (zero gradient everywhere). The paper visualizes this dynamic in Figure 2: panel (d) shows $-D(x_t, t)$ converging to zero everywhere as the generator converges to the ground-truth flow.
Why the JVP formulation doesn't collapse for high-dimensional data. A subtle concern: when the spatial dimension $n$ is large (e.g., $256 \times 256 \times 3 = 196,608$ for ImageNet images, or even larger in latent space), the discriminator projects an $n$-dimensional velocity onto a single scalar (the JVP). There exists an $(n-1)$-dimensional null space of velocities that produce the same JVP value—the generator could learn to output velocities in this null space that satisfy the discriminator but don't match the true velocity. The paper addresses this in Section 3.1:
"the discriminator, which projects an n-dimensional input to a scalar value, creates ambiguity because multiple
$v_t \in \mathbb{R}^n$can yield the same value. G may learn to exploit the null space, and D is then updated to counter this behavior. However, this causes slow convergence."
The solution is an optimal transport regularizer that encourages the generator to pick the minimum-norm solution among all velocities that satisfy the discriminator:
What it computes: the mean squared norm of the generator's predicted velocity, averaged over the batch and normalized by the dimension $n$. This penalizes large velocity magnitudes regardless of direction.
Why this form: this is the continuous-time limit of the discrete optimal transport regularization used in AFMs (Equation 10 in the paper, where the discrete version penalizes $\|G(x_s, s, t) - x_s\|^2_2 / (|t-s|)$). As $|t-s| \to 0$, the discrete regularizer converges to the continuous one (Equation 24). Among all velocity fields that satisfy the adversarial equilibrium, the minimum-norm solution is unique and corresponds to the optimal transport flow—the one that moves probability mass along the shortest paths. By encouraging the generator toward this solution, the regularizer disambiguates the null space and accelerates convergence.
Critical design choice: setting $\lambda_{\text{ot}} = 0$ during post-training. The paper sets $\lambda_{\text{ot}} = 0$ for all post-training experiments (Section 3.1):
"For post-training existing flow-matching models, we set
$\lambda_{\text{ot}} = 0$to completely eliminate the bias of the Euclidean norm."
The rationale is that the FM-pre-trained generator already produces velocities close to the ground-truth flow, so the null-space ambiguity is less severe—the generator starts near the optimal transport solution and doesn't need the regularizer to find it. Moreover, keeping $\lambda_{\text{ot}} > 0$ would bias the generator toward the Euclidean norm, which is precisely the inductive bias that CAFM is trying to escape. An ablation in Table 9 confirms this: $\lambda_{\text{ot}} = 0$ achieves FID 3.63 (SiT-XL/2 guidance-free) while $\lambda_{\text{ot}} = 0.01$ achieves FID 4.50—a significant degradation. For training from scratch, $\lambda_{\text{ot}}$ is needed and is gradually decreased over training (Section 4.3, Figure 5).
The Centering Penalty: Preventing Discriminator Logit Drift
The adversarial losses in Equations 20 and 21 only constrain the JVP—the directional derivative of $D$—but place no constraint on the absolute value of $D(x_t, t)$ itself. There is nothing preventing the discriminator's scalar potential from drifting to arbitrarily large positive or negative values while maintaining the correct derivatives. This drift can cause numerical instability.
The paper adds a centering penalty (Section 3.1, Equation 23):
What it computes: the mean squared value of the discriminator's scalar output, encouraging it to stay close to zero.
Why this form: this is the continuous-time analog of the centering penalty used in AFMs (Equation 13), which penalized $(D(x_t, t) + D(G(x_s, s, t), t))^2$. The quadratic penalty is simple and differentiable. The weight is set to $\lambda_{\text{cp}} = 0.001$ (Section 3.1), which is small enough not to interfere with the adversarial dynamics but large enough to prevent unbounded drift. At the equilibrium where the generator matches the true flow, the discriminator ideally becomes a flat zero potential everywhere—the centering penalty helps it converge to this flat-zero solution rather than a flat-nonzero one.
Complete Training Objectives
Putting the pieces together, the final objectives for CAFM (Equations 25-26) are:
What this means operationally: For each discriminator update, compute the JVP twice (once with real velocity, once with generated velocity), evaluate the least-squares contrastive loss, add the centering penalty on $D(x_t, t)$ itself, and backpropagate through the JVP into $D$'s parameters. For each generator update, compute the JVP once (with generated velocity only), evaluate the least-squares loss with the real target, add the optimal transport penalty if $\lambda_{\text{ot}} > 0$, and backpropagate through the JVP into $G$'s parameters.
The key differences from the discrete-time AFM objectives (Equations 14-15) are: (1) no $R_1$ or $R_2$ gradient penalties—the continuous-time JVP formulation mitigates vanishing gradients without them, (2) the centering penalty operates on $D(x_t, t)$ directly rather than on the sum of real and fake discriminator outputs, and (3) the generator's optimal transport regularizer is the continuous limit of the discrete version.
Practical Implementation: Efficient JVP Computation
The JVP-based discrimination requires computing directional derivatives of the discriminator network with respect to its input. The paper implements this efficiently using PyTorch's forward-mode automatic differentiation (Section 3.2, Algorithm 1).
Forward-mode AD (torch.func.jvp). Unlike standard backpropagation (reverse-mode AD), which computes gradients of a scalar output with respect to all inputs, forward-mode AD computes the directional derivative of all outputs with respect to a single input direction. This is exactly what we need: given the primal $(x_t, t)$ and the tangent $(\dot{x}_t, \dot{t})$, compute the scalar $D_{\text{jvp}}$ which is the directional derivative of $D$ along the tangent. Forward-mode AD computes both $D(x_t, t)$ and $D_{\text{jvp}}(x_t, t, \dot{x}_t, \dot{t})$ in a single forward pass through the network—the primal values flow through the network normally, and the tangent values are propagated alongside them using the chain rule.
Why forward-mode is more efficient here than reverse-mode: The discriminator has a scalar output but $n$-dimensional input (where $n$ can be hundreds of thousands). Reverse-mode AD would require one backward pass per output dimension to compute the full Jacobian, which would be prohibitive. Forward-mode AD computes the JVP—which is exactly the Jacobian-vector product we need—in a single forward pass, with cost comparable to a standard forward pass plus some constant-factor overhead for the tangent propagation.
Vectorizing map (torch.func.vmap). During discriminator updates, we need to compute the JVP twice at the same primal $(x_t, t)$: once with tangent $(\bar{v}_t, T)$ (real) and once with tangent $(G(x_t, t), T)$ (fake). Rather than running the discriminator forward twice, vmap vectorizes the JVP computation over the tangent dimension, computing both JVPs simultaneously (Algorithm 1, lines 18-21). The paper states that this works by "stacking" the two tangent vectors and applying jvp in a vectorized manner, which is more efficient on GPU hardware than sequential calls.
Compatibility with distributed training. The paper addresses a subtle implementation detail in Appendix F: the order of operations matters when combining JVP with distributed data parallel (DDP), fully sharded data parallel (FSDP), or gradient checkpointing. The correct pattern is ddp(jvp(D))—wrapping the JVP-computing module inside DDP—rather than jvp(ddp(D)), because applying JVP to a DDP wrapper would include incompatible gradient synchronization logic. For FSDP and gradient checkpointing, the paper finds it sufficient to wrap JVP and vmap only at the top-level submodules (rather than at every nn.Module as in prior work like rCM), reducing overhead.
Attention compatibility. The paper uses PyTorch's math fused kernel for scaled dot-product attention, which "supports both JVP and Vmap natively and is sufficient for image generation training" (Appendix F). This means no custom attention implementation is required.
Computational cost. For ImageNet SiT-XL/2 post-training, CAFM requires approximately $4.8 \times$ more wall-clock time per epoch than FM (Appendix H). This overhead comes from three sources: (1) the additional discriminator network, (2) the forward-mode AD through that network for JVP computation, and (3) the $N = 16$ discriminator updates per generator update. The paper considers this acceptable for post-training, where the total number of epochs is small (10 epochs).
The Post-Training Strategy
CAFM is primarily designed as a post-training method applied to already-trained flow-matching models (Section 3.3). This design choice is motivated by both efficiency and effectiveness.
Why post-training rather than training from scratch. The paper explicitly states the efficiency argument (Section 3.3):
"it is inherently less efficient than FMs due to the involvement of an extra discriminator network, the forward and backward computation of JVP, and the multiple steps of discriminator learning per generator update"
At approximately $4.8 \times$ the per-epoch cost of FM, training from scratch for the full 1400 epochs used for SiT-XL/2 would be prohibitively expensive. The post-training approach amortizes the cost: the expensive FM pre-training (1400 epochs for SiT, 600 for JiT) produces a good initialization, and only 10 epochs of CAFM fine-tuning are needed to achieve the reported gains.
The conceptual argument is deeper: both FM and CAFM target the same ground-truth flow in the infinite-capacity limit (as shown in Section 2.1 and Appendix C). The difference is purely in the generalization behavior under finite capacity. Starting from an FM-pre-trained model means the generator already approximately captures the correct velocity field; CAFM then "refines" the generalization—redirecting the model's finite capacity away from velocity errors that are perceptually irrelevant and toward errors that the discriminator identifies as important. The paper verifies this interpretation with a controlled experiment: applying additional FM training to the already-trained SiT model for 10 epochs (SiT+FM in Table 1) yields FID 8.64—essentially unchanged from the original 8.26 and within the error margin of random evaluation sampling. Only CAFM fine-tuning produces the dramatic improvement to 3.63. This proves that the gains are not from additional training iterations alone but specifically from the change of objective.
The discriminator warm-up phase. When post-training, the discriminator starts from scratch (random initialization) while the generator is already well-trained. This creates an initial imbalance: the discriminator is weak and provides uninformative gradients. The paper addresses this with a discriminator warm-up period (Sections 4.1): for the first 2 epochs on SiT and first 4 epochs on JiT, the generator is frozen and only the discriminator is updated. This allows the discriminator to adapt to the generator's current output distribution before the generator begins to change in response. Without this warm-up, the generator would receive noisy gradients from an undertrained discriminator and could drift from its good initialization.
Discriminator update frequency $N$. After warm-up, the discriminator is updated $N = 16$ times per generator update (Sections 4.1). This asymmetric schedule is common in GAN training: the discriminator needs more updates to stay near its local optimum because it has a simpler task (binary classification in JVP space) and converges faster, while the generator's task (predicting velocities) is more complex and benefits from more stable discriminator feedback. The paper sweeps $N \in \{8, 16, 32\}$ in an ablation (Table 8): $N = 8$ leads to divergence (FID 294.91), $N = 16$ achieves FID 3.63, and $N = 32$ achieves FID 3.68 (slightly worse, likely because the discriminator overpowers the generator). The value $N = 16$ is selected as the sweet spot.
Epoch counting. The paper notes (Section 4.1) that "epochs are measured as the combined number of images seen by both G and D throughout our experiments." This means one "epoch" of CAFM training involves $N = 16$ discriminator updates plus one generator update per batch, and the total number of images processed is $N+1$ times the batch size per iteration. The reported 10 epochs of CAFM post-training therefore correspond to substantially more discriminator updates than generator updates, but the generator sees only 10 passes through the dataset.
EMA decay. The paper uses an exponential moving average (EMA) of the generator weights with a short decay of 0.99 (Sections 4.1). This is standard practice in GAN training to stabilize the generator's output distribution against oscillations from the adversarial dynamics. The short decay (compared to typical EMA decay of 0.999 or 0.9999) reflects the small number of total post-training epochs: with only 10 generator passes through the data, a long decay would barely move the EMA weights from the initial FM-pre-trained values.
Post-training hyperparameters for SiT-XL/2: The generator architecture is completely unchanged from the pre-trained SiT-XL/2 checkpoint. The discriminator uses the same architecture and weight initialization as the generator, with modifications (LayerNorm → RMSNorm, prepended [CLS] token, projection layers for scalar output). Learning rate $1 \times 10^{-5}$ for both $G$ and $D$. Adam optimizer with $\beta = (0, 0.95)$ (note: $\beta_1 = 0$, not the typical $\beta_1 = 0.9$—this removes momentum along the gradient direction, which is common in GAN training to reduce oscillation). Batch size 256 (matching the original SiT training). No weight decay. $\lambda_{\text{ot}} = 0$. $\lambda_{\text{cp}} = 0.001$. EMA decay 0.99. Total 10 post-training epochs. TF32 precision. The discriminator is warmed up for 2 epochs (frozen generator, discriminator-only updates).
Post-training hyperparameters for JiT-H/16: The generator architecture is completely unchanged from the pre-trained JiT-H/16 checkpoint. The discriminator uses the same architecture and weight initialization as the generator. Because JiT already uses RMSNorm and has in-context class tokens, no normalization changes are needed; the discriminator simply takes the first class token and adds projection layers for scalar output. The generator's output is converted from x-prediction to v-prediction before being passed to the discriminator (JiT parameterizes the denoising prediction as the clean image $x$, not the velocity $v$). Learning rate $1 \times 10^{-5}$ for both $G$ and $D$. Adam with $\beta = (0, 0.95)$. Batch size 1024 (matching the original JiT training). No weight decay. $\lambda_{\text{ot}} = 0$. $\lambda_{\text{cp}} = 0.001$. EMA decay 0.99. Total 10 post-training epochs. Discriminator warm-up for 4 epochs. Timestep distribution $\log \text{Normal}(0.8, 0.8)$ reversed (JiT defines $x_0$ as noise rather than data, so the time convention is flipped).
Inference after post-training. No changes to the inference procedure. The paper uses the exact same evaluation code as the original models: Euler-Maruyama SDE with 250 integration steps for SiT (matching SiT's best configuration), Heun ODE with 50 steps for JiT. The discriminator is discarded after training.
Discriminator Architecture Design Choices
The discriminator shares the same backbone architecture as the generator (DiT transformer), but with three specific modifications:
RMSNorm instead of LayerNorm. The paper reports (Section 3.2 and Appendix G) that "switching LayerNorm to RMSNorm significantly improves training stability." Figure 12 shows the discriminator gradient norm during training: LayerNorm causes large spikes in the gradient norm, while RMSNorm produces smooth gradients. The authors hypothesize that this is related to the JVP computation, noting that "prior work involving JVP has also found that RMSNorm provides better stability." LayerNorm's mean-subtraction step may introduce discontinuities or near-discontinuities in the Jacobian that amplify during forward-mode differentiation. RMSNorm, which only divides by the root-mean-square without mean-centering, appears to produce smoother derivatives.
[CLS] token and projection layers for the discriminator logit. Following the practice from AFMs (Section 4.1), a learnable classification token is prepended to the input sequence of the discriminator. After the transformer processes the sequence, the output corresponding to this [CLS] token is extracted and passed through projection layers to produce the scalar output $D(x_t, t)$. This is analogous to how ViT models use [CLS] tokens for classification, but here the scalar is not a class probability—it's the discriminator's potential value.
Why not a separate discriminator architecture? The paper uses the same architecture as the generator (with these modifications) rather than a custom discriminator design. This is a practical choice: using the same architecture simplifies implementation (same initialization, same hyperparameters) and ensures the discriminator has comparable capacity to the generator. The paper reports (Section 3.2) that "CAFMs work well with standard transformers as both G and D," indicating that the architecture choice is not especially fragile.
No additional normalization on modulation. Some prior work involving JVP computation (e.g., terminal velocity matching) found that additional normalization on the modulation parameters (scale and shift in adaptive layer norm) was necessary for stability. The paper explicitly states (Section 3.2) that "unlike prior work, we do not find additional normalization on modulation necessary." This may be because the RMSNorm switch alone is sufficient to stabilize the gradients through the JVP.
For pixel-space models (JiT). The discriminator uses the same architecture but takes advantage of JiT's existing in-context class tokens: "we simply take the first class token and add projection layers for the discriminator output" (Section 4.1). This avoids the need to prepend an additional [CLS] token. The generator's output in JiT parameterizes the x-prediction (the clean image) rather than the velocity; the paper converts this to v-prediction before passing it to the discriminator by computing $v = G(x_t, t)$ using the relationship between x-prediction and v-prediction under the linear interpolation scheme.
For text-to-image models. In the text-to-image experiments (Appendix B), the discriminator follows the APT design: instead of prepending a [CLS] token at the input, a cross-attention layer is added at the last transformer layer that projects the visual features to a scalar logit. The paper notes that "this design allows most parts of the transformer to stay intact" since only the final layer is modified for the discriminator head.
Training from Scratch: Ablation Studies
Although post-training is the primary use case, the paper includes training-from-scratch experiments for completeness (Section 4.3). These experiments use the SiT-B/2 architecture (a smaller model than the XL variant) on ImageNet 256×256.
Hyperparameter schedule. The key finding from the training-from-scratch ablations (Figure 5) is that $\lambda_{\text{ot}}$ should decrease over training and $N$ should increase over training. Specifically:
-
Initial phase (epochs 0–50): The optimal transport weight is set to
$\lambda_{\text{ot}} = 4$(chosen by sweeping$\{0, 1, 2, 4, 8\}$; Figure 5b shows$\lambda_{\text{ot}} = 4$achieves FID 64.9 at epoch 50 vs. 74.3 for$\lambda_{\text{ot}} = 1$and 104.9 for$\lambda_{\text{ot}} = 0$). The discriminator update frequency is$N = 4$(chosen over$N = 1$; Figure 5a shows$N = 4$achieves FID 74.3 at epoch 50 vs. 82.1 for$N = 1$). -
Mid-training (epochs 160–400):
$\lambda_{\text{ot}}$is reduced to 1 at epoch 160. Figure 5c shows that$\lambda_{\text{ot}} = 1$continues improving to FID 25.8 at epoch 400 while$\lambda_{\text{ot}} = 4$plateaus at 29.6—reducing the optimal transport weight allows the adversarial objective to take over as the primary training signal. -
Late training (epochs 700–1000):
$N$is increased to 8 at epoch 700. Figure 5d shows that$N = 8$converges faster at this late stage (FID 23.2 vs. 24.0 for$N = 4$at epoch 950). The authors note that they also tried decreasing the learning rate, further decreasing$\lambda_{\text{ot}}$, and further increasing$N$to match post-training settings, but these "yield worse performance" within the 1000-epoch budget—suggesting that the post-training hyperparameters are only appropriate when the generator is already well-trained.
Convergence comparison. Figure 4 shows that CAFM training from scratch converges more slowly than FM under the same epoch count: after 400 epochs, FM achieves approximately FID 22.85 while CAFM is at approximately 25. After 1000 epochs, CAFM reaches 19.05—better than FM at 400 epochs but not obviously better than FM trained for the full 1000 epochs (the FM curve is only plotted to 400 epochs). The paper does not claim that training from scratch with CAFM outperforms FM when matched for total compute; the claim is specifically that post-training an FM-pre-trained model with CAFM produces large improvements over the FM baseline.
4. Key Insights and Innovations
Innovation 1: Recasting Test-Time Generalization as an Objective-Choice Problem Rather Than an Architecture or Manifold Problem
The most fundamental conceptual move in this paper is the recognition that the generalization failures of flow-matching models—the out-of-distribution samples, the reliance on guidance to produce acceptable images—are not primarily a problem of where the model learns (pixel space vs. latent space) or what manifold it operates on (Euclidean vs. Riemannian), but rather a problem of which loss function shapes the model's inductive biases under finite capacity. This reframing changes how the field should think about improving generative models.
Prior work had attacked the generalization problem from several angles. Representational autoencoders [83] moved the problem to a learned latent space, implicitly betting that the data manifold is more Euclidean there. Riemannian flow matching [7] required manually specifying the manifold geometry—an approach that works for known geometries but fails for general image datasets where the manifold structure is unknown a priori. Perceptual losses with fixed feature networks [44] replaced the Euclidean metric with a fixed learned metric, but introduced the generator-hacking vulnerability. Each of these approaches changed where or against what the model learns, but not the fundamental structure of the optimization.
The paper's key diagnostic insight, developed in Section 2.1 and formally justified in Appendix C, is that any loss of the form $d(a,b) = (a-b)^\top M (a-b)$ with strictly positive-definite $M$ converges to the same ground-truth flow in the infinite-capacity limit. The standard MSE ($M = I$) is just one admissible choice among infinitely many. The generalization problem under finite capacity is therefore not about finding a better target distribution—the target is fixed and correct—but about choosing the right metric $M$ that prioritizes perceptually meaningful velocity errors over irrelevant ones. This is a subtle but profound shift: it says the problem is not architectural or geometric but informational—the model needs to know which errors matter.
What makes this reframing genuinely novel is that it converts an ill-posed problem ("find the data manifold") into a well-posed one ("learn the metric jointly with the generator"). The data manifold is not specified in advance; it is discovered through adversarial competition. The discriminator learns $M(x_t, t)$ implicitly through its Jacobian $\partial D / \partial x_t$, which weights velocity errors differently at different points in space and time. This learned metric is adaptive—it responds to the generator's current weaknesses—and is grounded in the discriminator's ability to distinguish real from generated flows. This is a fundamental advance over fixed perceptual losses because it closes the exploitation loop: the discriminator is constantly retrained to detect the generator's latest failure modes, making hacking harder.
The significance of this reframing extends beyond the specific method. It suggests that the choice of loss function is the primary architectural decision in flow-based generative modeling under finite capacity, not the choice of network architecture, latent space, or ODE solver. This has implications for how the field allocates research effort: improving the loss function (making the metric better aligned with perceptual quality) may yield larger gains than improving the model architecture. The paper's ablation (Table 12) showing that the least-squares contrastive function outperforms non-saturating and hinge losses provides initial evidence that the specific form of the adversarial metric matters, opening a design space that prior work had not systematically explored for continuous-time flows.
Innovation 2: Resolving the Discrete-to-Continuous Gap in Adversarial Flow Training Through JVP-Space Discrimination
The paper's primary technical contribution is the specific mechanism that enables adversarial training in continuous time—discriminating in the Jacobian-vector product space of a learned scalar potential. However, the deeper innovation is not the mechanism itself but the diagnosis of why discrete-time adversarial flow training cannot be naively extended to continuous time, and the recognition that the solution requires exploiting a property (linearity of derivatives) that is simultaneously the reason naive approaches fail and the key to making continuous adversarial training work.
Prior work on adversarial flow models [43] operated in discrete time, training a generator $G(x_s, s, t)$ that transported samples between two timesteps and a discriminator that classified intermediate samples. The limitation was practical but fundamental: as the timestep interval $|t-s| \to 0$, training became increasingly unstable and the objective broke down. This meant AFMs were inherently discrete approximations—they could not access the continuous-time formulation that gives flow matching its flexibility (adaptive ODE solvers, consistent probability paths, theoretical connections to optimal transport).
The naive extension to continuous time would be to define $D(x_t, t, v_t)$—a discriminator that takes velocity as input—and train adversarially. The paper's Appendix D provides a rigorous diagnosis of why this fails: the fundamental problem is the mismatch between conditional and marginal targets when the discriminator is nonlinear. During training, only conditional velocities $\bar{v}_t$ (for specific $(x, z)$ pairs) are available, not the marginal velocity $v_t = \mathbb{E}[\bar{v}_t | x_t]$. For a nonlinear $D$, $\mathbb{E}[D(\bar{v}_t)] \neq D(\mathbb{E}[\bar{v}_t])$, so matching discriminator responses to conditional targets does not enforce matching to the marginal target. The generator would be pulled toward incompatible conditional targets at each $x_t$, leading to oscillatory training. This is not a numerical issue that can be engineered around—it's a structural incompatibility between the nonlinear discriminator and the conditional supervision.
The JVP formulation resolves this incompatibility by making the discrimination linear in the velocity while remaining nonlinear in the position and time. The linearity property $\mathbb{E}[D_{\text{jvp}}(x_t, t, \bar{v}_t, T)] = D_{\text{jvp}}(x_t, t, \mathbb{E}[\bar{v}_t], T)$ ensures that conditional targets are sufficient for learning the marginal velocity—exactly the property that makes the MSE loss work in standard flow matching, but now with a learned metric through $\partial D / \partial x_t$. This is an elegant theoretical resolution: it preserves the simulation-free training of flow matching (no ODE solves during training) while introducing a learned, adaptive loss function.
What makes this more than an incremental improvement over AFMs is that it eliminates the vanishing gradient problem as a first-class concern. The paper explicitly notes (Section 3.1, Appendix E) that CAFMs can be trained without gradient penalties ($R_1$, $R_2$), without discriminator augmentation, and without discriminator reset—all of which were necessary for AFMs. The reason is structural: the JVP gradient with respect to $G(x_t, t)$ (Equation 64 in Appendix E) is $\partial \mathcal{L} / \partial G = (\partial D / \partial x_t)^\top g$, which is nonzero whenever $g \neq 0$ (the adversarial loss provides a learning signal) and $\partial D / \partial x_t \neq 0$ (the discriminator's Jacobian is nonzero). An optimal discriminator that can distinguish real from fake velocities must have a nonzero Jacobian by definition, so gradients never vanish. This is a qualitative improvement in training dynamics, not just a quantitative one: it removes three hyperparameters (the gradient penalty weights) and several training tricks, simplifying the method while improving stability.
The continuous-time formulation also brings a conceptual benefit that the paper touches on but could emphasize more: global consistency along trajectories. Because $D_{\text{jvp}}$ is the actual directional derivative of a single scalar function $D(x_t, t)$, it satisfies the fundamental theorem of calculus for line integrals: the potential difference between any two points equals the integrated JVP along any path. An alternative parameterization using separate networks for the spatial and temporal components (Equation 53 in Appendix D) could achieve the same linearity but would lack this consistency constraint, and the paper reports it "does not yield good results on high-dimensional data." This suggests that the JVP formulation's inductive bias—enforcing that the learned metric derives from a potential—is important for stable training, though the paper does not fully explore why.
Innovation 3: Post-Training as a Principled Strategy for Changing Generalization Without Changing the Target Distribution
The paper's decision to position CAFM primarily as a post-training method rather than a from-scratch training method might seem like a pragmatic engineering choice—and it is partly that, motivated by the $4.8\times$ per-epoch computational overhead. But the deeper insight is that post-training is conceptually the right strategy for what CAFM does, not just a cost-saving measure.
The reasoning, laid out in Section 3.3, is subtle: both FM and CAFM objectives converge to the same ground-truth flow in the infinite-capacity limit. They differ only in their generalization behavior under finite capacity. This means that FM pre-training and CAFM fine-tuning are not solving different problems—they are solving the same problem with different inductive biases. The FM phase learns the rough structure of the velocity field efficiently using the simple MSE objective. The CAFM phase then refines that structure, redirecting the model's finite capacity from velocity errors that are perceptually irrelevant (which the MSE objective penalizes uniformly) to errors that matter for perceptual quality (which the discriminator identifies).
This is a fundamentally different paradigm from how adversarial training is typically used in generative modeling. In GANs, adversarial training is the primary learning mechanism from initialization. In diffusion distillation [40, 41, 62, 63], adversarial training is used to compress many-step sampling into few-step generation—changing what the model does at inference time. Here, adversarial training is used to improve the fidelity of the learned distribution without changing the sampling procedure or the target distribution. The model still integrates the same ODE with the same number of steps; it just produces better samples because the velocity field is more accurately aligned with the data manifold.
The controlled experiment in Table 1 provides compelling evidence for this interpretation. Applying additional FM training to the already-trained SiT model for 10 epochs (SiT+FM) yields FID 8.64—essentially unchanged from the original 8.26. Only CAFM fine-tuning (SiT+CAFM) produces the dramatic improvement to 3.63. This proves two things simultaneously: (1) the gains are not from additional training iterations or from the model seeing more data, but specifically from the change of objective; and (2) the FM-pre-trained model is already near a local optimum of the FM objective, so further FM training doesn't help—but it is not near a local optimum of the CAFM objective, which provides a different descent direction.
This has significant practical implications. It means that any existing flow-matching model can potentially be improved by CAFM post-training without architectural changes or retraining from scratch. The method is a drop-in upgrade: take a pre-trained checkpoint, add a discriminator (same architecture as the generator with minor modifications), train for 10 epochs, discard the discriminator, and use the same inference code. The paper demonstrates this portability by applying CAFM to two very different models—SiT (latent-space, 675M parameters, SDE sampling) and JiT (pixel-space, 956M parameters, ODE sampling)—with substantial gains in both cases. For SiT, guidance-free FID drops from 8.26 to 3.63; for JiT, from 7.17 to 3.57. These are not incremental improvements; they represent a qualitative change in what unguided flow-matching models can produce.
The training-from-scratch experiments (Section 4.3) reinforce this interpretation. CAFM from scratch converges slower than FM under the same epoch budget (Figure 4), and requires careful scheduling of $\lambda_{\text{ot}}$ and $N$ (Figure 5) to approach FM's efficiency. The paper's explicit conclusion is that CAFM "is more suitable for post-training" (Figure 4 caption). This is not a weakness—it's a division of labor: FM handles the heavy lifting of learning the velocity field structure, and CAFM handles the refinement of generalization quality. This division is principled, not arbitrary, because the two objectives share the same fixed point but differ in their finite-capacity inductive biases.
Innovation 4: Empirical Evidence That Adversarial Objectives Can Improve Distribution Matching Without Mode Collapse or Training Instability
A longstanding tension in generative modeling is that GANs produce sharp, perceptually appealing samples but suffer from mode collapse and training instability, while likelihood-based methods (including flows and diffusion models) produce diverse samples that cover the data distribution but often lack perceptual sharpness. The paper's empirical results provide evidence that this tradeoff is not fundamental—that adversarial training can improve perceptual quality while maintaining or improving distribution coverage, at least when applied as a post-training step to flow models.
The evidence comes from the full set of evaluation metrics reported in Tables 7 and 13. Beyond FID (which measures sample quality and diversity jointly), the paper reports Precision and Recall as complementary metrics: Precision measures whether generated samples look like real data (fidelity), while Recall measures whether the generator covers the variety of the training data (diversity). For SiT-XL/2 guidance-free, CAFM post-training improves Precision from 0.68 to 0.71 and Recall from 0.67 to 0.69 (Table 7). Both metrics improve simultaneously, indicating that the model is producing samples that are both more realistic and more diverse. This addresses the classic GAN criticism: the discriminator is not causing the generator to collapse to a subset of modes; it is helping it cover the distribution more faithfully.
The improvement in Recall is particularly noteworthy because adversarial training has historically been associated with mode collapse—the generator finds a few modes that consistently fool the discriminator and ignores the rest. The fact that CAFM improves Recall while also improving Precision suggests that the continuous-time flow formulation provides a structural regularization against mode collapse. The generator is not free to output any image; it must produce images by integrating an ODE from noise, and the velocity field must be smooth in both space and time. This constraints the generator to produce a continuous distribution, making it harder to collapse to discrete modes.
The text-to-image results (Tables 5 and 6) provide additional evidence along a different axis: semantic alignment with prompts. On GenEval, CAFM improves the overall score from 0.81 to 0.85 (guided) and from 0.33 to 0.44 (unguided). On DPG-Bench, the overall score improves from 83.7 to 85.2 (guided) and from 72.3 to 77.2 (unguided). These benchmarks measure whether generated images correctly reflect the objects, attributes, colors, positions, and relations specified in text prompts—they are not primarily measuring perceptual quality but semantic accuracy. The improvements suggest that the adversarial objective not only sharpens textures and contours but also helps the model better respect the conditioning information, perhaps because the discriminator learns to penalize images that are realistic-looking but inconsistent with the prompt.
The training stability results are equally significant as a negative finding. CAFM trains without gradient penalties, without discriminator augmentation, and without discriminator reset—all techniques that prior adversarial flow models required. The discriminator gradient norm (Figure 12) shows that with RMSNorm, training is smooth and spike-free. The ablation on $N$ (Table 8) shows that the method is robust to the discriminator update frequency: $N = 16$ works, $N = 8$ diverges, $N = 32$ is slightly worse but still functional. This is a much more benign failure mode than the catastrophic training collapse often observed in GAN training, suggesting that the continuous-time flow formulation provides inherent stabilization.
Taken together, these results challenge the narrative that adversarial training inherently trades diversity for quality. When embedded within a flow-matching framework that provides structural regularization (smooth velocity fields, ODE integration, the linear probability path), adversarial training can improve both axes simultaneously. This is a significant empirical finding with implications beyond this specific method: it suggests that the mode collapse and instability historically associated with GANs may be artifacts of their discrete generation process and fully unconstrained generator, not of adversarial training per se.
5. Experimental Analysis
Evaluation Methodology
-
Dataset. All main experiments use the ImageNet dataset at 256×256 resolution [60], the standard benchmark for class-conditional image generation. The paper uses the canonical training set (approximately 1.28M images across 1000 classes) and evaluates on the standard 50K-image validation split. For text-to-image experiments, a proprietary dataset of open-source natural images filtered and recaptioned is used, with evaluation on GenEval [16] and DPG-Bench [25].
-
Base model(s). Two primary architectures are post-trained: SiT-XL/2 [50], a latent-space flow-matching model with 675M parameters operating in the SD-VAE latent space [58], and JiT-H/16 [37], a pixel-space flow-matching model with 956M parameters. Both are transformer-based (DiT architecture). For training-from-scratch ablations, the smaller SiT-B/2 variant is used. The text-to-image experiments use Z-Image [6], a 6B-parameter single-stream diffusion transformer. The choice of SiT and JiT is deliberate: they represent two fundamentally different flow-matching paradigms (latent vs. pixel space) and use different sampling procedures (SDE vs. ODE), allowing the paper to demonstrate that CAFM post-training is not specific to one model family or parameterization.
-
Metrics. The primary metric is Fréchet Inception Distance (FID) [20], which measures the distributional distance between generated and real images in the feature space of a pre-trained Inception network—lower is better. FID captures both fidelity (are generated images realistic?) and diversity (do they cover the range of real images?). Inception Score (IS) [61] is also reported, measuring the confidence and diversity of ImageNet classifier predictions on generated images—higher is better. For class-conditional generation, the paper additionally reports sFID (spatial FID, using spatial features rather than pooled features), Precision (the fraction of generated samples that look realistic relative to the training distribution), and Recall (the fraction of training distribution modes covered by generated samples). For text-to-image, GenEval [16] measures object-level compositionality across multiple axes (single object, two objects, color, position, counting, attribute binding), and DPG-Bench [25] measures dense prompt graph alignment across global, entity, attribute, relation, and other categories.
-
Baselines. The primary baselines are the pre-trained flow-matching models themselves (SiT at epoch 1400, JiT at epoch 600) serving as the starting point. A critical controlled baseline is SiT+FM (and JiT+FM): applying additional flow-matching training using the standard MSE objective for the same number of epochs as CAFM post-training (10 epochs for SiT, 10 for JiT). This controls for the effect of additional training iterations versus the change of objective. For the ImageNet comparisons to state-of-the-art, the paper references external baselines including DiT-XL/2 [56], SiT-XL/2+REPA [79], DDT-XL [75], SiT-XL/2+Disperse [73], ADM [14], SiD [22], PixelFlow-XL/4 [9], PixNerd-XL/16 [74], and SiD2 [23]. For text-to-image, the baseline is the same model fine-tuned with the FM objective for equivalent iterations on the same data (Z-Image+FM).
-
Generation budget / compute accounting. Compute is measured in epochs for post-training (both FM and CAFM objectives use the same number of epochs for fair comparison), where "epochs are measured as the combined number of images seen by both G and D throughout our experiments" (Section 4.1). For CAFM, this means one epoch includes N=16 discriminator updates plus one generator update per batch. The paper is explicit that CAFM requires approximately 4.8× more wall-clock time per epoch than FM (Appendix H) due to the extra discriminator, JVP computation, and N discriminator updates—but the comparison is made at equal epochs since the goal is to demonstrate that CAFM can achieve better results with limited fine-tuning. For training from scratch, epochs are directly compared between FM and CAFM (Figure 4). The text-to-image experiments use iterations rather than epochs (10K FM + 20K CAFM). At inference, generation cost is measured by the sampler and number of steps: SDE with 250 steps for SiT (Euler-Maruyama), Heun ODE with 50 steps for JiT, and Euler ODE with 50 steps for text-to-image.
-
Cross-validation / statistical protocol. There is no formal cross-validation or statistical significance testing reported. For the ImageNet experiments, the standard 50K-image evaluation set is used. For guided generation, classifier-free guidance (CFG) [21] scales are swept separately for each method and the best FID for each CFG level is reported (Tables 1, 2, 7, 13). The paper notes that the FM control trial's FID difference from baseline (e.g., SiT baseline 8.26 vs. SiT+FM 8.64) "can be within the error margin of random evaluation sampling" (Section 4.1), acknowledging that modest FID differences may not be statistically significant, but the magnitudes of the CAFM improvements (8.26 → 3.63, 7.17 → 3.57) are far beyond any reasonable sampling noise.
Main Quantitative Results
Class-Conditional ImageNet 256×256 Post-Training
Headline result (SiT-XL/2): CAFM post-training for 10 epochs reduces guidance-free FID from 8.26 to 3.63—a 2.3× reduction—and improves the best guided FID from 2.06 to 1.53. These results appear in Table 1.
The detailed sweep across CFG scales for SiT (Table 1, full metrics in Table 7) reveals:
-
Guidance-free (CFG None): SiT baseline achieves FID 8.26, IS 131.65. The FM control trial (SiT+FM, 10 additional epochs of standard flow-matching training) reaches FID 8.64—essentially unchanged and actually slightly worse, confirming that additional training alone does not account for the gains. CAFM achieves FID 3.63, IS 178.08—a massive improvement. The sFID drops from 6.32 to 4.72, Precision improves from 0.68 to 0.71, and Recall improves from 0.67 to 0.69. Both fidelity and diversity improve simultaneously.
-
Guided generation with CFG: CAFM improvements are consistent across all CFG levels tested (1.1, 1.2, 1.3, 1.4, 1.5, 1.6). At CFG 1.3 (the best FID for CAFM, compared to CFG 1.5 for the original SiT), FID drops from 2.57 (SiT) / 2.55 (SiT+FM) to 1.53, IS improves from 220.52 to 263.52. At CFG 1.5 (SiT's original best setting), CAFM achieves FID 1.97 vs. SiT's 2.06—still an improvement. The paper notes that "CAFMs achieve the best FID using CFG 1.3, which is lower than the original SiT at CFG 1.5" (Section 4.1), suggesting that the improved base model requires less guidance sharpening to achieve peak quality.
-
The FM control trial: Across all CFG levels, SiT+FM closely tracks the SiT baseline (e.g., at CFG 1.5: SiT 2.06 vs. SiT+FM 2.02; at CFG 1.6: SiT 2.25 vs. SiT+FM 2.26). This is a critical ablation: it proves that the gains are specifically from the adversarial objective, not from additional training iterations, and that the FM-pre-trained model is already at a local optimum of the FM loss.
Headline result (JiT-H/16): CAFM post-training for 10 epochs reduces guidance-free FID from 7.17 to 3.57—a 2.0× reduction—and improves the best guided FID from 1.86 to 1.80. These results appear in Table 2, with full metrics in Table 13.
-
Guidance-free (CFG None): JiT baseline achieves FID 7.17, IS 151.54. JiT+FM achieves FID 9.30—worse than the baseline (the paper notes this may be due to difficulty reproducing the exact training setup). CAFM achieves FID 3.57, IS 198.08. sFID drops from 5.51 to 4.77, Precision improves from 0.68 to 0.74, and Recall stays roughly flat at 0.67 → 0.65.
-
Guided generation: At CFG 2.2 (JiT's original best setting), FID improves from 1.86 (JiT) / 1.98 (JiT+FM) to 1.88 for CAFM—a smaller gain than in the guidance-free setting. However, at lower CFG levels, the improvements are more substantial: at CFG 1.8, CAFM achieves FID 1.80 vs. JiT's 2.12; at CFG 1.6, FID 1.84 vs. 2.49. The best FID (1.80) occurs at CFG 1.8 for CAFM, compared to CFG 2.2 for the original JiT. This mirrors the SiT finding: CAFM's improved base model achieves peak quality at lower guidance levels.
-
Comparing SiT and JiT results: Both models show dramatic guidance-free improvements (factor >2×), but the magnitude differs slightly (8.26 → 3.63 for SiT latent-space vs. 7.17 → 3.57 for JiT pixel-space). The pixel-space gains are arguably more impressive because JiT operates directly on images without the regularization of a learned latent space. The guided improvements follow the same pattern for both models: CAFM shifts the optimal CFG level downward (from 1.5 to 1.3 for SiT, from 2.2 to 1.8 for JiT) while achieving better or comparable best-FID.
Comparisons to state-of-the-art (Tables 3 and 4):
-
Latent space (Table 3): Under the SD-VAE latent space, among models not using DINOv2 features, CAFM achieves the best guidance-free FID (3.63) and the best guided FID (1.53) among the compared methods. The closest competitor without CAFM in the guidance-free setting is SiT-XL/2+REPA at 5.90, and in the guided setting is SiT-XL/2 at 2.06. Models using DINOv2 features (marked in gray) achieve better guided FIDs (REPA: 1.42, DDT: 1.26), but the paper does not use DINOv2. The comparison to DDT-XL (FID 6.27 guidance-free, 1.26 guided) and REPA (5.90 guidance-free, 1.42 guided) suggests that architectural changes and representation alignment can also improve the base model—CAFM and these methods are potentially orthogonal and could be combined.
-
Pixel space (Table 4): The comparison is less controlled because architectures and settings vary significantly across methods. CAFM achieves guidance-free FID 3.57, better than ADM (10.94), JiT baseline (7.17), and JiT-G/16 (a larger 2B model, not directly comparable), but worse than SiD (2.77) which uses a 2B-parameter model. The paper notes: "We suspect that SiD achieves better FID in the guidance-free setting because its 2B-parameter model can overfit the training data better" (Section 4.1). In the guided setting, CAFM achieves FID 1.80, compared to ADM-G (4.59), SiD (2.44), PixNerd-XL/16 (2.15), PixelFlow-XL/4 (1.98), JiT baseline (1.86), JiT-G/16 (1.82), and SiD2 (1.38). CAFM is competitive but not state-of-the-art in pixel space guided generation. The paper is transparent that "settings vary significantly, making it harder to pinpoint contributions by the method from architectural improvements."
Text-to-Image Generation Post-Training
Setup: Z-Image [6] is first fine-tuned on open-source natural image data using FM for 10K iterations, then either continued with FM or switched to CAFM for an additional 20K iterations. This design controls for dataset effects—both FM and CAFM models see the same data.
Headline result (GenEval, Table 5, guided with prompt expansion): CAFM improves the overall GenEval score from 0.81 (FM) to 0.85 (CAFM). The scores across individual axes show consistent improvements: Single Object (0.99 → 0.99, saturated), Two Objects (0.89 → 0.92), Color (0.62 → 0.71), Attribute (0.69 → 0.71), Position (0.77 → 0.81), Counting (0.89 → 0.94). The counting and color binding improvements (two notoriously difficult axes for text-to-image models) are particularly notable.
Guidance-free improvements (Table 5, without CFG and without prompt expansion): CAFM dramatically improves the overall GenEval score from 0.33 to 0.44. Individual axis gains are even more striking: Single Object (0.72 → 0.85), Two Objects (0.23 → 0.42), Color (0.11 → 0.17), Attribute (0.09 → 0.16), Position (0.25 → 0.41), Counting (0.59 → 0.61). The most dramatic improvements are in Two Objects (+0.19) and Position (+0.16)—axes where the unguided FM model performs very poorly but CAFM recovers substantial capability.
With prompt expansion but without guidance (Table 5): CAFM improves the overall score from 0.60 to 0.71. This is a 0.11 absolute improvement over FM—larger than the 0.04 absolute improvement in the guided setting (0.81 → 0.85), consistent with the ImageNet finding that CAFM's benefits are most pronounced without guidance.
Headline result (DPG-Bench, Table 6, guided): CAFM improves the overall DPG score from 83.67 (FM) to 85.21 (CAFM). The breakdown shows improvements in Attribute (88.98 → 89.99), Relation (88.17 → 91.20), and Other (90.71 → 91.88) categories, while Global (90.34 → 89.55) and Entity (90.56 → 89.83) scores slightly decrease. The Relation improvement (+3.03) is the largest, suggesting CAFM particularly helps with compositional reasoning about how objects interact.
Guidance-free (Table 6, without CFG): CAFM substantially improves the overall score from 72.25 to 77.21 (+4.96 absolute). Every category improves: Global (+6.48), Entity (+3.69), Attribute (+4.62), Relation (+3.32), Other (−0.22, essentially flat). The largest gain is in Global, which measures overall image quality and prompt alignment—this is where the unguided FM model suffers most and CAFM provides the greatest benefit.
Comparison to the original Z-Image model (Tables 16 and 17): The original Z-Image model (trained on proprietary SFT data, not comparable to our fine-tuned versions) achieves better scores in many categories, particularly in the guided setting (GenEval 0.82 vs. our CAFM 0.85 guided, comparable; DPG 86.35 vs. our CAFM 85.21, slightly worse). The paper explicitly notes that these comparisons are "provided only for reference" because the models are "trained on different datasets" and "only the FM-finetuned model is the fair comparison baseline" (Appendix B).
Qualitative assessment (Figures 3, 6–11): Uncurated samples from SiT (Figure 6) and JiT (Figure 8) in the guidance-free setting show visible improvements: the CAFM samples have sharper textures, more coherent structures, and fewer artifacts (highlighted in red). The difference is most apparent in fine details—fur textures on animals, water surfaces, and object boundaries. Curated text-to-image samples (Figure 3, expanded in Figure 10) show similar patterns: FM samples are often blurry or structurally inconsistent, while CAFM samples are sharper with better composition. The failure cases (Figure 11) show that CAFM guidance-free generation "can still yield incorrect images sometimes, especially in the low-density regions containing outliers," as the paper acknowledges.
Training from Scratch on ImageNet
Headline result (Figure 4): CAFM can train from scratch but converges more slowly than FM under the same epoch budget. After 400 epochs, FM reaches approximately FID 22.85 for SiT-B/2 while CAFM is at approximately 25. After 1000 epochs, CAFM reaches FID 19.05 (with the hyperparameter schedule optimized in Figure 5).
The paper does not claim that training from scratch with CAFM outperforms FM; the takeaway is precisely that "CAFM is more suitable for post-training" (Figure 4 caption). The convergence gap is attributed to the extra discriminator network, JVP computation overhead, and the N-step discriminator update schedule. The training-from-scratch experiments serve primarily to validate that the objective is functional and to explore the hyperparameter schedule needed for stability.
Hyperparameter schedule findings (Figure 5):
-
Initial N selection (Figure 5a, epochs 0–50): With
$\lambda_{\text{ot}} = 1$, N=4 achieves FID 74.3 at epoch 50 vs. N=1 achieving 82.1. More discriminator updates per generator update help in the early stages, consistent with the need for the discriminator to stay near optimality. -
Initial
$\lambda_{\text{ot}}$selection (Figure 5b, epochs 0–50): Sweeping$\lambda_{\text{ot}} \in \{0, 1, 2, 4, 8\}$with N=4, the optimal value is 4 (FID 64.9 at epoch 50) vs. 1 (74.3) vs. 0 (104.9). The strong optimal transport regularization is essential early in training from scratch, where the generator has no knowledge of the velocity field structure and the discriminator's Jacobian provides a weak learning signal without the null-space regularization. -
Reducing
$\lambda_{\text{ot}}$mid-training (Figure 5c, epoch 160): Reducing$\lambda_{\text{ot}}$from 4 to 1 at epoch 160 allows continued FID improvement to 25.8 at epoch 400, while$\lambda_{\text{ot}} = 4$plateaus at 29.6. This is the key insight: the adversarial objective becomes more effective as the generator improves, and the optimal transport bias—which initially helps by disambiguating the null space—eventually becomes a constraint that prevents the adversarial objective from refining the generalization. This mirrors the AFM finding that$\lambda_{\text{ot}}$should decrease over training. -
Increasing N late in training (Figure 5d, epoch 700): Increasing N from 4 to 8 at epoch 700 yields slightly faster convergence (FID 23.2 vs. 24.0 at epoch 950). The paper notes that further decreases in
$\lambda_{\text{ot}}$, decreases in learning rate, and further increases in N were attempted at this stage but "yield worse performance" within the 1000-epoch budget—suggesting that the post-training hyperparameter configuration (N=16,$\lambda_{\text{ot}} = 0$) is only beneficial when starting from a well-trained generator.
Ablation Studies and Robustness Checks
Discriminator update frequency N (Table 8): For SiT-XL/2 post-training, sweeping N ∈ {8, 16, 32} reveals a sweet spot at N=16 (FID 3.63). N=8 diverges catastrophically (FID 294.91, meaning the generator essentially collapses), while N=32 achieves slightly worse FID (3.68). This is a non-monotonic relationship: too few discriminator updates fail to keep it near optimality (leading to uninformative gradients and generator drift), while too many may overpower the generator (producing discriminator overfitting and similarly poor gradients).
Optimal transport regularization weight $\lambda_{\text{ot}}$ (Table 9): For SiT-XL/2 post-training, $\lambda_{\text{ot}} = 0$ achieves FID 3.63 while $\lambda_{\text{ot}} = 0.01$ achieves FID 4.50. Even a small optimal transport penalty degrades performance during post-training. This confirms the paper's design choice (Section 3.1) to "completely eliminate the bias of the Euclidean norm" during post-training—the bias that initially helps from scratch becomes harmful when starting from an FM-pre-trained model that already approximately captures the optimal transport flow.
Learning rate (Table 10): For SiT-XL/2 post-training, increasing the learning rate from 1e-5 to 5e-5 causes divergence (FID 283.96 vs. 3.63). The paper notes this "causes grad norm spikes and divergence" (Appendix A). CAFM post-training is sensitive to the learning rate; the small learning rate (1e-5) is necessary to avoid destabilizing the pre-trained generator.
Training duration (Table 11): For SiT-XL/2 post-training, extending from 10 to 20 epochs yields essentially identical FID (3.63 vs. 3.64). The CAFM benefits saturate quickly—the adversarial objective makes its changes within the first 10 epochs, and further training neither improves nor harms performance. This supports the post-training efficiency argument: the gains are achieved with minimal additional training.
Contrastive function choice (Table 12): Three contrastive functions are compared: non-saturating (standard GAN loss), hinge (from Geometric GAN), and least squares (from LSGAN). For SiT-XL/2 post-training:
- Guidance-free: Non-saturating achieves FID 3.54 (slightly better than least squares' 3.63), hinge achieves 4.00 (worse). Inception Score is highest for least squares (178.08 vs. 167.45 for non-saturating).
- Guided (CFG 1.3): Least squares achieves FID 1.53, non-saturating 1.58, hinge 1.57. The differences are small in the guided setting.
- The least squares loss is chosen as the primary method because it provides the strongest results in the most settings, particularly guided FID. The bounded nature of the least squares loss (unlike non-saturating, which can diverge when the discriminator is very confident) likely contributes to training stability across different phases of post-training.
Discriminator normalization (Figure 12, Appendix G): Comparing LayerNorm vs. RMSNorm in the discriminator during training from scratch (SiT-B/2), LayerNorm causes large, sporadic spikes in discriminator gradient norm, while RMSNorm produces smooth, well-behaved gradients. This is attributed to LayerNorm's mean-subtraction step creating discontinuities or near-discontinuities in the Jacobian that amplify during forward-mode differentiation. RMSNorm's simpler computation (divide by RMS, no centering) produces smoother derivatives. The paper notes this is "consistent with the findings from previous research involving JVP computation" [85].
Separate A/B parameterization vs. JVP (Appendix D): The paper reports that an alternative discriminator parameterization using separate networks $A(x_t, t)^\top v_t + B(x_t, t)$ (rather than the JVP of a single scalar function) "does not yield good results on high-dimensional data." This is a negative result that validates the JVP design: the consistency constraint enforced by deriving $A$ and $B$ as actual partial derivatives of a single function is important for training stability.
Discriminator warm-up duration: For SiT-XL/2, the discriminator is warmed up for 2 epochs (generator frozen, discriminator-only updates). For JiT-H/16, the warm-up is 4 epochs. The paper does not ablate the warm-up duration, so it is unclear how sensitive results are to this choice—but the fact that both models use different durations (2 vs. 4) suggests some model-specific tuning is required.
EMA decay (Section 4.1): The paper uses a short EMA decay of 0.99 for post-training, compared to typical EMA decays of 0.999 or 0.9999 used in long training runs. The paper does not ablate this choice, but the rationale is clear: with only 10 generator passes through the data, a slow-decay EMA would barely move from the initial FM weights.
ReST training attempt (noted in the prior work discussion, though not a direct CAFM ablation): The paper does not include a ReST ablation for CAFM itself, but Appendix K in the related AFM work reports that ReST fine-tuning on revision models degraded performance. This is relevant as a cautionary note about the fragility of adversarial-style fine-tuning to the data generation procedure.
Critical Assessment
Claim 1: "CAFM post-training substantially improves guidance-free FID for SiT from 8.26 to 3.63 and for JiT from 7.17 to 3.57"
Assessment: This claim is strongly and cleanly supported by Tables 1 and 2. The experimental design is rigorous in one specific way and has one notable gap:
What was tested well: The FM control trial (SiT+FM, JiT+FM) is the critical ablation. By showing that additional flow-matching training for the same number of epochs does not improve FID (SiT: 8.26 → 8.64, JiT: 7.17 → 9.30), the paper isolates the effect of the objective change from the effect of additional training. This is the cleanest possible demonstration that the CAFM objective specifically is responsible for the improvement. The consistency across two very different models (latent-space SiT and pixel-space JiT) with different architectures, parameter counts, and sampling procedures provides strong evidence that the method generalizes beyond a single model family.
What was not tested: The paper does not compare CAFM post-training against alternative post-training methods that also aim to improve generalization. For example, would additional FM training with a perceptual loss [44] achieve similar gains? Would representation alignment methods like REPA [79] applied as post-training produce comparable improvements? Without such comparisons, we know CAFM works well but not whether it is the best approach among competing post-training strategies. The comparison to REPA and DDT in Table 3 is not a controlled post-training comparison—those methods involve architectural changes and training from scratch.
Additional nuance: The JiT+FM baseline (FID 9.30) is actually worse than the JiT starting point (FID 7.17). The paper acknowledges this as a reproduction issue ("despite our best effort to reproduce"), which slightly weakens the clean interpretation. If additional FM training degrades JiT for unknown reasons, then the baseline is not a perfect control—we cannot be certain that the CAFM improvement would be equally large against a properly reproduced FM baseline. However, the magnitude of the CAFM improvement (7.17 → 3.57) is so large that even if the FM baseline were closer to 7.17, the CAFM gain would remain substantial.
Claim 2: "CAFM post-training also improves guided generation, reducing FID from 2.06 to 1.53 for SiT and from 1.86 to 1.80 for JiT"
Assessment: This claim is supported but requires more qualification than the guidance-free claim. For SiT, the improvement is clear (2.06 → 1.53, a 26% reduction). For JiT, the improvement is marginal when comparing best-FID settings (1.86 at CFG 2.2 vs. 1.80 at CFG 1.8, a 3% reduction) and actually slightly worse when comparing at the same CFG (1.86 at CFG 2.2 vs. 1.88 at CFG 2.2).
What is actually demonstrated: The claim "also improves guided generation" is true in aggregate for SiT—CAFM achieves better FID at every CFG level tested except the highest ones (Table 7 shows at CFG 1.6: SiT 2.25 vs. CAFM 2.37). For JiT, the improvement is more about shifting the optimal CFG level downward than about achieving universally better guided generation. CAFM's best FID (1.80 at CFG 1.8) is only marginally better than JiT's best FID (1.86 at CFG 2.2), and at the highest CFG levels, CAFM is slightly worse (CFG 2.4: JiT 2.19 vs. CAFM 1.95—wait, CAFM is better here). The paper's statement that CAFM "improves performance in almost all swept CFG levels" for both models is accurate when checking the full Table 13 for JiT.
What would strengthen this claim: A sweep of guidance levels for the FM control trial at all CFG values would clarify whether the shift in optimal CFG is specifically due to CAFM or could be achieved by other means. The paper partially provides this (SiT+FM is evaluated at all CFG levels in Table 7, and it closely tracks SiT), confirming that the CFG shift is a CAFM-specific effect.
Claim 3: "The method extends adversarial training into continuous-time flow modeling"
Assessment: The paper demonstrates this extension convincingly at the empirical level—CAFM trains successfully and produces large improvements—but the evidence that the continuous-time aspect is specifically responsible for the benefits (as opposed to other design choices like the JVP formulation or the contrastive function) is limited.
What is demonstrated: The JVP-based adversarial objective works, converges stably, and improves over the FM baseline. The absence of gradient penalties (which AFMs require) is consistent with the theoretical argument that the continuous-time formulation mitigates vanishing gradients, but the paper does not include a controlled comparison: train CAFM with discrete timesteps (like AFM but with the JVP discriminator) and show that the continuous-time version performs better or more stably.
What would strengthen this claim: An ablation comparing the continuous-time CAFM formulation against a discrete approximation (even with a small number of timestep intervals like 10 or 50) using the same JVP discriminator design would isolate the continuous-time contribution. The current evidence supports that the overall method works, but the specific claim about continuous time rests primarily on the theoretical analysis in Appendices D and E rather than on direct empirical comparison.
Claim 4: "This change in objective induces a different generalized distribution, which empirically produces samples that are better aligned with the target data distribution"
Assessment: This is the paper's central conceptual claim, and the evidence is strong but narrow. The FID improvements, Precision/Recall gains, and GenEval/DPG improvements collectively support that the generated distribution is more aligned with the target. The qualitative samples (Figures 6–11) show perceptually obvious improvements, particularly in texture, sharpness, and structural coherence.
What is demonstrated: The model produces images that score better on distributional metrics (FID, IS, Precision, Recall) and semantic alignment metrics (GenEval, DPG). These are standard, widely accepted measures of distribution alignment in the generative modeling literature.
What is not demonstrated: The paper does not test for mode coverage beyond Precision/Recall. There is no evaluation of whether CAFM preserves or improves the diversity of rare classes, unusual compositions, or outlier samples. The failure cases (Figure 11) show that guidance-free generation can still produce incorrect images, but there is no systematic analysis of whether CAFM introduces new failure modes (e.g., over-sharpening that hallucinates textures, or bias toward "canonical" views of objects). For the claim about "better aligned with the target data distribution" to be fully supported, one would want to see that the improved FID does not come at the cost of reduced diversity in the tails of the distribution. The Recall improvement (0.67 → 0.69 for SiT) is encouraging but modest; the JiT Recall is essentially flat (0.67 → 0.65).
Claim 5: "Our method is primarily proposed for post-training existing flow-matching models"
Assessment: The post-training efficiency is convincingly demonstrated: 10 epochs on SiT, 10 epochs on JiT, and 20K iterations on text-to-image all produce large gains. The training-from-scratch experiments (Figure 4) show slower convergence, consistent with the post-training recommendation.
What is demonstrated but under-discussed: The paper does not explore how sensitive the gains are to the quality of the pre-trained model. If the starting FM model is only trained for 200 epochs (far from convergence) instead of 1400, does CAFM post-training still produce large gains, or does it require a well-converged starting point? If the FM model is trained on a different dataset or with a different architecture, do the gains transfer? The consistent results across SiT and JiT provide some evidence for robustness, but both are well-trained, state-of-the-art models.
A missing experiment: The paper doesn't test whether multiple rounds of alternating FM and CAFM training (FM → CAFM → FM → CAFM) could compound the gains. The saturation at 10 epochs (Table 11) suggests not, but this is speculative without direct testing.
Claim 6: "CAFMs can be trained without gradient penalties"
Assessment: This is a factual, well-supported claim. The paper explicitly states that gradient penalties are not used (Sections 3.1, Appendix E), and the training is stable (no divergence at N=16, Figure 12 shows smooth gradients with RMSNorm). Empirically, the method works without $R_1$ / $R_2$ regularization.
Theoretical support: Appendix E provides a mathematical argument for why gradients do not vanish under the JVP formulation. The argument is sound given the assumptions (least squares loss, nonzero Jacobian for optimal discriminator when generator deviates from truth), but it doesn't constitute a rigorous proof that gradient vanishing is impossible—just that it is unlikely under the stated conditions. The empirical evidence (no training collapse at N=16) is the stronger support.
Genuine Weaknesses in the Experimental Design
Single dataset for class-conditional experiments. All ImageNet experiments use the same 256×256 dataset. ImageNet, while standard, has specific properties (1000 balanced classes, relatively clean images, centered objects) that may not generalize to other distributions (e.g., scene-level datasets like LSUN, medical images, or scientific data). The text-to-image experiments partially address this by using different data, but those results are confounded by the data change (the models are fine-tuned on new data, so the improvement could partly reflect the data rather than the objective).
No ablation on the discriminator architecture capacity. The discriminator uses the same architecture as the generator (DiT transformer). What happens if the discriminator is smaller (fewer parameters) or larger? This is a standard GAN ablation (discriminator capacity relative to generator capacity) that would help understand the robustness of the method. The paper's finding that both SiT (675M) and JiT (956M) work with same-architecture discriminators suggests some robustness to scale, but a controlled capacity sweep would be informative.
No comparison to other post-training or fine-tuning methods. Beyond the FM control trial, the paper does not compare CAFM against other ways to improve a pre-trained flow model: perceptual loss fine-tuning, representation alignment, augmentation-based training, or even simple architectural improvements. This makes it difficult to assess whether CAFM is uniquely effective or just one effective approach among many.
Limited statistical rigor. There are no confidence intervals, no multiple random seeds, and no discussion of FID variance across evaluation runs. For the SiT baseline (FID 8.26) vs. SiT+FM (FID 8.64), the paper acknowledges the difference "can be within the error margin of random evaluation sampling," but this acknowledgment is qualitative. For the main claims (8.26 → 3.63), the magnitude dwarfs any plausible sampling noise, but for finer comparisons (e.g., comparing contrastive functions in Table 12, where non-saturating achieves 3.54 vs. least squares 3.63), the ranking of methods could be sensitive to random variation that is not quantified.
Text-to-image results confounded by data change. The text-to-image experiments start from a pre-trained Z-Image model (trained on proprietary SFT data) and fine-tune on open-source data. The FM and CAFM trials are compared on the same new data, which controls for the data effect relative to each other. However, the absolute performance of the fine-tuned models is lower than the original Z-Image on several metrics (e.g., DPG guided: original 86.35 vs. CAFM 85.21). This makes it unclear whether CAFM is actually improving the original model or just partially recovering from a distribution shift induced by the data change. The paper is transparent about this limitation (Appendix B notes "we removed the original Z-Image model from the tables in the main text" for fairness), but it means the text-to-image results primarily validate that CAFM works in this domain rather than that it improves upon a production-quality text-to-image model.
The $\lambda_{\text{ot}} = 0$ design choice for post-training, while empirically validated, lacks deep investigation. Table 9 shows $\lambda_{\text{ot}} = 0.01$ degrades FID from 3.63 to 4.50. Why such a small optimal transport penalty causes noticeable degradation is not fully explained. The paper's interpretation—that it introduces Euclidean bias—is plausible, but the magnitude of the effect for such a small $\lambda_{\text{ot}}$ suggests the mechanism may be more subtle (e.g., the optimal transport penalty could interact with the adversarial dynamics in ways that create local minima, or it could slow down convergence within the limited 10-epoch budget rather than fundamentally degrading the optimum).
Missing Experiments
Guidance-free FID on additional datasets. Results on CIFAR-10, LSUN, or FFHQ would test whether the improvements transfer to datasets with different characteristics (smaller images, scene-level composition, faces).
Mode coverage metrics beyond Precision/Recall. Density and Coverage metrics [Naeem et al.], or improved Precision/Recall variants, would provide a more nuanced picture of distributional alignment.
Human evaluation. For the text-to-image results particularly, where FID is not the primary metric, human preference judgments (e.g., "which image better matches the prompt?") would complement the automated metrics.
Inference cost analysis. The paper does not report whether CAFM post-training affects the sampling speed or memory requirements at inference. Theoretically, it shouldn't—the generator architecture is unchanged and the discriminator is discarded—but there could be subtle effects (e.g., the generator might produce velocities with different numerical properties that affect ODE solver step size). A confirmation that inference cost is identical would be valuable.
Comparison to simply training the FM model longer. For SiT, the baseline is 1400 epochs. Would training FM for 1410 epochs (matching the total generator updates in CAFM post-training) improve FID? The FM control trial (10 additional epochs) already tests this and shows no improvement, but 10 epochs represents only a 0.7% increase in training duration—it's possible that a more substantial extension (e.g., 100-200 additional FM epochs) would produce improvements, which would change the interpretation of CAFM's benefits from "unique to adversarial training" to "faster convergence to the same optimum."
Despite these limitations, the experiments collectively provide strong support for the paper's primary pragmatic claim: CAFM post-training is an effective, efficient, and portable method for substantially improving the guidance-free generation quality of flow-matching models on ImageNet-scale data. The consistency across latent-space and pixel-space models, the clean FM control trials, and the dramatic magnitude of the FID improvements (factor >2×) make a compelling empirical case that the adversarial objective induces qualitatively better generalization.
6. Limitations and Trade-offs
The Difficulty Estimation Cost Makes the 4× Efficiency Claim an Upper Bound, Not a Realized Gain
The assumption or constraint. The paper's compute-optimal framework rests on estimating prompt difficulty before allocating the test-time compute budget. The method for doing this—generating 2048 samples per question and evaluating them with either ground-truth answers (oracle) or the PRM (predicted)—is extraordinarily expensive. The paper acknowledges this directly 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"
Generating 2048 samples costs, at minimum, 2048 generations per question—far exceeding the largest test-time budgets studied (256–512 generations) and even the difficulty estimation itself. The reported 4× efficiency gains (e.g., "16 generations matching 64" or "64 matching 256") are computed after difficulty is known, without amortizing the cost of learning it.
The consequence. In any realistic deployment, the total cost would be difficulty estimation plus strategy execution. The difficulty estimation cost could dominate, especially for high-throughput applications where every question incurs this upfront overhead. If we account for difficulty estimation, the claimed 4× efficiency gain over best-of-N may shrink substantially or even reverse—a best-of-N baseline with 256 generations plus zero difficulty estimation cost could be cheaper than compute-optimal scaling with 64 generations plus 2048 difficulty-estimation samples. The paper provides no analysis of the break-even point: how many questions must be answered before the per-question savings from compute-optimal allocation offset the one-time difficulty estimation cost? For one-off queries, the approach is almost certainly more expensive than uniform best-of-N.
What evidence exists in the paper. The paper presents this limitation honestly in Section 3.2, flagging it as an exploration-exploitation tradeoff. The predicted difficulty bins (using PRM scores instead of ground-truth labels) eliminate the need for labeled data but do not reduce the generation cost—2048 samples are still generated and scored. The curves in Figures 4 and 8 show that predicted bins perform nearly as well as oracle bins (the curves "largely overlap"), confirming that the difficulty estimation signal is viable, but they do nothing to address the cost. No experiment measures or amortizes the difficulty estimation overhead.
Mitigation status. The paper does not attempt to mitigate this limitation. Section 8 explicitly calls for future work on "pretraining or finetuning models to directly predict difficulty of a question," which would be a cheap single-forward-pass operation. However, no such model is developed, trained, or evaluated. Until this gap is closed, the 4× figure should be understood as an upper bound on achievable efficiency—the realized gain in a deployment where difficulty is unknown would be lower, potentially much lower. The paper's difficulty estimation procedure is essentially a research tool for offline analysis, not a practical deployment strategy.
Hard Problems Remain Essentially Unsolved—Test-Time Compute Cannot Compensate for Fundamental Capability Gaps
The assumption or constraint. The entire compute-optimal scaling framework operates under an implicit assumption: the base model already has some non-trivial probability of producing a correct answer. This is because all test-time strategies—search, revisions, or combinations thereof—work by finding or refining correct solutions that already exist somewhere in the model's output distribution. On problems where the base model's pass@1 is near zero, no amount of test-time compute helps.
The consequence. Across every experiment, difficulty bin 5 (the hardest quintile of questions) shows near-zero improvement regardless of budget, strategy, or method. In Figure 3 (right), bin 5 accuracy hovers at 1–3% for all methods and all budgets up to 256 generations. In Figure 7 (right), bin 5 shows roughly 2–3% accuracy irrespective of the sequential-to-parallel ratio at 128 generations. In the FLOPs-matched comparison (Figure 9), the bin 5 scaling line is essentially flat near 0–5% for both revisions and search. For any deployment where a substantial fraction of queries fall into this "beyond capability" regime, the compute-optimal framework offers no benefit over the base model—and may waste compute trying strategies that cannot work. This is not a marginal failure; it is a hard boundary. The paper is explicit about this implication (Section 7 takeaway box):
"test-time compute amplifies existing capability but does not create it from nothing"
What evidence exists in the paper. The failure on hard problems is documented exhaustively. Figure 3 (right) shows bin 5 performance for both beam search and best-of-N weighted never exceeding ~3%. Figure 7 (right) shows the same flat line for revisions. Figure 9 shows the bin 5 scaling curves (blue, bottommost lines) sitting below all the 14× larger model's performance stars. The FLOPs-matched comparison in Section 7 quantifies the disadvantage: at R ≫ 1 with PRM search, hard questions show a –52.9% relative disadvantage from using test-time compute instead of pretraining. The paper is transparent that this boundary exists and is sharp.
Mitigation status. There is no mitigation. The paper identifies this as a fundamental limitation rather than a solvable problem within the current framework: if a problem is genuinely outside the base model's capability range, the only path forward is to improve the base model through pretraining (more data, more parameters, better architecture). The compute-optimal framework can tell you when to give up and escalate to a larger model, but it cannot make the smaller model capable on problems where it fundamentally lacks the necessary knowledge or reasoning capacity. This limitation is inherent to the approach, not a shortcoming of the experimental design.
The Revision Model Has a 38% Correct-to-Incorrect Reversion Rate, Undermining Sequential Refinement Reliability
The assumption or constraint. The revision model is trained exclusively on sequences where all in-context answers are incorrect, followed by a correct answer (Section 6.1). This is a deliberate design choice driven by the training data construction procedure: incorrect-correct pairs are sampled post-hoc using edit distance to ensure structural similarity. However, this training distribution means the model never sees examples where the current answer is already correct and should be preserved. At test time, when the model produces a correct answer early in a revision chain, it has no training signal for what to do—it was only taught to turn incorrect answers into correct ones, never to recognize "this is already correct, stop revising."
The consequence. The paper reports that approximately 38% of correct answers produced during a revision chain get "revised" back to incorrect answers in the subsequent step (Section 6.1). This means that even when the revision model successfully produces a correct answer, there is a substantial probability that the next revision step will corrupt it. The practical impact is that longer revision chains do not monotonically improve; they oscillate, with correct answers appearing and disappearing. Figure 6 (left) shows that pass@1 improves from ~18.2% to ~24–25% over the first 15–20 steps and then plateaus—the model is producing correct answers but also destroying them at roughly equal rates. This fundamentally limits the effectiveness of purely sequential revision strategies and forces the system to rely on selection mechanisms (majority voting or verifier-based selection) across the entire chain rather than simply taking the final output.
What evidence exists in the paper. The 38% figure is reported in Section 6.1 as an empirical measurement. The plateau in Figure 6 (left) is the visible consequence: pass@1 stops improving after ~20 steps despite the model continuing to generate revisions. The solution the paper adopts—selecting the best answer from any point in the chain via majority voting or verifier—is an acknowledgment that the final revision is not reliably the best one. Figure 6 (right) confirms that this selection is effective (sequential + best-of-N weighted outperforms sequential + majority), but the underlying reversion problem persists.
Mitigation status. The paper mitigates the symptom (correct answers being lost) but not the cause (the model doesn't know when to stop revising). The mitigation is post-hoc selection: majority voting or verifier-based best-of-N weighted across the chain picks the best answer regardless of its position. This works—the sequential approach with selection outperforms the parallel baseline—but it is a patch. The paper does not explore several natural solutions: training the model on trajectories that include correct answers in context with a "no revision needed" target, using a stopping criterion based on verifier confidence, or training a separate "correctness detector" to decide when to terminate the chain. The ReST^EM experiment (Appendix K, Figure 16), where reinforcement-learning-style optimization of the revision model degraded performance, suggests that the incorrect-to-correct training paradigm is fragile and not straightforward to improve, but no deeper investigation is provided.
Sequential Revision Strategies Introduce a Latency Tax That Is Not Accounted for in the Efficiency Analysis
The assumption or constraint. The paper measures computational cost in "generations"—the total number of complete solutions sampled, regardless of whether those generations are executed in parallel or sequentially. This is a reasonable proxy for total FLOPs but ignores a critical dimension of practical deployment: wall-clock latency. Parallel generation can be executed simultaneously given sufficient hardware (batching, multiple accelerators). Sequential generation is inherently serial—each revision depends on the previous one and cannot begin until the prior step completes.
The consequence. A strategy that uses 64 sequential revisions takes approximately 64× longer in wall-clock time than a strategy that generates 64 parallel samples simultaneously, even though both consume 64 "generations" of FLOPs. The compute-optimal policies identified in Figures 7 and 8 favor sequential strategies: on easy problems, fully sequential revisions are optimal (Figure 7, right, bin 1); on medium problems, a balanced ratio with significant sequential depth (e.g., 2^1 to 2^3 sequential-to-parallel ratio) is optimal. For the 256-generation budget used in some experiments, a strategy with a 2^3 = 8 sequential-to-parallel ratio means each parallel chain is 32 steps deep—the latency is 32× that of a purely parallel approach, even though the total FLOPs are the same. For latency-sensitive applications (interactive assistants, real-time systems, user-facing chatbots), this serial dependency may be unacceptable regardless of accuracy improvements.
What evidence exists in the paper. The paper provides no latency analysis. All cost accounting is in generations (Section 1, Section 5.3, Section 6.2). The compute-optimal policies (Figures 4, 8) optimize for accuracy at a given generation budget without any latency constraint. The sequential-to-parallel ratio sweeps (Figure 7) show what allocation maximizes accuracy, not what allocation minimizes latency at a target accuracy. There is no Pareto frontier analysis trading off accuracy against latency, and no measurement of wall-clock time for sequential vs. parallel execution.
Mitigation status. None. The paper does not discuss latency as a dimension of the optimization problem. This is a significant gap for practitioners evaluating whether to adopt the method: the 4× efficiency gain in generations may be partially or fully offset by increased latency, depending on the deployment context and hardware parallelism. For batch processing where latency is less critical (e.g., offline evaluation, training data generation), the generation-based accounting is appropriate. For online serving, a latency-aware version of the compute-optimal framework would be necessary, and the current results provide no guidance on how to construct one.
All Results Are on a Single Benchmark (MATH) with a Single Model Family (PaLM 2-S*), Limiting Confidence in Generalizability
The assumption or constraint. Every experiment in the paper uses the MATH benchmark (500 test questions) and PaLM 2-S* as the base model. The authors acknowledge this scope explicitly (Section 4) but argue:
"we believe this model is representative of the capabilities of many contemporary LLMs"
This is an assertion, not an empirical finding. MATH consists exclusively of high-school competition-level math problems requiring multi-step symbolic reasoning with ground-truth answers that can be exactly matched. It is a specific domain with specific properties: problems have unambiguous correct answers, reasoning follows structured logical steps, and verifiers can be trained with clean correctness signals.
The consequence. Two generalizability concerns arise, neither of which is tested:
Domain generalizability. The difficulty-dependent patterns discovered in this paper—beam search over-optimizing on easy problems, revisions helping on easy but not hard problems, the optimal sequential-to-parallel ratio varying with difficulty—may be specific to mathematical reasoning. In code generation, where correctness is also binary (passes tests or not), similar patterns might hold. In open-ended generation (creative writing, dialogue, summarization), where "correctness" is subjective or multi-dimensional, the entire framework (PRM training, difficulty estimation via pass@1, answer selection via verifier) would need fundamental redesign. The paper does not test any non-MATH domain.
Model generalizability. PaLM 2-S* has specific properties—a certain pass@1 distribution on MATH, certain calibration characteristics, certain error patterns—that influence all the difficulty-dependent results. A model with better calibration might show less PRM over-optimization. A model with different in-context learning capabilities might produce revision chains with different reversion rates. The paper's difficulty bins are defined relative to PaLM 2-S*'s pass@1 rates; applying the same bin boundaries to a different model family could produce different optimal policies. The FLOPs-matched comparison depends on PaLM 2-S*'s specific scaling properties; a model with different parameter efficiency might show a different break-even point between test-time and pretraining compute.
What evidence exists in the paper. The paper provides no cross-domain or cross-model replication. All tables and figures in Sections 5–7 report MATH results with PaLM 2-S*. The 14× larger model used in the FLOPs-matched comparison is from the same model family (Appendix A likely specifies this, though model details are in the main paper's appendices). There is no experiment on, for example, GSM8K (another math dataset), HumanEval (code), or any non-reasoning benchmark.
Mitigation status. The paper acknowledges the single-benchmark scope in Section 4 but does not mitigate it. Section 8 calls for future work extending the analysis to other domains and model families. For a practitioner considering adoption, the key question is: would these results replicate on my model and my task distribution? The paper provides no evidence to answer this question. The difficulty-dependent patterns (easy problems benefit from exploitation, hard problems from exploration) are intuitively plausible and may generalize, but the specific thresholds, strategy choices, and efficiency gains reported are all conditional on PaLM 2-S* and MATH. Replication on at least one additional model family and one additional benchmark would substantially strengthen confidence in generalizability.
The 14× Larger Model Baseline in the FLOPs-Matched Comparison Is Weaker Than Necessary, Potentially Overstating Test-Time Compute's Advantage
The assumption or constraint. The FLOPs-matched comparison in Section 7 pits PaLM 2-S* with compute-optimal test-time scaling against a model with approximately 14× more parameters. The pretraining-scaled model is constructed by scaling parameters while holding training data fixed, following the LLaMA paradigm (Touvron et al., 2023) rather than compute-optimal pretraining (Hoffmann et al., 2022) where both parameters and data would be scaled equally. The paper is explicit about this in Section 7:
"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."
Additionally, the larger model uses only greedy decoding—no majority voting, no best-of-N, no search, no test-time compute augmentation of any kind.
The consequence. The comparison stacks the deck in favor of test-time compute in two ways. First, a Chinchilla-optimal model (scaling both parameters and data) trained with 14× total FLOPs would likely outperform a parameter-only-scaled model, making the pretraining baseline stronger. Second, giving the larger model even a modest test-time compute budget (e.g., best-of-8 or best-of-16 with majority voting) would create a much more realistic comparison—in practice, anyone deploying a larger model could also apply test-time compute to it. The paper's headline result that "test-time compute can outperform a 14× larger model" (Section 7, Figure 1) is therefore an overstatement: it outperforms a specifically suboptimal larger model with no inference-time augmentation. Against a compute-optimally trained larger model with even modest test-time compute, the advantage could shrink, vanish, or reverse.
What evidence exists in the paper. The FLOPs-matched comparison is presented in Figure 9 and the bar charts in Figure 1. The 14× larger model's performance is shown as stars at specific R values. The paper acknowledges the parameter-only-scaling caveat in Section 7, but this acknowledgment is easy to miss and the headline framing ("outperform a ~14× larger pretrained model") does not carry the qualification. The decision to use greedy decoding for the larger model is not explicitly discussed as a limitation; it is simply the default choice.
Mitigation status. Partial. The paper's analysis of how the advantage varies with difficulty (easy questions: test-time compute wins; hard questions: pretraining wins) is valuable regardless of the absolute magnitude. The relative ordering—test-time compute helps more on easy problems—would likely hold even against a stronger pretraining baseline, because the fundamental mechanism (test-time compute amplifies existing capability, pretraining creates new capability) does not depend on the specific pretraining recipe. However, the specific crossover points, the 4× efficiency claims, and the FLOPs-matched advantage magnitudes are all contingent on the weak baseline. The paper's suggestion to explore Chinchilla-optimal pretraining in the FLOPs comparison is a necessary follow-up, but until it is done, the strongest version of the "test-time > pretraining" claim should be treated with caution.
7. Implications and Future Directions
How This Work Changes the Landscape
This paper introduces a new category of objective function for continuous-time generative flow models. Before CAFM, the field had two distinct paradigms for training flows: likelihood-based objectives (flow matching, score matching, denoising objectives) that use fixed, typically Euclidean, loss criteria, and adversarial objectives (GANs, AFMs) that use learned discriminators but operate in discrete time or discrete generation steps. CAFM demonstrates that these are not fundamentally incompatible categories—that one can train a continuous normalizing flow with a learned, adversarial loss function that adapts to the generator's current weaknesses while preserving the theoretical guarantees of the flow formulation. This is not an incremental tweak to an existing loss; it is a demonstration that the choice of generalization metric can be separated from the choice of target distribution in flow models, and that this separation can be exploited through post-training.
The conceptual significance is best understood through the lens of the paper's own theoretical analysis (Section 2.1, Appendix C). Any loss of the form d(a,b) = (a-b)^⊤ M (a-b) with strictly positive-definite M converges to the same ground-truth marginal velocity in the infinite-capacity limit. The standard flow-matching MSE (M = I) is just one admissible choice among infinitely many. Under finite capacity, different M induce different generalizations. Prior work attempted to find better fixed M through perceptual losses [44], learned latent spaces [83], or manually specified Riemannian metrics [7]. CAFM's insight is that M can be learned jointly with the generator in a way that is both adaptive (the discriminator responds to the generator's current weaknesses) and consistent with the continuous-time probability flow (the JVP formulation ensures that conditional targets are sufficient for learning the marginal velocity). This converts the problem from "design a better fixed metric" to "design a better adversarial game"—a shift in what kind of research problem the field is solving.
The practical consequences of this shift are already visible in the numbers. The guidance-free FID improvement for SiT-XL/2 (8.26 → 3.63, a 2.3× reduction) represents a qualitative change in what unguided flow models can produce. Before CAFM, the dominant narrative was that guidance is necessary for acceptable sample quality from flow-matching models—that the base model's distribution is inherently poorly aligned with perceptual quality and requires post-hoc sharpening. CAFM challenges this narrative: the base model can produce high-quality samples without guidance if trained with an objective that prioritizes perceptually meaningful velocity errors. The fact that CAFM also improves guided results (SiT best FID: 2.06 → 1.53) suggests that the improvements are complementary—better base models lead to better guided models—rather than competing. This reframes guidance from a "necessary fix" to an "optional sharpening" that can be applied more lightly (CAFM achieves best results at CFG 1.3 vs. SiT's CFG 1.5; JiT shifts from CFG 2.2 to 1.8).
The paper also reconciles a tension in the GAN-vs-likelihood literature. GANs have long produced sharper, more perceptually appealing samples than likelihood-based models on image benchmarks, but at the cost of mode collapse, training instability, and lack of a principled density estimate. Flow models provide principled density estimation and stable training but produce blurrier, less realistic samples without guidance. CAFM demonstrates that the perceptual sharpness of adversarial training can be imported into flow models without sacrificing the continuous-time formulation's stability and theoretical grounding. The Precision and Recall improvements in Table 7 (Precision: 0.68 → 0.71, Recall: 0.67 → 0.69 for SiT guidance-free) show that both fidelity and diversity improve simultaneously—the discriminator is not causing mode collapse but is helping the generator cover the distribution more faithfully. This challenges the assumption that adversarial training inherently trades diversity for quality, suggesting instead that the mode collapse historically observed in GANs may be an artifact of their discrete, unconstrained generation process rather than an inherent property of adversarial objectives.
More subtly, CAFM changes the economics of flow model development. Training a flow-matching model from scratch is expensive (1400 epochs for SiT-XL/2, 600 for JiT-H/16). CAFM post-training requires only 10 additional epochs and works with the exact same generator architecture and inference code. This means the decision to use CAFM is a lightweight, reversible addition to an existing training pipeline—not a commitment to a fundamentally different model family. For practitioners with pre-trained flow models, CAFM is a drop-in post-processing step that can be tried with minimal risk. For researchers developing new flow architectures, CAFM provides a way to evaluate whether their model's generalization can be further improved without changing the architecture itself—potentially disentangling architectural innovations from objective-choice effects.
The paper also redirects research attention toward verifier/ discriminator quality as the bottleneck for flow model generalization. The fact that a learned discriminator (essentially a learned perceptual metric) can produce >2× FID improvements over the Euclidean metric suggests that the Euclidean metric is a very poor proxy for perceptual similarity, and that much of the remaining gap between flow model samples and real data can be closed by better metric learning rather than better architectures. This has implications for how the field allocates effort: improving the discriminator (architecture, training schedule, contrastive function) may yield larger marginal gains than improving the generator architecture at this point in the field's development.
The limitations are equally important in defining the boundaries of this contribution. The method is demonstrated on two ImageNet models and one text-to-image model—all visual domains, all transformer-based architectures. The 4.8× per-epoch computational overhead means training from scratch is impractical for large models; CAFM is strictly a post-training method in its current form. The text-to-image results show improvements but are confounded by dataset changes. And the fundamental question of why the learned discriminator produces better generalization—what specific features it is sensitive to, how its internal representations differ from the MSE loss's implicit metric—remains uninvestigated. The paper establishes that adversarial post-training works; it does not deeply characterize what the discriminator learns.
Follow-Up Research This Work Enables
Characterizing what the discriminator learns versus what the MSE loss ignores. The paper hypothesizes that the discriminator is "more sensitive to perceptual details, e.g. texture, sharpness, contour" (Section 1), but provides no direct evidence. A natural follow-up would probe the discriminator's internal representations: does the discriminator's Jacobian ∂D/∂x_t align with perceptually meaningful image features (edges, textures, semantic boundaries)? One could compute the singular vectors of the learned metric M(x_t, t) = (∂D/∂x_t)^⊤(∂D/∂x_t) at various points along the flow and visualize what image perturbations are most heavily penalized versus ignored. Comparing these to the uniform metric of the MSE loss would reveal what structures the discriminator has learned to care about. A stronger version of this experiment would train a discriminator on a specific class of artifacts (e.g., blur, JPEG compression, GAN-generated images) and test whether it generalizes to penalize those artifacts in flow model outputs—determining whether the discriminator learns a general perceptual metric or merely memorizes the generator's specific failure modes at each stage of training.
Testing CAFM on non-image continuous flow models. The paper's experiments are exclusively visual, but the JVP formulation is domain-agnostic: it requires only that the generator outputs a vector field and that the discriminator can take the state x_t as input. Molecular dynamics, fluid simulation, robotic control, and audio synthesis all use flow-based or diffusion-based generative models that suffer from the same generalization problem (the MSE objective is unaware of domain-specific validity constraints). A concrete experiment would apply CAFM post-training to a flow model for molecular conformer generation (e.g., an E(3)-equivariant flow model on atomic coordinates), where "perceptual quality" corresponds to physical validity (bond lengths, angles, non-bonded interactions). The discriminator could be trained on the same coordinates without architectural changes beyond the RMSNorm switch. The key metric would be the fraction of generated conformers that pass basic sanity checks (no atomic clashes, reasonable bond lengths) compared to the FM baseline—a domain-specific analog of FID that tests whether the discriminator learns physically meaningful constraints.
Combining CAFM with representation alignment methods. Table 3 shows that CAFM and REPA both improve over the SiT baseline in the latent space, but via different mechanisms: CAFM changes the training objective, while REPA [79] adds an auxiliary representation-alignment loss that encourages the model's internal features to match those of a pre-trained vision encoder. These approaches are potentially orthogonal and composable. A direct experiment would post-train an already-REPA-trained SiT model with CAFM, or jointly apply both objectives during fine-tuning. The prediction: if the methods improve different aspects of generation (REPA improves feature structure, CAFM improves perceptual fidelity), the combination should outperform either alone. The metric would be the guidance-free FID of SiT-XL/2+REPA+CAFM compared to each individually. A negative result (no improvement, or worse performance) would suggest that the discriminator and the representation alignment loss are competing for the same model capacity or inducing conflicting generalizations.
Investigating the optimal discriminator architecture and capacity. The paper uses the same architecture as the generator for the discriminator, with minimal modifications (RMSNorm, CLS token, projection layers). There is no study of how discriminator capacity affects results. Standard GAN wisdom holds that the discriminator should not be too strong relative to the generator, but CAFM's N=16 update schedule already makes the discriminator much more heavily optimized. A systematic sweep of discriminator size—half the generator's parameters, equal, double—would reveal whether the current equal-capacity choice is optimal or whether a smaller discriminator would suffice (reducing the 4.8× overhead). The experiment would train SiT-B/2 from scratch with different discriminator sizes, measuring both final FID and training stability. A finding that a half-size discriminator achieves comparable results would make CAFM more practical; a finding that larger discriminators continue to improve would suggest that discriminator quality, not generator capacity, is the current bottleneck.
Latency-aware compute-optimal allocation for CAFM inference. The paper does not discuss inference cost beyond FLOPs, but CAFM's sequential discriminator updates during training suggest an interesting question: could the discriminator be used at inference time to improve sampling? The discriminator's JVP D_jvp(x_t, t, G(x_t, t), 1) provides a per-step estimate of whether the generator's predicted velocity is "realistic" at each point along the ODE trajectory. This could be used as an adaptive step-size controller: when the discriminator's score drops below a threshold, reduce the step size (the flow may be entering a region where the generator is unreliable); when it is high, increase the step size. This would be analogous to adaptive ODE solvers that use local error estimates, but with a learned, perceptually-aware error metric. A concrete experiment: compare CAFM samples generated with 250 fixed SDE steps versus a variable-step solver that targets a constant discriminator JVP value, measuring whether equal or better FID can be achieved with fewer average steps. The paper's own inference uses the same fixed-step solvers as the baseline models, so this extension would directly build on the discriminator that is already trained and discarded.
Stress-testing CAFM on out-of-distribution or adversarial inputs. The discriminator is trained to distinguish real from generated velocities, but only on the distribution of x_t encountered during training—mixtures of real data and noise along the linear probability path. What happens when CAFM is asked to generate from initial noise that is far from Gaussian, or when the ODE is integrated with unusually large step sizes that push x_t into regions the discriminator has never seen? The concern is that the discriminator, having learned to enforce perceptual quality near the data manifold, may produce nonsensical gradients far from it, causing the generator to behave pathologically. A systematic evaluation would measure FID when sampling with non-standard noise distributions (uniform, low-rank Gaussian) or with aggressive step-size schedules, comparing CAFM against the FM baseline. If CAFM is robust, it suggests the discriminator's metric generalizes beyond the training distribution; if it collapses, it reveals a brittleness that users need to be aware of—and that might be addressed by training the discriminator on a wider range of x_t distributions.
Practical Applications and Downstream Use Cases
Improving existing flow-matching image generators without retraining. The most direct application is post-training any pre-trained flow-matching image model with CAFM for 10 epochs to improve guidance-free generation quality. For SiT-XL/2, this reduces FID from 8.26 to 3.63—a change that is visually dramatic (Figures 6, 7) and requires no architectural changes, no additional data, and no modification to the inference pipeline. A deployed system currently using SiT-XL/2 with CFG 1.5 could either (a) continue using CFG and benefit from the FID improvement at the same guidance level (2.06 → 1.97 at CFG 1.5), or (b) reduce guidance to CFG 1.3 while achieving better FID (1.53) and potentially better diversity (lower guidance preserves more of the data distribution). The 10-epoch cost is modest: at 4.8× per-epoch wall-clock time relative to standard training, 10 CAFM epochs cost roughly as much wall-clock time as 48 FM epochs—a one-time fine-tuning cost that pays off permanently in improved sample quality for all future inference.
Reducing or eliminating guidance dependence in text-to-image models. The text-to-image results (Tables 5, 6) demonstrate that CAFM substantially improves guidance-free generation: on GenEval, the overall score without CFG or prompt expansion improves from 0.33 to 0.44, with the Two Objects subscore jumping from 0.23 to 0.42. While these numbers are not production-quality (guided generation still substantially outperforms), they represent a significant step toward text-to-image models that produce diverse, prompt-aligned images without guidance—which would be valuable for applications where distribution fidelity matters more than peak perceptual quality, such as synthetic data generation for downstream model training, creative exploration tools that benefit from diversity, or scientific visualization where guidance-induced biases could be misleading. A practical deployment could use CAFM without guidance as a "diverse sampling" mode and add guidance only when a specific high-quality sample is requested, providing users with a fidelity-diversity tradeoff that is not available from standard flow-matching models.
Cost-efficient improvement of any Transformer-based flow model. The paper demonstrates CAFM on three architectures (SiT, JiT, Z-Image) with parameter counts ranging from 675M to 6B, in both latent and pixel space, using both SDE and ODE samplers. The consistent gains across this diverse set suggest that CAFM is not architecture-specific but is a general post-training recipe for transformer-based flow models. An organization with an existing flow-matching pipeline can adopt CAFM by: switching the discriminator's normalization to RMSNorm, adding a CLS token and projection head to the existing architecture, and running the Algorithm 1 training loop with the hyperparameters from Table 14. The discriminator is discarded after training, so there is no inference-time overhead. The main practical barrier is the 4.8× per-epoch training cost and the JVP implementation complexity (Appendix F discusses DDP/FSDP/gradient-checkpointing compatibility), but the paper provides a reference PyTorch implementation that can be adapted.
When to Prefer This Method
The paper explicitly positions CAFM as a post-training method for existing flow-matching models, not as a replacement for flow-matching pre-training or as an alternative to guidance. The decision points are:
Prefer CAFM post-training when:
- You have a pre-trained flow-matching model (particularly transformer-based, in visual domains) whose guidance-free generation quality is unsatisfactory, and you want to improve it without changing architectures, retraining from scratch, or modifying the inference pipeline.
- You need improved sample quality at lower guidance levels—CAFM shifts the optimal CFG scale downward (SiT: 1.5 → 1.3; JiT: 2.2 → 1.8) while achieving better or comparable FID, providing better diversity at a given quality level.
- You can afford approximately
48FM-epoch-equivalents of additional training compute (10 CAFM epochs at4.8×cost) as a one-time fine-tuning step. - You want to improve both fidelity and diversity simultaneously—the Precision and Recall improvements in Table 7 show CAFM does not trade one for the other.
Prefer standard flow matching (without CAFM post-training) when:
- You are training from scratch and compute budget is the primary constraint—CAFM from scratch converges more slowly than FM (Figure 4) and requires careful hyperparameter scheduling (Figure 5).
- Your model uses LayerNorm in the discriminator and you cannot switch to RMSNorm—Figure 12 shows LayerNorm causes gradient norm spikes during JVP computation that may destabilize training.
- Your inference pipeline is already achieving acceptable quality with guidance, and the additional training complexity (JVP implementation, distributed training compatibility, discriminator warm-up) is not justified by the marginal guided-generation improvements (JiT guided FID: 1.86 → 1.80, a modest gain).
Conditional on domain: The paper's results are exclusively in image generation (class-conditional and text-to-image). There is no evidence for or against CAFM's effectiveness in non-visual flow models, and the method's reliance on a discriminator that learns perceptual quality through adversarial training may depend on properties specific to natural images (hierarchical structure, translational equivariance, the existence of a "natural image manifold"). Extrapolation to other domains should be treated as speculative until tested.
