Intrinsic Distribution Matching

AI summary

Overview

Research area: Generative computer vision — specifically one-step and few-step text-to-image generation via distribution matching distillation (DMD), with a focus on long-text rendering fidelity.

Technical level: Advanced. The paper leans on score-based diffusion theory, reverse-KL gradients, a Gaussian-plus-residual score decomposition, and Wasserstein-style moment matching. The core idea is intuitive, but the motivation section requires comfort with score functions and rectified flow.

Scope: The paper proposes Intrinsic Distribution Matching (IDM), which adds an explicit feature-moment loss on top of internal DMD supervision, and validates it through full-parameter and LoRA tuning of Z-Image-Turbo 6B and Qwen-Image 20B at 1 and 2 network evaluations.

What This Paper Is About

Diffusion and flow-matching models produce strong images but require many network evaluations, which is expensive. One-step and few-step generators cut that cost, but they have almost no opportunity to refine, so fine details — especially long passages of legible text — are easily lost. Distribution matching distillation addresses this by pushing the generated distribution toward a reference distribution using a difference of scores, and this paper asks whether adding an explicit constraint on feature statistics (means and per-coordinate standard deviations) makes that learned supervision better.

Key Contributions

1. Intrinsic distributional supervision. IDM augments internal DMD with an explicit loss on the means and per-coordinate standard deviations of the generator's own exponential-moving-average (EMA) features — no separate discriminator, score network, or external feature encoder.

2. A distribution-gradient explanation. Building on the Gaussian-plus-residual score representation, the authors decompose the reverse-KL field into a Gaussian moment component and a non-Gaussian residual, which identifies the statistics that can be constrained explicitly and explains why moments alone are insufficient.

3. Few-step training and evaluation at scale. Full-parameter and LoRA training (ranks 32 and 64) on 6B and 20B backbones, with ablations on the added moment loss, the choice of EMA features versus an external encoder, and velocity-versus-endpoint regression.

4. A long-text rendering result. At 2 NFE, IDM reaches 94.25% on LongText-CN and 90.20% on LongText-EN, the highest among compared methods at this budget, with the code released.

Main Findings

• Best long-text scores at 2 NFE. IDM (Z-Image) achieves 94.25% on LongText-CN and 90.20% on LongText-EN, improving over Z-Image-Turbo++ by 2.63 and 0.32 percentage points and over APEX by 3.87 and 1.44 points respectively.

• Competitive general quality at similar latency. With full-parameter tuning at 2 NFE, GenEval is 0.78 on Z-Image and 0.91 on Qwen-Image. On Qwen-Image under the same setting, APEX scores 0.90; FID is 6.34 for IDM versus 6.44 for APEX, with latencies of 0.44 and 0.45 seconds — so the quality gain comes at comparable sampling cost.

• A single checkpoint serves both budgets. At 1 NFE, Qwen-Image IDM reaches 0.90 GenEval at 0.35 seconds per sample, using the same trained checkpoint.

• Extra evaluations still help. Full-parameter checkpoints improve from 1 to 2 NFE: Z-Image FID drops from 7.13 to 6.71 and Qwen-Image FID from 6.77 to 6.34, each gaining 0.01 GenEval.

• LoRA retains most of the benefit. Rank-64 LoRA on Qwen-Image trains 0.4B parameters and reaches 0.90 GenEval at 2 NFE, versus 0.91 under full tuning.

• Ablations support the moment loss. Controlled ablations show gains from the added moment loss, and generator EMA features perform competitively with an external DINOv2 encoder across the five reported metrics.

• Throughput. IDM (Z-Image) reports 5.18 samples/s at 2 NFE and 5.41 at 1 NFE (batch of 10, single A100, BF16); IDM (Qwen-Image) reports 3.28 and 3.44 samples/s respectively.

Methodology in Plain English

The starting point is rectified flow: a network predicts a velocity that moves noise toward data, and at one network evaluation the generator is just the noise minus that predicted velocity. To train a fast generator, the authors use distribution matching distillation, which needs two scores — one describing the reference (real) distribution and one describing what the generator currently produces. Instead of training a separate "fake score" network to estimate the second, they follow APEX and carve two internal velocity branches out of a single model by feeding it the condition vector and its negation. The difference between these branches, converted through a score–velocity identity, gives the correction direction that DMD applies to the samples.

The new part comes from a mathematical observation: the reverse-KL field can be split exactly into a piece determined only by means and covariances, plus a leftover "non-Gaussian residual" that captures everything else. That split says two things at once — means and spread are things you can constrain directly, and they are not enough on their own, since two distributions can share moments and still look different. So the authors keep the full DMD correction and add a second, complementary loss.

