💻 computer science

Representation Distribution Matching for One-Step Visual Generation

This paper introduces Improved Representation Distribution Matching (iRDM), a one-step visual generation paradigm that achieves state-of-the-art performance on ImageNet and enhances multi-step models like FLUX.2 by optimizing distribution matching with large batch sizes and a robust, multi-encoder evaluation metric to prevent gaming.

Original authors: Lan Feng, Wuyang Li, Eloi Zablocki, Matthieu Cord, Alexandre Alahi

Published 2026-07-03

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Lan Feng, Wuyang Li, Eloi Zablocki, Matthieu Cord, Alexandre Alahi

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Idea: From Slow Cooking to Instant Noodles

Imagine you want to create a perfect painting.

Old Way (Diffusion Models): This is like a chef slowly cooking a complex stew. They start with a pot of noise (random static) and add ingredients step-by-step, stirring and tasting for 20 or 30 minutes (steps) until the flavor is just right. It takes a long time to make one bowl.
The Goal: The authors want to make a "magic instant noodle" that tastes exactly as good as the slow-cooked stew but is ready in one single step.

The paper introduces a new method called iRDM (improved Representation Distribution Matching) that achieves this. It trains a model to generate a high-quality image instantly, without needing a "teacher" model to guide it step-by-step.

How They Did It: The Two "Knobs" on the Machine

The authors realized that previous attempts to make these "instant" generators failed because they were turning the wrong knobs. They identified two main "knobs" (design axes) that control the quality:

1. The "How to Compare" Knob (The Comparison Axis)

To teach the generator, you have to compare its output to real photos.

The Old Mistake: Previous methods tried to compare images using a blurry, low-resolution ruler (like the Fréchet distance) or a ruler that only looked at a tiny sample of the real photos. It was like trying to judge a whole orchestra by listening to just one violin player for a split second.
The Fix (The Nyström Attraction): The authors realized that the classic "Maximum Mean Discrepancy" (MMD) metric was actually great, but people were using it wrong. They fixed it by:
- Freezing the Reference: Instead of looking at a random few real photos every time, they took the entire dataset of 1.28 million images, compressed it into a perfect "summary map" (called a Nyström reference), and froze it. This is like having a perfect, unchanging blueprint of what "real" looks like.
- Big Batches: They realized the generator needs to look at thousands of images at once (batches of 2,000+) to get a clear picture. Small batches are too noisy, like trying to hear a whisper in a crowded room.

2. The "What to Measure" Knob (The Representation Axis)

Once you have a way to compare, you need to decide what features to measure.

The Trap (Gaming the System): If you only use one "judge" (a single AI encoder) to grade the images, the generator will cheat. It will learn to trick that specific judge. It's like a student who memorizes the answers to one specific teacher's test but can't actually do the math. The student gets a perfect score, but the images still look fake to humans.
The Fix (The Balanced Panel): The authors used a "panel of judges" (14 different AI encoders). They didn't just average their scores; they used a special "Lagrangian controller" (a smart balancing system).
- The Analogy: Imagine a bucket of water held together by wooden staves (the judges). The water level (the quality) is only as high as the shortest stave. If one judge says the image is fake, the whole image is considered fake. The system forces the generator to satisfy the hardest judge, not just the easiest one. This prevents cheating.

The Results: The "Magic Instant Noodle"

By combining these fixes, they created iRDM. Here is what happened:

On ImageNet (Standard Pictures): They took an existing model and turned it into a one-step generator. It became the best one-step generator in the world according to their new, hard-to-cheat metric (SWr14).
- The Metric: They created a new test called SWr14. It's like a "human preference simulator" that looks at images through 14 different lenses. The real photos score a perfect 1.0. Their model scored 1.30, which is incredibly close to reality, beating all other one-step models.
- Human Taste: When they asked a separate AI (PickScore) to pick between their images and the old best models, humans (via the AI proxy) preferred the new model 71% of the time.
On FLUX.2 (Text-to-Image): They took a famous, high-quality 4-step model called FLUX.2 and "post-trained" it to be a 1-step model.
- The Surprise: The new 1-step version was actually better than the original 4-step version at following instructions (GenEval score went from 0.794 to 0.826).
- Speed: They did this training in about 90 hours on powerful GPUs.

Why This Matters (According to the Paper)

No Cheating: The biggest breakthrough is that they stopped the models from "gaming" the system. By using a diverse panel of frozen encoders and a balanced optimization strategy, the model can't just trick one specific metric; it has to actually look real.
No Teacher Needed: Unlike other fast methods that need a slow, complex teacher model to guide them, this method learns by directly comparing its output to a frozen map of reality.
The "One-Step" Reality: They proved that you don't need 20 steps to make a great image. You just need the right way to measure "realness."

Summary Analogy

Imagine you are trying to teach a robot to draw a perfect apple.

Old Way: You show the robot a picture, then a slightly blurry version, then a blurrier one, and ask it to guess what the original was, step by step.
The Paper's Way: You give the robot a perfect, frozen "Apple Blueprint" (the Nyström reference) and a panel of 14 expert art critics (the encoders). You tell the robot: "Draw an apple. If any of the 14 critics say it looks fake, you fail. If you satisfy the pickiest critic, you win."
The Result: The robot learns to draw a perfect apple in a single, instant stroke, and it's so good that even the pickiest critic can't tell it apart from a real photo.

Technical Summary: Representation Distribution Matching for One-Step Visual Generation

Problem Statement

Generative modeling fundamentally aims to match the output distribution of a generator to the data distribution. While diffusion and flow models achieve this implicitly by learning to reverse a noising process over many steps, recent approaches seek to do so explicitly and directly in a single network evaluation. These "one-step" generators aim to match generated and real feature distributions under frozen pretrained encoders without relying on online teachers, adversaries, or simulated trajectories.

However, existing one-step methods fix their design choices jointly, leaving it unclear which specific components drive quality. Furthermore, classical distribution matching objectives like Maximum Mean Discrepancy (MMD) were previously dismissed as too weak to train competitive generators, and single-encoder evaluation metrics are prone to "gaming," where a generator optimizes for a specific encoder's score while producing visibly fake images.

Methodology: Representation Distribution Matching (RDM)

The paper formalizes Representation Distribution Matching (RDM) as a paradigm defined by two design axes: how distributions are compared and the representations in which they are compared. Through controlled studies along these axes, the authors propose Improved RDM (iRDM).

1. The Comparison Axis

Re-evaluating MMD: The authors demonstrate that MMD is not inherently weak but was previously badly estimated. They propose a specific estimation strategy:
- Exact Within-Batch Repulsion: The repulsion term (preventing mode collapse) is calculated exactly over the generated batch ( $B \times B$ kernel sum).
- Nyström Attraction: The attraction term (pulling toward data) is estimated against the full training set using a Nyström kernel mean embedding. Instead of resampling the reference at every step (which introduces noise), the entire training set (e.g., 1.28M ImageNet images) is compressed once into a frozen set of 4,096 landmarks. This provides a zero-variance summary of the data distribution.
Batch Size: The generated batch size is identified as a critical variable. While larger batches improve the estimate, they reduce the number of optimizer steps for a fixed compute budget. The authors find an optimum above 2,048 (specifically 5,120 for ImageNet and 10,240 for FLUX), far exceeding standard practice. Gradient caching is used to manage memory.
Conditional Tasks: For text-to-image, the method matches the joint law of caption and image features rather than the image marginal alone. This ensures prompt fidelity is part of the objective.

2. The Representation Axis

The Problem of Single Encoders: A single encoder can be "gamed." A generator can drive the distance below the real data score on that specific encoder while remaining visibly fake (e.g., fixing some classes but leaving artifacts in others).
Diverse Encoder Battery: To prevent gaming, iRDM matches against a diverse panel of 14 pretrained encoders (including DINOv2, SigLIP, CLIP, etc.).
Constrained Optimization: Instead of uniform weighting, the authors use a proportional Lagrangian controller (based on Stooke et al., 2020). This controller dynamically adjusts weights: it upweights encoders that are "violating" the real-data floor (hardest to satisfy) and downweights those that are already satisfied or overfitted. This ensures the generator cannot ignore difficult feature spaces.

3. Evaluation Metric: SWr14

To evaluate performance without gaming the training loss, the authors introduce SWr14, a Sliced-Wasserstein distance averaged over the 14 encoders.

It is independent of the training loss (which uses MMD).
Real validation data is scaled to a score of 1.0 by construction.
It is harder to game than Fréchet distance or MMD.

Key Contributions

A Unifying Framework (RDM): The paper formalizes one-step generation as a design space governed by two axes (comparison and representation), allowing the quality ceiling of methods to be traced to specific design choices rather than headline ideas.
A Simple Recipe at the State of the Art (iRDM): By varying axes in isolation, the authors identify the optimal configuration: exact within-batch repulsion, Nyström attraction to a frozen full-data reference, large fresh batches, joint image-text matching for conditional tasks, and constrained optimization for encoder balancing.
A Metric Resistant to Gaming: The introduction of SWr14, an optimal-transport metric over 14 encoders that shares no machinery with the training loss, providing a robust evaluation standard.

Results

ImageNet-256 Generation

SWr14 Score: iRDM achieves a state-of-the-art SWr14 of 1.30, significantly outperforming the previous best one-step generator (pMF-H FD-SIM at 2.05) and approaching the real data floor of 1.0.
Human Preference: Evaluated via PickScore (a learned human-preference proxy not used in training), iRDM is preferred over the prior best one-step generator (FD-SIM) on 71.2% of matched samples. Notably, it is the first one-step model to surpass held-out real photographs on PickScore (63.6% win rate).

Text-to-Image Post-Training (FLUX.2)

The recipe was applied to post-train the four-step FLUX.2 [klein] (4B parameters) into a one-step generator.
GenEval: The one-step iRDM model surpasses the four-step teacher, achieving 0.826 vs. 0.794.
PickScore: The one-step model achieves 22.76 vs. 22.58 for the four-step version.
Efficiency: These results were achieved in approximately 90 H200 GPU-hours without an online teacher or reward model.

Significance and Claims

The paper claims that Representation Distribution Matching is a viable, teacher-free paradigm for high-quality one-step generation. The significance lies in:

Demystifying MMD: Showing that MMD, when estimated correctly (Nyström attraction + exact repulsion) with large batches, is a powerful objective.
Preventing Reward Hacking: Demonstrating that single-encoder matching is insufficient and that a balanced, constrained optimization over a diverse encoder battery is necessary to ensure genuine realism.
Efficiency: Proving that a four-step diffusion model can be distilled into a one-step model that not only matches but surpasses the teacher in alignment (GenEval) and human preference (PickScore) using a simple distribution matching objective, without the complexity of adversarial training or trajectory distillation.

The authors modestly acknowledge that a gap remains between the best one-step generator (SWr14 1.30) and a fresh real draw (1.0), suggesting future work in multi-scale kernels and richer conditional couplings.