Representation Distribution Matching for One-Step Visual Generation

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/.

• The paper introduces representation distribution matching (RDM) as a unifying paradigm that enables one-step visual generation with high distribution fidelity.
• It leverages frozen pretrained encoder spaces, employing Nyström-MMD and multi-encoder constraints to stabilize training and mitigate mode collapse.
• Empirical results on ImageNet show superior SW14 and PickScore metrics, demonstrating practical improvements for rapid, high-fidelity image synthesis.

Representation Distribution Matching for One-Step Visual Generation

Introduction

This work formalizes Representation Distribution Matching (RDM) as a unifying paradigm for one-step visual generative modeling, eschewing the conventional need for online teachers, adversarial losses, or intricate multi-step denoising trajectories. The method leverages frozen pretrained encoder feature spaces to directly align the distributions of generated and reference images. The study deconstructs this design space along two axes—discrepancy estimation and representation selection—and provides a rigorous empirical breakdown, identifying optimal strategies for each. The proposed instantiation, improved RDM (iRDM), achieves state-of-the-art distributional fidelity among one-step models on ImageNet, as measured by the multi-encoder Sliced-Wasserstein ($SW_{14}$) metric.

The RDM Paradigm and Its Design Axes

RDM comprises any approach that directly matches the feature distribution of generated and reference (typically real) data using a set of frozen encoders. The generator becomes inherently one-step: every sample is produced in a single forward pass. This contrasts with diffusion and flow models, which approximate the same goal via iterative denoising or invertible mappings.

The paper identifies two orthogonal axes that fully parameterize RDM methods:

1. Discrepancy Estimation: How to compare distributions in feature space. Options include maximum mean discrepancy (MMD), Fréchet distance, sliced Wasserstein distance, and variations in reference/sample pairing or batch sizing.
2. Representation Selection: Which feature spaces to deploy, i.e., the choice and weighting of pretrained encoders. Strategies range from single-encoder matching (e.g., CLIP) to balanced batteries of diverse transformers, vision-LLMs, and self-supervised learners.

This decomposition enables a systematic attribution of limitations or gains to specific design choices, clarifying claims in related works.

Discrepancy Estimation: MMD and Nyström Landmarks

The authors thoroughly reevaluate MMD, finding that prior dismissals reflected estimation issues rather than an intrinsic deficiency. With high-quality kernel estimates and sufficient batch size, MMD emerges as a robust, scalable objective. Key findings:

• Nyström Mean Embeddings: The real/reference distribution is condensed via $k$-means into a set of 4096 landmark embeddings per encoder. These act as a frozen attraction target for the generator, ensuring stability and minimizing noise.
• Exact Within-Batch Repulsion: The generator side employs per-batch pairwise repulsion, which is crucial for avoiding mode collapse.
• Large Fresh Batches: Generator batches far exceeding prior practice (optimal $>2048$) are necessary to stabilize MMD estimation, made feasible by gradient caching. This departs from traditional diffusion batch regimes.
• Conditional (Text-to-Image) Tasks: The loss matches the joint distribution over image and text features, incorporating alignment objectives naturally—improving prompt fidelity as quantified by GenEval.

Comparisons to alternatives (Fréchet, sliced Wasserstein, random Fourier features, drifting fields) demonstrate that Nyström-MMD dominates in the regime relevant for images and conditioning, both in terms of convergence and fidelity to the real manifold.

Representation Selection: Multi-Encoder Constraint

A major insight is the inadequacy of single-encoder supervision. While a generator can be tuned to outperform real data in the metric induced by any single encoder (FID, CLIP, DINOv2, etc.), this results in visible artifacts and "gaming" without genuine realism. The proposed remedy:

• Diverse Encoder Battery: Training integrates ten frozen encoders spanning supervised, self-supervised, and multimodal backbones (Inception, ConvNeXt, MAE, CLIP, DINOv3, etc.). Evaluation averages over 14, with four held out for generalization.
• Proportional Lagrange Control: Encoder weights are dynamically adjusted based on constraint satisfaction, upweighting those hardest to satisfy and downweighting overfit cases. This constrained optimization mitigates blind spots and enforces a "weakest-link" criterion analogous to known perceptual models.
• Robust Evaluation ($SW_{14}$): The multi-encoder Sliced-Wasserstein ratio, not used in training, is robust to reward hacking and more aligned with human preference measures such as PickScore.

This approach is theoretically grounded: with a diverse and sufficiently expressive encoder set, the aggregate MMD becomes characteristic and vanishes only at the real distribution.

Empirical Results

ImageNet One-Step Generation

iRDM substantially improves over previous one-step generative models:

• $SW_{14}$ Metric: Achieves 1.30 (real baseline 1.0), outperforming all released baselines, including pMF-H FD-SIM (2.05). iRDM sets the benchmark on 9/14 encoders.
• PickScore: Human-preference proxy agrees; iRDM surpasses the strongest prior generator on 71.2% of class-matched samples and is the first to exceed held-out real photo references on this metric for one-step models.
• Ablations: The advantage derives from (1) Nyström-attraction to frozen reference, (2) exact within-batch repulsion, (3) large fresh batches, (4) constrained multi-encoder matching—not from the particular architectural choice.

Text-to-Image Post-Training (FLUX.2 [klein])

Applying iRDM to post-train the four-step FLUX.2 model results in a one-step generator that:

• GenEval (COCO prompts): Increases from 0.794 (four-step) to 0.826 (one-step, joint loss).
• PickScore: Increases from 22.58 to 22.76.
• Conditional Matching: Only the joint image-text objective achieves these gains; marginal (image-only) feature matching fails to improve compositional or attribute alignment.

Additional Findings

• The proportional Lagrange scheme significantly improves the worst-case encoder match, compared to uniform weighting.
• No single loss or encoder is optimal in all regimes; empirical rankings confirm the superiority of Nyström-MMD for training and $SW_{14}$ for evaluation.

Implications and Future Directions

Practically, iRDM enables efficient, high-fidelity one-step generation with compute budgets competitive with diffusion or flow-based distillation, but with direct control over distributional matching. This is immediately valuable for applications requiring rapid synthesis or high-throughput evaluation, such as data augmentation, generative simulation, and large-scale preference-driven alignment.

Theoretically, the decomposition clarifies lingering questions on the limitations of current metrics, metric gaming, and content drift in generative modeling. It establishes that minimizing distances in a single feature space is neither necessary nor sufficient for realism, and that carefully architected multi-space objectives are robust.

Intersection with reward model alignment and direct human feedback is an obvious path forward. Additionally, the RDM framework is straightforwardly applicable wherever pretrained, domain-specific encoders exist—thus extending to speech, video, 3D, or cross-modal settings.

The main outstanding gap is to further decrease $SW_{14}$ toward the real data floor. Directions include more expressive or task-tailored encoder sets, multi-scale kernel aggregation, and richer conditional objectives.

Conclusion

The paper rigorously establishes representation distribution matching as the operative design principle behind one-step visual generation, resolving it into independently optimizable axes of discrepancy estimation and representation selection. With the iRDM instantiation, it is shown that careful discrepancy estimation (Nyström-MMD, joint conditioning, large batches) and constrained, multi-encoder supervision are pivotal for high performance. The model sets a new one-step state of the art on core distributional and preference metrics. Further gains will likely come from scaling representation diversity, refining conditional objectives, and porting the paradigm across data modalities.

---

Reference: "Representation Distribution Matching for One-Step Visual Generation" (2607.02375)