OBS-Diff: Accurate Pruning For Diffusion Models in One-Shot

本文提出了 OBS-Diff，一种新颖的无需训练的单次剪枝框架，通过适配现代扩散模型架构的 OBS 算法、基于误差累积视角的 timestep 感知 Hessian 构建以及高效的组级顺序剪枝策略，实现了大规模文生图扩散模型的准确压缩与推理加速。

要点

这篇论文介绍了一种名为 OBS-Diff 的新方法，它的核心任务是给庞大的“文生图”AI 模型（比如 Stable Diffusion 3）做**“瘦身”，而且是在不重新训练**的情况下，一次性完成。

为了让你轻松理解，我们可以把整个过程想象成**“给一位正在画连环画的超级画家做手术”**。

1. 背景：为什么需要“瘦身”？

现在的文生图 AI 就像一位才华横溢但极其臃肿的超级画家。

• 优点：他画得极好，能根据文字描述画出惊人的图片。
• 缺点：他太胖了（参数量巨大，几十亿甚至上百亿），画画时需要巨大的房间（显存）和极长的时间（计算成本）。普通人根本请不起他，也跑不动他。

科学家想把他变瘦（剪枝/Pruning），去掉那些不重要的笔触（权重），让他变快、变小。但以前的方法有两个大麻烦：

1. 太慢：以前的方法像“边剪边练”，剪一刀就要重新学半天，成本太高。
2. 不懂行：以前的方法（比如直接剪掉数值小的）像是一个不懂画画的理发师，只敢剪头发，不敢动眉毛，结果剪完画家画歪了，或者画出来的东西全是乱码。

2. OBS-Diff 的三大绝招

OBS-Diff 就像一位精通解剖学且懂画画的顶级外科医生，它用了三个巧妙的策略来解决上述问题：

绝招一：时间感知手术刀（Timestep-Aware Hessian）

比喻：画画的“起稿”比“上色”更重要。

• 原理：AI 画画不是一下子完成的，而是像倒放视频一样，从一团模糊的噪点开始，一步步“去噪”变成清晰的图。
  • 早期步骤（起稿）：决定了画面的整体构图和骨架。如果这里剪错了，后面画得再细也是歪的。
  • 晚期步骤（上色）：只是修饰细节。
• OBS-Diff 的做法：以前的方法对所有步骤一视同仁。OBS-Diff 给早期步骤戴上了“放大镜”，认为它们最重要。它在决定剪掉哪根神经时，会问：“这根神经在起稿阶段重要吗？”如果重要，就坚决保留；如果不重要，就大胆剪掉。
• 效果：保证了画家在起稿时手不抖，画出来的图结构依然稳固。

绝招二：分组流水线（Module Packages）

比喻：不要一个一个剪，要“打包”剪。

• 原理：如果要把画家身上的几亿根神经（权重）一根根检查、剪掉、再检查，那得剪到猴年马月。
• OBS-Diff 的做法：它把画家的身体分成几个**“大包裹”**（比如把画眼睛、画鼻子、画嘴巴的神经打包）。
  • 它先让画家画一张图，同时收集这一整包神经的数据。
  • 然后一次性把这包里所有不重要的神经都剪掉。
  • 再换下一包。
• 效果：大大加快了“体检”和“手术”的速度，让整个过程在几分钟内就能完成，而不需要几天。

绝招三：全能手术刀（支持多种剪法）

比喻：既能剪碎发，也能剪掉整条胳膊。

• 原理：以前的方法只能剪“碎发”（非结构化剪枝，随机剪几个点），或者只能剪“整条胳膊”（结构化剪枝，剪掉整个模块），很难兼顾。
• OBS-Diff 的做法：它非常灵活。
  • 非结构化：像修剪草坪，哪里草高了剪哪里（适合通用加速）。
  • 半结构化：像剪出整齐的 2:4 发型（适合特定硬件加速）。
  • 结构化：直接剪掉整个“画眼睛的模块”或“画背景的模块”（适合大幅减小模型体积）。
• 效果：不管你想怎么剪，它都能剪得漂亮，不会让画家“残疾”。

3. 结果如何？

实验证明，OBS-Diff 的效果惊人：

