智猩猩AI整理 

编辑：六六 

大规模文本到图像扩散模型虽功能强大，却受制于高昂的计算成本。鉴于扩散模型的迭代去噪特性，现有的一步式网络剪枝方法难以直接应用于此类模型。

为填补这一空白，西湖大学提出OBS-Diff——一种创新的一步式、无需训练的剪枝框架，能够对大规模文本到图像扩散模型进行准确且无需训练的压缩。

大量实验表明，OBS-Diff 为无需训练的扩散模型剪枝确立了新的性能标杆。该方法在保持高视觉质量的同时实现了推理加速，并且在多种稀疏度水平与稀疏模式下均优于其他逐层剪枝方法。

• 论文标题：OBS-Diff: Accurate Pruning For Diffusion Models in One-Shot

• 论文链接：https://arxiv.org/abs/2510.06751

• 代码链接：https://github.com/Alrightlone/OBS-Diff

01 方法

图 1  所提出的OBS-Diff框架在MMDiT架构上的应用示意图。

图 1 所示OBS-Diff 框架将模型划分为序列化的 “模块包（Module Packages）”，以分摊校准成本。在每个模块组内，采用一种新颖的时序感知 Hessian 矩阵构建方式，优先处理早期去噪步骤，从而实现对当前模块包内所有目标层的同步剪枝。

1. 时序感知的 Hessian 矩阵构造

扩散模型的去噪过程由离散时间步参数化，具有迭代特性。剪枝产生的误差在该过程中影响不均匀：早期步骤引入的误差会向后续步骤传播并累积，导致最终输出的偏差显著增大。因此，鲁棒的剪枝策略必须优先保证这些关键早期阶段的网络功能不受损。

研究重新构建了逐层优化问题，以最小化加权重构误差为目标，该误差赋予早期、高影响度步骤更高的权重：

  是一个依赖于时间步的权重，此策略确保权重在推理开始时最高，并平滑衰减。

通过引入此加权方案，对 Optimal Brain Surgeon（OBS）框架进行了适配。用于捕捉此加权损失二阶信息的Hessian矩阵，现计算为所有推理步骤的加权和：

该矩阵被称为时序感知 Hessian 矩阵。它概括了参数在生成过程中不断变化的重要性。因此，从其逆矩阵导出的显著性分数，对在去噪过程早期形成阶段至关重要的权重更为敏感，从而能获得更忠实保持性能的剪枝后模型。

2. 模块包：组级顺序剪枝策略

该方法基于两个概念构建。基本单元是指在单次前向传播中具有相互独立输入的一组层（例如，查询、键和值的投影层），允许并行处理。模块包由一个或多个基本单元组成，这些单元被集体剪枝和校准。

本框架顺序处理这些模块包。对于每个包，首先执行组级数据收集阶段：在校准数据集上运行一次完整的去噪轨迹，同时使用前向钩子收集该包内所有模块的输入统计量。随后，使用各自对应的时序感知 Hessian 矩阵同时对所有模块进行剪枝。

网络状态在包与包之间顺序更新，但在单个包的数据收集过程中保持不变。这在更粗的组级粒度上保留了顺序校准的原则，从而使该过程在计算上可行。此策略极大地减少了校准运行的次数，其主要权衡在于需要同时存储多个 Hessian 矩阵导致内存占用增加。

3. 拓展至半结构化与结构化剪枝

OBS-Diff 框架的一个关键优势在于其适应性。在聚焦于非结构化剪枝的同时，该框架可以轻松扩展到半结构化与结构化稀疏模式。

半结构化剪枝。 对于像2:4这样的半结构化稀疏模式，扩展是直接的。在每个包含四个权重的块中，直接剪除其中具有最低单权重 OBS-Diff 显著性得分的两个权重，从而高效地创建对硬件友好的模型。

结构化剪枝。对于前馈网络（FFN）层的结构化剪枝，通过聚合其相关权重的显著性来评估一个神经元的重要性。整个神经元的显著性  及其对应的权重更新量如公式所示，其中得分最低的神经元将被移除。

同理，针对多头注意力模块执行完整注意力头的剪枝。该方法受 SlimGPT启发，可量化每个注意力头的显著性。

