WaDi: Weight Direction-aware Distillation for One-step Image Synthesis

1. Introduction

• In U-Net distillation, direction changes average 2.2% (STD 2.1%) versus norm changes of 0.1% (STD 0.2%), or 22× and 10× larger, respectively.The introduction reports a similar trend in DiT architectures.Source: §1 · p. 2, Figure 2 · p. 2
• Retaining 30% of the directional residual matrix’s rank recovers 93% of its information, indicating that directional changes have a low-rank structure.The residual is the difference between one-step and multi-step direction matrices.Source: §1 · p. 2, Figure 2 · p. 2
• Replacing DMD2’s norm changes FID by +0.7 with unchanged CLIP, while replacing its direction changes FID by +241.3 and CLIP by -0.18.The authors suggest teacher initialization and weight decay constrain norm drift, leaving distillation to adjust direction and reduce representational discrepancy.Source: §1 · p. 2, Figure 2 · p. 2
• Full fine-tuning and LoRA-based fine-tuning [16] jointly update norm and direction, increasing optimization coupling; both also face slow convergence [9] [21], instability [12] [13], and overfitting.The introduction identifies the mismatch between small norm changes and substantial direction changes as a source of optimization difficulty.Source: §1 · pp. 2–3
• WaDi achieves state-of-the-art FID on COCO 2014 and COCO 2017, outperforming all existing one-step generation methods with only about 10% of U-Net parameters trainable.The paper also reports improved inference efficiency and image quality across its evaluations.Source: §1 · p. 3
• LoRaD models pretrained-weight directions with learnable low-rank rotation matrices and is integrated into VSD to form the WaDi one-step text-to-image distillation framework.Parameterizing rotation angles as a product of two low-rank matrices further reduces the number of learnable parameters.Source: §1 · p. 3
• WaDi is applied to controllable generation, relation inversion, high-resolution synthesis, and image customization, demonstrating acceleration and broad applicability across downstream tasks.The paper reports qualitative and quantitative results on these tasks.Source: §1 · p. 3

Diffusion models

• Diffusion models [15] [17] [18] excel at image generation, but pixel-space computation is costly; LDM moves denoising to latent space.Text-guided methods [30] [46] remain slow because generation uses multiple steps.Source: §2 · p. 3
• DiT [45] replaces the commonly used U-Net with a Transformer for better scalability, yet iterative denoising remains slow.Text-to-image generation advances through DiT [10].Source: §2 · p. 3

Diffusion model acceleration

• Training-free acceleration uses caching [28] [42] [53] or high-order solvers [33] [37] [38], but its acceleration effects are limited.Examples include Faster Diffusion [28], DDIM, and DPM-Solver [37] [38].Source: §2 · p. 3
• Training-based acceleration includes consistency distillation, progressive distillation, diffusion-GAN distillation, and VSD, each using a distinct strategy.CD [25] [39] [40] learns trajectory consistency but can reduce fidelity; PD reduces steps in stages with training overhead; diffusion-GAN [23] [31] [41] targets fidelity; VSD [8] [43] [76] aligns distributions with dual teachers. SwiftBrush [43], SwiftBrushv2 [8], DMD, and DMD2 develop VSD toward image-free one-step, ensembling, regression-loss improvement, and few-step generation, respectively.Source: §2 · p. 3

Motivation

• Because common FT or LoRA methods can make optimization harder, WaDi uses LoRaD to model directional rotations, which are more influential in distillation.The authors identify directional weight changes as generally more influential than other changes during distillation.Source: §2 · p. 3

3.1. Preliminary

• VSD alternates updates to a fake model trained with the diffusion objective and a one-step generator trained to synthesize images from input noise.The VSD objective uses real and fake models, a time-dependent weighting term, and a one-step generator conditioned on prompts.Source: §3.1 · p. 4

3.2. Low-rank Rotation of Weight Direction

