Source-linked AI summary
High-Fidelity Two-Step Image Generation via Teacher-Aligned End-to-End Distillation
Dongyang Liu, Ruoyi Du, David Liu, Dengyang Jiang, Liangchen Li, Qilong Wu, Zhen Li, Steven C. H. Hoi, Hongsheng Li, Peng Gao
TL;DR
Two-step diffusion generation remains difficult because each denoising step must span a large noise-to-data interval. Z-Image Turbo++ combines teacher-aligned adversarial learning, step-specific parameters, and end-to-end iterative regularization, substantially narrowing the quality gap with 8-step generation while using only two steps.
Problem
Two-step diffusion distillation remains challenging because each step must cover a large noise-to-data interval, making the learning target crucial for stable training.
Method
Z-Image Turbo++ combines teacher-generated adversarial targets, independent parameters for each denoising step, and end-to-end training with iterative regularization.
Results
The model substantially narrows the gap between 2-step and 8-step generation while preserving most of the teacher’s visual quality and benchmark performance.
Takeaways & Limitations
Carefully tailored distillation strategies can improve the quality-efficiency trade-off of few-step image generation.
Takeaways & Limitations
Step-decoupled parameterization increases storage requirements, which can limit deployment on low-memory devices.
Abstract
from arXiv · showhide
Few-step diffusion distillation has become increasingly mature for 4-8-step generation, yet pushing further to 2 steps remains challenging. In this work, we introduce Z-Image Turbo++, a high-quality 2-step image generation model distilled from the 8-step Z-Image Turbo teacher. Our method addresses the central bottlenecks of increased task difficulty and limited model capacity in 2-step generation through three simple but effective design choices tailored to this regime. First, we propose Distribution-Aligned Adversarial Learning, which uses teacher-generated images rather than external real images as real samples for GAN training, providing a more attainable and informative adversarial target. Second, we adopt Step-Decoupled Parameterization, assigning independent model parameters to the two denoising steps to better match their distinct capacity demands. Third, we perform End-to-End Training with Iterative Regularization, allowing the first step to receive gradients from final image quality while preserving a meaningful intermediate generation through an explicit step-1 loss. Together, these designs substantially narrow the quality gap between 2-step and 8-step generation in both qualitative and quantitative evaluations, highlighting the potential of carefully tailored distillation strategies for improving the quality-efficiency trade-off in few-step generation.
1 Introduction
Z-Image Turbo++ targets the severe quality degradation of naive 2-step diffusion distillation by addressing optimization difficulty and limited capacity under step specialization. Its teacher-aligned adversarial training, step-decoupled parameterization, and end-to-end iterative regularization narrow the gap to 8-step generation while using only two denoising steps.
- Motivation: Diffusion sampling typically requires 40–100 neural network evaluations, motivating few-step distillation for more efficient image generation.Few-step methods and released models have established successful compression to 4–8 steps.
- Problem: Naively reducing existing methods to 2 steps causes severe performance degradation, although 2 steps retain iterative structure while maximizing inference efficiency.Single-step generation remains too challenging for satisfactory quality, whereas 4–8 steps leave efficiency gains available.
- Challenges: Two-step generation is difficult because each denoising step spans a large noise-to-data interval and the steps perform sharply different, demanding roles.The first step transforms pure noise into a meaningful intermediate state, while the second refines it into a clean image; shared parameters become restrictive.
- Contribution: Z-Image Turbo++ is distilled from the 8-step Z-Image Turbo teacher and specialized for high-quality 2-step image generation.The method is designed around the two core challenges identified for this regime.
- Result: The proposed designs narrow the gap between 2-step and 8-step generation, preserving most of the teacher’s visual quality and benchmark performance with only two denoising steps.The design set includes adversarial teacher alignment and step-specific models to address optimization and capacity limitations.
- Method: Using 8-step teacher generations as GAN real samples, rather than external real images, provides a more stable, attainable optimization path with more informative gradients.The closer teacher–student distributional alignment improves training stability and final quality.
2 Related Work
Prior work accelerates diffusion models through faster samplers, model-level efficiency techniques, and student–teacher step distillation. Related approaches include consistency- and distribution-matching methods, adversarial distillation, and the Z-Image-Turbo lineage.
- Diffusion Model Acceleration: Diffusion acceleration reduces sampling cost through advanced ODE solvers, pruning, quantization, caching, and step-distillation methods.ODE-based methods reduce sampling steps without retraining, while other approaches modify the model or train a student generator.
- Few-Step Distillation: Few-step distillation includes iterative step halving, trajectory self-consistency, and distribution matching between student and teacher outputs.Representative families include Progressive Distillation, Consistency Models, Distribution Matching Distillation, and Decoupled DMD.
- GAN-Based Distillation: GAN-based distillation progresses from external visual discriminators to teacher-derived generative features and combines adversarial objectives with progressive distillation.ADD uses lightweight discriminator heads with a pretrained visual feature extractor, whereas LADD performs adversarial distillation directly in latent space using diffusion-teacher features.
- Z-Image: Z-Image uses a unified S3-DiT stream for text, image, and latent tokens, while Z-Image-Turbo combines Decoupled DMD and DMDR for an 8-step model.The 6B-parameter Z-Image demonstrates strong photorealistic generation, and its Turbo variant targets high visual fidelity with practical inference efficiency.
3 Preliminary
The paper builds on flow matching, where a learned velocity field transports noise to data through ODE integration, and on few-step distillation methods including DMD and GAN objectives. Decoupled DMD identifies CFG augmentation as the main driver of few-step conversion, while distribution matching regularizes training and suppresses artifacts.
- Flow Matching: Flow matching defines t = 0 as pure noise and t = 1 as clean data, then trains a neural network to predict the velocity field between them.Generation integrates the learned field from t = 0 to t = 1 with a numerical ODE solver.
- Distribution Matching Distillation: DMD trains a few-step student generator by minimizing the integral KL divergence between student and teacher distributions.Its practical gradient uses a frozen real score model and a concurrently trained fake score model.
- Decoupled DMD: Decoupled DMD attributes few-step conversion primarily to CFG Augmentation, while Distribution Matching mainly regularizes training and suppresses artifacts.This insight motivates schedule design for 4–8-step distillation and supports Z-Image-Turbo's 8-step generation capability.
- GAN Objectives: GAN distillation uses a discriminator to distinguish real from fake images and provide an adversarial gradient to the generator.The paper follows the common architecture of a frozen cloned multi-step diffusion model with lightweight discriminator heads.
4 Method
The method distills an 8-step teacher into a 2-step generator using distribution-aligned adversarial learning, step-decoupled parameters, and end-to-end training with iterative regularization. These choices target distribution mismatch, per-step capacity demands, and the need to preserve meaningful intermediate representations.
- Overall Pipeline: The two-phase pipeline assumes an established 8-step teacher and distills it into a 2-step generator using three synergistic techniques.Phase 1 treats few-step teacher preparation as a prerequisite; Phase 2 applies the proposed distillation designs.
- Distribution-Aligned Adversarial Learning: Distribution-aligned adversarial learning uses 8-step teacher outputs as discriminator real samples instead of external real images.Teacher outputs better match the student’s target distribution, avoiding artifacts caused by distributional differences between diffusion outputs and natural photographs.
- Distribution-Aligned Adversarial Learning: Teacher-generated real samples improve training stability and final generation quality, while external real images produce pronounced artifacts.With teacher outputs, GAN loss initially rises and plateaus; with external images, it remains substantially higher and continues growing.
- Step-Decoupled Parameterization: Step-decoupled parameterization independently trains two step-specific models initialized from the 8-step teacher, addressing the larger per-step burden in 2-step generation.The approach substantially decreases the step-2 generator GAN loss, but doubles parameter count and can increase device-side storage demands.
- End-to-End Training with Iterative Regularization: End-to-end training propagates final-output gradients through both steps, while an explicit step-1 loss preserves a meaningful intermediate representation.Removing the step-1 loss surges the step-2 GAN loss and visibly degrades generation quality, making the constraint essential for stable training.
5 Experiments
Experiments show that Z-Image Turbo++ produces high-quality two-step images and approaches the 8-step teacher across four benchmarks. Ablations support teacher-aligned GAN targets and step-specific parameters, while dense text and complex secondary objects remain challenging.
- Qualitative Results: With only two inference steps, Z-Image Turbo++ generates rich details, sharp textures, photorealistic images, and strong text rendering.Qualitative comparisons indicate that it preserves most of the original Z-Image model’s image quality and appearance.
- Quantitative Results: Benchmark trends are non-uniform: OneIGBench and LongTextBench separate model variants more clearly, whereas GenEval and DPGBench produce more mixed comparisons.The per-step LoRA variant outperforms the 8-step teacher on GenEval and DPGBench despite performing substantially worse on OneIGBench and LongTextBench.
- Quantitative Results: Z-Image Turbo++ outperforms TwinFlow, DMD2, and direct two-step Z-Image Turbo inference while approaching the 8-step Turbo baseline on nearly all metrics.The evaluation uses DPGBench, GenEval, OneIGBench, and LongTextBench with prompts taken directly from official benchmark sets without enhancement.
- Limitations: A consistent text-generation gap remains on LongText-CN, LongText-EN, and OneIG-Text, despite visual fidelity and primary-object generation largely matching the 8-step baseline.Dense text and secondary or underspecified objects in complex scenes are less reliable.
- Ablation Summary: Replacing external real images with 8-step teacher-generated images removes GAN-training artifacts and produces more natural outputs.External high-quality data causes training instability and systematic artifacts.
- Ablation Summary: Independent weights for the two denoising steps lower the generator GAN loss, indicating that the resulting two-step distribution is harder to distinguish from the 8-step teacher.This ablation supports parameter decoupling as a useful design choice for two-step generation.
6 Conclusion and Limitations
Z-Image Turbo++ is a 2-step image generation model distilled from the 8-step Z-Image Turbo teacher. It combines three tailored techniques to address stability, capacity, and knowledge-preservation challenges, with additional parameter storage as its main limitation.
- Contribution: Z-Image Turbo++ is distilled from the 8-step Z-Image Turbo teacher into a 2-step image generation model.The paper presents this model as its central contribution.
- Contribution: Distribution-Aligned Adversarial Learning, Step-Decoupled Parameterization, and End-to-End Training with Iterative Regularization target the 2-step regime’s stability, capacity, and knowledge-preservation challenges.The method combines all three design choices to address these challenges.
- Limitations: Additional parameter storage is the model’s main limitation.The conclusion explicitly identifies parameter storage overhead as the primary limitation.
A Pseudo-code for Memory-Efficient Training
The pseudo-code contrasts naive two-step generator training with a memory-efficient update using detached intermediate activations and inherited gradients. The implementation preserves separate step losses while controlling gradient transfer into step 0.
- Naive generator training: Naive training computes GAN and DMD losses for both denoising steps before backpropagating their sum.The step-0 and step-1 losses each include GAN_LOSS and DMD_LOSS, followed by (step0_loss + step1_loss).backward().
- Memory-efficient implementation: The inherited gradient is scaled by inherit_weight before backpropagation to avoid gradient explosion when passing the full gradient directly to step 0.The implementation sets inherit_weight=0.1 by default and backpropagates total_step0_loss = step0_loss + inherit_weight * inherit_loss.
- Memory-efficient implementation: The memory-efficient update detaches step 0’s velocity prediction while keeping it differentiable for receiving step-1 gradients.step0_v_prediction is cloned with detach(), then requires_grad_(True), allowing step1_loss.backward() to produce an inherited gradient.
B More Experimental Details · B.1 Experiment Settings
The experiments use a fixed Adam-based training setup with teacher-generated samples, distributed execution, and explicit optimization, loss-weight, and benchmark-reporting details. Training runs for 20,000 iterations with exponential moving averaging and takes approximately 80 hours on 16 H100 GPUs.
- B.1 Experiment Settings: All trainable models use Adam with learning rate 1 × 10^-5, β1 = 0.0, β2 = 0.9, and no weight decay.These settings apply throughout the trainable-model experiments.
- B.1 Experiment Settings: Following TTUR, the generator and guidance model are updated at a 1:5 frequency ratio.The ratio follows the TTUR strategy in DMD2 [34].
- B.1 Experiment Settings: Training lasts 20,000 iterations, while generator parameters use exponential moving averaging with decay rate 0.99.The moving average is maintained throughout full training.
- B.1 Experiment Settings: Samples are pre-generated from the 8-step teacher to define the teacher distribution used in subsequent training.This distribution is prepared before optimization begins.
- B.1 Experiment Settings: Experiments run on 16 H100 GPUs with global batch size 64, and the complete training process takes approximately 80 hours.These figures describe the stated implementation setup and runtime.
- B.1 Experiment Settings: The section specifies loss weights for DMD, diffusion, GAN, and inherit losses, and reports detailed results on DPG-Bench and GenEval.DMD and guidance-model diffusion losses are each weighted 1 × 10^-2; GAN generator and discriminator losses are each 1 × 10^-3; inherit loss is 0.1.