跳转至

Importance-Aware OBS Pruning for Diffusion Models

会议: NeurIPS 2026(录用信息来自任务元数据)
arXiv: 2607.20048v3
领域: 模型压缩
关键词: 扩散模型、参数剪枝、空间重要性、最优脑外科、免训练压缩

一句话总结

本文把提示词相关的空间重要性注入 OBS 的逐层重建目标,在不微调的情况下改善高稀疏率扩散模型的主体保真度,但全局掩码系数消融与给出的精确公式存在未解释的一致性问题。

研究背景与动机

文本到图像扩散模型不仅参数量大,还要反复运行去噪网络。OBS-Diff 已经把 Optimal Brain Surgeon(OBS,最优脑外科)的二阶参数剪枝推广到扩散模型:逐层收集不同去噪时刻的激活,估计移除一个权重后的重建代价,再解析地调整其他权重进行补偿。这条路线不需要重新训练,但平均重建误差未必对应人眼最在意的生成质量。

同样大小的特征误差,落在背景纹理与落在人物脸部、物体边缘上的后果可能完全不同。标准逐层重建目标没有显式区分这些位置,因此在高稀疏率下,图像可能仍然与提示词大致相关,却出现主体变形或细节缺失。作者进一步指出,扩散模型原始任务损失在单个校准批次、单个时刻上未必处于驻点;不过这与逐层教师重建目标是否为精确二次函数是两回事,不能混为一谈。

本文没有重新设计剪枝求解器,而是改变求解器需要保护的误差分布:先找出提示词影响较强的空间位置,再让这些位置更多地参与曲率统计。核心 idea:用空间重要性过滤校准激活,将均匀的逐层重建改为内容相关的重建,再沿用 OBS 的权重排序与解析补偿。

方法详解

整体框架

输入是预训练扩散模型、校准提示词和目标稀疏率,输出是剪枝后的参数模型。校准阶段运行条件与无条件去噪分支,提取 CFG 空间信号,并把相应位置的权重施加到各层激活上;随后用加权激活构建 Hessian,进行 OBS 排序、移除与补偿。

这是离线校准和参数修改流程,不是训练一个新网络。稠密层输出提供重建参考,而非人工标注监督;剪枝完成后的常规生成只使用剪枝模型,不需要在每次生成时重新计算剪枝重要性图。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    P["稠密模型与校准提示词"] --> A["CFG空间重要性"]
    A --> B["重要性加权Hessian"]
    P -.->|稠密层输出作重建参考| B
    B --> C["OBS排序与补偿"]
    C --> D["剪枝参数模型"]
    Q["推理提示词与噪声"] --> E["常规去噪生成"]
    D --> E

关键设计

1. CFG空间重要性:用条件响应确定需要优先保护的位置

CFG 即 classifier-free guidance(无分类器引导)。这里的重要性不是条件预测本身,也不是文本交叉注意力,而是同一噪声状态下条件与无条件预测的差异幅度。作者先对逐通道差值取绝对值,再沿通道维平均,因此相反符号的响应不会彼此抵消,最终得到二维空间图。

\[ \tilde{M}_{t}(i,j)=\frac{1}{C}\sum_{c=1}^{C}\left|\epsilon_{\theta}(x_{t},t,y)_{c,i,j}-\epsilon_{\theta}(x_{t},t,\varnothing)_{c,i,j}\right|. \]

这里的通道下标与文本条件分别用不同符号表示。接着,对每个样本、每个去噪时刻独立做空间 min–max 归一化;它不是跨批次归一化,也不是把各时刻的响应强度保留在统一尺度上。

\[ M_{t}=\frac{\tilde{M}_{t}-\min(\tilde{M}_{t})}{\max(\tilde{M}_{t})-\min(\tilde{M}_{t})+\epsilon},\qquad A_{t}=\lambda M_{t}. \]

小常数仅出现在原文的归一化分母中。原文明确使用 \(A_{t}=\lambda M_{t}\),没有添加恒为一的背景项;低响应位置可以接近零,从而几乎不参与目标。CFG 反映的是提示词对预测的影响,不是经过验证的人类显著性真值,也不保证所有重要的小物体都被突出。

