Amortized Moment Matching for Visual Generation

Unknown

Conceptio › Archive › arXiv CS

arXiv CSopen access

Amortized Moment Matching for Visual Generation

Unknown · 2026 · arxiv_cs

arXiv CS · Papers · License: Open Access · 2026

Open Source ↗Direct PDF ↓

machine learning, deep learning, neural networks

𝜇 fake, Σfake

Preprint 𝑥 fake

𝑥 real

A MORTIZED M OMENT M ATCHING FOR V ISUAL Polynomial Denoiser Projection (Theorem 3.1) G ENERATION Wenze Liu1,∗ Xintao Wang2 1 2

Pengfei Wan2

Xiangyu Yue1,†

MMLab, CUHK Kling Team, Kuaishou Technology Noise

A BSTRACT

arXiv:2607.26860v1 [cs.LG] 29 Jul 2026

𝜇 real, Σreal

Noise

Generator𝜃

FD-loss

Batch B

Generator𝜃

We propose amortized moment matching, utilizing neural networks Population N to learn data Amortizer update moments as distributional training signals. By casting diffusion denoisers through Real/Fake batch B … Batch Bframework polynomial projections, we establish a general … for moment amortizaCondition (class, text, …) tion, revealing that an n-th degree projection explicitly identifies data moments up to order n + 1. Derived from the tractable affine case, (a) weFD-loss instantiate the Amortized Fréchet Distance (AMFD) loss. Unlike FD-loss which relies on explicit marginal moment calculations, AMFD is able to dynamically learn conditional moments via an alternating, matrix-free optimization Noise pipeline that effortlessly Noise Batch B scales to high-dimensional data. When operating on global representation features, AMFD serves as a powerful post-training objective; empirically, its neural Generator𝜃 FD-loss Generator𝜃 formulation yields more robust training dynamics than exact statistical matching, 6 substantially surpassing the FD baseline on the FDr metric and Nachieving supePopulation rior one-step generation on ImageNet. Furthermore, it unlocks direct exploration Amortizer update … two moments can idenwithin native generative spaces, suggesting thatB the first Batch Real/Fake Batch B … tify target distributions only in spaces with strong semantics. Finally, when scaled to text-to-image generation, the condition-aware nature(a)ofFD-loss AMFD unlocks massive gains in instruction-following capabilities, enabling our one-step models to outperform their multi-step FLUX.2 [klein] 4B teachers on the GenEval benchmark while achieving on-par performance on PickScore. Code and checkpoints are available at https://github.com/poppuppy/amfd.

I NTRODUCTION

Distribution matching is a natural paradigm for generative modeling, aiming to directly align the generated distribution with the target distribution. Classical approaches adopt moment matching and maximum mean discrepancy (MMD) objectives (Sriperumbudur et al., 2010; Gretton et al., 2012; Li et al., 2015; Dziugaite et al., 2015; Mathiasen & Hvilshøj, 2020). However, they struggled to become competitive for visual generation, largely due to challenges in finite-sample estimation, kernel selection, and optimization stability (Li et al., 2017). In contrast, dominant paradigms such as diffusion and flow matching (Sohl-Dickstein et al., 2015; Song & Ermon, 2019; Ho et al., 2020; Liu et al., 2022a; Lipman et al., 2022), autoregressive models (Van Den Oord et al., 2016; Razavi et al., 2019; Esser et al., 2021; Sun et al., 2024), GANs (Goodfellow et al., 2014; Karras et al., 2019), and normalizing flows (Dinh et al., 2016; Kingma & Dhariwal, 2018; ∗ †

Work done at Kuaishou Technology. Corresponding author.

𝑒 ! = 𝑎" + 𝑎# 𝑥 + 𝑎$ 𝑥 $ + ⋯ + 𝑎% 𝑥 % + ⋯ 1 1 2! 1! Optimal diffusion denoiser: 1

1 𝑛!

𝒗∗ 𝒙, 𝑡 = 𝐴" + 𝐴# 𝒙⊗# + ⋯ + 𝐴% 𝒙⊗% + ⋯ mean cov

n+1th moment

Figure 1: Diffusion denoiser polynomial projection. The Maclaurin series expands ex into scalar polynomial terms; analogously, expanding the optimal diffusion denoiser via tensor powers yields highdimensional coefficients. We show that these coefficients explicitly map to the underlying distributional moments, where A0 corresponds to the mean, A1 to the covariance, and An to the (n + 1)-th moment.

Neura Amortiz

(b) AMFD

Neura Amortiz

(b) AMFD

Population N

Amortizer update Real/Fake batch B Condition (class, text, …)

… …

Batch B

Preprint

(a) FD-loss

Noise FD-loss

Generator step

Batch B

Generator𝜃

Neural Amortizer𝜓

Population N

… …

Batch B

Poly diff loss

(b) AMFD loss

Noise Generator𝜃

Neural Amortizer𝜓

Real/Fake Batch B

(a) FD-loss

Neural Amortizer𝜓

AMFD loss Amortizer step Poly Diff loss

(b) AMFD loss

Figure 2: Comparison between FD-loss and our AMFD loss. (a) FD-loss explicitly computes marginal distribution statistics over a large population, propagating gradients through mini-batches. (b) Our AMFD loss uses neural amortizers which can leverage neural generalization to naturally encode conditions. Through an alternating training scheme, these amortizers dynamically maintain the conditional mean and covariance statistics, which the generator directly queries to obtain gradient signals for updating. !

𝑎" + 𝑎# 𝑥 + 𝑎$ 𝑥 + ⋯ + 𝑎% 𝑥 + ⋯ Zhai et 𝑒al.,=2024) circumvent explicit distribution matching by adopting particle-level generative rules, allowing distributional 1 1 alignment 1 to emerge implicitly. 1

2! reveals that 𝑛! the historical weakness of direct distribution matching 1! Recent progress, however, Optimal diffusion denoiser: stemmed∗ largely from the ⊗# underlying data ⊗%spaces. Once shifted to semantically rich representation 𝒙, 𝑡 = 𝐴" matching + 𝐴# 𝒙 emerges + ⋯ + 𝐴%as𝒙an +⋯ spaces,𝒗distribution effective paradigm. For instance, Drifting models (Deng thevolve generated distributions, while the FD-loss (Yang et al., et al., 2026) utilize kernel forces to mean cov n+1 moment 2026) directly optimizes mean and covariance matching with a large population. Advancing the perspective of moment matching, we propose amortized moment matching. Rather than estimating moments purely from empirical batches or fixed global references, we amortize the statistics themselves: a neural amortizer predicts moment operators as dynamic functions of the conditioning signal and/or the target representation space. Our theoretical motivation stems from diffusion models, where a denoising model learns to approximate the conditional expectation of the clean data or the velocity field given the noisy observation. As illustrated in Figure 1, we analyze this denoiser function through polynomial projections: the affine projection recovers mean and covariance information, while higher-degree projections correspond to amortizing higher-order moments. This formulation offers a generalized route to moment amortization.

While high-order moment tensors grow exponentially in size, we instantiate a tractable and powerful second-order case termed the Amortized Fréchet Distance (AMFD) loss. When operating on global representation features, AMFD serves as an effective post-training objective that can enhance generative models and transform multi-step models into one-step generators. Crucially, in contrast to FD-loss (Yang et al., 2026) which computes aggregated marginal moments, AMFD injects the conditioning signal directly into the amortizer. This aligns naturally with modern diffusion models, where denoisers operate as conditional functions of labels or text prompts, enabling dynamic, condition-aware distribution matching. As illustrated in Figure 2(b), AMFD formulates this as a fully deep-learning-native pipeline driven by an alternating optimization between the generator and the neural amortizer. To avoid the prohibitive complexity of explicit covariance matrices, we parameterize the affine function via Jacobian-Vector-Product (JVP) actions. This matrix-free efficiency effortlessly scales to high dimensions, which additionally unlocks the exploration of whether matching merely the first two moments provides sufficient generative signals within native generative spaces (e.g., raw pixels, VAE/RAE latents (Rombach et al., 2022; Zheng et al., 2025)), where FD-loss is computationally intractable. Ultimately, AMFD strategically trades strict empirical exactness for superior conditional generalization and scalability. We evaluate AMFD in diverse visual generation tasks. First, in ImageNet-256×256 post-training, its neural formulation yields more robust optimization dynamics than the FD-loss baseline, achieving state-of-the-art FDr6 scores of 1.79 on JiT-H/16 and 1.75 on pMF-H/16. Second, exploration within native generative spaces suggests that the first two moments can identify target distributions only in spaces with strong semantics, such as the RAE space. Finally, when scaled to text-to-image generation with pixel-space PixelGen (Ma et al., 2026) and latent-space FLUX.2 [klein] 4B (Black Forest Labs, 2026), the condition-aware nature of AMFD unlocks massive gains in instruction-following capabilities. Notably, our one-step post-trained FLUX models significantly surpass their multi-step and few-step teachers on the GenEval benchmark, while maintaining on-par on PickScore. 2

Preprint

Figure 3: One-step text-to-image samples generated by our AMFD post-trained FLUX.2 [klein] 4B. Our main contributions are summarized as follows: • For the first time, we establish a rigorous connection between diffusion denoisers and data moments by proving that an n-th degree polynomial projection of a denoiser explicitly identifies moments up to order n + 1. Grounded in this theory, we introduce amortized moment matching, a scalable framework that learns and matches distributional moments via neural amortizers. • We instantiate the tractable second-order case as the Amortized Fréchet Distance (AMFD) loss. Unlike FD-loss that is restricted to unconditional marginals, AMFD leverages neural generalization to efficiently match fine-grained conditional moments—a capability crucial for scaling to complex text-to-image (T2I) generation. • Empirically, AMFD achieves state-of-the-art FDr6 for one-step generation on ImageNet256×256, reaching 1.79 (JiT-H/16) and 1.75 (pMF-H/16). Furthermore, when extended to text-toimage generation, our one-step FLUX.2 models surpass their multi-step teachers on the GenEval benchmark (0.846 vs. 0.794) while maintaining on-par performance on PickScore.

P RELIMINARIES

2.1

F R ÉCHET D ISTANCE

Let Pr and Pg denote the real and generated data distributions over the image space X . Given a frozen representation encoder ϕ : X → RD , we aim to compare the push-forward feature distributions ϕ# Pr and ϕ# Pg . We define their respective first two moments for i ∈ {r, g} as: µi = EX∼Pi [ϕ(X)], Σi = CovX∼Pi [ϕ(X)], (1) where X denotes the random variable in X . The squared Fréchet distance, which measures the discrepancy between their Gaussian approximations, is defined as: 1/2 1/2 d2FD (Pg , Pr ; ϕ) = ∥µg − µr ∥22 + Tr Σg + Σr − 2(Σ1/2 . (2) g Σr Σg ) When ϕ is Inception-v3 (Szegedy et al., 2016), this recovers the widely used Fréchet Inception Distance (FID) (Heusel et al., 2017). Recently, the FD-loss (Yang et al., 2026) proposed directly optimizing Eq. 2 across multiple representation spaces to train generative models. 2.2

D IFFUSION M ODELS

While our theoretical framework applies broadly to diffusion models, we adopt the flow matching formulation with linear interpolants (Liu et al., 2022a; Lipman et al., 2022; Albergo et al., 2023) for clarity. Consider a base distribution X0 ∼ P0 = N (0, I) and a target conditional data distribution X1 ∼ Pr (· | c). The linear interpolant is defined as Xt = (1 − t)X0 + tX1 for t ∈ [0, 1]. Flow matching trains a neural vector field vθ by minimizing the regression objective: h i 2 LFM (θ) = E ∥vθ (Xt , t, c) − (X1 − X0 )∥2 , 3

(3)

Preprint

which yields the population-level optimum: v ⋆ (x, t, c) = E[X1 − X0 | Xt = x, c].

(4)

This optimal vector field drives the probability flow ordinary differential equation (ODE) dXt = v ⋆ (Xt , t, c), dt

(5)

which deterministically transports the prior distribution P0 to the target distribution Pr (· | c).

A MORTIZED M OMENT M ATCHING : T HEORETICAL M OTIVATION

While the flow-matching objective Eq. 3 is conventionally interpreted as learning a transport vector field, we demonstrate from a distributional perspective that its optimal vector field (Eq. 4) fundamentally encodes conditional data statistics. Specifically, by analyzing the orthogonal projection of this denoiser onto polynomial function classes, we reveal a natural moment hierarchy. This connection allows us to isolate the affine case to derive a computationally tractable second-order matching objective. We formalize this insight below (see Appendix B for the rigorous formulation and proof). Theorem 3.1 (Denoising Projection Moment Hierarchy, Informal). Under the linear interpolant Xt = tX1 + (1 − t)X0 , fix a condition c and time t ∈ (0, 1). Let Pn denote the space of vectorvalued polynomials in Xt of degree at most n. The optimal degree-n polynomial denoiser vn⋆ = arg min E ∥v(Xt , t, c) − (X1 − X0 )∥2 | c , (6) v∈Pn

is exactly the orthogonal projection of the unconstrained flow-matching velocity v ⋆ onto Pn . Crucially, vn⋆ satisfies the normal equations, meaning it perfectly matches the moment tensors of the target velocity up to order n: E vn⋆ (Xt , t, c) ⊗ Xt⊗k | c = E (X1 − X0 ) ⊗ Xt⊗k | c := Gtk (P | c), ∀k = 0, . . . , n. (7) Writing the data moments as Mj (P | c) = E[X1⊗j | c], these matched tensors Gtk obey a triangular recurrence relation: Gtk (P | c) = tk Mk+1 (P | c) + Fkt (M0 , . . . , Mk ),

(8)

where Fkt is a deterministic function depending solely on t, the Gaussian moments of X0 , and the lower-order data moments {Mj }kj=0 . Consequently, optimizing the degree-n denoising objective uniquely determines the conditional data moments up to order n + 1. This theorem reveals that diffusion models inherently learn a hierarchy of amortized statistics, where the coefficients of the degree-zero, affine, and higher-degree terms correspond to the mean, covariance, and higher-order tensors, respectively. Since matching high-order tensors is computationally prohibitive in high-dimensional spaces, we isolate the affine case. This provides a tractable yet effective second-order instantiation, formalized as follows. Corollary 3.2 (Affine Projection Identifies Conditional Covariance). Under the setting of Theorem 3.1, assume X1 | c has a finite second moment. Let µ(c) = E[X1 | c] and Σ(c) = Cov(X1 | c). Then the degree-zero projection is exactly the conditional mean: v0⋆ (z, t, c) = µ(c),

(9)

and the degree-one projection admits the affine form: v1⋆ (z, t, c) = µ(c) + Rt (c) z − tµ(c) , where the affine operator is given by −1 Rt (c) = tΣ(c) − (1 − t)I t2 Σ(c) + (1 − t)2 I .

(10)

(11)

Preprint

T HE A MORTIZED F R ÉCHET D ISTANCE L OSS

Motivated by the insight that affine denoising projections uniquely encode conditional covariances, we propose the Amortized Fréchet Distance (AMFD) loss. We instantiate this matching process within a pre-trained representation space, as adopted by Yang et al. (2026). Below we detail the parameterization, followed by the full alternating training algorithm. 4.1

A MORTIZER PARAMETERIZATION

ℓ Let {ϕℓ }L ℓ=1 be a set of frozen representation encoders. For a data branch b ∈ {r, g}, let Xr = ϕℓ (Xr ) and Xgℓ = ϕℓ (Gθ (z, c)) denote the real and generated features in RDℓ . To capture their statistics, AMFD employs a neural amortizer that predicts a conditional mean µℓb (c) ∈ RDℓ and an ℓ affine operator, which we parameterize residually as Rb,t (c) = −I +tAℓb,t (c) for numerical stability.

For a target feature X1ℓ ∈ {Xrℓ , Xgℓ } and noise X0ℓ ∼ N (0, I), we construct the interpolant Xtℓ = tX1ℓ + (1 − t)X0ℓ . Substituting the residual parameterization, our amortized affine denoiser takes the form: ℓ v̂bℓ (Xtℓ , t, c) = µℓb (c) + Rb,t (c)X̃tℓ = µℓb (c) − X̃tℓ + tAℓb,t (c)X̃tℓ , (12) where X̃tℓ = Xtℓ − tµℓb (c) is the time-centered feature. Materializing Aℓb,t (c) explicitly would incur a prohibitive O(Dℓ2 ) memory cost in high-dimensional representations. We circumvent this by implicitly defining the operator via its JVP action on arbitrary ℓ vectors u ∈ RDℓ . Given a differentiable neural mapping fb,ψ (s, t, c) : RDℓ → RDℓ governed by an auxiliary input s, the exact operator action is efficiently computed as: Aℓb,t (c)u =

ℓ ∂fb,ψ (s, t, c) ∂s

ℓ u = JVPs fb,ψ (s, t, c) s=0 (u).

(13)

s=0

Detailed configurations for the JVP operation are provided in Appendix D.2 and D.4. 4.2

T RAINING A LGORITHM

Figure 2 (b) outlines the AMFD training procedure, which relies on an alternating optimization scheme (detailed alogrithm see Appendix C.1). The amortizers learn to accurately estimate the conditional real and generated statistics, while the generator is trained to minimize the discrepancy between them. To ensure unbiased scale aggregation across diverse representation spaces and prevent high-dimensional features from disproportionately dominating the objective, all per-encoder mean, covariance, and amortizer losses employ the normalized inner product ⟨a, b⟩ℓ = D1ℓ ⟨a, b⟩ and its induced norm ∥a∥2ℓ = ⟨a, a⟩ℓ . Amortizer update. For branch b ∈ {r, g}, define the target feature as ℓ hr , b = r, ℓ X1 = sg[hℓg ], b = g, where sg[·] denotes the stop-gradient operation. Given a sampled time t ∼ U (0, 1) and noise X0ℓ ∼ N (0, I), we form the linear interpolant Xtℓ and train the affine denoiser (Eq. 12) via: h i h i 2 2 Lℓb (ψ) = E v̂bℓ (Xtℓ , t, c) − (X1ℓ − X0ℓ ) ℓ + E µℓb (c) − X1ℓ ℓ . (14) At the population optimum, the diffusion matching term natively identifies the conditional mean and covariance functions. The explicit mean-regression term is theoretically redundant but practically serves as a regularization factor that consistently improves empirical convergence. The total amortizer objective sums over all branches and encoders: Lamort (ψ) =

L X

Lℓr (ψ) + Lℓg (ψ) .

ℓ=1

(15)

Preprint

Generator mean loss. The generator should modify its features along the direction that explicitly minimizes the mean discrepancy µℓg (c) − µℓr (c). We optimize this using the first-variation objective: Lℓµ (θ) = sg µℓg (c) − µℓr (c) , hℓg ℓ . (16) Crucially, gradients are only routed through the generated features hℓg , while the amortizer-predicted means act as frozen empirical targets. Generator covariance loss. For a given encoder ℓ, we define the centered generated feature as uℓ = hℓg − sg[µℓg (c)]. Ideally, the generator would directly minimize the quadratic discrepancy between the two conditional covariance matrices. Though the amortizer-predicted affine operator Aℓt (c) maps uniquely to Σℓ (c), directly substituting Aℓt (c) for Σℓ (c) is poorly conditioned. Instead, we project the discrepancy through the local Jacobian of the mapping Aℓt 7→ Σℓ evaluated ℓ at the generated branch. This yields the Jacobi operator Jg,t (c)u = (1 + t)u − t2 Aℓg,t (c)u, which defines the local pullback force: ℓ ℓ δaℓJ = Jg,t (c) Aℓg,t (c) − Aℓr,t (c) Jg,t (c)uℓ . (17) We then compute the generator covariance loss as: Lℓcov (θ) = ωcov (t) uℓ , δaℓJ ℓ ,

(18)

where ωcov (t) is a scaling coefficient (derivation detailed in Appendix C.2) required to approximately match the weighting of the FD objective (Eq. 2). During the generator update, the amortizers are frozen, but the internal JVP operator actions (Eq. 13) are dynamically tethered to uℓ , ensuring exact gradient flow to the generator. Notably, ignoring this local geometry yields a non-Jacobi variant δaℓnoJ = (Aℓg,t (c) − Aℓr,t (c))uℓ , which we empirically found to be significantly less effective, thereby validating our Jacobi-adjusted formulation. Multi-representation normalization. Aggregating losses across multiple heterogeneous encoders often produces generator gradients with highly disproportionate magnitudes. Similar to FDloss (Yang et al., 2026), we normalize each loss via a per-sample gradient proxy, thereby stabilizing the optimization process. For encoder ℓ, the proxy vector is defined as: ℓ gproxy = sg µℓg (c) − µℓr (c) + 2ωcov (t)δaℓJ . (19) The normalized per-encoder generator loss is subsequently computed as:   ℓ ℓ (θ) (θ) + L L cov µ , Lℓnorm (θ) = E  P 2 ρ Dℓ ℓ g + ϵ proxy,j j=1

(20)

where ρ and ϵ serve as scaling hyperparameters. The final unified generator objective aggregates these normalized losses across all L representation encoders: Lgen (θ) =

L X

Lℓnorm (θ).

(21)

ℓ=1

E XPERIMENTS

In this section, we empirically evaluate our proposed AMFD framework. First, we isolate and ablate the specific design choices of the amortizers. Second, we assess AMFD in ImageNet post-training and extend our study to native generative spaces. Finally, we investigate the transferability of the AMFD objective to text-to-image post-training across both pixel and latent spaces. 5.1

A MORTIZER I MPLEMENTATION

We instantiate one amortizer for each representation encoder, yielding a total of L amortizers with randomly initialized weights. For global feature representations, the amortizer is an AdaLNconditioned MLP (Li et al., 2024); in text-to-image models, a cross-attention module is integrated into every AdaLN block to process text conditioning. For native generative spaces, the amortizer 6

Preprint

Variant

MSEµ MSEΣ

AMFD-C 5.75 2.18 AMFD-U 0.0106 0.114

Figure 4: Estimation accuracy across interpolation time t.

Figure 5: Estimation accuracy across amortizer MLP sizes.

Table 1: Estimation accuracy with (AMFD-C) and without (AMFD-U) class conditioning (MSE ×10−3 ).

Pixel SD-VAE VA-VAE RAE FID = 188.61 FID = 117.07 FID = 84.82 FID = 3.13 > > > FDr6 = 150.17 FDr6 = 94.09 FDr6 = 76.48 FDr6 = 8.88

Figure 6: Impact of pretraining epochs on JiT-B/16 post-training.

Figure 7: Impact of pretraining epochs on JiT-L/16 post-training.

Figure 8: AMFD in native generative spaces. >

mirrors the generator’s architecture. We refer to this condition-aware configuration as AMFD-C (Conditional). For unconditional generation, denoted as AMFD-U (Unconditional), we retain the same network architecture but replace the conditioning input with a fixed dummy condition. The real and fake branches share identical amortizer parameters, relying exclusively on the time step t to distinguish the specific prediction task. More details are provided in Appendix D.3 and D.4. 5.2

A MORTIZER A BLATIONS

To isolate estimation accuracy, we train the amortizer solely on real Inception-v3 (Szegedy et al., 2016) features. We evaluate its performance using MSEµ and MSEΣ to measure the discrepancy against the analytical ground truth (details in Appendix D.3). Effect of interpolation time t. Unlike standard diffusion training that learns over the entire time continuum, Eq. 11 shows that the target moments can be perfectly identified at any fixed t ∈ (0, 1). By sweeping t across this interval, we observe that smaller t values yield consistently lower amortization errors for both the mean and covariance (Figure 4). Scaling model capacity. We evaluate estimation accuracy across varying MLP capacities. As depicted in Figure 5, increasing model size steadily reduces approximation error. However, as MSEµ plateaus and MSEΣ gains diminish at the c2048d16a4 scale, we adopt it as our default. Effect of class conditioning. As shown in Table 1, the unconditional setting (AMFD-U) achieves considerably higher estimation accuracy than its conditional counterpart (AMFD-C). This gap likely stems from the inherent simplicity of modeling a single global marginal distribution, compounded by the sparse per-class data density that limits statistical efficiency during conditional matching. 5.3

I MAGE N ET P OST-T RAINING

Superior dynamics and results when matching in global representation spaces. We evaluate AMFD for ImageNet post-training using both the multi-step JiT and one-step pMF models. As shown in Table 2, AMFD-U substantially improves the FDr6 score over the FD-loss baseline across both architectures (see Appendix E for full FDr metrics and Appendix F for qualitative samples). The table also reveals that AMFD-C slightly underperforms AMFD-U in final generation quality. We attribute this to two intertwined factors. First, amortizing unconditional marginal moments is an inherently easier optimization task than fitting fine-grained conditional moments, a difference in estimation accuracy explicitly corroborated in Table 1. Second, the discrete class labels in ImageNet represent a relatively simple semantic condition; in such low-complexity conditional spaces, the 7

Preprint

Table 2: System-level comparison for class-conditional generation on ImageNet 256×256. † indicates using CFG interval. FD-loss baseline and AMFD employ SIM (SigLIP, Inception, MAE). NFE

Params

FDr6 ↓

FID↓

IS↑

Precision↑

Recall↑

1.00

1.68

232.2

0.75

0.66

discrete discrete latent latent latent latent latent latent latent latent latent latent latent pixel pixel pixel pixel pixel pixel pixel pixel pixel pixel pixel pixel pixel

10 × 2 256 × 2×4 250 × 2 256 × 2×100 50 × 2† 256 × 2×100 256 × 2×100 250 × 2† 250 × 2† 250 × 2 250 × 2 250 × 2† 50 × 2† 100 × 2 50 × 2 × 2† 1 1 1 50 × 2 × 2† 1 1 1 50 × 2 × 2† 1 1 1

2B 1.1B 675M 478M 1.9B 942M 478M 685M 675M 675M 675M 676M 839M 1.0B 131M 131M 131M 131M 459M 459M 459M 459M 953M 953M 953M 953M

6.70 3.57 8.44 6.68 6.13 5.61 5.49 4.64 5.45 4.57 5.70 3.04 3.26 5.01 15.65 5.53 4.75 3.91 10.73 3.24 2.67 2.02 7.66 2.65 2.15 1.79

1.97 1.01 2.12 1.80 1.68 1.56 1.39 1.54 1.42 1.42 1.26 1.17 1.16 2.10 3.71 1.00 0.95 0.95 2.59 0.77 0.87 0.85 1.97 0.75 0.85 0.83

304.6 281.9 256.7 293.4 274.1 299.5 306.2 302.9 306.1 294.3 309.3 298.3 261.0 318.8 269.0 344.6 319.4 325.2 288.5 317.3 325.1 319.8 296.0 313.0 328.1 312.2

0.82 0.77 0.81 0.80 0.80 0.80 0.81 0.78 0.80 0.80 0.79 0.79 0.77 0.81 0.81 0.78 0.77 0.77 0.79 0.77 0.76 0.77 0.78 0.76 0.76 0.76

0.59 0.68 0.60 0.60 0.62 0.62 0.62 0.62 0.65 0.64 0.66 0.66 0.67 0.59 0.50 0.60 0.67 0.60 0.59 0.66 0.69 0.65 0.63 0.66 0.69 0.67

latent pixel latent latent pixel pixel pixel pixel pixel pixel pixel pixel pixel pixel pixel pixel

1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1

463M 465M 610M 610M 118M 118M 118M 118M 410M 410M 410M 410M 935M 935M 935M 935M

10.92 10.51 8.39 2.45 13.70 3.50 3.94 3.43 9.09 2.09 2.25 2.01 6.87 1.89 1.93 1.75

1.53 1.43 1.82 0.76 3.31 0.85 0.95 0.92 2.72 0.78 0.88 0.86 2.29 0.77 0.86 0.85

257.2 305.8 278.9 301.3 254.6 331.4 315.9 310.1 261.7 309.2 321.8 306.7 267.2 310.1 323.0 307.3

0.79 0.81 0.78 0.77 0.81 0.77 0.77 0.77 0.81 0.76 0.76 0.76 0.80 0.77 0.76 0.76

0.63 0.60 0.63 0.67 0.52 0.65 0.66 0.65 0.56 0.67 0.68 0.68 0.59 0.68 0.69 0.68

Method

Space

50k validation images

Multi-step models VAR-d30 (Tian et al., 2024) BAR-L (Yu et al., 2026) SiT-XL/2 (Ma et al., 2024) MAR-L (Li et al., 2024) FlowAR-H (Ren et al., 2024) MAR-H (Li et al., 2024) MAR-L, DeTok (Yang et al., 2025) REG (Wu et al., 2026) SiT-XL/2-REPA (Yu et al., 2024) LightningDiT (Yao et al., 2025) DDT-XL (Wang et al., 2026) REPA-E (Leng et al., 2025) RAE-XL (Zheng et al., 2025) PixNerd-XL (Wang et al., 2025) JiT-B (Li & He, 2026) + FD-loss (Yang et al., 2026) + AMFD-C (Ours) + AMFD-U (Ours) JiT-L (Li & He, 2026) + FD-loss (Yang et al., 2026) + AMFD-C (Ours) + AMFD-U (Ours) JiT-H (Li & He, 2026) + FD-loss (Yang et al., 2026) + AMFD-C (Ours) + AMFD-U (Ours) One-step models Drift-L (Deng et al., 2026) Drift-L (Deng et al., 2026) iMF-XL (Geng et al., 2026) + FD-loss (Yang et al., 2026) pMF-B (Lu et al., 2026) + FD (Yang et al., 2026) + AMFD-C (Ours) + AMFD-U (Ours) pMF-L (Lu et al., 2026) + FD-loss (Yang et al., 2026) + AMFD-C (Ours) + AMFD-U (Ours) pMF-H (Lu et al., 2026) + FD-loss (Yang et al., 2026) + AMFD-C (Ours) + AMFD-U (Ours)

marginal benefits of conditional matching are outweighed by the increased estimation difficulty. We further study the robustness of AMFD against FD-loss when the base model is far from a one-step generative regime by analyzing post-training convergence across varying pre-training checkpoints. Notably, while FD-loss completely diverges on early-stage checkpoints, both AMFD variants consistently converge to competitive FIDs (Figures 6 and 7). This demonstrates that the primary driver of robustness is the neural formulation itself. Furthermore, AMFD-C achieves strong FID scores with even fewer pre-training epochs than AMFD-U. This indicates that, despite being harder to fit perfectly, precise class-aware supervision provides further training stabilization. Exploration in native generative space. Empowered by the space complexity efficiency of AMFD, we further investigate whether matching solely the first and second moments within native generative spaces is sufficient to capture the target data distribution. Our evaluation spans four distinct environments with increasing semantic density: raw pixel space, SD-VAE (Rombach et al., 2022) latent space, VA-VAE (Yao et al., 2025) latent space, and RAE (Zheng et al., 2025) latent space. As detailed in Figure 8, both qualitative visualizations and quantitative metrics (FID and FDr6 ) reveal that generative performance improves progressively as the semantic strength of the space increases. Notably, while the lower-semantic spaces struggle to produce coherent structures, the RAE latent space successfully yields reasonable visual quality with a competitive FID of 3.13. This contrast 8

Preprint

Table 3: GenEval and PickScore for text-to-image post-training. Higher is better for all metrics. Row Method

NFE

Single Obj. Two Obj. Counting Colors Position Color Attr. Overall PickScore

1 2 3 4

Pixel space, multi-step PixelGen 25×2×2 FD-SIM (Yang et al., 2026) 1 AMFD-U-SIM (Ours) 1 AMFD-C-SIM (Ours) 1

0.994 0.984 0.975 0.969

0.886 0.795 0.831 0.864

0.578 0.416 0.469 0.513

0.902 0.753 0.769 0.846

0.708 0.535 0.658 0.723

0.698 0.415 0.578 0.645

0.794 0.650 0.713 0.760

21.02 20.33 20.68 20.80

5 6 7 8 9

VAE latent space, multi-step FLUX.2 [klein] 4B Base FD-SIM (Yang et al., 2026) AMFD-U-SIM (Ours) AMFD-C-SIM (Ours) AMFD-C-10 enc. (Ours)

50×2 1 1 1 1

1.000 1.000 0.997 1.000 1.000

0.864 0.934 0.917 0.957 0.962

0.759 0.653 0.644 0.750 0.766

0.899 0.824 0.827 0.910 0.886

0.505 0.448 0.523 0.733 0.720

0.642 0.598 0.590 0.650 0.685

0.778 0.743 0.749 0.833 0.836

21.50 21.31 21.37 21.41 21.46

10 11 12 13 14 15 16

VAE latent space, few-step FLUX.2 [klein] 4B DMD2 (Yin et al., 2025) FD-SIM (Yang et al., 2026) iRDM (Feng et al., 2026) AMFD-U-SIM (Ours) AMFD-C-SIM (Ours) AMFD-C-10 enc. (Ours)

4 1 1 1 1 1 1

0.994 0.997 1.000 0.994 0.994 0.997 1.000

0.904 0.894 0.944 0.924 0.929 0.955 0.957

0.791 0.806 0.716 0.756 0.741 0.791 0.778

0.880 0.864 0.878 0.923 0.902 0.923 0.920

0.575 0.603 0.618 0.650 0.638 0.670 0.680

0.623 0.660 0.658 0.708 0.678 0.740 0.733

0.794 0.804 0.802 0.826 0.813 0.846 0.845

21.85 21.62 21.82 21.77 21.85 21.82

suggests that matching first- and second-order statistics becomes progressively more effective as the representation space encodes richer and more structured semantics. 5.4

T EXT- TO -I MAGE P OST-T RAINING

We evaluate AMFD on text-to-image generation at 512×512 resolution using the GenEval (Ghosh et al., 2023) (instruction following) and PickScore (Kirstain et al., 2023) (human preference) benchmarks across diverse architectures: the pixel-space PixelGen (Ma et al., 2026), the latent-space FLUX.2 [klein] 4B Base (Black Forest Labs, 2026), and its 4-step distilled counterpart. As shown in Table 3, our AMFD establishes a new paradigm with several pivotal insights: Inherent stability in complex T2I. Even without conditioning, AMFD-U-SIM consistently outperforms the FD-SIM baseline across all architectures on both GenEval and PickScore (Rows 2 vs. 3, 6 vs. 7, and 12 vs. 14). This confirms that our amortized formulation inherently provides better training dynamics in highly complex text-to-image scenarios. The critical role of conditioning. When post-training using a few-step teacher (which inherently operates closer to a one-step generator and already possesses robust conditional alignment), AMFDU-SIM performs remarkably well (Row 14). However, when using a standard multi-step diffusion model, condition-aware matching becomes absolutely crucial. For instance, in the multi-step FLUX.2 Base, AMFD-C-SIM dramatically boosts the GenEval score from 0.749 to 0.833 compared to its unconditional counterpart (Rows 7 vs. 8), unlocking similarly massive gains in PixelGen (0.713 to 0.760, Rows 3 vs. 4). This completely validates the necessity of condition-aware matching for standard diffusion distillation—a stark contrast to our ImageNet findings, driven by the vastly higher semantic complexity of text prompts compared to class labels. Sufficiency of the first two moments. FD-loss (Yang et al., 2026) has empirically established that matching the first two statistical moments yields strong generative performance. We further corroborate this by comparing our AMFD (Rows 15 & 16) with the recent iRDM (Feng et al., 2026) (Row 13), which theoretically matches all statistical moments via an MMD criterion. By achieving superior GenEval scores and equal or better PickScores, our results further demonstrate that, when empowered by neural amortization, aligning merely the first two moments already provides abundant supervisory signals. Effectiveness of the compact SIM set. Expanding the representation set from SIM to 10 encoders (Table 4, similar to Feng et al. (2026)) yields comparable performance across both FLUX.2 [klein] models (Rows 8 vs. 9; 15 vs. 16). While standard metrics may not capture all real-world visual nuances, this suggests the lightweight SIM ensemble already provides sufficient features for strong text-to-image alignment, offering a practical balance of efficiency and quality. 9

Preprint

Ultimately, these combined advantages allow our one-step AMFD models to significantly surpass both the multi-step FLUX.2 [klein] 4B Base teacher and its 4-step distilled variant on GenEval, while achieving highly competitive, on-par performance on PickScore (e.g., matching the 21.85 PickScore of the 4-step teacher). This demonstrates the formidable capability of our method. Qualitative results are in Figure 3 and Appendix Figure 11.

L IMITATIONS

While effective, our study presents several limitations. First, although our theory accommodates arbitrary-order moments, AMFD is restricted to matching only the first two moments, leaving higher-order statistics uncaptured. Second, our reliance on off-the-shelf representation encoders, which are typically pretrained at low resolutions, bottlenecks the direct applicability to highresolution generation. Finally, AMFD serves as a post-training objective, while it is currently unable to train generative models entirely from scratch.

C ONCLUSION

In this work, we introduced amortized moment matching, a principled framework that leverages neural amortizers to provide scalable distributional training signals for generative models. By analyzing diffusion denoisers through the lens of polynomial projections, we established that they inherently learn a hierarchy of conditional statistics. We instantiated the tractable second-order case as the Amortized Fréchet Distance (AMFD) loss, a post-training paradigm designed to elevate the synthesis quality of existing models and convert multi-step models into efficient one-step generators. By efficiently matching conditional means and covariance operators, AMFD yields superior one-step generation performance in ImageNet post-training. Furthermore, it scales seamlessly to complex text-to-image architectures across both pixel and VAE latent spaces, demonstrating strong instruction-following proficiencies and high-fidelity generation.

R EFERENCES Michael S Albergo, Nicholas M Boffi, and Eric Vanden-Eijnden. Stochastic interpolants: A unifying framework for flows and diffusions. arXiv preprint arXiv:2303.08797, 2023. Black Forest Labs. FLUX.2 [klein]: Towards Interactive Visual Intelligence. https://bfl. ai/blog/flux2-klein-towards-interactive-visual-intelligence, 2026. Model weights: https://huggingface.co/collections/black-forest-labs/ flux2. Nicholas M Boffi, Michael S Albergo, and Eric Vanden-Eijnden. Flow map matching. arXiv preprint arXiv:2406.07507, 2024. Daniel Bolya, Po-Yao Huang, Peize Sun, Jang Hyun Cho, Andrea Madotto, Chen Wei, Tengyu Ma, Jiale Zhi, Jathushan Rajasegaran, Hanoona Bangalath, et al. Perception encoder: The best visual embeddings are not at the output of the network. Advances in Neural Information Processing Systems, 38:60884–60937, 2026. Benjamin Boys, Mark Girolami, Jakiw Pidstrigach, Sebastian Reich, Alan Mosca, and O Deniz Akyildiz. Tweedie moment projected diffusions for inverse problems. arXiv preprint arXiv:2310.06721, 2023. Jiuhai Chen, Zhiyang Xu, Xichen Pan, Yushi Hu, Can Qin, Tom Goldstein, Lifu Huang, Tianyi Zhou, Saining Xie, Silvio Savarese, et al. Blip3-o: A family of fully open unified multimodal models-architecture, training and dataset. arXiv preprint arXiv:2505.09568, 2025a. Junying Chen, Zhenyang Cai, Pengcheng Chen, Shunian Chen, Ke Ji, Xidong Wang, Yunjin Yang, and Benyou Wang. Sharegpt-4o-image: Aligning multimodal models with gpt-4o-level image generation. arXiv preprint arXiv:2506.18095, 2025b. Mingyang Deng, He Li, Tianhong Li, Yilun Du, and Kaiming He. Generative modeling via drifting. arXiv preprint arXiv:2602.04770, 2026. 10

Preprint

Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. Density estimation using real nvp. arXiv preprint arXiv:1605.08803, 2016. Gintare Karolina Dziugaite, Daniel M Roy, and Zoubin Ghahramani. Training generative neural networks via maximum mean discrepancy optimization. arXiv preprint arXiv:1505.03906, 2015. Patrick Esser, Robin Rombach, and Bjorn Ommer. Taming transformers for high-resolution image synthesis. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 12873–12883, 2021. David Fan, Shengbang Tong, Jiachen Zhu, Koustuv Sinha, Zhuang Liu, Xinlei Chen, Michael Rabbat, Nicolas Ballas, Yann LeCun, Amir Bar, et al. Scaling language-free visual representation learning. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 370–382, 2025. Lan Feng, Wuyang Li, Eloi Zablocki, Matthieu Cord, and Alexandre Alahi. Representation distribution matching for one-step visual generation. arXiv preprint arXiv:2607.02375, 2026. Enrico Fini, Mustafa Shukor, Xiujun Li, Philipp Dufter, Michal Klein, David Haldimann, Sai Aitharaju, Victor G Turrisi da Costa, Louis Béthune, Zhe Gan, et al. Multimodal autoregressive pre-training of large vision encoders. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 9641–9654, 2025. Stephanie Fu, Netanel Tamir, Shobhita Sundaram, Lucy Chai, Richard Zhang, Tali Dekel, and Phillip Isola. Dreamsim: Learning new dimensions of human visual similarity using synthetic data. arXiv preprint arXiv:2306.09344, 2023. Zhengyang Geng, Mingyang Deng, Xingjian Bai, J Zico Kolter, and Kaiming He. Mean flows for one-step generative modeling. arXiv preprint arXiv:2505.13447, 2025. Zhengyang Geng, Yiyang Lu, Zongze Wu, Eli Shechtman, J Zico Kolter, and Kaiming He. Improved mean flows: On the challenges of fastforward generative models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 30467–30476, 2026. Dhruba Ghosh, Hannaneh Hajishirzi, and Ludwig Schmidt. Geneval: An object-focused framework for evaluating text-to-image alignment. Advances in Neural Information Processing Systems, 36: 52132–52152, 2023. Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. Advances in neural information processing systems, 27, 2014. Arthur Gretton, Karsten M Borgwardt, Malte J Rasch, Bernhard Schölkopf, and Alexander Smola. A kernel two-sample test. The journal of machine learning research, 13(1):723–773, 2012. Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked autoencoders are scalable vision learners. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 16000–16009, 2022. Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge to a local nash equilibrium. Advances in neural information processing systems, 30, 2017. Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840–6851, 2020. Tero Karras, Samuli Laine, and Timo Aila. A style-based generator architecture for generative adversarial networks. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 4401–4410, 2019. Tero Karras, Miika Aittala, Jaakko Lehtinen, Janne Hellsten, Timo Aila, and Samuli Laine. Analyzing and improving the training dynamics of diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 24174–24184, 2024. 11

Preprint

Durk P Kingma and Prafulla Dhariwal. Glow: Generative flow with invertible 1x1 convolutions. Advances in neural information processing systems, 31, 2018. Yuval Kirstain, Adam Polyak, Uriel Singer, Shahbuland Matiana, Joe Penna, and Omer Levy. Picka-pic: An open dataset of user preferences for text-to-image generation. Advances in neural information processing systems, 36:36652–36663, 2023. Xingjian Leng, Jaskirat Singh, Yunzhong Hou, Zhenchang Xing, Saining Xie, and Liang Zheng. Repa-e: Unlocking vae for end-to-end tuning of latent diffusion transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 18262–18272, 2025. Chun-Liang Li, Wei-Cheng Chang, Yu Cheng, Yiming Yang, and Barnabás Póczos. Mmd gan: Towards deeper understanding of moment matching network. Advances in neural information processing systems, 30, 2017. Tianhong Li and Kaiming He. Back to basics: Let denoising generative models denoise. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 36115– 36125, 2026. Tianhong Li, Yonglong Tian, He Li, Mingyang Deng, and Kaiming He. Autoregressive image generation without vector quantization. Advances in Neural Information Processing Systems, 37: 56424–56445, 2024. Yujia Li, Kevin Swersky, and Rich Zemel. Generative moment matching networks. In International conference on machine learning, pp. 1718–1727. PMLR, 2015. Shanchuan Lin, Xin Xia, Yuxi Ren, Ceyuan Yang, Xuefeng Xiao, and Lu Jiang. Diffusion adversarial post-training for one-step video generation. arXiv preprint arXiv:2501.08316, 2025. Shanchuan Lin, Ceyuan Yang, Zhijie Lin, Hao Chen, and Haoqi Fan. Continuous adversarial flow models. arXiv preprint arXiv:2604.11521, 2026. Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXiv:2210.02747, 2022. Wenze Liu and Xiangyu Yue. Learning to integrate diffusion odes by averaging the derivatives. arXiv preprint arXiv:2505.14502, 2025. Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprint arXiv:2209.03003, 2022a. Zhuang Liu, Hanzi Mao, Chao-Yuan Wu, Christoph Feichtenhofer, Trevor Darrell, and Saining Xie. A convnet for the 2020s. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 11976–11986, 2022b. Ilya Loshchilov and Frank Hutter. arXiv:1711.05101, 2017.

Decoupled weight decay regularization.

arXiv preprint

Yiyang Lu, Susie Lu, Qiao Sun, Hanhong Zhao, Zhicheng Jiang, Xianbang Wang, Tianhong Li, Zhengyang Geng, and Kaiming He. One-step latent-free image generation with pixel mean flows. arXiv preprint arXiv:2601.22158, 2026. Eric Luhman and Troy Luhman. Knowledge distillation in iterative generative models for improved sampling speed. arXiv preprint arXiv:2101.02388, 2021. Nanye Ma, Mark Goldstein, Michael S Albergo, Nicholas M Boffi, Eric Vanden-Eijnden, and Saining Xie. Sit: Exploring flow and diffusion-based generative models with scalable interpolant transformers. In European Conference on Computer Vision, pp. 23–40. Springer, 2024. Zehong Ma, Ruihan Xu, and Shiliang Zhang. Pixelgen: Pixel diffusion beats latent diffusion with perceptual loss. arXiv preprint arXiv:2602.02493, 2026. Hila Manor and Tomer Michaeli. On the posterior distribution in denoising: Application to uncertainty quantification. In International Conference on Learning Representations, volume 2024, pp. 49233–49263, 2024. 12

Preprint

Alexander Mathiasen and Frederik Hvilshøj. Backpropagating through fr\’echet inception distance. arXiv preprint arXiv:2009.14075, 2020. Maxime Oquab, Timothée Darcet, Théo Moutakanni, Huy Vo, Marc Szafraniec, Vasil Khalidov, Pierre Fernandez, Daniel Haziza, Francisco Massa, Alaaeldin El-Nouby, et al. Dinov2: Learning robust visual features without supervision. arXiv preprint arXiv:2304.07193, 2023. Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International conference on machine learning, pp. 8748–8763. PmLR, 2021. Ali Razavi, Aaron Van den Oord, and Oriol Vinyals. Generating diverse high-fidelity images with vq-vae-2. Advances in neural information processing systems, 32, 2019. Sucheng Ren, Qihang Yu, Ju He, Xiaohui Shen, Alan Yuille, and Liang-Chieh Chen. Flowar: Scalewise autoregressive image generation meets flow matching. arXiv preprint arXiv:2412.15205, 2024. Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. Highresolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 10684–10695, 2022. Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusion models. arXiv preprint arXiv:2202.00512, 2022. Axel Sauer, Kashyap Chitta, Jens Müller, and Andreas Geiger. Projected gans converge faster. Advances in Neural Information Processing Systems, 34:17480–17492, 2021. Axel Sauer, Frederic Boesel, Tim Dockhorn, Andreas Blattmann, Patrick Esser, and Robin Rombach. Fast high-resolution image synthesis with latent adversarial diffusion distillation. In SIGGRAPH Asia 2024 Conference Papers, pp. 1–11, 2024. Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International conference on machine learning, pp. 2256–2265. pmlr, 2015. Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems, 32, 2019. Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. arXiv preprint arXiv:2303.01469, 2023. Bharath K Sriperumbudur, Arthur Gretton, Kenji Fukumizu, Bernhard Schölkopf, and Gert RG Lanckriet. Hilbert space embeddings and metrics on probability measures. The Journal of Machine Learning Research, 11:1517–1561, 2010. Peize Sun, Yi Jiang, Shoufa Chen, Shilong Zhang, Bingyue Peng, Ping Luo, and Zehuan Yuan. Autoregressive model beats diffusion: Llama for scalable image generation. arXiv preprint arXiv:2406.06525, 2024. Peng Sun, Zhenglin Cheng, Deyuan Liu, Jun Xie, Xinyi Shang, and Tao Lin. Three-body scattering for generative modeling. arXiv preprint arXiv:2607.18198, 2026. Christian Szegedy, Vincent Vanhoucke, Sergey Ioffe, Jon Shlens, and Zbigniew Wojna. Rethinking the inception architecture for computer vision. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 2818–2826, 2016. Keyu Tian, Yi Jiang, Zehuan Yuan, Bingyue Peng, and Liwei Wang. Visual autoregressive modeling: Scalable image generation via next-scale prediction. Advances in neural information processing systems, 37:84839–84865, 2024. 13

Preprint

Michael Tschannen, Alexey Gritsenko, Xiao Wang, Muhammad Ferjad Naeem, Ibrahim Alabdulmohsin, Nikhil Parthasarathy, Talfan Evans, Lucas Beyer, Ye Xia, Basil Mustafa, et al. Siglip 2: Multilingual vision-language encoders with improved semantic understanding, localization, and dense features. arXiv preprint arXiv:2502.14786, 2025. Aäron Van Den Oord, Nal Kalchbrenner, and Koray Kavukcuoglu. Pixel recurrent neural networks. In International conference on machine learning, pp. 1747–1756. PMLR, 2016. Shuai Wang, Ziteng Gao, Chenhui Zhu, Weilin Huang, and Limin Wang. Pixnerd: Pixel neural field diffusion. arXiv preprint arXiv:2507.23268, 2025. Shuai Wang, Zhi Tian, Weilin Huang, and Limin Wang. Ddt: Decoupled diffusion transformer. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 40633–40642, 2026. Zhengyi Wang, Cheng Lu, Yikai Wang, Fan Bao, Chongxuan Li, Hang Su, and Jun Zhu. Prolificdreamer: High-fidelity and diverse text-to-3d generation with variational score distillation. Advances in Neural Information Processing Systems, 36:8406–8441, 2023. Ge Wu, Shen Zhang, Ruijing Shi, Shanghua Gao, Zhenyuan Chen, Lei Wang, Zhaowei Chen, Hongcheng Gao, Yao Tang, Ming-Ming Cheng, et al. Representation entanglement for generation: Training diffusion transformers is much easier than you think. Advances in Neural Information Processing Systems, 38:7714–7743, 2026. Jiawei Yang, Tianhong Li, Lijie Fan, Yonglong Tian, and Yue Wang. Latent denoising makes good visual tokenizers. arXiv preprint arXiv:2507.15856, 2025. Jiawei Yang, Zhengyang Geng, Xuan Ju, Yonglong Tian, and Yue Wang. Representation fr\’echet loss for visual generation. arXiv preprint arXiv:2604.28190, 2026. Jingfeng Yao, Bin Yang, and Xinggang Wang. Reconstruction vs. generation: Taming optimization dilemma in latent diffusion models. In Proceedings of the Computer Vision and Pattern Recognition Conference, pp. 15703–15712, 2025. Junyan Ye, Dongzhi Jiang, Zihao Wang, Leqi Zhu, Zhenghao Hu, Zilong Huang, Jun He, Zhiyuan Yan, Jinghua Yu, Hongsheng Li, et al. Echo-4o: Harnessing the power of gpt-4o synthetic images for improved image generation. arXiv preprint arXiv:2508.09987, 2025. Tianwei Yin, Michaël Gharbi, Richard Zhang, Eli Shechtman, Fredo Durand, William T Freeman, and Taesung Park. One-step diffusion with distribution matching distillation. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 6613–6623, 2024. Tianwei Yin, Michaël Gharbi, Taesung Park, Richard Zhang, Eli Shechtman, Fredo Durand, and Bill Freeman. Improved distribution matching distillation for fast image synthesis. Advances in Neural Information Processing Systems, 37:47455–47487, 2025. Qihang Yu, Qihao Liu, Ju He, Xinyang Zhang, Yang Liu, Liang-Chieh Chen, and Xi Chen. Autoregressive image generation with masked bit modeling. arXiv preprint arXiv:2602.09024, 2026. Sihyun Yu, Sangkyung Kwak, Huiwon Jang, Jongheon Jeong, Jonathan Huang, Jinwoo Shin, and Saining Xie. Representation alignment for generation: Training diffusion transformers is easier than you think. arXiv preprint arXiv:2410.06940, 2024. Shuangfei Zhai, Ruixiang Zhang, Preetum Nakkiran, David Berthelot, Jiatao Gu, Huangjie Zheng, Tianrong Chen, Miguel Angel Bautista, Navdeep Jaitly, and Josh Susskind. Normalizing flows are capable generative models. arXiv preprint arXiv:2412.06329, 2024. Richard Zhang, Phillip Isola, Alexei A Efros, Eli Shechtman, and Oliver Wang. The unreasonable effectiveness of deep features as a perceptual metric. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 586–595, 2018. Boyang Zheng, Nanye Ma, Shengbang Tong, and Saining Xie. Diffusion transformers with representation autoencoders. arXiv preprint arXiv:2510.11690, 2025. 14

Preprint

Linqi Zhou, Stefano Ermon, and Jiaming Song. Inductive moment matching. arXiv preprint arXiv:2503.07565, 2025a. Linqi Zhou, Mathias Parger, Ayaan Haque, and Jiaming Song. Terminal velocity matching. arXiv preprint arXiv:2511.19797, 2025b.

Preprint

A PPENDIX A

R ELATED W ORK

A.1

D IFFUSION AND F LOW M ATCHING

Diffusion models (Sohl-Dickstein et al., 2015; Song & Ermon, 2019; Ho et al., 2020) and flow matching (Liu et al., 2022a; Lipman et al., 2022) typically learn unrestricted denoising functions or velocity fields to iteratively transport samples from a prior to the data distribution. From a statistical perspective, classical Tweedie identities connect the optimal Gaussian denoiser to posterior moments conditioned on a particular noisy observation (Manor & Michaeli, 2024; Boys et al., 2023). Our perspective differs fundamentally: by restricting the denoising regression to polynomial function classes, we study its projection coefficients. Crucially, these coefficients depend on the conditional data distribution, rather than on an individual noised sample, and form a hierarchy that identifies increasingly high-order data moments. Accordingly, we do not use the learned denoiser for generative sampling. Instead, we utilize these amortized coefficients as explicit distributional training signals for a separate generator. A.2

M OMENT M ATCHING IN R EPRESENTATION S PACES

Moment matching provides a direct approach to generative modeling by comparing statistics of real and generated distributions. Traditional methods rely on nonparametric frameworks, optimizing empirical metrics such as kernel mean embeddings and maximum mean discrepancy (MMD) from finite sample sets (Sriperumbudur et al., 2010; Gretton et al., 2012; Li et al., 2015; Dziugaite et al., 2015; Li et al., 2017). A complementary line of work demonstrates that pretrained visual encoders provide highly effective spaces for generative supervision. Perceptual objectives (Zhang et al., 2018) compare paired images using frozen deep features, projected GANs (Sauer et al., 2021) construct discriminators on top of pretrained representations, representation alignment methods (Yu et al., 2024) supervise internal generative features, and representation encoders (Zheng et al., 2025) train generative models natively in representation spaces. Recent work has begun to directly optimize discrepancies between real and generated distributions in representation spaces. Drifting Models (Deng et al., 2026) construct a kernel-induced field that evolves the generated distribution during training. FD-loss (Yang et al., 2026) demonstrates that representation-space means and covariances can be explicitly optimized by decoupling population-level moment estimation from generator minibatches. Several concurrent works have also explored representation-space distribution matching to enable strong one-step generation. Specifically, Representation Distribution Matching (RDM) (Feng et al., 2026) utilizes MMD estimators, while Three-Body Scattering Modeling (TBSM) (Sun et al., 2026) employs an online tracker to learn an energy-distance descent field. Rather than estimating empirical marginal statistics or learning sample-dependent transport fields, AMFD amortizes condition-level moment coefficients derived from polynomial projections of a diffusion denoiser. A.3

D ISTILLATION AND P OST-T RAINING FOR O NE -S TEP G ENERATION

State-of-the-art generative paradigms, such as diffusion models (Sohl-Dickstein et al., 2015; Song & Ermon, 2019; Ho et al., 2020) and flow matching (Liu et al., 2022a; Lipman et al., 2022), achieve unprecedented synthesis quality but typically rely on slow, iterative generation processes. To address this, generative distillation and post-training techniques have been extensively developed to compress these iterative processes into efficient one-step generators. Existing approaches can be broadly categorized into several directions. Direct distillation methods (Luhman & Luhman, 2021; Salimans & Ho, 2022) explicitly matches the input-output or intermediate mappings of a multistep teacher. Trajectory consistency methods (Song et al., 2023; Boffi et al., 2024; Liu & Yue, 2025; Geng et al., 2025; Zhou et al., 2025a) enforce consistency along the teacher’s deterministic ODE trajectory. GAN-based post-training methods (Sauer et al., 2024; Lin et al., 2025) leverage adversarial objectives to directly align the one-step generated distribution with the real data distribution. Furthermore, variational score distillation (VSD) (Wang et al., 2023; Yin et al., 2024) optimize the one-step generator by estimating and matching the score functions of the target and generated distributions. Although AMFD shares the alternating optimization structure of adversarial and VSD-based post-training, its amortizers serve a fundamentally different statistical role. A 16

Preprint

GAN discriminator learns to classify real versus generated samples, and VSD-style networks estimate the full score fields of intermediate noisy distributions. In contrast, AMFD neural amortizers directly regress specific statistical targets: conditional means and covariance-dependent operators. Consequently, AMFD bypasses traditional adversarial dynamics and the necessity to reproduce a fixed teacher trajectory, offering a principled, condition-aware approach to one-step generation.

P ROOFS

In this section, we present the formal statement and proof of Theorem 3.1, alongside the proof of Corollary 3.2. Theorem B.1 (Denoising Projection Moment Hierarchy, Formal Statement). Fix a condition c and t ∈ (0, 1). Let X1 ∼ P (· | c) and X0 ∼ N (0, I) be independent variables. Define the interpolant Z = tX1 + (1 − t)X0 and the target velocity V = X1 − X0 , assuming E[∥V ∥2 | c] < ∞. Let H = L2 (PZ|c ; RD ) be the Hilbert space of square-integrable vector fields, wherein the unconstrained flow-matching solution is given by v ⋆ = E[V | Z, c]. For any integer n ≥ 0, let Hn ⊂ H be the subspace of vector-valued polynomials in Z up to degree n. The optimal degree-n polynomial denoiser, (22) vn⋆ = arg min E ∥v(Z) − V ∥2 | c , v∈Hn

satisfies vn⋆ = Πn v ⋆ , where Πn is the orthogonal projection onto Hn . If the polynomial space S ⋆ H n is dense in H, these projections converge strongly to v , admitting the orthogonal decomn≥0 P ∞ ⋆ position v ⋆ = n=0 ∆n in L2 (PZ|c ), where ∆0 = v0⋆ and ∆n = vn⋆ − vn−1 . Furthermore, by the first-order optimality over Hn , vn⋆ satisfies the normal equations, thereby perfectly matching the target moment tensors of the unconstrained velocity: E vn⋆ (Z) ⊗ Z ⊗k | c = Gtk (P | c) := E V ⊗ Z ⊗k | c , k = 0, . . . , n. (23) Writing the j-th conditional data moment as Mj (P | c) = E[X1⊗j | c], the target tensors Gtk obey the triangular recurrence relation: Gtk (P | c) = tk Mk+1 (P | c) + Fkt (M0 , . . . , Mk ),

(24)

where Fkt is a deterministic tensor-valued function depending solely on t, the known Gaussian moments of X0 , and the lower-order data moments. Consequently, the optimal degree-n polynomial denoiser uniquely determines the conditional data moments up to order n + 1. Proof. We begin by establishing the orthogonal projection property and the resulting decomposition. For an arbitrary v ∈ H, we decompose the residual as: v(Z) − V = v(Z) − v ⋆ (Z) + v ⋆ (Z) − V .

(25)

By definition, v ⋆ (Z) = E[V | Z, c], which immediately implies E[V − v ⋆ (Z) | Z, c] = 0. Consequently, the cross term vanishes in expectation: h ⊤ ⋆ i E v(Z) − v ⋆ (Z) v (Z) − V | c = 0. (26) This decouples the mean squared error into: E ∥v(Z) − V ∥2 | c = E ∥v(Z) − v ⋆ (Z)∥2 | c + E ∥v ⋆ (Z) − V ∥2 | c .

(27)

Since the second term is independent of v, minimizing the flow-matching objective over Hn is mathematically equivalent to minimizing the distance ∥v − v ⋆ ∥2H over v ∈ Hn . Thus, the minimizer is exactly the orthogonal projection: vn⋆ = Πn v ⋆ . ⋆ To establish the orthogonality of the increments ∆n = vn⋆ − vn−1 , consider any h ∈ Hn−1 . Since vn⋆ = Πn v ⋆ and Hn−1 ⊂ Hn , the projection residual v ⋆ − vn⋆ is orthogonal to h. Similarly, ⋆ ⋆ v ⋆ − vn−1 ⊥ h. Subtracting these equations yields ⟨vn⋆ − vn−1 , h⟩H = 0, implying ∆n ⊥ Hn−1 . Because ∆m ∈ Hm ⊂ Hn−1 for any m < n, we obtain ⟨∆n , ∆m ⟩H = 0. Under the assumption

Preprint

that the polynomial spaceP is dense in H, the orthogonal projections converge strongly, yielding the ∞ valid infinite series v ⋆ = n=0 ∆n in L2 (PZ|c ). Having established the projection properties, we now turn to the moment hierarchy and the triangular recurrence relation. Since Hn is spanned by the tensor polynomials {Z ⊗k }nk=0 , the condition that the residual (vn⋆ − V ) is orthogonal to Hn in L2 (PZ|c ) is equivalent to testing against the basis elements. Thus, for any k ∈ {0, . . . , n}: ⟨vn⋆ − V, z 7→ z ⊗k ⟩H = 0 =⇒ E vn⋆ (Z) − V ⊗ Z ⊗k | c = 0. (28) Rearranging this yields the moment matching property: E[vn⋆ (Z) ⊗ Z ⊗k | c] = E[V ⊗ Z ⊗k | c] = Gtk (P | c). To establish the triangular recurrence, we explicitly expand Gtk (P | c). Recalling V = X1 − X0 and Z = tX1 + (1 − t)X0 , we write: h ⊗k i Gtk (P | c) = E (X1 − X0 ) ⊗ tX1 + (1 − t)X0 |c . (29) By the non-commutative binomial expansion of the tensor power, the highest-order term of X1 inside the parentheses is tk X1⊗k . Distributing (X1 − X0 ) into this expansion, the term with the ⊗(k+1) highest tensor power of X1 is generated solely by X1 ⊗ (tk X1⊗k ) = tk X1 . All other terms in the expansion consist of tensor products mixing at most k copies of X1 with various copies of X0 . Because X0 ∼ N (0, I) is independent of X1 given c, the expectations of these mixed terms factorize. The expectations of the X0 components yield standard Gaussian moments (which are known constants), and the expectations of the X1 components yield lower-order data moments Mj (P | c) for j ≤ k. Gathering all these lower-order terms into a single deterministic tensor-valued function Fkt (M0 , . . . , Mk ), we obtain: h i ⊗(k+1) Gtk (P | c) = tk E X1 | c + Fkt (M0 , . . . , Mk ) = tk Mk+1 (P | c) + Fkt (M0 , . . . , Mk ). (30) Since t ∈ (0, 1), the coefficient tk is strictly positive. Therefore, knowing Gtk (P | c) (which is perfectly matched by the degree-n projection for k ≤ n) and the previously determined lower-order moments allows iterative back-substitution to uniquely solve for Mk+1 (P | c). This concludes the proof. Corollary B.2 (Affine Component Identifies a Covariance Operator). Under the setting of Theorem 3.1, assume X1 | c has a finite second moment. Let µ(c) = E[X1 | c] and Σ(c) = Cov(X1 | c). Then the degree-zero projection is exactly the conditional mean: v0⋆ (z, t, c) = µ(c),

(31)

(32)

(33)

For a fixed t ∈ (0, 1), the mapping Σ(c) 7→ Rt (c) is rigorously injective. Therefore, the affine component of the denoising velocity serves as an exact, injective transformation of the covariance. Proof. By standard least-squares theory, the optimal degree-zero predictor is simply the conditional expectation of the target: v0⋆ (z, t, c) = E[V | c] = µ(c). (34) Similarly, the optimal degree-one (affine) predictor is given by the exact population ordinary least squares (OLS) solution: v1⋆ (z, t, c) = E[V | c] + Cov(V, Z | c) Cov(Z | c)−1 z − E[Z | c] . (35) 18

Preprint

Algorithm 1 Alternating Training Pipeline for AMFD L Require: Generator Gθ , frozen encoders {ϕℓ }L ℓ=1 , amortizers {ψℓ }ℓ=1 1: repeat 2: Sample real data (xr , c) and noise z; generate xg = Gθ (z, c) 3: Compute feature representations hℓr = ϕℓ (xr ) and hℓg = ϕℓ (xg ) for all ℓ = 1, . . . , L 4: Update amortizers by minimizing Lamort (ψ) (Eq. 15), with generated features hℓg detached 5: Freeze amortizers and update generator Gθ by minimizing Lgen (θ) (Eq. 21) 6: until convergence

Recall that Z = tX1 + (1 − t)X0 and V = X1 − X0 . Taking conditional expectations yields E[Z | c] = tµ(c) and E[V | c] = µ(c). Utilizing the independence of X1 and X0 ∼ N (0, I), we directly evaluate the covariance matrices: Cov(Z | c) = t2 Σ(c) + (1 − t)2 I,

Cov(V, Z | c) = tΣ(c) − (1 − t)I.

(36)

Substituting these moments into the OLS formula directly yields v1⋆ (z, t, c) = µ(c) + Rt (c) z − tµ(c) , where Rt (c) = Cov(V, Z | c) Cov(Z | c)−1 . Finally, to establish injectivity, we isolate Σ from Rt . Right-multiplying Rt by (t2 Σ + (1 − t)2 I) and grouping the terms involving Σ yields: t(I − tRt )Σ = (1 − t) I + (1 − t)Rt . (37) Observe that the matrix factor I − tRt = (1 − t)(t2 Σ + (1 − t)2 I)−1 is strictly positive definite, hence invertible. Left-multiplying by its inverse yields the explicit closed-form expression for Σ(c): Σ(c) =

−1 (1 − t) I − tRt (c) I + (1 − t)Rt (c) . t

(38)

Thus, Rt uniquely determines the covariance matrix.

A DDITIONAL D ETAILS ON AMFD T RAINING

C.1

AMFD T RAINING A LGORITHM

Algorithm 1 outlines the complete alternating optimization procedure for the AMFD framework. C.2

D ERIVATION OF THE JACOBI C OVARIANCE W EIGHTING

In this section, we derive the weighting factor ωcov (t) introduced in Section 4.2. Our objective is to calibrate the covariance loss such that its gradient approximately matches the standard Fréchet Distance force, 12 dΣ, at the isotropic equilibrium point Σ = I. Crucially, at Σ = I, the operators Rt , At , and the Jacobi operator Jt = I − tRt all reduce to scalar multiples of the identity matrix. Because these operators commute at this specific point, we can compute their differentials using standard scalar calculus. By the proof of Corollary B.2, the covariance is expressed as: Σ=

1−t (I − tRt )−1 (I + (1 − t)Rt ). t

(39)

Evaluating the differential dΣ with respect to Rt around the commuting point yields: dΣ =

1−t 1 − t −2 (I − tRt )−2 (1 − t)(I − tRt ) − (−t)(I + (1 − t)Rt ) dRt = Jt dRt . (40) t t

Using our residual parameterization Rt = −I + tAt , we substitute dRt = tdAt to obtain: dΣ = (1 − t)Jt−2 dAt . 19

(41)

Preprint

Our applied generator force is defined as δaJ = Jt dAt Jt . Since all matrices commute at Σ = I, this simplifies to δaJ = Jt2 dAt , which implies dAt = Jt−2 δaJ . Substituting this back into the dΣ equation provides the direct relationship: dΣ = (1 − t)Jt−4 δaJ .

(42)

To compute the explicit scalar value, we evaluate Jt (I) = I − tRt (I) using Eq. 11: 1−t 2t − 1 I = 2 I. Jt (I) = I − t 2 2 t + (1 − t) t + (1 − t)2

(43)

Plugging the scalar value of Jt (I)−4 into Eq. 42 yields: t2 + (1 − t)2 dΣ = (1 − t)3

4 δaJ .

(44)

To precisely recover the target FD force 21 dΣ, the loss weighting must be: 4 1 t2 + (1 − t)2 . ωcov (t) = 2 (1 − t)3

(45)

I MPLEMENTATION D ETAILS

In this section, we provide comprehensive implementation details for our experiments. Specifically, we specify the representation encoders utilized (§D.1), introduce a custom analytical JVP implementation designed to accelerate amortizer training (§D.2), outline the experimental settings for our amortizer ablation studies (§D.3), and provide the default amortizer settings for generation tasks (§D.4). Furthermore, we supply explicit data curation procedures and hyperparameter configurations for both ImageNet (§D.5) and text-to-image (§D.6) post-training environments. D.1

R EPRESENTATION E NCODERS

Table 4 summarizes the representation encoders utilized across our experimental settings. For the pooling operations, cls denotes the class token, avg indicates global average pooling over spatial or patch tokens, and attn refers to the attention-based pooling head. During amortizer training, we apply per-dimension normalization to all extracted features. This ensures that the inputs to the amortizer maintain zero mean and unit variance, facilitating stable neural network optimization. Table 4: Detailed configuration of the representation encoders. We list the model names and architectures, the used weight checkpoints, the feature extraction pooling operations, the input image resolutions, and the resulting feature dimensions. Model

Checkpoint

Arch.

Pooling

Resolution

Dim

Inception-v3 (Szegedy et al., 2016) ConvNeXt-v2 (Liu et al., 2022b) SigLIP2 (Tschannen et al., 2025) MAE (He et al., 2022) CLIP (Radford et al., 2021) DINOv2 (Oquab et al., 2023) PE-Core (Bolya et al., 2026) AIM-v2 (Fini et al., 2025) Web-SSL (Fan et al., 2025) DreamSim (Fu et al., 2023)

torch-fidelity inception v3 convnextv2 base.fcmae ft in22k in1k vit so400m patch16 siglip 256.v2 webli vit large patch16 224.mae vit large patch14 clip 224.openai vit large patch14 dinov2.lvd142m vit pe core large patch14 336.fb aimv2 huge patch14 224.apple pt webssl-dino1b-full2b-224 DINO + CLIP + OpenCLIP ensemble

CNN CNN ViT ViT ViT ViT ViT ViT ViT ViT

avg avg cls/attn cls cls cls attn avg cls cls

299 224 224 224 256 224 224 224 224 224

2048 1024 1152 1024 1024 1024 1024 1536 1536 1792

D.2

C USTOM JVP I MPLEMENTATION FOR MLP A MORTIZER

Evaluating the AMFD objective in Eq. 18 requires frequent JVP computations through the amortization network. In practice, relying on PyTorch’s generic forward-mode automatic differentiation (AD) for our MLP amortizer resulted in severe GPU underutilization, emerging as a primary training 20

Preprint

Algorithm 2 Manual JVP Forward Pass for MLP Require: Tangent vector v, condition label c, time t, amortizer gθ . v. Ensure: Exact input-direction JVP ẋout = ∂gθ (s,t,c) ∂s s=0

1: x ← 0 2: ẋ ← v 3: (x, ẋ) ← LinearJVP(x, ẋ; Win , bin ) ▷ x = Win x + bin , ẋ = Win ẋ 4: % Conditioning path has zero tangent w.r.t. s and is evaluated once. 5: y ← SiLU TimeEmbed(t) + ClassEmbed(c) 6: (γ, β, α) ← AdaLN0 (y) 7: for i = 0, . . . , N − 1 do 8: if i > 0 and i mod K = 0 then 9: (γ, β, α) ← AdaLNi/K (y) ▷ K is the AdaLN switch frequency 10: end if 11: % RMSNorm JVP, sharing the primal normalization statistics. 12: m ← Mean(x2 ) 13: r ← (m + ϵ)−1/2 14: xn ← x · r 15: ṙ ← −r3 · Mean(x ⊙ ẋ) 16: ẋn ← ẋ · r + x · ṙ 17: if RMSNorm has affine weight wn then 18: xn ← xn ⊙ wn 19: ẋn ← ẋn ⊙ wn 20: end if 21: % AdaLN modulation. Shift has no input tangent. 22: h ← xn ⊙ (1 + γ) + β 23: ḣ ← ẋn ⊙ (1 + γ) 24: % SwiGLU MLP JVP. 25: (u, u̇) ← LinearJVP(h, ḣ; W1 , b1 ) 26: (u1 , u2 ) ← Chunk(u) 27: (u̇1 , u̇2 ) ← Chunk(u̇) 28: σ ← Sigmoid(u1 ) 29: s ← SiLU(u1 ) = u1 ⊙ σ 30: ṡ ← u̇1 ⊙ σ ⊙ 1 + u1 ⊙ (1 − σ) 31: z ← s ⊙ u2 32: ż ← ṡ ⊙ u2 + s ⊙ u̇2 33: (o, ȯ) ← LinearJVP(z, ż; W2 , b2 ) 34: % Residual gated update. 35: x←x+o⊙α 36: ẋ ← ẋ + ȯ ⊙ α 37: end for 38: % Final AdaLN + RMSNorm + linear projection. 39: (γf , βf ) ← FinalAdaLN(y) 40: (xn , ẋn ) ← RMSNormJVP(x, ẋ) 41: x ← xn ⊙ (1 + γf ) + βf 42: ẋ ← ẋn ⊙ (1 + γf ) 43: ( , ẋout ) ← LinearJVP(x, ẋ; Wout , bout ) 44: return ẋout

bottleneck. To circumvent this, we implement an exact analytical JVP forward pass tailored specifically to our architecture, as outlined in Algorithm 2. By jointly propagating primal activations and tangent vectors across linear layers, RMSNorm, SiLU/SwiGLU activations, and AdaLN-modulated residual blocks, our custom implementation maximally reuses shared intermediate tensors. This approach strictly preserves the generic AD outputs up to numerical precision while yielding substantial gains in GPU utilization and overall training throughput. 21

Preprint

D.3

A MORTIZER A BLATION D ETAILS

Amortizer architecture details. The amortizer is parameterized as an AdaLN-conditioned MLP (Li et al., 2024). A grouped AdaLN mechanism is used to minimize conditioning overhead: rather than computing AdaLN parameters at every layer, residual blocks are grouped to share a single modulation projection. Our notation c[channels]d[depth]a[AdaLN-blocks] reflects this structure; for instance, c2048d16a4 applies AdaLN only 4 times across its 16 layers. Evaluation metrics and protocol. To quantify the amortization accuracy, we evaluate its estimates of the conditional mean and the full covariance. Let µc ∈ RD and Σc ∈ RD×D denote the empirical mean and covariance under condition c, and let µ̂c and Σ̂c be their respective estimates. We define the mean estimation error as 1 X 2 MSEµ = ∥µ̂c − µc ∥2 . (46) |C|D c∈C

To evaluate the covariance estimates, we materialize the complete covariance operator by applying it to the canonical basis vectors. Because the learned operator is not explicitly constrained to be symmetric, we extract its symmetric component 1 Σ̂sym = Σ̂c + Σ̂⊤ , (47) c c 2 and define the covariance estimation error as 2 1 X sym Σ̂c − Σc . (48) MSEΣ = 2 |C|D F c∈C

For our ImageNet experiments, C denotes the set of evaluated classes. We compute the mean estimation error MSEµ across all 1,000 classes (i.e., |C| = 1000). However, materializing and evaluating the full D × D covariance operators across all classes is computationally demanding. Therefore, we evaluate MSEΣ on a subset of the 100 most populated classes (|C| = 100), each containing exactly 1,300 images. This subset corresponds to class indices 0–42, 44–50, 52–61, 63–97, 99–102, and 104. For both metrics, we report the macro-average over their respective evaluated classes. In contrast, for the unconditional (marginal) experiments, we compute both MSEµ and MSEΣ directly over the global data distribution, treating the entire dataset as a single entity without class-wise partitioning. The symmetric covariance error is inherently better aligned with our generator objective than the MSE of the raw, potentially asymmetric operator. Specifically, the covariance term in the generator loss is evaluated through quadratic forms of the type u⊤ (Σ̂g − Σ̂r )u. For any skew-symmetric matrix K, u⊤ Ku = 0; hence, the antisymmetric component of the estimated operator has no impact on this generator objective. Consequently, MSEΣ precisely measures the covariance component that can influence the generator update, properly avoiding penalties on antisymmetric estimation artifacts that are invisible to the generator loss. Training configurations. Table 5 details the default hyperparameter configurations for these ablation studies. Table 5: Default hyperparameters for the amortizer ablation studies. representation encoder amortizer architecture amortization variant task conditioning time (t) optimizer learning rate weight decay precision batch Size training steps evaluation metrics

Inception-v3 avg AdaLN MLP AMFD-C (conditional) real: t = 2 for mean, t ∈ (0, 1) for cov; generated: −t AdamW, β1 = 0.9, β2 = 0.95 1e-4 0 bf16 1024 20,000 MSEµ , MSEΣ

Preprint

D.4

D EFAULT A MORTIZER C ONFIGURATIONS FOR I MAGE G ENERATION

Unless otherwise specified, the following settings apply universally to all amortizer training procedures: • Architecture & initialization: Parameterized by a c2048d16a4 MLP with random initialization. • Condition Injection: The external conditioning signal is explicitly injected into the amortizer network, corresponding to the AMFD-C configuration. • Optimization: We adopt the same optimizer and global batch size as the generator, applying a constant learning rate of 1e-4 without weight decay. • Diffusion parameters: The interpolation time is fixed at t = 0.25. Given the lightweight architecture of the MLP amortizer, we independently draw 4 noise samples per data input to expand the effective batch size and fully utilize hardware computational capacity. • Task-conditioning time: For the real branch, it is set to t = 2 for mean prediction and t ∈ (0, 1) for covariance prediction. The generated branch utilizes the corresponding negative time steps (−t). • JVP Calculation: We employ a custom analytical implementation for the Jacobian-vector product (JVP), as detailed in Appendix D.2. • Multi-Representation Normalization: To compute the normalized generator loss (Eq. 20), we consistently fix ρ = 1 and ϵ = 0.01, aligning with the multi-representation FD-loss formulation in Yang et al. (2026). Notably, this normalization is bypassed during single-representation training. For our default MLP amortizers, we utilize the custom analytical JVP implementation. For experiments conducted in native generative spaces where the amortizers adopt Transformer architectures, we opt to employ PyTorch’s functional API (torch.func.jvp and vmap). While our analytical MLP JVP inherently avoids distributed synchronization issues, the PyTorch JVP engine can conflict with the asynchronous communication hooks of Distributed Data Parallel (DDP). To ensure a unified, robust codebase compatible across all amortizer architectures, we systematically structure our distributed training pipeline by wrapping the JVP operation strictly inside the DDP module (i.e., DDP(jvp(Amortizer))). This uniform architectural arrangement natively circumvents synchronization failures during forward-mode differentiation, similar to the strategy noted in Lin et al. (2026). Moreover, we empirically observed that computing JVPs through LayerNorm leads to severe training instability. Conversely, RMSNorm provides significantly better stability under such operations. While the majority of our adopted amortizer architectures natively employ RMSNorm, the SiT (Ma et al., 2024) model used in our native space experiments defaults to LayerNorm. To ensure stable training, we explicitly replace all LayerNorm layers with RMSNorm exclusively within the SiT amortizer, while leaving the SiT generator entirely unmodified to preserve its exact pretrained formulation. This issue has also been reported by Zhou et al. (2025b) and Lin et al. (2026). D.5

I MAGE N ET P OST-T RAINING

Training settings. Table 6 details the comprehensive hyperparameter configurations for ImageNet post-training, covering both the representation-feature approaches (employing JiT and pMF) and the investigation across various native generative spaces. Evaluation. We measure the model performance using the standard FID alongside the recently introduced FDr6 metric (Yang et al., 2026). D.6

T EXT- TO -I MAGE P OST-T RAINING

Training settings. We evaluate AMFD in both the pixel space using PixelGen (Ma et al., 2026) and the VAE latent space using FLUX.2 [klein] 4B Base (Black Forest Labs, 2026) and FLUX.2 [klein] 4B (Black Forest Labs, 2026). i) To construct the training corpus for the PixelGen experiments, we aggregate subsets from BLIP3o-60k (Chen et al., 2025a) (58, 859 samples), ShareGPT4o-Image (Chen et al., 2025b) (45, 717 samples), and Echo-4o-Image (Ye et al., 2025) (excluding 23

Preprint

Table 6: Configurations for ImageNet-256×256 post-training. The top section details post-training using representation features, while the bottom section specifies settings for four native generative spaces (pixel, SD-VAE (Rombach et al., 2022), VA-VAE (Yao et al., 2025), and RAE (Zheng et al., 2025)). For native spaces, we adopt JiT (Li & He, 2026), SiT (Ma et al., 2024), LightningDiT (Yao et al., 2025), and DiTDH as architectures for both the generator and the amortizer. Representation feature spaces

JiT

pMF

model sizes initialization AMFD-Incep. AMFD-SIM batch size optimizer weight decay learning rate lr schedule total epochs (Table 2) total epochs (Figure 6 & 7) warmup epochs precision dropout augmentation EMA decay

B, L, H

B, L, H official pretrained weights Inception-v3 avg SigLIP2 cls + Inception-v3 avg + MAE cls 1024 AdamW (Loshchilov & Hutter, 2017), β1 = 0.9, β2 = 0.95 0 1e-5 1e-6 cosine 100 20 1 bf16 mixed 0 center crop, horizontal flip EDM2-style (Karras et al., 2024)

Native generative spaces

Pixel

generator model amortizer model initialization batch size optimizer weight decay learning rate lr schedule total epochs (Figure 8) warmup epochs precision dropout augmentation EMA decay amortizer noise samples amortizer JVP calculation

JiT-L/16 JiT-B/16

SD-VAE

VA-VAE

RAE

SiT-XL/2 LightningDiT-XL/1 DiTDH -XL/2 SiT-B/2 LightningDiT-B/1 DiTDH -B/2 official pretrained weights 1024 AdamW (Loshchilov & Hutter, 2017), β1 = 0.9, β2 = 0.95 0 1e-5 cosine 20 1 bf16 mixed 0 center crop, horizontal flip EDM2-style (Karras et al., 2024) 1 torch.func.jvp

the multi-reference split, 105, 506 samples). This curation yields a total of 210, 082 training samples. During post-training, the additional transformer blocks introduced for text conditioning remain frozen. ii) For the FLUX.2 [klein] 4B models, our data preparation protocol and the selection of the ten representation encoders generally follow the settings introduced by iRDM (Feng et al., 2026). However, as DINOv3 is currently inaccessible, we substitute DINOv3-L with DINOv2-L. The reference dataset is partitioned into two distinct subsets: a perception block based on natural COCO captions, and a composition block based on detector-verified GenEval prompts. All reference images are generated offline at a 512 × 512 resolution using the four-step FLUX.2 [klein] 4B teacher. Unlike iRDM, which truncates captions to 48 tokens, we preserve the original 512 text token length of FLUX.2 [klein]. For the perception block, we generate 24 candidates for each of the 82, 783 COCO train2014 captions and retain the top three scored by PickScore. This exactly matches the iRDM dataset scale, yielding 248, 349 samples. For the composition block, given the 553 GenEval prompts, we sample 150 seeds (expanding up to 1, 000 seeds if necessary) to retain up to 100 detector-verified generations per prompt. Notably, two specific prompts (one targeting position and one color attribute) yielded zero valid samples even after exploring 1, 000 seeds. Consequently, our final composition block comprises 53, 504 samples, closely approximating iRDM’s 53, 800. In total, the FLUX.2 [klein] training corpus contains 301, 853 samples. The complete training configurations for PixelGen and FLUX.2 [klein] experiments are summarized in Table 7. 24

Preprint

Table 7: Hyperparameter configurations for text-to-image post-training. The top section details post-training in pixel space (PixelGen), while the bottom section specifies settings for VAE space (FLUX.2 [klein]). Pixel-space generation

PixelGen

original sampling steps initialization base resolution trainable parameters text encoder (frozen) text tokens AMFD-SIM training data batch size optimizer learning rate total steps warmup steps precision EMA decay amortizer architecture

25 × 2 × 2 official checkpoint 512 × 512 image generator only Qwen3-1.7B 128 SIM (Inception avg, SigLIP2 cls, MAE cls) public datasets (§D.6) 1024 AdamW (Loshchilov & Hutter, 2017), β1 = 0.9, β2 = 0.95 1e-5 5000 500 bf16 none c2048d16a4 MLP w/ 8-head cross-attn per adaLN

Latent-space generation

FLUX.2 [klein] 4B Base

original sampling steps initialization base resolution trainable parameters text encoder (frozen) text tokens AMFD-SIM AMFD-10 enc. training data batch size optimizer learning rate total steps warmup Steps precision EMA decay amortizer architecture

50 × 2

FLUX.2 [klein] 4B

4 official pretrained weights 512 × 512 image generator only Qwen3-4B 512 SIM (Inception avg, SigLIP2 cls, MAE cls) 10 encoders in Table 4, SigLIP2 using attn prepared following Feng et al. (2026) 1024 AdamW (Loshchilov & Hutter, 2017), β1 = 0.9, β2 = 0.95 5e-6 1500 150 bf16 none c2048d16a4 MLP w/ 8-head cross-attn per adaLN

