WaDi:用“权重方向旋转”提升单步扩散图像生成质量

WaDi 用低秩旋转建模权重方向,让扩散模型实现高质量单步生成。

原文标题:CVPR 2026 | 当LoRA遇上RoPE!WaDi:面向单步图像生成的权重方向感知蒸馏

原文作者:机器之心

冷月清谈:

南开大学 PCA Lab 提出 WaDi,一种面向单步文本到图像生成的扩散蒸馏框架。作者通过分析多步教师模型与单步学生模型的权重变化发现,蒸馏过程中权重范数变化很小,真正影响生成质量的是权重方向变化;进一步的 SVD 分析显示,这种方向变化具有明显低秩结构。基于这一观察,团队设计了 LoRaD 模块,用低秩旋转矩阵直接建模权重方向,而非像全量微调或 LoRA 那样同时耦合改变范数和方向。WaDi 将 LoRaD 集成进 VSD 蒸馏流程,在 COCO 2014、COCO 2017 上取得领先 FID,并且可训练参数仅约为 U-Net/DiT 全量参数的 10%。实验还展示了它在 ControlNet、关系反演、高分辨率合成、DreamBooth 定制化等任务中的加速和泛化能力,推理时间可大幅缩短,同时保持较好的图像质量。

怜星夜思:

1、如果蒸馏里真正重要的是“权重方向”而不是“权重大小”,这会不会影响以后 LoRA/微调方法的设计思路?
2、单步生成质量越来越好之后,多步扩散模型还有必要存在吗?
3、WaDi 在 ControlNet、DreamBooth 这类下游任务上也能用,这说明它更像通用加速器吗?
4、只训练约 10% 参数就能取得不错效果,这对普通开发者或小团队有实际意义吗?

原文内容


本文作者王雷,南开大学,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 的训练目标可以表示为:



其中图片为高斯噪声,图片为时间步图片处的潜变量,图片表示用于引导图像生成的条件(如提示词),图片为由参数图片参数化的模型预测的噪声。


变分得分蒸馏(VSD)最初用于文本到 3D 生成,以解决过饱和和多样性不足的问题,后被扩展至 2D 文本到图像生成,应用于 Swiftbrush、DMD、DMD2 和 SiD 等方法,实现了单步生成。VSD 的训练目标为:



其中图片是时间相关的权重项,图片是由图片参数化的真实模型,图片是由图片参数化的虚假模型,图片是以图片为输入噪声参数化的单步生成器,图片。此外,图片利用公式 (1) 进行训练。VSD 交替更新图片图片直至收敛。


2.权重方向低秩旋转


图 3. (左)权重方向低秩旋转(LoRaD)模块的详细结构。LoRaD 使用可学习的低秩旋转角来旋转预训练权重方向。(右)权重方向感知蒸馏(WaDi)框架概览。


分析多步 U-Net 模型与其单步对应模型之间的权重变化,揭示了显著的方向偏移,而范数变化相对较小。受此启发,团队提出权重方向低秩旋转(LoRaD)(见图 3 左),通过学习仅改变方向的旋转来更新权重。此外,团队观察到方向变化具有低秩结构(见图 2 (b))。为利用这一特性并降低全秩建模的开销(全秩建模引入的额外参数相当于原始权重的 50%),团队采用 LoRA 的低秩分解策略。从 2D 情形(图片)出发,给定权重向量图片,应用 2D 旋转矩阵:



其中图片是旋转后的权重向量。受旋转位置编码(RoPE)的启发,RoPE 将 2D 情形推广至任意偶数维度图片,团队对预训练权重矩阵图片的每列应用不同的旋转矩阵:




其中旋转矩阵图片




其中图片


注:旋转不影响范数,因此无需显式分离范数矩阵。


考虑到公式 (5) 中图片的稀疏性,矩阵 - 向量乘法图片可高效计算为:



