PaperlayerView exact revision on arXiv

Source-linked AI summary

Representation Distribution Matching for One-Step Visual Generation

Lan Feng, Wuyang Li, Eloi Zablocki, Matthieu Cord, Alexandre Alahi

arXiv:2607.02375v1cs.CV

TL;DR

One-step image generation lacks clear evidence about which distribution-comparison and representation choices produce genuinely realistic samples. This paper formalizes and tests those choices, yielding iRDM, which achieves state-of-the-art one-step ImageNet generation and improves a four-step FLUX.2 model when reduced to one step.

  • Problem

    Existing RDM methods conflate comparison and representation choices, while single-encoder feature distances can report near-real scores for visibly fake images.

  • Method

    The paper varies RDM’s two design axes and combines well-estimated MMD, large fresh batches, and balanced matching across diverse frozen encoders.

  • Results

    iRDM reaches one-step ImageNet state of the art at SWr14 1.30 and post-trains four-step FLUX.2 into a one-step model surpassing it on GenEval and PickScore.

  • Takeaways & Limitations

    RDM quality can be traced to specific comparison and representation choices, while SWr14 provides an evaluation metric designed to resist objective-specific gaming.

  • Takeaways & Limitations

    The best one-step generator remains measurably short of a fresh real draw, with SWr14 1.30 versus a real-data floor of 1.

Abstract

from arXiv · show

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

1 INTRODUCTION

Representation Distribution Matching (RDM) directly matches generated and real feature distributions under frozen pretrained encoders to enable one-step generation without an online teacher, adversary, or simulated trajectory. Controlled studies identify comparison estimation and representation choice as the two governing design axes, whose preferred choices combine into iRDM.

  • RDM framework: RDM generates images in one network evaluation by directly matching generated and real distributions in a frozen pretrained encoder’s feature space.It requires no online teacher, adversary, or trajectory to simulate.
  • Design axes: Controlled studies vary comparison and representation choices separately, showing that these two axes govern quality in existing one-step generators.This isolates the effect of specific design choices rather than attributing quality to a method’s headline idea.
  • Comparison: The MMD becomes competitive when estimated with a structured feature space, enough samples on each side, and a frozen Nyström reference built from the full 1.28M-image training set.The reference is compressed once and then kept fixed rather than moving during training.
  • Representations and evaluation: A single encoder can be gamed, so iRDM matches a diverse encoder battery and evaluates with SWr14, a relative Sliced-Wasserstein distance averaged over 14 pretrained encoders.Because SWr14 is not used as the training objective, improvements are intended to rule out reward hacking.
  • iRDM: iRDM combines the preferred choices into a one-step recipe built around accurate MMD estimation, exact within-batch repulsion, Nyström attraction to a frozen full-data reference, large fresh batches, and a balanced encoder battery.The contribution passage also identifies joint image-text training for text-to-image tasks and constrained optimization for balancing the encoder battery.

2 RELATED WORK

Related work frames one-step and few-step generation as reducing the per-denoising-step cost of diffusion and flow models. It also connects generation-as-distribution-matching methods to fixed feature spaces and frozen pretrained encoders.

  • One-step and few-step generation: Diffusion and flow models incur inference cost at each denoising step, motivating methods that reduce the number of steps.Step reduction can distill a pretrained teacher or remove it entirely.
  • One-step and few-step generation: Distillation-based step reduction matches a teacher’s trajectory, score, or moments.
  • Matching distributions in fixed feature spaces: Generation as distribution matching extends the GAN program through fixed-kernel moment matching, adversarial kernels, and sliced Wasserstein generators.
  • Matching distributions in fixed feature spaces: The feature space has shifted toward frozen pretrained encoders, which are used for perceptual losses.

3 REPRESENTATION DISTRIBUTION MATCHING AND ITS DESIGN SPACE

RDM trains one-step generators by matching generated and real feature distributions under frozen encoders, with its design space organized around distribution comparison and representation choice. The preferred design uses exact within-batch repulsion, frozen Nyström attraction, large fresh batches, and a balanced encoder battery.

  • RDM formulation: RDM aligns generated and real feature distributions under a frozen encoder, making one-step generation a property of the output-distribution constraint.The same objective can post-train a few-step sampler by treating its final output as the generator output.
  • Design axes: The design space has two axes: how feature distributions are compared and which representations define the feature spaces.Comparison includes the discrepancy, estimator, reference distribution, and conditional joint law; representation choice includes encoder selection and weighting.
  • Estimator design: Exact within-batch repulsion prevents collapse, while frozen Nyström attraction pulls generated features toward a full-training-set reference at O(Bm) cost.The repulsion uses the exact B × B kernel sum; the reference mean embedding is precomputed over n = 1.28M training images and frozen.
  • Estimator comparison: MMD Nyström stays sharpest across batch sizes in the spiral diagnostic, whereas MMD exact degrades as its per-batch reference thins.At the largest batch, MMD exact and MMD Nyström lock onto the spiral; random features remain diffuse, sliced-Wasserstein remains loose, and drifting collapses.
  • Batch-size choice: Quality climbs with the number of generated samples before flattening into a broad optimum; the adopted values are N=5120 for ImageNet and N=10240 for FLUX post-training.The smallest batch is noise-dominated and regresses despite more optimizer steps.
  • iRDM: iRDM combines exact repulsion, frozen Nyström attraction, large fresh batches, conditional image-text matching, and a diverse encoder battery balanced by constrained optimization.The reference is precomputed once per encoder, while each step generates and encodes a fresh batch.

