用神经网络动态学习数据统计量,提升图像生成质量与指令遵循能力
Amortized Moment Matching for Visual Generation

- 通过多项式投影将扩散模型转换为可学习的统计量匹配框架
- 在ImageNet上比弗雷歇距离基线提升FDr⁶指标,一步生成效果更优
- 适用于文本到图像生成,显著增强指令跟随能力,适合高阶生成任务
我们提出了一种摊销矩匹配方法,利用神经网络学习数据矩作为分布训练信号。通过多项式投影将扩散去噪器建模,建立了一个通用的矩摊销框架,表明n阶投影可显式识别至n+1阶的数据矩。基于可解析的仿射情形,我们实例化了摊销弗雷歇距离(AMFD)损失。与依赖显式边缘矩计算的FD损失不同,AMFD可通过交替、无矩阵优化流程动态学习条件矩,轻松扩展至高维数据。当作用于全局表示特征时,AMFD可作为强大的后训练目标;实验证明其神经形式相比精确统计匹配更具鲁棒性,在FDr⁶指标上显著超越FD基线,并在ImageNet上实现更优的一步生成效果。此外,它支持在原生生成空间中的直接探索,表明前两个矩仅在语义强的空间中能识别目标分布。当扩展至文生图任务,AMFD的条件感知特性带来巨大收益,使我们的一步模型在GenEval基准上超越其多步教师模型FLUX.2 [klein] 4B,同时在PickScore上表现相当。代码与检查点已公开于https://github.com/poppuppy/amfd。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。