Source-linked AI summary
WaDi: Weight Direction-aware Distillation for One-step Image Synthesis
Lei Wang, Yang Cheng, Senmao Li, Ge Wu, Yaxing Wang, Jian Yang
TL;DR
Slow multi-step diffusion inference motivates one-step distillation, but the relevant weight changes and optimization challenges require better understanding. WaDi analyzes these changes, introduces low-rank directional rotations through LoRaD within VSD, and achieves strong image-quality, efficiency, and downstream-task results with about 10% of U-Net parameters trainable.
Problem
Multi-step diffusion inference is slow, while existing fine-tuning approaches jointly update norm and direction, increasing optimization difficulty.
Method
WaDi integrates LoRaD, which uses learnable low-rank rotation matrices to model weight-direction changes, into Variational Score Distillation.
Results
WaDi achieves state-of-the-art FID scores on COCO 2014 and COCO 2017 using about 10% of U-Net parameters as trainable components.
Takeaways & Limitations
The distilled one-step model generalizes to controllable generation, relation inversion, high-resolution synthesis, and image customization.
Takeaways & Limitations
The DreamBooth experiment is illustrative rather than a comprehensive study of diffusion fine-tuning.
Abstract
from arXiv · showhide
Despite the impressive performance of diffusion models such as Stable Diffusion (SD) in image generation, their slow inference limits practical deployment. Recent works accelerate inference by distilling multi-step diffusion into one-step generators. To better understand the distillation mechanism, we analyze U-Net/DiT weight changes between one-step students and their multi-step teacher counterparts. Our analysis reveals that changes in weight direction significantly exceed those in weight norm, highlighting it as the key factor during distillation. Motivated by this insight, we propose the Low-rank Rotation of weight Direction (LoRaD), a parameter-efficient adapter tailored to one-step diffusion distillation. LoRaD is designed to model these structured directional changes using learnable low-rank rotation matrices. We further integrate LoRaD into Variational Score Distillation (VSD), resulting in Weight Direction-aware Distillation (WaDi)-a novel one-step distillation framework. WaDi achieves state-of-the-art FID scores on COCO 2014 and COCO 2017 while using only approximately 10% of the trainable parameters of the U-Net/DiT. Furthermore, the distilled one-step model demonstrates strong versatility and scalability, generalizing well to various downstream tasks such as controllable generation, relation inversion, and high-resolution synthesis.
1. Introduction
Diffusion distillation targets faster one-step generation, but the analysis finds that weight-direction changes dominate weight-norm changes. WaDi addresses this with low-rank directional rotations integrated into VSD, improving efficiency and image quality.
- Multi-step diffusion models generate high-quality images but incur high computational cost and slow inference, motivating one-step distillation.
- Weight direction varies more than weight norm between one-step students and their multi-step teachers.In U-Net architectures, weight norm changes have mean 0.1% and standard deviation 0.2% across layers.
- 30% of the residual direction matrix rank recovers 93% of its information, indicating that directional changes are low-rank and structured.
- Replacing the norm changes DMD2 by +0.7 FID with unchanged CLIP, whereas replacing direction causes +241.3 FID and -0.18 CLIP.
- LoRaD models directional adjustments with learnable low-rank rotation matrices and is integrated into VSD as WaDi.The framework optimizes only weight direction and uses about 10% of U-Net parameters as trainable components.
- WaDi achieves state-of-the-art FID scores on COCO 2014 and COCO 2017 while improving inference efficiency and image quality.It is also evaluated on downstream tasks including controllable generation, relation inversion, high-resolution synthesis, and image customization.
2. Related Work
Related acceleration methods reduce diffusion sampling costs through training-free or training-based strategies, but training-based approaches can make optimization difficult. WaDi focuses on directional rotations to address this difficulty.
- Latent diffusion and DiT improve diffusion efficiency and scalability, but iterative denoising remains slow.
- Training-based acceleration includes consistency, progressive, diffusion-GAN, and variational score distillation, each with distinct fidelity or training-cost trade-offs.Consistency distillation can suffer low fidelity, while progressive distillation introduces significant training overhead.
- Existing training-based methods commonly use full fine-tuning or LoRA, which can increase optimization difficulty.
- Directional changes are generally more influential in distillation, motivating WaDi's use of LoRaD to model directional rotations.
3. Method
The method analyzes directional weight changes in diffusion distillation, introduces LoRaD to rotate pretrained weight directions with low-rank parameters, and integrates it into VSD as WaDi.
- LoRaD design: LoRaD rotates pretrained weight directions with learnable rotation matrices while preserving their norms.Rotations are applied independently across paired dimensions of the weight space.
- LoRaD design: The rotation angles are represented as the product of two low-rank matrices to reduce trainable parameters.This exploits the observed low-rank structure of directional changes rather than modeling a full-rank update.
- LoRaD design: For a weight matrix W, LoRaD applies parallel 2 × 2 rotations to odd-even row pairs across the d-dimensional space.The resulting rotated matrix is computed using block-diagonal rotations and element-wise operations.
- Framework overview: WaDi builds on VSD by adding LoRaD, a direction-aware module motivated by the finding that distillation changes weight directions more than norms.The framework uses a pretrained diffusion model as teacher, a trainable fake model, and a one-step generator.
- WaDi optimization: WaDi applies LoRaD to both the one-step generator and fake model, then alternates their optimization to improve alignment with the teacher distribution.The student uses a high-rank rotation, whereas the fake model uses a low-rank rotation for adaptive guidance.
4. Experiment
WaDi is evaluated on COCO benchmarks and downstream tasks, where it combines strong image quality with substantial parameter and inference efficiency. Ablations examine adapter type and rank, while qualitative and user-study results assess versatility and perceptual quality.
- Comparison with State-of-the-Art Methods: WaDi achieves the best FID and Recall across SD 1.5, SD 2.1, and PixArt-α on COCO 2014, while ranking first or second in CLIP and Precision.Only 9.74%, 10.92%, and 13.30% of parameters are trainable for SD 1.5, SD 2.1, and PixArt-α, respectively.
- Qualitative and User-Study Results: Qualitative comparisons report more coherent, prompt-aligned images, while a 57-participant user study evaluates zero-shot generation and downstream tasks against existing baselines.The qualitative examples emphasize structure, stylistic fidelity, subject attributes, and contextual details.
- Downstream Tasks: 86.26% lower inference time with ControlNet preserves image quality, spatial-condition fidelity, and prompt adherence comparable to ControlNet.WaDi is also reported to reduce Reversion inference time by 88.89% while maintaining quality close to the original multi-step method.
- Downstream Tasks: In DreamBooth, LoRaD is compared with full fine-tuning and LoRA as an illustrative parameter-efficiency experiment rather than a comprehensive diffusion fine-tuning study.The comparison addresses vanilla DreamBooth overfitting and LoRA’s improved balance between subject capture and prompt sensitivity.
- Ablation Studies: LoRaD attains FID 20.86 and CLIP 0.31 on COCO 2017 with 83.8M trainable parameters, approximately 90% fewer than full fine-tuning.Its direction mean is 2.89%, compared with 2.21% for FT and at most 0.92% for LoRA/DoRA variants.
- Ablation Studies: Increasing student rank from setting A to C reduces FID from 13.64 to 10.79, but further increasing rank to D degrades FID to 12.75 and CLIP to 0.31.Fake-model rank changes FID more than CLIP, and setting C provides the reported capacity–performance trade-off.
5. Conclusion
WaDi frames weight-direction changes as the key mechanism in one-step distillation and introduces LoRaD to model those adjustments efficiently. Experiments show improved image quality and inference speed, while the distilled model generalizes to downstream tasks.
- WaDi identifies weight direction as the key mechanism in one-step diffusion distillation, whereas norm changes play a comparatively smaller role.
- LoRaD models directional adjustments through a parameter-efficient module tailored to one-step text-to-image distillation.
- WaDi significantly outperforms DMD, SiD-LSG, and SwiftBrush in image quality and inference speed.
- The distilled model adapts to diverse downstream tasks, demonstrating generalization and practical applicability.