这个接口还允许换成 Canny 边缘图,或由 YOLOv8 检测对象区域后保留相应位置的 CFG 响应。前者偏向边界与纹理,后者偏向检测到的对象;检测器在这里提供空间选择,不是额外的训练监督。不同信号改变的是保护目标,而非 OBS 求解器本身。

2. 重要性加权Hessian:改变误差度量,而不是在最终分数上乘掩码

作者把重要性图广播到输出通道,对重建残差逐位置相乘,再平方求和。时间权重 \(\alpha_{t}\) 沿用 OBS-Diff 的对数递减设置,用于强调较早的去噪步骤;它与空间图负责不同维度的加权。

\[ \mathcal{J}(\hat{W}_{l})=\sum_{t=1}^{T}\alpha_{t}\left\|A_{t}\odot\left(\hat{W}_{l}X_{l,t}-W_{l}X_{l,t}\right)\right\|_{F}^{2},\qquad \operatorname{Sparsity}(\hat{W}_{l})=S. \]

对在各空间 token 上独立、沿通道作用的线性层,位置权重可以移到输入激活上。把过滤后的激活送入原有二阶统计即可,因此改动发生在曲率估计阶段,而不是用一张图直接给每个权重指定保留概率。

\[ X'_{l,t}=A_{t}\odot X_{l,t},\qquad H_{l,\mathrm{imp}}=2\sum_{t=1}^{T}\alpha_{t}\mathbb{E}\left[X'_{l,t}X_{l,t}^{\prime\top}\right]. \]

重要细节是平方:残差先乘空间权重再取平方范数,所以在展开的 Hessian 中,每个位置的贡献乘的是 \(A_{t}(i,j)^{2}\),不是 \(A_{t}(i,j)\)。可把一个输出行对应的输入 token 向量写成下面的等价展开;这是笔记对原文目标的代数展开,不是论文新增的另一套公式。

\[ H_{l,\mathrm{imp}}=2\sum_{t=1}^{T}\alpha_{t}\mathbb{E}\left[\sum_{i,j}A_{t}(i,j)^{2}x_{l,t,i,j}x_{l,t,i,j}^{\top}\right]. \]

空间图由此改变输入协方差的相对组成。与重要位置共同激活的通道方向会在目标中占更大比重,但最终保留哪一个参数仍取决于权重大小及逆 Hessian 中的相关性,不能简化为“前景激活大就必然保留”。权重过滤与通道线性映射可交换的推导也不能不加说明地推广到任意跨空间算子。

3. OBS排序与补偿:沿用解析求解,并检查全局尺度的实际作用

对一个输出行,OBS 设某个权重必须变为零,允许其余权重一起变化,以最小化二次重建误差。对应的最小目标增量就是 saliency(剪枝重要性分数),分数较小的参数先被移除;解析更新则利用通道之间的相关性,把该权重造成的误差分摊给其他权重。

\[ S_{l,q}=\frac{w_{q}^{2}}{2[H_{l,\mathrm{imp}}^{-1}]_{qq}},\qquad \delta w=-\frac{w_{q}}{[H_{l,\mathrm{imp}}^{-1}]_{qq}}H_{l,\mathrm{imp}}^{-1}\mathbf{e}_{q}. \]

这两式分别对应正文的排序公式与附录 C 的补偿推导,并以相应 Hessian 可逆为前提。实际剪枝沿用 OBS-Diff 的流程及基于 Cholesky 的逆 Hessian 更新;结构化版本沿用基线的聚合策略,本文没有提出一个独立的新结构化求解器。

附录 C 明确说明,在冻结校准激活、以稠密层为教师时,\(\hat{W}_{l}=W_{l}\) 处的残差与一阶梯度都为零。这个逐层代理目标是关于待剪权重的精确二次函数;不能把原始扩散任务的非驻点问题解释成该代理目标也必然存在被忽略的一阶项。真正的近似在于代理重建与最终生成质量之间的联系,以及具体剪枝流程对联合最优解的近似。