Evaluation. We evaluate our models on the GenEval (Ghosh et al., 2023) and PickScore (Kirstain et al., 2023) benchmarks. For PickScore, we report the mean score over the 499 prompts from the Pick-a-Pic test set. Regarding the baseline results in Table 3, the GenEval metrics for FLUX.2 [klein] 4B and its post-training methods DMD2 and iRDM are directly cited from the iRDM paper (Feng et al., 2026). For PickScore, since the specific prompt set evaluated in their paper is not publicly available, we directly cite the baseline scores reported in their official codebase, which are based on the Pick-a-Pic test set. The PickScore for DMD2 is left blank, as it is not available in the repository.

A DDITIONAL M ETRICS ON I MAGE N ET P OST- TRAINING

Table 8 reports the FDr (Yang et al., 2026) metrics evaluated across six diverse pre-trained representation encoders for ImageNet post-training.

A DDITIONAL V ISUALIZATIONS

Class-conditional generation on ImageNet. We provide additional generated samples on ImageNet 256×256 using post-trained JiT-H/16 and pMF-H/16 in Figure 9 and Figure 10, respectively. In each figure, the columns from left to right display samples from base models post-trained with 25

Preprint

Table 8: Additional FDr metrics for one-step post-training on ImageNet-256×256. Lower is better for all metrics. FDr ↓ Inception

