Sparse Layer Sharpness-Aware Minimization for Efficient Fine-Tuning
This work addresses the efficiency bottleneck of SAM for practitioners in machine learning fine-tuning, offering a method that reduces computational overhead while maintaining performance, though it is incremental as it builds directly on SAM.
The paper tackles the high computational cost of Sharpness-Aware Minimization (SAM) in fine-tuning by proposing SL-SAM, which uses a sparse layer selection method to reduce active parameters in backpropagation, achieving comparable performance to state-of-the-art baselines while activating only 47%, 22%, and 21% of parameters in vision, moderate, and large language model tasks, respectively.
Sharpness-aware minimization (SAM) seeks the minima with a flat loss landscape to improve the generalization performance in machine learning tasks, including fine-tuning. However, its extra parameter perturbation step doubles the computation cost, which becomes the bottleneck of SAM in the practical implementation. In this work, we propose an approach SL-SAM to break this bottleneck by introducing the sparse technique to layers. Our key innovation is to frame the dynamic selection of layers for both the gradient ascent (perturbation) and descent (update) steps as a multi-armed bandit problem. At the beginning of each iteration, SL-SAM samples a part of the layers of the model according to the gradient norm to participate in the backpropagation of the following parameter perturbation and update steps, thereby reducing the computation complexity. We then provide the analysis to guarantee the convergence of SL-SAM. In the experiments of fine-tuning models in several tasks, SL-SAM achieves the performances comparable to the state-of-the-art baselines, including a \#1 rank on LLM fine-tuning. Meanwhile, SL-SAM significantly reduces the ratio of active parameters in backpropagation compared to vanilla SAM (SL-SAM activates 47\%, 22\% and 21\% parameters on the vision, moderate and large language model respectively while vanilla SAM always activates 100\%), verifying the efficiency of our proposed algorithm.