• LoRaD models the observed directional shifts with learnable low-rank rotations, changing weight directions without changing their norms.The low-rank decomposition is adopted to reduce parameter overhead compared with full-rank modeling, which adds parameters equivalent to 50% of the original weights; block-diagonal rotations can be applied across odd-even row pairs using element-wise operations.Source: §3.2 · pp. 4–5

3.3. Weight Direction-aware Distillation

• WaDi integrates LoRaD into VSD: a teacher-initialized fake model approximates the pretrained teacher’s distribution while a teacher-initialized one-step student learns image synthesis.LoRaD uses a high-rank rotation for the student and a low-rank rotation for the fake model; their parameters are optimized alternately.Source: §3.3 · p. 5

4.1. Experimental Setup

• WaDi is evaluated zero-shot on COCO 2014 (30k samples) and COCO 2017 (5k), using FID, CLIP score, precision, recall, and HPSv2.Training uses 1.4M JourneyDB prompts, teacher-initialized student weights, and 2 epochs; student and fake-model learning rates are 1e-4 and 1e-2.Source: §4.1 · pp. 5–6

4.2. Comparison with State-of-the-Art Methods

• WaDi achieves the best FID and Recall across SD 1.5, SD 2.1, and PixArt-α, ranks first or second in CLIP and Precision, and trains 9.74%, 10.92%, and 13.30% of parameters, respectively.The comparison uses COCO 2014 and reports results across three backbones; the authors attribute the improvements to LoRaD’s low-rank rotation reparameterization.Source: Table 1 · p. 6, §4.2 · p. 6

4.3. Downstream Tasks

• WaDi reduces ControlNet inference time by 86.26% while preserving image quality, spatial-condition following, and prompt adherence comparable to ControlNet.The comparison is shown in Figure 5.Source: §4.3 · p. 6, Figure 5 · p. 7

4.5. Ablation Studies

• On COCO 2017, LoRaD obtains FID 20.86 and CLIP score 0.31 with 83.8M trainable parameters, ~31% fewer than LoRA/DoRA and ~90% fewer than FT.LoRaD has direction mean 2.89%, versus 2.21% for FT and ≤0.92% for LoRA/DoRA variants; it parameterizes low-rank orthogonal rotations that preserve weight norms.Source: Table 2 · p. 7, §4.5 · p. 8

4.3. Downstream Tasks

• Integrating WaDi into Reversion [22] reduces inference time by 88.89% while producing high-fidelity images aligned with relational prompts and close in quality to the original multi-step method.The comparison is shown in Figure 6.Source: Figure 6 · p. 7, §4.3 · p. 7
• In DreamBooth, LoRaD maintains subject fidelity while following prompts, balancing the overfitting of vanilla fine-tuning against LoRA’s degradation of identity and image fidelity [16].The paper presents this as an illustrative example, not a comprehensive diffusion fine-tuning study.Source: §4.3 · pp. 7–8, Figure 7 · p. 8

4.5. Ablation Studies

• On COCO 2014, raising student rank from setting A to C lowers FID from 13.64 to 10.79, while settings C versus D yield FID 10.79 versus 12.75 and CLIP 0.30 versus 0.31.Changing fake-model rank across settings C, E, and F affects FID while leaving CLIP largely stable; setting C offers a favorable capacity–performance trade-off.Source: Table 3 · p. 8, §4.5 · p. 8

4.4. User Study

• In a user study with 57 participants covering zero-shot generation and downstream tasks, WaDi results demonstrate superiority over existing baselines.The results are presented in Figure 9.Source: Figure 9 · p. 8, §4.4 · p. 8

• Weight-direction changes are a key distillation mechanism, while weight-norm changes play a comparatively smaller role.Source: §5 · p. 8
• LoRaD models directional adjustments in WaDi with a parameter-efficient module.Source: §5 · p. 8
• WaDi significantly outperforms DMD, SiD-LSG, and SwiftBrush in image quality and inference speed.Source: §5 · p. 8
• WaDi’s distilled model adapts to a wide range of downstream tasks, demonstrating generalization and practical applicability.Source: §5 · p. 8