CVPR 2026|当LoRA遇上RoPE!WaDi:面向单步图像生成的权重方向感知蒸馏
本文作者王雷,南开大学,PCA Lab成员
尽管扩散模型(如 Stable Diffusion,SD)在图像生成领域表现出色,但其缓慢的推理速度限制了实际部署。近期工作通过将多步扩散蒸馏为单步生成器来加速推理。
为了更好地理解蒸馏机制, 南开大学 PCA Lab 团队 分析了单步学生模型与多步教师模型之间 U-Net/DiT 权重的变化规律。分析表明,权重方向上的变化显著超过权重范数上的变化,这揭示了方向是蒸馏过程中的关键因素。受此启发,团队提出了 权重方向低秩旋转 (Low-rank Rotation of weight Direction,LoRaD)—— 一种专为单步扩散蒸馏设计的参数高效适配器。
LoRaD 通过可学习的低秩旋转矩阵对预训练权重的方向进行建模。团队进一步将 LoRaD 集成到变分得分蒸馏(Variational Score Distillation,VSD)中,提出了 权重方向感知蒸馏 (Weight Direction-aware Distillation,WaDi)—— 一种新颖的单步蒸馏框架。
WaDi 在 COCO 2014 和 COCO 2017 上取得了最先进的 FID 分数,而可训练参数仅占 U-Net/DiT 全量参数的约 10%。此外,蒸馏后的单步模型展现出强大的通用性和可扩展性,能够良好地泛化到可控生成、关系反演、高分辨率合成等多种下游任务。
论文链接:https://arxiv.org/abs/2603.08258
代码链接:https://github.com/gudaochangsheng/WaDi
Project: https://gudaochangsheng.github.io/WaDi-Page/
Demo: https://huggingface.co/spaces/gudaochangsheng/WaDi-1.5
讲解视频: https://www.youtube.com/watch?v=j6CuQxynJcA
图 1. 使用我们提出的方法 WaDi(即 SD 2.1)一步生成的图像。 引言
扩散模型(DMs)在图像生成领域受到了广泛关注,在文本到图像生成、文本到视频生成以及图像到视频生成等任务中均有广泛应用。然而,扩散模型依赖多步采样,导致计算成本高、推理速度慢。
为此,近期蒸馏方法将采样步数压缩至数步甚至一步。有趣的是,在蒸馏过程中,团队发现权重范数在各层间保持相对稳定,而在将权重重参数化为范数与方向时,方向则呈现出更大的变化幅度。
受权重重参数化的启发,团队采用类似的分解方式来分析扩散蒸馏中的权重变化。为此,团队研究了最先进(SOTA)单步模型(如 DMD2 和 Pixart-α DMD)与其对应多步模型(如 SD 1.5 和 Pixart-α)之间的权重更新。
如图 2 (a) 所示,在基于 U-Net 的架构中,各层权重范数几乎保持稳定,均值和标准差(STD)分别约为 0.1% 和 0.2%。相比之下,权重方向的变化则明显更大,均值为 2.2%,标准差为 2.1%,对应为范数变化的 22 倍和 10 倍。在基于 DiT 的架构中也观察到类似规律(见图 2 (a) 右)。
这些观察表明, 权重方向可能携带了蒸馏中更丰富、更敏感的信息。
此外,方向上的变化是否具有结构化规律?为此,团队对残差矩阵(单步与多步方向矩阵之差)进行奇异值分解(SVD),发现仅保留 30% 的秩即可恢复 93% 的信息,突显了其低秩本质(见图 2 (b))。
图 2. 我们方法的动机分析。(a) 一步学生模型与教师模型之间的权重范数和方向差异。更多细节和补充示例见补充材料 E。(b) DMD2 残差矩阵的 SVD 分析。(c) 将一步模型的范数替换为多步模型的范数影响很小①④);替换方向会严重降低生成质量②⑤)。(d) 与 (c) 对应的定性示例。(e) LoRaD 示意图。
为了量化这两个分量的影响,团队通过有选择地将单步模型的范数或方向替换为多步模型的对应值,进行受控消融实验(见图 2 (d))。如图 2 (c) 所示,替换范数对性能影响微乎其微(如 DMD2:FID +0.7,CLIP 不变),而替换方向则导致严重退化(如 DMD2:FID +241.3,CLIP -0.18)。
这些发现表明, 方向重建是蒸馏中性能提升的核心因素,而范数变化的影响相对次要。 一种可能的解释是:用教师权重初始化学生模型对齐了初始范数,训练过程中的权重衰减进一步约束了范数漂移;而蒸馏信号则主要通过调整权重方向来减少表征差异。
综合来看,这些结果表明 方向重建是蒸馏性能提升的核心驱动因素。
上述蒸馏方法大致可分为两类: 全量微调(FT) 和 基于低秩适配(LoRA)的微调。 然而,二者在优化范数和方向时均直接更新模型参数,导致范数和方向的变化相互耦合,增加了优化难度。此外,FT 和 LoRA 均面临收敛慢、不稳定和过拟合等问题,进一步增加了优化的复杂性。
为此,团队提出了权重方向低秩旋转(LoRaD)(见图 2 (e)),通过可学习的旋转矩阵调整预训练权重的方向。鉴于方向变化的结构化特性(即低秩特性),旋转角度被参数化为两个低秩矩阵的乘积,以进一步减少可学习参数数量。团队将 LoRaD 集成到变分得分蒸馏(VSD)中,提出了权重方向感知蒸馏(WaDi),一种新颖的单步文本到图像蒸馏框架。
在 COCO 2014 和 COCO 2017 数据集上的实验表明,WaDi 取得了最先进的 FID 分数,超越所有现有单步生成方法。这一成果仅通过优化方向实现,将蒸馏难度降低,同时 U-Net 可训练参数仅约占 10%,极大提升了参数效率。
此外,团队将 WaDi 应用于可控生成、关系反演、高分辨率合成和图像定制化等下游任务,展示了其加速能力和广泛适用性。本文贡献总结如下:
团队对多步与单步生成模型之间 U-Net 权重变化进行了深入分析,将权重方向调整确定为单步蒸馏的关键驱动因素,为高效蒸馏提供了新的理论视角;
团队提出了一种新颖的单步文本到图像蒸馏框架 WaDi,采用 LoRaD 通过低秩旋转建模权重方向,有效引导学生模型对齐教师分布;
WaDi 在 COCO 数据集和多个下游任务上进行了评估,定性和定量结果均表明 WaDi 在显著提升推理效率的同时取得了实质性的图像质量提升。
相关工作
扩散模型
扩散模型在图像生成领域表现卓越,但像素空间计算开销大。为提升效率,Rombach 等人提出了潜在扩散模型(LDM),将去噪过程迁移至潜在空间。然而,现有基于文本引导的方法由于多步生成仍然较慢。尽管大多数方法采用 U-Net 骨干,扩散 Transformer(DiT)以 Transformer 替代 U-Net 以获得更好的可扩展性,推动了文本到图像生成的进步。尽管有所改进,迭代去噪仍是一个缓慢的过程。近期,许多加速方法相继出现。
扩散模型加速
现有加速方法可分为无训练和基于训练两类。无训练加速方法主要通过缓存减少冗余计算,或采用高阶求解器减少采样步数。然而这两类方法的加速效果有限,因此基于训练的方法受到了更多关注。
基于训练的加速方法大致可分为四类:一致性蒸馏(CD)、渐进蒸馏(PD)、扩散 - GAN 蒸馏和变分得分蒸馏(VSD)。CD 在轨迹层面学习一致性以加速采样,但图像保真度往往较低。PD 分阶段降低步数,引入显著训练开销。扩散 - GAN 蒸馏(如 Diffusion2GAN)通过将多步扩散蒸馏为 GAN 来提升保真度。VSD 采用双教师策略实现分布对齐,SwiftBrush 实现了单步无图像生成,SwiftBrushv2 利用模型集成进一步改善,DMD 引入回归损失进一步提升性能,DMD2 将 VSD 扩展至少步生成,并支撑了近期文本到视频加速框架。
然而,现有基于训练的方法通常使用 FT 或 LoRA,这可能增加优化难度。团队发现方向变化在蒸馏中通常更具影响力,因此提出 WaDi,利用 LoRaD 专注于建模方向旋转。
方法
团队首先简要回顾变分得分蒸馏(VSD),它是本工作的基础。受权重方向变化在蒸馏中发挥关键作用这一观察的启发,紧接着引入权重方向低秩旋转(LoRaD)模块(更多理论说明见补充材料 D)。最后,团队将 LoRaD 集成到 VSD 中,形成我们提出的蒸馏框架 —— 权重方向感知蒸馏(WaDi)。
1.预备知识
潜在扩散模型(LDM)在低维潜在空间中执行扩散过程,提升了计算效率。LDM 的训练目标可以表示为: