Representation Distribution Matching for One-Step Visual Generation

2026-07-02

AI Summary

This work addresses the limitations of distribution matching and representation selection in single-step image generation by proposing a novel paradigm that directly aligns the distributions of real and generated images within the frozen feature space of pretrained encoders. Systematic investigation reveals that the classical Maximum Mean Discrepancy (MMD), when properly estimated, is highly competitive; large-batch training (batch size > 2048) is crucial for performance; and balanced matching across multiple encoders significantly enhances generalization. The proposed method achieves state-of-the-art results in single-step generation on ImageNet (SW↓r14 = 1.30) and attains a 71.2% win rate in human preference evaluations via PickScore. Furthermore, it successfully distills the four-step FLUX.2 model into a stronger single-step variant that surpasses the original in both GenEval and PickScore metrics, requiring only 90 H200 GPU-hours.

Abstract

We elucidate the design space of Representation Distribution Matching (RDM), our name for the paradigm that trains a one-step image generator by matching generated and reference feature distributions under frozen pretrained encoders. We identify two design axes, how the distributions are compared and the representations they are compared in, and controlled studies along them yield three findings. First, the classical MMD, which could not train convincing generators a decade ago, becomes a strong and scalable objective once estimated right. Second, the generated batch is then the operative variable, with an optimum above 2048, far beyond customary batch sizes. Third, any single representation can be gamed, driven below the real score while images stay visibly fake, so we match against a balanced battery of encoders and evaluate with SW_r14, a Sliced-Wasserstein distance over 14 encoders that is independent of the training loss and resists gaming. Combining the preferred choices yields improved RDM (iRDM): it sets the one-step state of the art on ImageNet at SW_r14 1.30, corroborated by PickScore, a human-preference proxy our objective never optimizes, which prefers it over the prior best one-step generator on 71.2% of matched samples. The same recipe post-trains the four-step FLUX.2 [klein] into a one-step generator, surpassing the four-step version on GenEval, 0.826 to 0.794, and on PickScore, 22.76 to 22.58, in 90 H200 GPU-hours. Project page: https://alan-lanfeng.github.io/rdm/.

Problem

• one-step generation
• representation distribution matching
• feature distribution
• image generation
• distribution comparison

Innovation

• Representation Distribution Matching
• one-step generation
• MMD
• large batch training
• multi-encoder evaluation