[论文解读] Make Sharpness-Aware Minimization Stronger: A Sparsified Perturbation Approach

Peng Mi, Li Shen+5 more|arXiv (Cornell University)|Oct 11, 2022

Advanced Neural Network Applications被引用 16

一句话总结

本文提出稀疏SAM（SSAM），一种更高效的锐度感知最小化（SAM）变体，通过二值掩码而非完整参数扰动实现稀疏扰动。通过使用Fisher信息或动态稀疏训练识别关键参数进行扰动，SSAM在达到最高50%稀疏度且计算成本减半的情况下，实现了与SAM相当或更优的泛化性能，同时保持相同的收敛速率$O(\log T/\sqrt{T})$。该方法在CIFAR-10、CIFAR-100和ImageNet上得到验证，实现了更低训练开销下的最先进性能。

ABSTRACT

Deep neural networks often suffer from poor generalization caused by complex and non-convex loss landscapes. One of the popular solutions is Sharpness-Aware Minimization (SAM), which smooths the loss landscape via minimizing the maximized change of training loss when adding a perturbation to the weight. However, we find the indiscriminate perturbation of SAM on all parameters is suboptimal, which also results in excessive computation, i.e., double the overhead of common optimizers like Stochastic Gradient Descent (SGD). In this paper, we propose an efficient and effective training scheme coined as Sparse SAM (SSAM), which achieves sparse perturbation by a binary mask. To obtain the sparse mask, we provide two solutions which are based onFisher information and dynamic sparse training, respectively. In addition, we theoretically prove that SSAM can converge at the same rate as SAM, i.e., $O(\log T/\sqrt{T})$. Sparse SAM not only has the potential for training acceleration but also smooths the loss landscape effectively. Extensive experimental results on CIFAR10, CIFAR100, and ImageNet-1K confirm the superior efficiency of our method to SAM, and the performance is preserved or even better with a perturbation of merely 50% sparsity. Code is availiable at https://github.com/Mi-Peng/Sparse-Sharpness-Aware-Minimization.

研究动机与目标

• 为解决锐度感知最小化（SAM）的高计算成本问题，其因完整参数扰动导致训练时间翻倍。
• 探究SAM中对所有模型参数进行无差别扰动是否必要或次优。
• 开发一种稀疏扰动机制，在降低计算开销的同时保持泛化性能。
• 从理论上证明所提方法的收敛速率与SAM相同，确保理论鲁棒性。
• 在多种模型和数据集上验证该方法，确认其效率与性能的提升。

提出的方法

• 提出稀疏SAM（SSAM），一种SAM的变体，通过可学习或数据驱动的二值掩码仅对部分模型参数施加扰动。
• 提出两种掩码生成策略：SSAM-F利用Fisher信息选择高影响力参数进行扰动；SSAM-D利用动态稀疏训练联合优化参数与稀疏掩码。
• 仅对掩码参数应用扰动，减少每步的梯度计算量，相比标准SAM将计算成本降低约50%。
• 采用两步优化过程：首先在掩码子集上计算扰动方向；其次使用扰动后的梯度更新模型。
• 理论分析证明，在非凸随机设置下，SSAM的收敛速率与SAM相同，为$O(\log T/\sqrt{T})$。
• 实验中使用2:4稀疏运算支持，尽管硬件加速依赖未来进展。

实验结果

研究问题

• RQ1 SAM中的完整参数扰动是否对有效损失景观平滑与泛化性能是必要的？
• RQ2 仅对部分参数施加稀疏扰动，是否能保持或优于完整SAM的泛化性能？
• RQ3 使用Fisher信息或动态稀疏选择是否比随机或均匀选择更有效、更高效？
• RQ4 在稀疏扰动下，SAM的收敛速率是否可保持，以确保理论鲁棒性？
• RQ5 在图像分类基准中，可在不牺牲模型准确率的前提下实现多大程度的稀疏度？

主要发现

• 在CIFAR-10和CIFAR-100上，SSAM在扰动掩码达到50%稀疏度时，测试准确率与SAM相当或更优。
• 在ResNet18和WideResNet28-10上，SSAM在50%稀疏度下优于SAM，在CIFAR-100上实现84.20%的top-1准确率（$\rho = 0.1$）。
• SSAM-F使用$N_F = 128$样本时，在CIFAR-10上达到96.84%准确率，与完整SAM相当，同时相比全批量估计将Fisher信息计算成本降低90%。
• 消融研究证实，随机或基于尖锐权重的掩码会降低性能，验证了基于Fisher信息与动态稀疏策略的有效性。
• 掩码更新频率显著影响性能：间隔越长，准确率越低，表明频繁更新掩码对最优收敛至关重要。
• SSAM-D通过动态掩码更新在ImageNet上使用ResNet50实现77.25% top-1准确率，与SAM性能相当，同时降低计算负载。