4 EXPERIMENTS

Experiments show that iRDM achieves state-of-the-art one-step ImageNet distributional quality and human-preference results, while post-training FLUX.2 [klein] makes one-step generation outperform its four-step start. Ablations attribute gains to joint image-text matching, adaptive encoder weighting, and properly estimated kernel MMD.

  • ImageNet generation: SWr14 1.30 makes iRDM the ImageNet-256 state of the art, below every released generator and best on nine of fourteen encoders.SWr14 averages floor-normalized Sliced-Wasserstein ratios over 14 encoders; approximately 1 matches a fresh real draw.
  • ImageNet generation: 71.2% of matched pairs favor iRDM over its pMF-H FD-SIM start, with PickScore increasing from 20.61 to 20.96.PickScore is an off-objective learned human-preference measure; iRDM also wins against RAE-XL and REPA-E SiT-XL on 75.7% and 73.2% of classes.
  • FLUX.2 [klein] post-training: 0.826 GenEval and 22.76 PickScore let one-step FLUX.2 [klein] surpass its four-step version at 0.794 and 22.58, respectively.The post-training used about 90 H200 GPU-hours and a frozen reference collected once from the four-step teacher.
  • FLUX.2 [klein] post-training: 0.826 versus 0.801 shows that joint image-text matching improves overall GenEval over the marginal image-only variant.The largest reported gaps occur on two-object prompts, 0.924 versus 0.899, and attribute binding, 0.708 versus 0.608.
  • Weighting ablation: 1.88 versus 1.90 shows a small mean advantage for gated proportional weighting over uniform weighting, while the worst encoder improves from 4.06 to 3.49.The controller allocates budget to the worst encoder, PE-Core, and gates out three encoders already at their floors.
  • Training-distance ablation: mmdx ≻ mmd rff ≻ mmd exact ≻ fd ≻ sw ≻ drifting is the ranking across both neutral evaluation distances in the DINOv2 training-distance ablation.Nyström MMD moves the feature distribution closest to real, whereas Sliced-Wasserstein is among the least effective training losses despite being an effective judge.

5 CONCLUSION · A EXTENDED RELATED WORK · B ENCODER PANEL

The paper frames representation distribution matching as a two-axis design space, favoring accurate scalable comparison and balanced multi-encoder representations. It situates these choices within related work, defines the encoder panel, and identifies remaining performance and design gaps.

  • 5 CONCLUSION: RDM instances vary along two axes: how generated and real feature distributions are compared, and which representations are compared.Varying each axis separately yields a preferred design with an associated mechanism.
  • 5 CONCLUSION: Exact within-batch repulsion, Nyström attraction to a frozen reference, and large fresh batches make classical MMD a strong objective.Conditional tasks match the joint image-text law rather than only the image marginal.
  • 5 CONCLUSION: No single encoder is sufficient because its score can fall below the real score while images remain visibly fake; balanced multi-encoder matching addresses this failure.The encoder battery is maintained through constrained optimization.
  • 5 CONCLUSION: SWr14 1.30 remains above the real-data floor of 1, leaving a measurable gap and motivating multi-scale kernels, task-specific panels, richer conditional couplings, and transfer beyond images.The proposed recipe uses a frozen reference and a single network evaluation when transferred to other modalities with pretrained encoders.
  • A EXTENDED RELATED WORK: Nyström methods use data-dependent landmarks and can dominate random Fourier features when the kernel spectrum decays quickly, while preserving full kernel-mean-embedding estimation rates with fewer landmarks.This summarizes the related-work context for scalable kernel estimators.
  • A EXTENDED RELATED WORK: Single-encoder metrics inherit encoder blind spots: scores vary with resizing, can improve through fringe features without quality gains, and can re-rank models when encoders change.The cited metric family includes FID, KID, CMMD, and feature-space precision and recall.
  • A EXTENDED RELATED WORK: Few-step text-to-image systems are commonly distilled from multi-step teachers and then steered toward human taste using learned preference rewards.The related work describes adversarial and score-based distribution matching, followed by policy-gradient or direct-preference optimization.
  • B ENCODER PANEL: The fourteen-encoder panel freezes released backbones, uses pooled image embeddings at listed resolutions, assigns ten encoders to training, and holds four out for evaluation.MMDr14 averages the fourteen encoder distances without feature normalization across diverse training paradigms.

C BATCH-SIZE SWEEP

The section presents a batch-size sweep summarized in Table 6 and Figure 4. It also defines the MMDR14 evaluation reported for released ImageNet-256 generators and iRDM across 14 vision encoders.

  • C BATCH-SIZE SWEEP: Table 6 tabulates the batch-size sweep plotted in Figure 4.
  • C BATCH-SIZE SWEEP: MMDR14 is the arithmetic mean of MMD-RFF distance ratios across 14 vision encoders.
  • C BATCH-SIZE SWEEP: The validation baseline is MMDR = 1 by definition, using real held-out data.