ConvNeXt

DINOv2

MAE

SigLIP

CLIP

FID ↓

FDr6 ↓

50K validation images

1.00

1.68

1.00

JiT-B/16 + FD-loss + AMFD-C + AMFD-U

SIM SIM SIM

0.56 0.56

1.14 1.19

5.38 4.64

2.16 1.28

5.77 3.26

13.47 12.54

3.71 1.00 0.95 0.95

15.65 5.53 4.75 3.91

JiT-L/16 + FD-loss + AMFD-C + AMFD-U

SIM SIM SIM

1.54 0.52 0.51

3.49 0.88 0.92

6.10 2.73 2.42

8.07 0.62 0.31

19.37 3.33 1.94

25.82 7.94 6.03

2.59 0.77 0.87 0.85

10.73 3.24 2.67 2.02

JiT-H/16 + FD-loss + AMFD-C + AMFD-U

SIM SIM SIM

1.18 0.45 0.51 0.49

2.52 0.86 0.84 0.79

4.28 2.10 2.12 1.88

5.65 0.43 0.38 0.21

11.91 1.68 2.73 1.73

20.40 10.37 6.30 5.65

1.97 0.75 0.85 0.83

7.66 2.65 2.15 1.79

pMF-B/16 + FD-loss + AMFD-C + AMFD-U