That second loss compares generated and reference samples in the feature space of the generator's own EMA copy. Features are mean-pooled last-layer activations; the loss matches the feature means and the per-coordinate standard deviations (a squared 2-Wasserstein distance between diagonal Gaussians over features). Pools of up to K = 16384 detached model and reference features ("queues") supply stable statistics, with model moments combining the current batch with the model queue. The total objective is the RCGM base loss plus a weighted score loss (λ_score = 1) and the moment loss (λ_pop = 0.5), with the EMA updated as θ̄ ← 0.99 θ̄ + 0.01 θ. Feature extraction and queues exist only during training; inference is a single forward pass.

Training used BF16, FlashAttention4, and FSDP2 with global batch size 64: 15K steps on Z-Image-Turbo 6B and 10K steps for Qwen-Image 20B LoRA, on 32× H100 80GB and 8× A100 80GB GPUs. Reference data combined ShareGPT-4o, BLIP-3o, and samples synthesized with Qwen-Image-2512 and Z-Image-Turbo — 600K from FLUX-Reason-6M prompts and 5M from text-rendering prompts, versus the 200K poster-prompt samples used by APEX.

Why This Matters

Impact on research. The paper offers a clean theoretical justification for why a moment constraint is a sensible addition to DMD rather than a replacement, and it shows a way to get auxiliary distributional supervision without a discriminator or an external encoder. That removes a whole extra network from the training stack while keeping the statistical signal, which is useful for anyone distilling large backbones under compute constraints.

Real-world applications:

• Fast on-device or API-served image generation where per-image latency and GPU cost dominate.
• Poster, packaging, and signage creation, where long strings of text must be rendered correctly in one shot.
• E-commerce and advertising pipelines needing high-volume, low-latency text-in-image asset generation.
• Multilingual design tools, given the separate Chinese and English long-text evaluation.

Industry relevance. Both backbones (Z-Image-Turbo 6B and Qwen-Image 20B) are large production-scale models, and the paper reports LoRA results that train only 0.4B parameters while reaching 0.90 GenEval. That matters for teams that fine-tune rather than pretrain, and the reported throughput/latency numbers on A100 hardware give a concrete picture of serving economics.

Future Directions

• Going beyond pooled means and per-coordinate standard deviations: the decomposition implies the residual carries higher-order structure, so richer feature statistics or conditional (per-prompt) matching instead of pooled statistics are natural extensions.
• Testing whether the moment loss transfers to other distillation recipes and architectures beyond the internal velocity-branch construction borrowed from APEX.
• Pushing below the reported budgets or into other fine-grained conditional tasks where one-shot rendering is the bottleneck — the paper motivates long text specifically, but the same argument applies to other detail-dense prompts.
• Understanding the interaction between the moment loss and the queue-based statistics: queue entries come from earlier model and EMA states, so the moments only approximate current population statistics, and the authors leave the exactness of this approximation as an open design question.

Target Audience

Researchers and engineers working on diffusion/flow-model acceleration, distillation, and few-step image generation — particularly those already familiar with DMD-style training who want to know what an explicit statistical constraint buys them. It is also relevant to practitioners deploying large text-to-image backbones who care about text rendering fidelity and serving latency, and to readers interested in the theoretical link between score differences and moment matching. Beginners will find the motivation sections dense, though the high-level idea and the experimental tables are readable without the derivations.

Authors’ abstract

One-step and few-step text-to-image generation must accurately render the details specified in the prompt despite limited opportunities for refinement. Distribution matching distillation (DMD) uses score differences to bring the generated distribution closer to a reference distribution. Following prior work, we estimate this difference with internal velocity branches, avoiding a separate fake-score network. We ask whether an explicit moment constraint improves this learned score supervision. A decomposition of the reverse-KL field into a Gaussian moment component and a non-Gaussian residual identifies statistics that can be constrained explicitly. We introduce Intrinsic Distribution Matching (IDM), which retains the full DMD correction and matches feature means and per-coordinate standard deviations using the generator's exponential-moving-average (EMA) copy, requiring no external feature encoder. Experiments use full-parameter and low-rank adaptation (LoRA) tuning on Z-Image-Turbo 6B and Qwen-Image 20B. At two network evaluations, IDM reaches 94.25% and 90.20% on Chinese and English LongText-Bench, the highest scores among the compared methods at this budget. Controlled ablations show gains from the added moment loss, while generator EMA features perform competitively with an external DINOv2 encoder across the five reported metrics. Code is available at https://github.com/inclusionAI/IDM.