ZO-SAM: Zero-Order Sharpness-Aware Minimization for Efficient Sparse Training¶
会议: CVPR2025
arXiv: 2603.13115
代码: 待确认
领域: others
关键词: 稀疏训练, 零阶优化, Sharpness-Aware Minimization, 梯度方差, 模型压缩
一句话总结¶
提出 ZO-SAM，将零阶优化策略性地整合到 SAM 的扰动步骤中，仅需一次反向传播即可获得 SAM 的平坦最小值优势，在稀疏训练场景下将计算开销减半的同时提升精度和鲁棒性。
研究背景与动机¶
- 稀疏神经网络通过大幅减少参数量来降低计算和内存开销，但现有稀疏训练方法在高稀疏度时面临梯度信号混乱噪声大的核心问题，严重阻碍收敛和泛化
- SAM（Sharpness-Aware Minimization）通过引导模型走向平坦最小值来降低梯度方差、改善泛化，但每步需要两次完整反向传播，计算开销翻倍
- 在稀疏训练场景下，计算效率本身就是核心诉求，SAM 的双倍开销直接削弱其可行性
- 核心问题：能否在保留 SAM 平坦最小值优势的同时，消除双次反向传播的冗余计算？
方法详解¶
ZO-SAM 框架设计¶
ZO-SAM 的核心思想是在 SAM 的两步过程中选择性使用零阶/一阶梯度：
- 扰动步骤（零阶替代）：用 Random Gradient Estimation (RGE) 替代反向传播计算扰动方向
- \(\hat{\nabla}\mathcal{L}(\theta) = \frac{1}{m}\sum_{i=1}^{m}\frac{\mathcal{L}(\theta+\delta u_i)-\mathcal{L}(\theta-\delta u_i)}{2\delta}u_i\)
- 只需前向传播（函数评估），无需反向传播
-
扰动步骤的目的是确定扰动方向，对梯度精度要求较低，容忍近似误差
-
梯度更新步骤（保留一阶精度）：在扰动后的参数点 \(\theta+\epsilon\) 处执行完整反向传播
- \(\theta \leftarrow \theta - \eta\nabla\mathcal{L}(\theta+\epsilon)\)
- 更新步骤直接影响参数优化质量，需要精确梯度
设计选择的合理性¶
- 选择 RGE 而非 CGE（逐坐标估计）：RGE 仅需 \(m \ll d\) 次函数评估，CGE 需 \(d\) 次，大规模模型中不可行
- RGE 的随机方向采样提供更全局的损失面探索，有助于避免尖锐最小值
- 零阶估计放在第一步（扰动）而非第二步（更新）：扰动步骤容错性高，更新步骤需要精度
计算开销分析¶
- 标准 SAM：每步需 2 次完整反向传播 → 开销约为 SGD 的 2 倍
- ZO-SAM：第一步仅需 2m 次前向传播（函数评估），第二步 1 次反向传播 → 总开销约 SGD 的 1 + 2m·(前向/反向比)
- 当 m 较小时（论文中通常取 m=1），ZO-SAM 开销远低于 SAM，实测吞吐量约为 SAM 的 1.6 倍
实验关键数据¶
ResNet-32 on CIFAR-10/100（稀疏度 90%/95%/98%）¶
- ZO-SAM 在所有 7 种稀疏训练方法上均带来一致提升
- CIFAR-10 提升范围：0.38%～2.31%（ResNet-32）
- CIFAR-100 提升范围：0.45%～2.54%（ResNet-32）
- 最大提升出现在 RigL + ZO-SAM：CIFAR-100 90% 稀疏度提升 2.54%
DeiT on ImageNet-1K¶
- DeiT-Tiny 50% 稀疏度：最大提升 1.14%（SViTE → SNIP+ZO-SAM）
- DeiT-Small 70% 稀疏度：最大提升 1.17%（RigL+ZO-SAM）
与 SAM 系列效率对比（ResNet-32, 90% 稀疏度, MEST）¶
| 方法 | CIFAR-10 | 吞吐量(img/s) | 相对效率 |
|---|---|---|---|
| SAM | 93.77% | 2704 | 47.67% |
| GSAM | 93.72% | 2701 | 47.60% |
| ZO-SAM | 93.50% | 4349 | 76.67% |
鲁棒性（CIFAR-10-C, SNIP, 90%）¶
- ZO-SAM 在腐蚀数据集上精度提升 3.10%，Δ 最小，鲁棒性最强
收敛速度¶
- 达到 90% 精度的 epoch 数：ZO-SAM 70 epochs vs SGD 104 epochs（95% 稀疏度下 88 vs 131）
- 收敛速度优于所有 SAM 变体（ESAM 75/92, LookSAM 79/94, GSAM 84/113）
损失面可视化¶
- 无 ZO-SAM 时高稀疏度损失面呈窄而陡的 basin
- 加入 ZO-SAM 后损失面变为宽而平坦的 basin，表明梯度稳定性改善
特征图对比¶
- 在 ResNet-32 的浅层(3)/中层(17)/深层(31) 均观察到更清晰、更集中的特征激活
- 基线方法的特征图呈现散乱或模糊模式，梯度方差大
亮点¶
- 精妙的混合策略：不是简单替换所有梯度为零阶，而是根据 SAM 两步各自的精度需求选择性使用，兼顾效率与质量
- 即插即用：可与任意稀疏训练方法（静态/动态）组合，均获得一致提升，无需修改基线方法的训练流程
- 多维度验证充分：覆盖 CNN/Transformer、多数据集、多稀疏度、损失面可视化、收敛速度、特征图、鲁棒性
- 实用性强：吐量达 SGD 的76.67%，远超 SAM 的47.67%，在资源受限场景下真正可用
局限性¶
- 精度略低于完整 SAM（~0.27% on CIFAR-10），用效率换取了少量精度
- 超参数 \(m\)（RGE 采样方向数）和 \(\delta\)（步长）的选择缺乏深入分析和敏感性实验
- 仅在分类任务上验证，未涉及检测、分割等下游任务
- 理论分析（收敛保证）相对薄弱，缺少零阶估计在 SAM 框架下的正式收敛证明
- RGE 在极高维参数空间的近似质量可能下降，论文未讨论模型规模的上限
- 未与其他零阶优化方法（如 ZO-SGD、ZO-Adam）在 SAM 框架下的整合进行对比
- 实验仅覆盖视觉分类骨干，未验证在 NLP/多模态等其他领域的通用性
评分¶
- 新颖性: ⭐⭐⭐（将零阶优化融入 SAM 的想法直觉合理但并不意外）
- 实验充分度: ⭐⭐⭐⭐⭐（大量方法/架构/数据集/指标的全面覆盖）
- 写作质量: ⭐⭐⭐⭐（动机清晰、逻辑流畅）
- 价值: ⭐⭐⭐⭐（稀疏训练实践中有直接应用价值）