其中图片表示逐元素乘法。该实现利用了旋转矩阵的稀疏性,仅通过逐元素运算完成计算,从而显著降低了计算成本。


此外,由于公式 (5) 和 (6) 中的旋转矩阵由独立的图片子矩阵构成块对角结构,整个计算可以高效地实现为多个图片旋转在奇偶索引对上的并行应用。如图 3 左所示,团队将预训练权重矩阵图片图片维空间拆分为图片个子空间,对每个子空间独立旋转。通过分离图片的奇偶行,定义:



得到两个矩阵图片图片


每对奇偶行上的并行图片旋转可以紧凑地表示为:



其中图片是旋转后的权重矩阵,图片是可学习的旋转角度参数矩阵。为进一步减少可训练参数数量,受 LoRA 启发,团队对图片应用低秩分解:



其中图片图片是秩为图片的低秩参数矩阵。最终,公式 (8) 可以改写为:



3.权重方向感知蒸馏


为了充分利用蒸馏中观察到的方向特性,团队将 LoRaD 集成到 VSD 中,得到一个方向感知蒸馏框架,称为权重方向感知蒸馏(WaDi)。如图 3 右所示,WaDi 采用预训练扩散模型图片作为教师(真实模型),并引入可训练的虚假模型图片(从图片初始化)来近似教师分布。最终学生模型(单步生成器)图片同样从图片初始化,经训练后能够单步合成高质量图像。算法详见补充材料 F.3。


为增强与真实分布的对齐,团队对学生模型和虚假模型均应用 LoRaD。具体而言,单步生成器图片引入高秩旋转矩阵图片以更好地拟合教师;虚假模型图片采用低秩旋转矩阵图片以提供自适应引导。最终,团队交替优化图片图片以共同提升生成质量。


相应地,WaDi 的训练目标可以从公式 (2) 改写为:



虚假模型图片的训练目标可以从公式 (1) 改写为:



实验


1.实验设置


评估数据集与指标。团队在 COCO 2014 和 COCO 2017 数据集上系统评估 WaDi 的零样本文本到图像生成能力,分别随机采样 30k 和 5k 张图像。为全面评估生成质量,团队使用 Fréchet Inception Distance(FID)衡量图像保真度,CLIP 分数评估文本 - 图像语义对齐。FID 使用 Inception V3 作为特征提取器,CLIP 分数基于 ViT-G/14 模型。团队还采用精确率和召回率评估保真度与多样性,并在 Human Preference Score v2(HPSv2)基准上评估文本 - 图像对齐质量。详见补充材料 G.1。


实现细节。遵循先前方法,WaDi 中的学生模型采用与教师相同的架构,并用教师权重初始化。WaDi 在从 JourneyDB 数据集中采样的 140 万条提示词上进行训练。训练过程中,学生模型的学习率(LR)设为 1e-4,虚假模型的学习率设为 1e-2。团队使用 AdamW 作为优化器,批大小为 128(每 GPU 16 张)。无分类器引导(CFG)系数设为 1.5,训练进行 2 个 epoch。团队在三种不同骨干上蒸馏学生模型,分别为 SD 1.5、SD 2.1 和 PixArt-α(256×256)。对于 SD 1.5 和 SD 2.1,学生模型的 LoRaD 秩设为 256,而对于 PixArt-α 设为 128。所有虚假模型的 LoRaD 秩统一设为 32。详见补充材料 F.1。


2.与最先进方法的比较


定量结果。团队在 COCO 2014 数据集上与三种骨干(SD 1.5、SD 2.1、PixArt-α)的 SOTA 零样本单步生成方法全面评估 WaDi。为保证公平比较并考虑计算约束,团队遵循 TiUE 的设置,统一使用 140 万条提示词复现 WaDi、DMD2、SiD-LSG 和 SwiftBrushv2。如表 1 所示,WaDi 在所有骨干上均取得最佳 FID 和召回率分数,表现出优越的保真度和多样性。在 CLIP 和精确率上也排名第一或第二,显示出强大的文本 - 图像对齐和感知质量。