D KERNEL-MMD EVALUATION

MMDr14 is a training-aligned kernel-MMD cross-check of SWr14, computed across the full released field using a 14-encoder panel. Its rankings broadly agree with SWr14, though mid-field models are reordered.

  • Metric definition: MMDr14 averages per-encoder RFF-MMD ratios against real training data across 14 encoders.Real validation scores 1 by construction.
  • Metric definition: MMDr14 serves as a training-aligned kernel-MMD cross-check of the primary SWr14 metric.The evaluation covers the full released field and uses the 14-encoder panel.
  • Metric comparison: MMDr14 broadly preserves SWr14’s model ordering, with some reordering among mid-field models.Because the training loss is itself a kernel MMD, individual encoders can be pushed below the real score.

E TEXT-TO-IMAGE POST-TRAINING DETAILS · E.1 REFERENCE CURATION

The text-to-image post-training reference is collected once from the four-step FLUX.2 [klein] teacher, frozen thereafter, and combines perception-focused and composition-verified generations. These blocks are embedded with caption features under the training encoders and compressed into reusable Nyström references.

  • E.1 REFERENCE CURATION: The frozen reference contains roughly 300K image-caption pairs from two independently curated teacher-generation blocks: perception and composition.The four-step FLUX.2 [klein] teacher is queried only during reference collection, never during post-training.
  • E.1 REFERENCE CURATION: 248,349 perception pairs cover all 82,783 COCO train2014 images, retaining the three highest-PickScore candidates from 24 teacher draws per caption.Each retained image stays paired with the caption that generated it.
  • E.1 REFERENCE CURATION: Composition curation retains only teacher generations certified correct by the standard GenEval Mask2Former detector across 553 prompts.The teacher is sampled at 150 seeds per prompt, with additional sampling when fewer than 100 correct generations are available.
  • E.1 REFERENCE CURATION: Teacher correctness ranges from 98% on single-object prompts to 40% on attribute binding and 33% on position.The filter therefore most strongly reshapes binding and spatial prompts, where the one-step model later improves.
  • E.1 REFERENCE CURATION: The joint reference concatenates each image feature with its caption’s frozen SigLIP2 text embedding under the ten training encoders.Features are compared with one Gaussian kernel using bandwidth 0.25 of the median-heuristic value and text weight 1.
  • E.1 REFERENCE CURATION: 8192-landmark Nyström references are constructed once per encoder to compress the embedded reference distribution.This compression follows the joint image-text embedding and kernel construction.

E.2 DMD2 BASELINE

The DMD2 baseline adapts the original three-network distillation method to klein’s flow-matching parameterization, requiring a shifted timestep distribution. Its best LAION-prompt configuration reaches GenEval 0.804 and PickScore 22.36.

  • Method: DMD2 distills the four-step FLUX.2 [klein] teacher into a one-step student using a generator, trainable critic, and frozen teacher.The DMD gradient is the teacher-minus-critic score difference.
  • Timestep schedule: Uniform timestep sampling reaches low noise, where the teacher collapses to the mode-averaged posterior mean and harms the generator.The klein parameterization instead requires mapping uniform draws through its signal-to-noise shift, using empirical µ ≈2.03.
  • Configuration and result: GenEval 0.804 and PickScore 22.36 are achieved by the best LAION-prompt configuration at its 500-step peak.Broader LAION captions outperform COCO-only training at every checkpoint and reduce a 4.7-point collapse to a 1.1-point plateau.

F ONE-STEP TEXT-TO-IMAGE SAMPLES

Figure 7 presents additional single-step iRDM generations from the post-trained four-step FLUX.2 [klein], with each 512 × 512 image produced in one network evaluation.

  • F ONE-STEP TEXT-TO-IMAGE SAMPLES: Figure 7 shows additional single-step iRDM generations from post-trained four-step FLUX.2 [klein], each a 512 × 512 image produced in one network evaluation.These examples demonstrate the one-step output format of the post-trained model.

G QUALITATIVE COMPARISON

Qualitative comparisons place uncurated iRDM samples beside the strongest external one-step baseline across five ImageNet-256 classes, with the methods appearing close by eye despite iRDM’s better SW_r14 score. The post-trained one-step FLUX.2 [klein] also produces 512 × 512 samples with one network evaluation.

  • ImageNet-256 comparison: iRDM and pMF-H FD-SIM are compared using uncurated samples across five ImageNet-256 classes, without cherry-picking.The classes span a bird, animal coat, deformable garment, rigid man-made object, and natural landscape.
  • FLUX.2 [klein]: The post-trained one-step FLUX.2 [klein] generates 512 × 512 text-to-image samples with one network evaluation each.Figure 7 shows single-step generations from the post-trained model.
  • ImageNet-256 comparison: SW_r14 is 1.30 for iRDM versus 2.05 for pMF-H FD-SIM, although the samples are close by eye.Figure 8 presents ten additional draws per method under each class label.
Loading 2607.02375v1…