还有一个影响复现的重要检查:若掩码形状固定,所有层与时刻共享同一个正的全局系数,那么只改变该系数不会改变空间位置之间的比例。以下是由论文公式直接得到的尺度关系;其中 \(H_{0}\) 表示同一重要性图在系数为一时的 Hessian,不是未加权 OBS 的 Hessian。

\[ H_{\lambda}=\lambda^{2}H_{0},\qquad H_{\lambda}^{-1}=\lambda^{-2}H_{0}^{-1},\qquad S_{\lambda,q}=\lambda^{2}S_{0,q},\qquad \delta w_{\lambda}=\delta w_{0}\quad(\lambda>0). \]

因此,精确公式下所有分数统一缩放,排序不变,补偿更新中的尺度也抵消。主文却报告不同正系数产生不同质量,并把它解释成语义引导强弱;这个解释不能仅由已给出的方程支持。全文没有给出能解释该差异的阻尼或稳定化定义,不应替作者补造实现机制。\(\lambda=0\) 更会令该加权 Hessian 为零,使通常的逆矩阵公式不可用,而不是自动退化为未加权 OBS。

损失函数 / 训练策略

方法不进行微调或梯度训练。校准数据与稠密层输出只用于构造局部重建目标,后续权重调整是 OBS 解析补偿,不是新增的优化训练轮次。

校准采用 GCC3M 的 1,000 个样本,批大小为 2,运行 10 个去噪步骤;测试运行 25 步。校准 CFG scale 对 SD3-Medium 为 7.0、对 PixArt-Σ 为 4.5,测试两者均为 7.0;默认掩码系数为 1.0。

原文未在全文中完整定义时间权重的精确实现、各层空间映射的工程细节及逆矩阵数值处理。复现时应核实这些环节,而不是把“保持原流程”视为实现已经完全披露。

实验关键数据

主实验

下表选取正文表 1、2 的关键非结构化剪枝结果,所有指标越高越好。CLIP Score 是整体图文对齐代理,ImageReward 是学习到的人类偏好代理,MUSIQ 是不直接依赖文本的图像质量指标。

模型 稀疏率 方法 CLIP Score ImageReward MUSIQ
SD3-Medium 稠密 原模型 32.26 0.98 72.68
SD3-Medium 30% OBS-Diff 32.25 0.96 72.53
SD3-Medium 30% 本文 CFG 32.24 0.98 72.41
SD3-Medium 50% OBS-Diff 32.16 0.71 65.98
SD3-Medium 50% 本文 CFG 32.20 0.76 66.30
PixArt-Σ 稠密 原模型 31.86 0.94 71.11
PixArt-Σ 60% OBS-Diff 31.44 0.49 65.24
PixArt-Σ 60% 本文 CFG 31.54 0.52 67.13

数据范围存在原文冲突:实验设置称测试使用 MS-COCO 2017 验证集的 1,000 个提示词,但表 1 标为 GCC3M 结果,表 2、5 的标题也提到 GCC3M 样本。这里保留原表数值,不擅自把表格全部重新归属为某一测试数据集。

消融实验

下表来自表 6,使用 PixArt-Σ、60% 稀疏率。它是论文报告的结果,不代表前述全局尺度一致性问题已经解决。

方法 掩码系数 CLIP Score ImageReward
OBS-Diff 不适用 31.44 0.49
本文 0.05 31.54 0.52
本文 0.1 31.60 0.54
本文 0.3 31.58 0.55
本文 0.5 31.57 0.53
本文 1.0 31.54 0.52
本文 2.0 31.56 0.51
本文 3.0 31.49 0.50

默认设置为 1.0,但表中 CLIP 最佳的是 0.1,ImageReward 最佳的是 0.3。因此不能把默认设置称为消融得到的最优值;更不能据此声称精确目标的排序对全局正系数有可解释的敏感性。