值得注意的是,SD 1.5、SD 2.1 和 PixArt-α 可训练参数分别仅占 9.74%、10.92% 和 13.30%,突显了 WaDi 的参数高效性。这些改进源于团队提出的 LoRaD,通过低秩旋转重参数化权重更新,实现稳定高效的蒸馏。


表 1. WaDi 与其他方法在零样本 COCO 2014 结果上的定量比较。∗ 表示我们复现的结果,≀ 表示使用官方预训练模型得到的结果。「-」表示未知。最佳和次佳分数分别用粗体和下划线标出。「Image-free」指在没有真实图像监督的情况下进行训练。


定性结果。图 4 展示了 WaDi 与 SOTA 单步生成方法在 SD 1.5 和 SD 2.1 骨干上的定性比较。在不同提示词下,WaDi 一致地生成视觉连贯、语义对齐的结果。例如,在第一、二行中,WaDi 更好地保留了结构和风格保真度,捕捉到清晰的细节和鲜艳的色彩,无伪影或失真。在第三、四行中,它能准确跟随涉及特定主体(如 sphynx cat、corgi、shiba inu)和场景(如剧院、服装)的提示词,而其他方法往往漏掉关键属性或生成不真实的形状。值得注意的是,在最后一行,WaDi 生成了空间构图和背景细节一致的复杂场景(如狗看电视),展示了相对于其他基线更优越的整体理解能力。详见补充材料 G.5。


图 4. 与其他方法的定性比较,其中 ∗ 表示团队复现的结果。


3.下游任务


可控生成。ControlNet 是一种广泛使用的可控生成模型,通过将空间条件整合到 SD 中实现精细控制。如图 5 所示,将 WaDi 应用于 ControlNet 可显著提升推理效率,将推理时间缩短 86.26%,同时保持图像质量,忠实遵循空间条件,提示词遵从度与 ControlNet 相当。


图 5. 使用或不使用 WaDi 的 ControlNet [74] 质量结果。


关系反演。Reversion 是首个在 SD 中通过关系提示词引导特定对象关系合成的方法。将 WaDi 集成到 Reversion 中可显著加速推理。如图 6 所示,WaDi 将推理时间缩短 88.89%,生成与关系提示词对齐的高保真图像,质量接近原始多步 Reversion。


图 6. 使用或不使用 WaDi 的 Reversion [22] 质量结果。


图像定制化。Dreambooth 是一种开创性的个性化文本到图像框架,通过对 U-Net 进行微调将目标主体绑定到稀有标记。为增强参数效率,团队将 LoRaD 集成到 Dreambooth 中,并与 Dreambooth(FT)和 LoRA 进行比较。如图 7 所示,原始 DreamBooth 会捕捉主体但记忆训练图像,降低提示词敏感性。LoRA 缓解了过拟合,但降低了主体保真度和图像质量。相比之下,LoRaD 在保持提示词遵从度的同时维持了主体保真度,实现了更好的平衡。团队将此 DreamBooth 实验仅作为说明性示例,而非扩散微调的全面研究。


图 7. 使用或不使用 LoRaD 的 DreamBooth 质量结果。


4.用户研究


为评估图像质量和文本 - 图像对齐,团队邀请 57 名参与者进行了用户研究,涵盖零样本生成和下游任务。如图 9 所示,结果清楚地表明团队的方法优于现有基线。详见补充材料 F.5。


图 9. 与其他方法相比的用户研究结果。


5.消融研究