• 剪得多：即使剪掉了 50% 甚至 70% 的参数（相当于把画家体重减半），他画出来的图依然清晰、漂亮，和没剪之前几乎没区别。
• 剪得快：整个过程不需要重新训练，“一刀切” 完成。
• 对比强烈：其他方法（比如直接剪掉数值小的）在剪掉 50% 后，画出来的图可能变成一团乱麻（全是噪点或扭曲的人脸），而 OBS-Diff 依然能画出“从头发里长出彩色花朵”这样复杂的画面。

总结

OBS-Diff 就像是给庞大的 AI 画家做了一次精准的“微创整形手术”。
它不再盲目地乱剪，而是懂得画画的规律（时间步权重），懂得批量操作（分组策略），并且什么类型的剪法都会。最终，它让这位超级画家变得轻便、快速，却依然才华横溢，让普通人的电脑也能跑得动最顶尖的 AI 绘画模型。

技术摘要

这是一篇发表于 ICLR 2026 的会议论文，题为 OBS-Diff: Accurate Pruning for Diffusion Models in One-Shot。该论文提出了一种针对大规模文本到图像（Text-to-Image）扩散模型的单次、无需训练（One-shot, Training-free）的剪枝框架。

以下是对该论文的详细技术总结：

1. 研究背景与问题 (Problem)

• 计算成本高昂：大规模扩散模型（如 Stable Diffusion 3/3.5, Flux）参数量巨大（数十亿），导致推理计算和内存需求极高，限制了其广泛应用。
• 现有方法的局限性：
  • 缺乏通用性：现有的剪枝方法多针对特定架构（如 U-Net），难以适应现代扩散模型（如 MMDiT）的复杂结构。
  • 迭代特性挑战：扩散模型具有迭代去噪的特性，参数在多个去噪步（Timesteps）中共享。传统的基于单次前向传播的剪枝方法（如 LLM 领域的 SparseGPT）无法直接应用，因为早期步长的误差会在后续步骤中累积并放大。
  • 依赖训练/微调：许多现有方法需要昂贵的梯度计算或剪枝后的微调（Fine-tuning），增加了部署成本。
  • 非结构化剪枝缺失：针对大规模扩散模型的非结构化（Unstructured）和半结构化（Semi-structured）剪枝研究较少。

2. 核心方法论 (Methodology)

OBS-Diff 框架旨在通过改进经典的**最优脑外科医生（Optimal Brain Surgeon, OBS）**算法，解决上述问题。其核心包含三个关键创新：

A. 时间步感知的 Hessian 构建 (Timestep-Aware Hessian Construction)

• 问题洞察：扩散模型的去噪过程是迭代的。早期去噪步（小 $t$）引入的误差会传播并累积到最终输出，而后期步长的误差影响较小。
• 解决方案：
  • 重新定义剪枝目标，最小化加权的重构误差，而非均匀误差。
  • 提出一种对数递减（Logarithmic-decrease）的权重调度方案 $\alpha_t$，赋予早期去噪步更高的权重（$\alpha_1 > \alpha_2 > \dots > \alpha_T$）。
  • 构建时间步感知的 Hessian 矩阵：$H_l = 2 \sum_{t=1}^T \alpha_t E[X_{l,t}X_{l,t}^T]$。这使得剪枝算法能识别并保留对早期关键生成步骤至关重要的权重，从而抑制误差累积。

B. 组序剪枝策略与模块包 (Group-wise Sequential Pruning & Module Packages)

• 问题洞察：在迭代模型中，逐层（Layer-wise）校准需要多次执行完整的去噪轨迹，计算成本 prohibitively high（高得无法接受）。
• 解决方案：
  • 引入**“模块包”（Module Packages）**概念：将模型划分为若干组（Package），每组包含输入相互独立的层（Basic Units）。
  • 组序处理：
    1. 数据收集：对每个模块包，仅执行一次完整的去噪轨迹，利用 Hook 同时收集包内所有层的激活值统计信息。
    2. 并行剪枝：基于收集的数据构建 Hessian 矩阵，同时剪枝包内所有层。
    3. 状态更新：处理完一个包后更新网络状态，再进入下一个包。
  • 优势：大幅减少了校准所需的完整前向传播次数，在显存开销和计算时间之间取得了平衡。

