Amortized Moment Matching for Visual Generation

2026-07-29

AI Summary

This work addresses the limitations of traditional moment-matching methods, which rely on explicit marginal moment computations and suffer from poor scalability to high-dimensional data and training instability. The authors propose an amortized moment-matching framework that, for the first time, integrates amortized inference into moment matching by employing neural networks to dynamically learn conditional moments as distributional training signals. The diffusion denoiser is modeled as a polynomial projection, circumventing explicit matrix operations. Using an alternating optimization strategy, the method effectively identifies target distributions in high-dimensional semantic spaces using only low-order moments. Experiments demonstrate superior single-step generation performance on ImageNet, with FDr⁶ significantly outperforming FD baselines, and competitive results in text-to-image synthesis, where the single-step generator surpasses the multi-step FLUX.2 teacher model in GenEval while achieving comparable PickScore.

Abstract

We propose amortized moment matching, utilizing neural networks to learn data moments as distributional training signals. By casting diffusion denoisers through polynomial projections, we establish a general framework for moment amortization, revealing that an $n$-th degree projection explicitly identifies data moments up to order $n+1$. Derived from the tractable affine case, we instantiate the Amortized Fré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 pipeline that effortlessly scales to high-dimensional data. When operating on global representation features, AMFD serves as a powerful post-training objective; empirically, its neural formulation yields more robust training dynamics than exact statistical matching, substantially surpassing the FD baseline on the FDr$^6$ metric and achieving superior one-step generation on ImageNet. Furthermore, it unlocks direct exploration within native generative spaces, suggesting that the first two moments can identify target distributions only in spaces with strong semantics. Finally, when scaled to text-to-image generation, the condition-aware nature of 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.

Problem

• moment matching
• visual generation
• conditional moments
• high-dimensional data
• distributional training signals

Innovation

• amortized moment matching
• AMFD
• diffusion denoiser
• conditional moments
• one-step generation