Summary (Overview)

• Intrinsic Distribution Matching (IDM) is a novel training method for one-step and few-step text-to-image generation that augments internal Distribution Matching Distillation (DMD) with an explicit feature-moment matching loss, eliminating the need for external feature encoders or separate fake-score networks.
• The method decomposes the reverse-KL field into a Gaussian moment component (determined by means and covariances) and a non-Gaussian residual component, providing theoretical motivation for why explicit moment constraints complement learned score supervision.
• IDM matches feature means and per-coordinate standard deviations in the generator's exponential-moving-average (EMA) feature space, using a squared 2-Wasserstein distance between diagonal Gaussians as the population loss.
• At 2 network function evaluations (NFE), IDM achieves state-of-the-art scores of 94.25% and 90.20% on Chinese and English LongText-Bench respectively, surpassing all compared methods at this budget.
• Controlled ablations demonstrate that generator EMA features perform competitively with an external DINOv2 encoder across five reported metrics, while requiring no additional architectural components.

Introduction and Theoretical Foundation

Background

Diffusion and flow-matching models produce high-quality images but require repeated network evaluations, imposing substantial inference costs. One-step and few-step generation reduces this cost by learning a much shorter sampling process, but must accurately render prompt details—particularly challenging for long-text rendering where many characters must remain legible and correct within a coherent image.

Distribution Matching Distillation (DMD) trains a fast generator to match a reference distribution by using score differences:

• A pretrained teacher supplies the reference score
• An auxiliary fake-score network estimates the score of the current generated distribution
• Their difference guides updates to generated samples

Recent work (TwinFlow, APEX) obtains this score difference from internally separated velocity branches of the generator, avoiding a separate fake-score network. IDM adopts this internal construction and asks whether the generator can also supply an explicit constraint on statistics across samples.

Theoretical Foundation: Rectified Flow

Let $p_z = \mathcal{N}(\mathbf{0}, \mathbf{I}_d)$ be the noise distribution and $p_r$ the reference data distribution. Rectified flow uses the path:

$$\mathbf{x}_{r,t} = (1-t)\mathbf{x}_r + t\epsilon$$

with independently sampled $\mathbf{x}_r \sim p_r$ and $\epsilon \sim p_z$. A network $F_\theta(\mathbf{x}_{r,t}, t, \mathbf{c})$ regresses the pair velocity $\epsilon - \mathbf{x}_r$, where $\mathbf{c} \in \mathcal{C}$ is a text-conditioning vector.

The score-velocity identity connects scores and velocities:

$$\mathbf{s}_t(\mathbf{x}_t) = -\frac{\mathbf{x}_t + (1-t)\mathbf{v}_t^*(\mathbf{x}_t)}{t} \tag{1}$$

Thus, a velocity estimate also provides a score estimate, and subtracting two velocities at the same input gives a scaled score difference.

Methodology

Reverse-KL Field Decomposition

For differentiable positive densities, the reverse KL divergence at fixed condition $\mathbf{c}$ is:

$$\mathcal{D}_{\mathrm{KL}}(p_\theta \| p_r) = \int p_\theta(\mathbf{x}|\mathbf{c}) \log \frac{p_\theta(\mathbf{x}|\mathbf{c})}{p_r(\mathbf{x}|\mathbf{c})} d\mathbf{x} \tag{2}$$

with sample-space distribution gradient:

$$\mathbf{g}_{\mathrm{KL}}(\mathbf{x}|\mathbf{c}) = \mathbf{s}_\theta(\mathbf{x}|\mathbf{c}) - \mathbf{s}_r(\mathbf{x}|\mathbf{c}) \tag{3}$$

The score of any distribution $p$ decomposes exactly as:

$$\mathbf{s}_p(\mathbf{x}) = \underbrace{\mathbf{s}_{q_p}(\mathbf{x})}_{= -\Sigma_p^{-1}(\mathbf{x} - \mu_p)} + \underbrace{\boldsymbol{\eta}_p(\mathbf{x})}_{\text{non-Gaussian score residual}} \tag{6}$$

where $q_p := \mathcal{N}(\mu_p, \Sigma_p)$ is the moment-matched Gaussian. Subtracting decompositions for model and reference gives:

$$\mathbf{g}_{\mathrm{KL}}(\mathbf{x}|\mathbf{c}) = \underbrace{\left(\mathbf{s}_{q_\theta}(\mathbf{x}|\mathbf{c}) - \mathbf{s}_{q_r}(\mathbf{x}|\mathbf{c})\right)}_{\mathbf{g}_{\mathrm{KL}}^{\mathrm{G}}: \text{ Gaussian moment component}} + \underbrace{\left(\boldsymbol{\eta}_\theta(\mathbf{x}|\mathbf{c}) - \boldsymbol{\eta}_r(\mathbf{x}|\mathbf{c})\right)}_{\mathbf{g}_{\mathrm{KL}}^{\mathrm{nG}}: \text{ non-Gaussian residual component}} \tag{7}$$

Key theoretical result: Under integrability and boundary conditions, score residuals satisfy orthogonality properties:

$$\mathbb{E}_p[\boldsymbol{\eta}_p(\mathbf{x})] = \mathbf{0}, \qquad \mathbb{E}_p[\boldsymbol{\eta}_p(\mathbf{x})\mathbf{x}^\top] = \mathbf{0}_{d \times d} \tag{8}$$

DMD Loss

The DMD correction supervises velocity and endpoint predictions. The internal velocity branches are obtained by negating the condition vector:

$$\widehat{\mathbf{v}}_{r,t}(\mathbf{x}_t|\mathbf{c}) := F_\theta(\mathbf{x}_t, t, \mathbf{c}), \quad \widehat{\mathbf{v}}_{\theta,t}(\mathbf{x}_t|\mathbf{c}) := F_\theta(\mathbf{x}_t, t, -\mathbf{c}) \tag{10}$$

The DMD targets are constructed with stop-gradient operations:

$$\mathbf{v}_t^{\text{DMD}} := \text{sg}\left(\mathbf{F}_\theta(\mathbf{x}_t, t, \mathbf{c}) - \widehat{\mathbf{d}}_{t'}^{\text{DMD}}\right), \quad \mathbf{x}_0^{\text{DMD}} := \text{sg}\left(\mathbf{x}_t - t\mathbf{v}_t^{\text{DMD}}\right) \tag{11}$$

The combined squared regression loss:

$$\mathcal{L}_{\text{score}}(\theta) := \mathbb{E}\left[\underbrace{\|F_\theta(x_t, t, c) - v_t^{\text{DMD}}\|_2^2}_{\ell_v(t)} + \underbrace{\|\hat{x}_0(x_t, t, c) - x_0^{\text{DMD}}\|_2^2}_{\ell_x(t)}\right] \tag{12}$$

Feature-Moment Loss

The population loss matches feature moments in the EMA feature space:

$$\mathcal{L}_{\mathrm{pop}} = \|\mu_\theta^\phi - \mu_r^\phi\|_2^2 + \sum_{j=1}^{d_\phi}\left(\sqrt{\nu_{\theta,j}^\phi} - \sqrt{\nu_{r,j}^\phi}\right)^2 \tag{13}$$

This is the squared 2-Wasserstein distance between diagonal Gaussians $\mathcal{N}(\mu_i^\phi, \mathrm{diag}(\nu_i^\phi))$.

Training Objective

$$\mathcal{L}_{\mathrm{IDM}}(\theta) = \mathcal{L}_{\mathrm{base}}(\theta) + \lambda_{\mathrm{score}}\mathcal{L}_{\mathrm{score}}(\theta) + \lambda_{\mathrm{pop}}\widehat{\mathcal{L}}_{\mathrm{pop}} \tag{14}$$

with default weights $\lambda_{\mathrm{score}} = 1$ and $\lambda_{\mathrm{pop}} = 0.5$.

Empirical Validation / Results

Experimental Setup

• Backbones: Z-Image-Turbo 6B and Qwen-Image 20B
• Tuning: Full-parameter and LoRA (ranks 32 and 64)
• Training: 15K steps on Z-Image-Turbo, 10K steps for Qwen-Image LoRA
• Hardware: 32 H100 80GB and 8 A100 80GB GPUs
• Data: ShareGPT-4o, BLIP-3o, plus synthesized samples (600K from FLUX-Reason-6M prompts, 5M from text-rendering prompts)

Key Results

Table 1: System-level comparison (excerpts)

LongText-Bench Results (2 NFE): IDM (Z-Image) achieves 94.25% on LongText-CN and 90.20% on LongText-EN—the highest among compared methods, improving over Z-Image-Turbo++ by 2.63 and 0.32 percentage points, and over APEX by 3.87 and 1.44 points respectively.

Ablation Results

Table 5: Overall ablation results (2 NFE)

Key ablation findings:

• Adding $\mathcal{L}_{\mathrm{pop}}$ to the shared DMD control raises GenEval from 77.46 to 77.81 (+0.35 points)
• EMA features reach 77.81 GenEval versus 77.75 for DINOv2—nearly identical without external encoder
• Equal-weight v+x regression reaches 77.81 GenEval, versus 76.64 for v-only and 77.26 for x-only

Theoretical and Practical Implications

Theoretical Implications

The decomposition of the reverse-KL field into Gaussian moment and non-Gaussian residual components provides a principled justification for combining learned score supervision with explicit moment constraints. The key insight is that:

1. Moments alone do not determine the full distribution—the residual component explains why moment matching alone is insufficient
2. The full DMD correction contains both components, but learned score estimates and finite training may leave discrepancies
3. Explicit moment constraints address statistics that can be constrained directly while the learned correction handles the residual structure

Practical Implications

• Architecture efficiency: IDM requires no separate discriminator, score model, or external feature encoder—all training signals come from the generator and its EMA copy
• Scalability: Demonstrated on backbones up to 20B parameters with both full-parameter and LoRA tuning
• Sampling flexibility: The same trained checkpoint supports both 1-NFE and 2-NFE sampling without retraining
• Text rendering: Significant improvements in long-text fidelity, which is critical for real-world applications like poster generation, menus, and document illustration

Conclusion

Intrinsic Distribution Matching (IDM) augments internal DMD with explicit feature-moment matching in the generator's EMA feature space. The reverse-KL field decomposition motivates retaining full score-difference supervision while additionally matching feature means and standard deviations. Controlled ablations show consistent gains over the internal DMD baseline from EMA feature moments, with performance competitive to DINOv2 features without requiring an external encoder.

At 2 NFE, IDM achieves state-of-the-art LongText-Bench scores of 94.25% and 90.20% on Chinese and English subsets respectively, demonstrating that using the generator's predictions and representations together is an effective strategy for few-step text-to-image generation.

Future directions could include:

• Extending moment matching to higher-order statistics or full feature distributions
• Exploring adaptive weighting schemes for the population loss
• Applying IDM to other generative modalities and architectures
• Investigating the interaction between moment matching and other distillation objectives