表 2 在 COCO 2017 的 VSD 损失下比较了五种适配器类型。LoRaD 以最少 83.8M 可训练参数(比 LoRA/DoRA 少约 31%,比 FT 少约 90%)取得最低 FID(20.86)和竞争性 CLIP 分数(0.31)。它还取得最高的方向均值(2.89,而 FT 为 2.21%,LoRA/DoRA 变体为 ≤0.92%),表明在紧凑参数化下具有更广泛、更有效的更新方向空间。与 DoRA 和 DoRA(冻结范数)不同,后两者通过 LoRA 式加性更新对归一化权重后跟动态重归一化进行方向优化,LoRaD 直接将预训练权重参数化为低秩正交旋转,保留范数并纯在方向空间中操作。总体而言,LoRaD 展示了良好的质量 - 效率权衡。


表 2. 在 COCO 2017 数据集上,WaDi(SD 1.5,VSD 损失)中适配器类型影响的消融实验。「NM」和「DM」分别表示所有层的范数均值和方向均值。


团队在 COCO 2014 上进行了秩配置的消融研究。如表 3 所示,团队得出三个主要观察:1)增大学生秩可持续改善性能。将秩从设置 A 提升到 C,FID 从 13.64 降至 10.79,表明更高的秩使学生能够更好地捕捉教师分布,提升生成质量。2)超过阈值后继续增大秩会导致收益递减。对比设置 C 和 D,进一步增大秩导致 FID 退化(12.75 vs. 10.79),CLIP 也从 0.31 降至 0.30,表明过大的秩可能导致过拟合。3)虚假模型的秩对保真度的影响大于对齐度。调整虚假模型秩(设置 C、E、F)会改变 FID 但 CLIP 基本稳定,表明保真度对容量更敏感而对齐度较为稳健。


综上,设置 C 在模型容量和性能之间取得了良好的权衡,与图 8 的定性结果一致。详见补充材料 G.2、G.4。


表 3. 在 COCO 2014 数据集上,WaDi(SD 1.5,VSD 损失)中秩影响的消融实验。


图 8. 不同设置下的一步图像生成结果。


结论


本文提出了权重方向感知蒸馏(WaDi),一种高效的单步文本到图像蒸馏框架。通过对多步与单步模型权重变化的深入分析,团队发现权重方向的变化是蒸馏中的关键机制,而范数变化的作用相对次要。


基于这一洞见,团队引入权重方向低秩旋转(LoRaD)模块,以参数高效的方式建模方向调整。


大量实验表明,WaDi 在图像质量和推理速度上显著优于现有单步方法 —— 包括 DMD、SiD-LSG 和 SwiftBrush。此外,蒸馏后的模型可无缝适配多种下游任务,展示了强大的泛化能力和实际适用性。本工作为高效扩散模型蒸馏提供了新颖的理论视角和实践方案。




© THE END 

转载请联系本公众号获得授权

投稿或寻求报道:liyazhou@jiqizhixin.com

从优化角度看,这个发现有价值。范数稳定意味着预训练模型已经提供了较好的尺度先验,蒸馏阶段主要需要改变表示方向来对齐教师分布。如果这个规律能在更多架构和任务上复现,未来的参数高效微调可能会从加性更新转向几何约束式更新。

1 个赞

有意义,不过别误会成“家用显卡随便训”。它省的是可训练参数,不等于训练完全白菜价。对普通用户来说,最现实的收益是别人训好后,我们生成图更快、等图时间更短。

2 个赞

关于“WaDi 算不算通用加速器”,我倾向于说它有这个潜力,但还不能直接下结论。文章里的 ControlNet、Reversion、DreamBooth 示例说明迁移性不错,不过要看更多真实工作流,比如 ComfyUI 复杂节点、IP-Adapter、多 ControlNet 叠加时表现如何。

3 个赞

针对“权重方向是不是更重要”这个问题,我觉得影响会挺大。以前很多微调方法默认是在参数空间里加一个增量,但 WaDi 这类工作等于提醒大家:也许不该只看参数变了多少,而要看它往哪个方向转了。以后可能会出现更多专门约束方向、角度、子空间的适配器。

2 个赞