SIM SIM SIM

0.57 0.55

0.81 0.83

4.34 3.87

2.20 1.69

6.98 5.80

8.78 7.86

3.31 0.85 0.95 0.92

13.70 3.50 3.94 3.43

pMF-L/16 + FD-loss + AMFD-C + AMFD-U

SIM SIM SIM

1.62 0.47 0.52 0.51

1.36 0.57 0.53 0.55

6.70 2.21 2.38 2.14

9.72 0.56 0.67 0.47

20.34 3.03 4.07 3.36

14.81 5.68 5.35 5.03

2.72 0.78 0.88 0.86

9.09 2.09 2.25 2.01

pMF-H/16 + FD-loss + AMFD-C + AMFD-U

SIM SIM SIM

1.37 0.46 0.51 0.50

1.15 0.57 0.58 0.60

5.43 1.74 1.91 1.74

6.25 0.35 0.41 0.27

15.33 2.46 3.47 2.86

11.68 5.77 4.71 4.54

2.29 0.77 0.86 0.85

6.87 1.89 1.93 1.75

Method

Repr.

FD-loss, our conditional AMFD (AMFD-C), and our unconditional AMFD (AMFD-U). To ensure a direct visual comparison, identically positioned samples across the three columns are generated using the exact same initial noise. Text-to-image generation. Figure 11 showcases additional text-to-image samples produced by the FLUX.2 [klein] 4B model post-trained with our AMFD loss. Ordered column by column from left to right (and top to bottom within each column), the corresponding text prompts are as follows: (1) A striking close-up portrait of a young woman with vitiligo wearing a simple cobalt silk blouse, natural curls framing her face, confident calm expression, pale peach studio background, soft beauty lighting, truthful unretouched skin detail. (2) A large bold word AMFD painted in glowing anime bubble letters across the center of a night festival stage, fireworks behind it, colorful spotlights, dynamic celebratory composition. (3) A silver electric concept coupe parked on a mirror-like salt flat at blue hour, seamless aerodynamic body, thin white running lights, distant mountains doubled in shallow water, minimal luxury automotive photography. (4) A poetic close-up portrait of a young adult Chinese woman standing among tall silver grass on a coastal hill, cream scarf, hair lifted by sea wind, overcast horizon softened behind, introspective expression, muted cinematic color and fine film grain. (5) A sculptural ceramic vase filled with one branch of white magnolia, placed on a travertine pedestal against a muted terracotta wall, clean side light and subtle shadows, refined design editorial. (6) A macro still life of a cracked blue ceramic bowl repaired with gold kintsugi, placed on dark linen with white flowers. 26

Preprint

(7) A delicate watercolor-style scene of a sailboat crossing a lake at sunrise, reeds, ducks, pale orange sky, calm composition. (8) A tiny hummingbird hovering beside a red trumpet flower, iridescent throat flashing emerald and magenta, wings rendered as delicate motion, dew-covered garden in soft morning bokeh, macro wildlife photograph. (9) A curious alpaca standing at the edge of a high Andean village, woven red tassel on its halter, terraced green mountains and low clouds behind, bright documentary travel photograph. (10) A secluded tropical cove viewed through limestone arches, luminous aquamarine water, a narrow white beach and one wooden boat, humid morning haze, realistic Southeast Asian travel photograph. (11) A jewel-toned fruit tart beneath a glass bakery dome, concentric raspberries, kiwi and mango with a mirror glaze, soft reflections on a walnut counter, elegant close food photography. (12) A proud Friesian horse galloping along a wide Atlantic beach, black mane streaming, hooves scattering silver surf, overcast sky and distant dunes, dramatic equine photography. (13) A massive green word FRESH placed above a farm stand with fruit crates, morning sunlight, clean commercial photography. (14) A red electric motorcycle parked in a narrow alley after rain, neon signs reflected in wet pavement, steam rising from vents, realistic urban night photography, low angle lens. (15) A delicate close-up portrait of a young Chinese ink painter in a bright studio, subtle charcoal marks on one fingertip and cheek, long dark hair loosely tied, rice paper bokeh, introspective gaze, soft high-key natural photography. (16) A small blue delivery van parked beside a quiet seaside road under a bright but soft afternoon sky. The van has rounded corners, clean windows, slightly worn tires, and a simple roof rack carrying two neatly tied wooden crates. The road curves gently along the coast, with a low stone barrier, pale grass, and calm water visible beyond it. A few distant sailboats appear near the horizon, and the sunlight gives the van a mild highlight along its side panel. The composition should feel clear and pleasant, like a travel photograph from a slow coastal route. Focus on the van’s shape, paint surface, glass reflections, tire texture, road markings, stone barrier, sea color, and open air. The mood should be relaxed, bright, and realistic, with simple forms and a clean sense of place. The camera frames the van from a relaxed roadside angle, giving the rounded body, roof rack, crates, tires, glass, road curve, stone barrier, grass, sea, and horizon a clear readable structure. Colors stay bright and gentle, with blue paint, pale road, warm grass, soft sky, and calm water creating a clean coastal mood. (17) A cinematic portrait of a weathered polar expedition pilot inside an icy cockpit, frost on the glass, amber instrument lights reflecting on the face, shallow depth of field, realistic skin texture, dramatic cold atmosphere. (18) A green sea turtle drifting above a colorful coral garden, small reef fish surrounding its shell, clear sunbeams descending from the surface, richly detailed but natural underwater photography. (19) A black cat sitting on the hood of a vintage cream convertible parked beside a lavender field, sunset sky, chrome reflections, playful fashion campaign mood. (20) The Namib Desert at sunrise seen from a high dune, sweeping orange ridges casting blue shadows, a line of gemsbok tracks crossing untouched sand, minimal and precise fine-art landscape photograph. (21) A sunlit close-up portrait of a weathered Greek fisherman at a harbor, white beard, knitted navy cap, deep smile lines, turquoise boat paint softly visible behind him, direct honest gaze, Mediterranean color, high-detail environmental portrait. (22) A colorful cereal box on a breakfast table, the front label clearly reads PIXEL CRUNCH in playful letters, morning sunlight, realistic product photography.

Preprint

Figure 9: Uncurated samples from JiT-H post-trained with FD-loss (left), our AMFD-C (middle), and AMFD-U (right), generated using identical initial noise.

Preprint

Figure 10: Uncurated samples from pMF-H post-trained with FD-loss (left), our AMFD-C (middle), and AMFD-U (right), generated using identical initial noise.

Preprint

Figure 11: Additional text-to-image samples generated by the FLUX.2 [klein] 4B model post-trained with AMFD.

Record · ID 411062 · SHA-256 2a26645659b1acb4
Retrieved via Conceptio — every document is proof-bundled with source, license, and retrieval metadata.