Intrinsic Distribution Matching

IDM keeps the full internal DMD correction and adds a feature-moment loss on generator EMA features, improving 2-NFE LongText scores on Z-Image-Turbo and Qwen-Image.

One-step and few-step text-to-image generation must render prompt details with few refinement steps. DMD uses score differences between generated and reference distributions but relies on learned score estimates and finite training. The paper asks whether an explicit moment constraint can improve this learned score supervision.

The paper decomposes the reverse-KL score field into a Gaussian moment component determined by means and covariances and a non-Gaussian residual. It retains the full internal DMD correction estimated from velocity branches and adds a feature-moment loss. The loss matches means and per-coordinate standard deviations of features from the generator's EMA copy for generated and reference samples, without an external feature encoder. Training combines a base RCGM loss, the DMD score loss, and the feature-moment loss; EMA parameters are updated with 0.99/0.01 and feature queues are detached. Features and queues are used only in training; one-step inference uses one forward pass, and the checkpoint also supports few-step sampling.

• IDM on Z-Image-Turbo at two network evaluations reaches 94.25% and 90.20% on Chinese and English LongText-Bench, the highest scores among the compared methods at this budget. §1
• The same 2-NFE result improves over Z-Image-Turbo++ by 2.63 and 0.32 percentage points on LongText-CN and LongText-EN. §1
• In Table 5, the full IDM row reports OneIG 53.06, GenEval 77.81, DPG-Bench 86.82, LongText-CN 94.25, and LongText-EN 90.20. §5.3
• Removing the DMD score loss lowers LongText-CN to 93.33 and LongText-EN to 89.62; removing the feature-moment loss gives 93.87 and 89.79; removing both gives 92.70 and 88.88. §5.3
• In the feature-map ablation, EMA features score 94.25 and 90.20 on LongText-CN and LongText-EN, versus DINOv2 94.13 and 90.09, teacher features 93.79 and 89.81, and live self-features 93.95 and 89.95. §5.3
• System-level Table 1 reports IDM Qwen-Image at 2 NFE with FID 6.34, CLIP 28.56, and GenEval 0.91; IDM Z-Image at 2 NFE reports FID 6.71, CLIP 28.22, and GenEval 0.78. §5

• The feature-moment loss is described as a pooled feature-space moment regularizer motivated by the Gaussian component, not an estimate of it.
• Practical neural branches and mixed objectives do not guarantee recovery of exact marginal velocities, so the DMD estimate is approximate.
• Queue staleness can introduce bias, and the paper states it does not establish that the feature-moment loss has lower empirical parameter-gradient variance than DMD.
• All controlled ablations use IDM (Z-Image) with full-parameter tuning at 2 NFE, leaving other settings open.

• The reverse-KL distribution gradient is s_theta - s_r, and the descent direction is s_r - s_theta. §3.1
• The score of p decomposes exactly into the score of its moment-matched Gaussian plus a non-Gaussian score residual. §3.3
• The reverse-KL field is the sum of a Gaussian moment component depending only on means and covariances and a non-Gaussian residual component; equal moments can leave a nonzero residual. §3.3
• Under stated integrability and boundary conditions, the residual is orthogonal to constant and linear functions under its own distribution. §3.3
• The feature-moment loss compares generated and reference samples in the feature space of the generator's EMA copy. §4.2
• The generator objective is L_base + lambda_score L_score + lambda_pop L_pop_hat, with lambda weights fixed in Section 5.1. §4.3
• The EMA is updated as bar theta <- 0.99 bar theta + 0.01 theta; queued features are detached and receive no gradients, while current features keep their computation graph. §4.3
• Feature extraction and queues are used only during training; one-step inference is one forward pass and the same checkpoint supports few-step sampling. §4.3