计算始于输出投影层的完整Hessian矩阵  ，通过聚合其各个权重的重要性来得到第  个头的总显著性  。显著性计算公式如下：

但 MMDiT 的联合注意力机制存在独特挑战。共享注意力头处理拼接后的多模态输入，却接入独立的模态专属输出支路。该结构使同一组注意力头对应两种不同的重要性排序（各模态对应一种），而 OBS-Diff 在分离后处理两组输出投影矩阵。为解决该问题，采用互逆排序融合（Reciprocal Rank Fusion, RRF）将两种排序融合为单一决策序列：

随后，使用完整的Hessian矩阵  ，更新整个输出投影层的权重。

02 评估

1. 非结构化剪枝结果

表 1  文本到图像扩散模型上非结构化剪枝方法的定量比较。各项指标的最佳结果以加粗形式标出。

表 1 中的结果展示了 OBS-Diff 在 CLIP 分数和 ImageReward 指标上的优越性。

关于FID指标的有趣现象——剪枝后的模型偶尔能够超越原始密集模型，可以认为 FID 在此可能不是评估不同剪枝方法的非常可靠的指标。

关于 CLIP 分数，OBS-Diff 在绝大多数测试案例中都是性能最佳的方法，与密集模型相比仅显示出轻微的下降。

最值得注意的是，OBS-Diff 在所有基准测试的 ImageReward 指标上持续领先，表明其输出与人类审美偏好具有更优的一致性。

该方法在高稀疏度水平下的优越性最为显著。如图 2 所示，基线方法在高稀疏度下生成的图像经常完全损坏并存在严重伪影，而 OBS-Diff 则能持续生成高质量且连贯的结果。

图 2  在SD3-Medium模型上对不同非结构化剪枝方法的定性比较。所有生成图像的分辨率均为512 × 512。

2. 半结构化剪枝结果

2:4 半结构化剪枝的结果展示于表 2。OBS-Diff 在CLIP分数与ImageReward上均大幅超过最强基线方法。这凸显了本方法在对硬件友好的稀疏度约束下，仍能有效保持高层次语义一致性与视觉保真度。

表 2  在 Stable Diffusion 3.5-Large 模型上应用半结构化剪枝的性能表现。剪枝操作应用于第3至第25个MMDiT模块。

3. 结构化剪枝结果

结果展示于表 4 与表 3。基线  剪枝方法即使在较低稀疏度下也出现了严重的性能退化，与此形成鲜明对比的是，OBS-Diff 方法表现出显著的鲁棒性，这种鲁棒性一直持续到 30% 的稀疏度。这些发现突显了 OBS-Diff 在进行激进的结构化剪枝时，仍能有效保留关键模型结构的卓越能力。

为了与现有技术进行基准比较，本研究将 OBS-Diff 与当前先进的文本到图像扩散模型结构化剪枝框架 EcoDiff 进行了直接对比。如表 3 和表 4 所示，证实了 OBS-Diff 能够有效地泛化至不同的模型架构，在 MMDiT 与 U-Net 两种骨干网络上均优于基线方法。

表 3  SDXL（U-Net）模型在不同稀疏度水平下的结构化剪枝性能。TFLOP 指标代表整个U-Net单次前向传播的理论计算成本。

表 4   Stable Diffusion 3.5-Large 模型在不同稀疏度水平下的结构化剪枝性能。剪枝过程排除了第一个和最后一个 Transformer 模块。TFLOPs指标代表整个 Transformer 进行单次前向传播的理论计算成本。

4. 实际耗时对比

为量化实际效率增益，测量了 SD3.5-Large 模型中单个 MMDiT 模块执行一次前向传播的实际耗时。测试在单张 NVIDIA 4090 GPU 上进行，批次大小为 4，分辨率为 1024×1024。

表 5 显示，两种剪枝方法均有效降低了推理延迟。2:4半结构化剪枝实现了 1.23 倍的加速，而结构化剪枝方法在 30% 稀疏度下达到了 1.31 倍的加速。这些结果验证了应用 OBS-Diff 剪枝技术所带来的切实可行的实际加速效益。

表 5  不同稀疏方案下单个 MMDiT 模块的实际推理耗时（毫秒）与加速比。