C. 多粒度剪枝支持 (Multi-granularity Support)

OBS-Diff 具有高度灵活性，支持三种剪枝粒度：

1. 非结构化剪枝：直接应用 OBS 算法移除单个权重。
2. 半结构化剪枝 (N:M)：例如 2:4 稀疏模式，在每 4 个权重中移除 2 个重要性最低的。
3. 结构化剪枝：
   • FFN 神经元：聚合神经元所有权重的显著性得分进行移除。
   • MHA 注意力头：针对 MMDiT 架构中多头注意力的特殊结构（多模态输入共享头但输出路径分离），提出使用**倒数秩融合（Reciprocal Rank Fusion, RRF）**技术，融合不同模态路径下的头部重要性排名，统一决定哪些头被剪除。

3. 主要贡献 (Key Contributions)

1. 首个针对扩散模型的通用 OBS 框架：成功将经典 OBS 算法适配到现代扩散模型（如 MMDiT），支持非结构化、半结构化和结构化剪枝。
2. 时间步感知机制：首次从误差累积角度提出时间步感知的 Hessian 构建方法，通过加权机制优先保护早期去噪步骤，显著提升了高稀疏度下的生成质量。
3. 高效的校准策略：提出的“模块包”组序剪枝策略，解决了迭代模型校准成本过高的问题，实现了单次、无需训练的剪枝。
4. SOTA 性能：在多个主流模型（SD v2.1, SD 3-Medium, SD 3.5-Large, Flux.1-dev）上验证了方法的有效性。

4. 实验结果 (Results)

• 评估指标：FID (Frechet Inception Distance), CLIP Score, ImageReward (人类偏好对齐度)。
• 非结构化剪枝：
  • 在 SD 3-Medium 和 SD 3.5-Large 上，OBS-Diff 在 50%-70% 的稀疏度下，ImageReward 和 CLIP Score 均显著优于基线（Magnitude, Wanda, DSnoT）。
  • 在高稀疏度（如 60%）下，基线方法（如 Wanda, DSnoT）性能崩溃（ImageReward 为负或极低），而 OBS-Diff 仍能保持高质量的图像生成。
  • 定性结果显示，OBS-Diff 生成的图像细节丰富、结构完整，而基线方法常出现严重的伪影或结构破坏。
• 结构化剪枝：
  • 在 SDXL (U-Net) 和 SD 3.5 (MMDiT) 上，OBS-Diff 在 15%-30% 的稀疏度下，FID 和 ImageReward 几乎与稠密模型持平。
  • 相比之下，L1-norm 基线在 15% 稀疏度时 FID 即发生灾难性下降（从 31 升至 158+）。
• 效率：
  • 无需训练：整个剪枝过程无需梯度更新。
  • 速度快：在单张 RTX 4090 上，2B 参数的 SD 3-Medium 模型剪枝耗时不到 15 分钟。
  • 加速效果：半结构化（2:4）剪枝带来 1.23x 加速，结构化（30%）剪枝带来 1.31x 加速。
• 鲁棒性：
  • 对不同的 CFG 尺度、采样步数（15/25/50 步）和采样器（Euler, Heun, DPM++ 等）表现出良好的泛化性。
  • 校准数据集（100 个提示词）与评估数据集（MS-COCO）分布不同（OOD），模型未出现过拟合。

5. 意义与影响 (Significance)

• 理论突破：将经典的二阶优化理论（OBS）成功迁移并适配到具有迭代特性的扩散模型中，填补了该领域在训练-free 剪枝方面的空白。
• 实际应用价值：提供了一种低成本、高效率的模型压缩方案，使得在消费级硬件（如单张 4090）上部署和运行超大规模扩散模型成为可能，同时保持了极高的生成质量。
• 通用性：该方法不仅适用于当前的 MMDiT 架构，其核心思想（时间步加权、组序校准）也可推广至视频生成、Diffusion LLMs 及其他迭代式生成模型。

总结：OBS-Diff 通过引入时间步感知的误差累积机制和高效的组序校准策略，解决了扩散模型剪枝中的核心难点，实现了在无需微调的情况下，对大规模扩散模型进行高精度的压缩，是目前该领域的 State-of-the-Art 方法。