关键发现

  • 改善不是逐指标、逐设置的全面胜出:SD3-Medium 在 30% 稀疏率下,本文 CLIP 与 MUSIQ 都略低于 OBS-Diff,仅 ImageReward 更高;高稀疏率的优势更明确。
  • 信号选择存在取舍:表 3 的 PixArt-Σ、60% 设置中,CFG+Detector 的 MUSIQ 为 67.27,高于 CFG 的 67.13,但 ImageReward 为 0.50,低于 CFG 的 0.52。Canny 的 ImageReward 为 0.53,却没有取得最佳 MUSIQ。
  • 结构化剪枝也有支持:表 5 的 PixArt-Σ、40% 设置中,ImageReward 从 OBS-Diff 的 -0.04 提高到本文的 0.30,MUSIQ 从 65.17 提高到 68.64;不过没有对应的真实推理延迟表。
  • 类别定向校准不保证所有指标上升:表 4 的 Cat 设置中,定向版本 MUSIQ 为 71.64,高于通用版本的 69.52,但 CLIP 从 32.66 降至 32.28,ImageReward 从 0.41 降至 0.34。

附录 A 的人工研究采用五级图像质量评分,下表对应表 7 的 PixArt-Σ、60% 稀疏率、20 张图像评估集。它是一个选定的小规模研究,不能外推为所有模型、提示词或随机种子的平均效果。

方法 CLIP Score ImageReward MUSIQ 人工评分
OBS-Diff 31.90 0.72 69.83 2.59
本文 31.06 0.78 72.93 3.97

人工评分明显偏好本文,但 CLIP 恰好更低。表 8 还报告,在这个评估集上,CLIP、ImageReward、MUSIQ 的逐样本偏好与人类偏好一致率分别为 25.00%、50.00%、70.00%;这说明自动指标各有盲区,不构成普遍有效的指标校准结论。

亮点与洞察

  • 把感知偏好放进重建误差的空间度量,比在最终排名上附加一个启发式分数更有结构性。它保留了 OBS 的解析补偿,同时改变参数必须共同保护的激活分布。
  • 内容引导并不等同于更高的整体对齐分数。人工研究与 CLIP 的反向结果提醒压缩评估需要关注对象结构,而不能只看全局嵌入相似度。
  • 校准分布与空间图是两个不同的控制接口:前者决定模型经常看到什么概念,后者决定这些样本内部哪些位置更重要。二者可用于应用定制,但也需要检查未被重点保护的内容。

局限与展望

  • 作者承认 CFG 图在抽象提示词、小物体和复杂场景中可能噪声较大或覆盖不足;用检测器也会继承检测遗漏与类别偏差。
  • 全局系数消融与精确方程的尺度不变性未对齐,归一化分母的小常数并不能解释后续全局缩放为何改变精确排序。需要明确实现与实验条件,而非猜测阻尼机制。
  • 原始任务的非驻点动机与附录中的精确二次代理需要区分。本文没有证明更小的加权层误差必然带来更好的最终生成质量。
  • 校准与测试数据集表述冲突,且小规模人工研究缺少足够的跨模型、跨随机种子证据。正文对 SD3 的 60% 描述也超出了对应主表和图示展示的最高 50% 范围。
  • 非结构化稀疏率不是实测加速比:实际延迟收益依赖硬件或稀疏算子支持,本文没有展示真实推理提速。微调恢复、更丰富的生成任务与非目标类别鲁棒性仍待研究。

相关工作与启发

  • 与 OBS-Diff 对比:继承时间加权、二阶排序与补偿流程,主要变化是空间重要性过滤后的曲率统计。因此贡献是感知相关的目标修改,而不是全新剪枝优化器。
  • 与 SparseGPT、逐层 OBS 对比:共享激活协方差与相关权重补偿的思想;本文强调扩散校准中的时空内容差异,不能把局部重建的精确性误当作最终任务的精确保证。
  • 与 importance-based token merging 对比:都可利用 CFG 判断重要区域,但 token merging 改变中间表示的计算,本文修改参数。两者理论上可叠加,本文没有给出联合方法的实测收益。

评分

  • 新颖性: 3/5。空间感知目标是清晰、有针对性的扩展,但求解框架主要继承 OBS。
  • 实验充分度: 3/5。覆盖两种骨干、不同稀疏形式和引导信号,但数据标注冲突、人工样本量与缺少延迟结果限制结论。
  • 写作质量: 2/5。核心流程易懂,但全局系数解释与精确公式不一致,部分实验叙述超出表格范围。
  • 价值: 4/5。提供可复用的内容感知压缩接口,工程复现与数学一致性仍需核查。