Z-Image-Turbo-SDA is a highly efficient LoKr (Low-Rank Kronecker Product) adapter designed to rescue the "Diversity Collapse" problem in few-step distilled Flow Matching / Diffusion models.
By applying our novel Semantic Directional Alignment (SDA) Loss, this LoKr adapter recovers 70.2% of the original teacher model's compositional diversity (LPIPS) while perfectly maintaining the extreme sharpness and fast inference speed of the 8-step baseline model.
📖 The Problem: Diversity Collapse in Distillation
When distilling a 50-step diffusion model (Teacher) into an 8-step model (Student), the student often learns to shortcut the manifold, converging to the "mean" of the distribution. As a result, changing the random noise seed yields almost identical compositions (same poses, same layouts).
Standard SFT (Supervised Fine-Tuning) fails to fix this because forcing the student to match the teacher's absolute velocity ($v$) destroys the rectified straight-line trajectory, leading to blurry, gray, or scrambled images.
Instead of matching raw velocity, we fine-tuned this model using a custom 4-pillar physics-based architecture, backed by the parameter-efficient LoKr structure:
1. $x_0$-Space Cosine Loss (Direction Without Magnitude)
Instead of matching the raw velocity $v$, we project predictions into the clean $x_0$ space ($\hat{x}_0 = z - v$). We then apply Spatial Low-Pass Filtering (AvgPool2d, $k=8$) to prevent the model from cheating via high-frequency noise. Finally, we use Standard Cosine Loss to align the semantic direction of the student's variance with the teacher's, completely freeing the student to maintain its high-contrast magnitude.
2. Self-Reference SFT Anchor (The "Thermostat" Effect)
To prevent the diversity push from destroying the 8-step trajectory, we introduce a continuous anchor:
sft_loss = Huber(v_student, v_baseline_detached, delta=0.08)
This acts as an elastic tether. If the model strays too far to change the composition, the Huber loss provides a constant (L1) pull to polish the high-frequency details without violently snapping the model back to the collapsed mean.
3. Asymmetric Adapter Masking (The Breakthrough)
If a Pre-trained Training Adapter is used, we apply it asymmetrically:
Diversity Forward (Adapter OFF): Exposes the raw collapsed manifold of the student, generating massive, accurate gradients for composition correction.
SFT Forward (Adapter ON): Provides an indestructible quality baseline, locking in the sharpness.
4. Extreme-Low-Noise Skip (div_skip)
At the final inference step (e.g., $t \approx 8.9$), macro-composition is already decided. Pushing diversity here only causes high-frequency jitter (blur). We completely cut off the Diversity Loss at the lowest noise step, utilizing pure SFT to guarantee razor-sharp final pixels.
💻 How to Use
This adapter is trained as a LoKr (LyCORIS ecosystem). Modern versions of diffusers coupled with peft natively support loading LoKr weights seamlessly.
python
1from diffusers import DiffusionPipeline
2import torch
34# 1. Load the base 8-step distilled model5pipeline = DiffusionPipeline.from_pretrained(6"Tongyi-MAI/Z-Image-Turbo",7 torch_dtype=torch.float16
8).to("cuda")910# 2. Load the SDA Diversity LoKr Adapter11# (Ensure you have the latest `diffusers` and `peft` installed)12pipeline.load_lora_weights("F16/z-image-turbo-sda", adapter_name="sda_diversity")1314prompt ="A lone traveler standing on a mountain peak, epic fantasy lighting"1516# 3. Generate diverse images with different seeds!17# You will now get completely different poses, camera angles, and layouts across seeds.18for seed in[42,123,777,999]:19 generator = torch.Generator(device="cuda").manual_seed(seed)20 image = pipeline(21 prompt=prompt,22 num_inference_steps=8,23 guidance_scale=1,24 generator=generator
25).images[0]26 image.save(f"traveler_seed_{seed}.png")
📊 Quantitative Evaluation
Evaluated across complex prompts and 16 random seeds per prompt at Step 2500:
Model Variant
Avg LPIPS (Perceptual Diversity) $\uparrow$
Pixel StdDev (Magnitude Variance) $\downarrow$
8-Step Baseline (Collapsed)
0.564
0.190
8-Step + SDA LoKr (Ours)
0.691(+0.127)
0.209(Pristine Sharpness)
50-Step Teacher
0.745
0.243
Why this matters:
The SDA LoKr successfully decoupled structural diversity from destructive noise. It bumped the macro-compositional diversity (LPIPS) close to the 50-step teacher, while strictly constraining the Pixel StdDev to prevent the "darkening/blurring" effect typical in continuous ODE fine-tuning.
💻 Implementation & Usage
The core loss mechanism can be easily integrated into any Diffusers / Flow Matching training loop.
Pseudo Code
python
1import torch
2import torch.nn.functional as F
34defcompute_sda_loss(student_model, teacher_model, z1, z2, latents, t_continuous, v_self_detached):5# 1. Teacher Predictions (No Grad)6with torch.no_grad():7 x0_T_z1 = z1 - teacher_model(z1, t_continuous)8 x0_T_z2 = z2 - teacher_model(z2, t_continuous)910# 2. Student Baseline (from detached SFT forward, Adapter is ON)11 x0_S_z1 = z1 - v_self_detached
1213# 3. Asymmetric Adapter Masking: Adapter-OFF for Diversity14# Temporarily disable the pre-trained assistant adapter to expose 15# the raw collapsed manifold and generate massive diversity gradients.16 student_model.assistant_adapter.is_active =False17try:18# Student forward with Adapter OFF19 x0_S_z2 = z2 - student_model(z2, t_continuous)20finally:21# Crucial: Turn it back ON for subsequent SFT / validation steps22 student_model.assistant_adapter.is_active =True2324# Calculate Deltas25 delta_T = x0_T_z2 - x0_T_z1
26 delta_S = x0_S_z2 - x0_S_z1
2728# 4. Spatial Low-Pass Filter (Force Macro-Composition changes, prevent high-freq cheating)29 pooled_T = F.avg_pool2d(delta_T, kernel_size=8).view(delta_T.size(0),-1)30 pooled_S = F.avg_pool2d(delta_S, kernel_size=8).view(delta_S.size(0),-1)3132# 5. Standard Cosine Loss (Direction alignment without magnitude penalty)33 cos_sim = F.cosine_similarity(pooled_S, pooled_T, dim=-1)34 div_loss =(1.0- cos_sim).mean()3536return div_loss
3738# --- Training Loop Integration ---3940# SFT Forward (Adapter ON) -> Protects pristine image quality41with torch.no_grad():42 v_self = student_baseline(xt_z1, t)4344# Self-Reference Huber Loss (Elastic anchor to preserve 8-step trajectory)45sft_loss = F.huber_loss(v_pol_z1, v_self.detach(), delta=0.08)4647# Time-Segmented Diversity (Enabled only for t > div_skip_threshold)48# Skip Index 7 (extremely low noise) to lock in final pixel sharpness49if t > LOW_NOISE_THRESHOLD:50 div_loss = compute_sda_loss(student_model, teacher_model, z1, z2, latents, t, v_pol_z1.detach())51 loss = sft_loss + diversity_lambda * div_loss
52else:53 loss = sft_loss # Pure SFT
Timesteps: Strict discrete scheduler sampling (indices 1-6) with div_skip at extremely low noise (index 7).
❓ FAQ (Frequently Asked Questions)
Q1: Can I use SDA LoRA alongside other Z-Image-Turbo (ZIT) LoRAs?
A: Yes, but with potential compatibility risks. You will need to balance the weights manually.
The Physics: The core mechanism of SDA LoRA involves rotating the model's velocity field. Because the prediction direction has been intentionally deviated from the original rectified trajectory to achieve diversity, it may clash with other ZIT LoRAs that expect the original "straight-line" path.
The Conflict: When stacking multiple LoRAs, you might encounter structural artifacts or blurring due to conflicting directional gradients.
The Solution: Try lowering the inference weight of the SDA LoRA (e.g., to 0.5 ~ 0.7). Note the inherent trade-off: reducing the weight weakens the "semantic push," causing diversity to regress back toward the collapsed mean. You will need to find the "sweet spot" based on your specific needs.
Q2: Why is the diversity more dramatic in simple prompts compared to complex ones?
A: This depends on the "topological constraints" the prompt places on the composition space.
Simple Prompts (e.g., "A cat sitting"): SDA can unleash maximum diversity here because the prompt imposes few restrictions on the latent space. The "cat" has the freedom to sit, lie down, or look in different directions.
Complex Prompts (e.g., "A cat on the left, looking at a red cup on the right"): These prompts act like a rigid "blueprint," strictly anchoring elements to specific coordinates. Since the text encoder's guidance is extremely strong, the SDA LoRA cannot easily flip the composition without violating the prompt's instructions. Consequently, longer, more descriptive prompts will naturally show less spatial variation.
⚠️ Developer Notes & Limitations
SDA is a result of my personal study and experimental journey into the underlying mechanics of Flow Matching models. As such, it is far from perfect.
The current v1 release was derived from our best internal checkpoint (div_27_2500, which achieved up to 70% diversity recovery), followed by an additional 360 steps of quality-focused fine-tuning.
A Regrettable Trade-off (Anatomy vs. Diversity): Maintaining rigid structures like limbs and hands requires the model to stay close to its original distilled manifold. To achieve better anatomical stability, I had to compromise on the "semantic push." As a result, the composition diversity recovery regressed from its peak of 70% back to 58% in this final version.
Residual Issues: Despite sacrificing 12% of the diversity, the model may still produce occasional anatomical errors in complex poses—a common "Achilles' heel" for few-step distilled models.
Solving the zero-sum game between "macro-compositional diversity" and "micro-anatomical rigidity" within an 8-step inference framework remains an open challenge. I hope to find a more elegant solution in the future, and I welcome the community to build upon this exploration.
🤝 Citation
If you use this model or the SDA methodology in your research, please cite:
