Sparse Layer Sharpness-Aware Minimization for Efficient Fine-Tuning

TL;DR

This paper introduces SL-SAM, a sparse layer approach to sharpness-aware minimization that reduces computational costs in fine-tuning large models while maintaining high performance, by dynamically selecting layers via a multi-armed bandit strategy.

Contribution

SL-SAM innovatively applies sparsity and multi-armed bandit algorithms to optimize layer selection in SAM, significantly reducing computation during fine-tuning without sacrificing accuracy.

Findings

SL-SAM achieves comparable or better performance than state-of-the-art methods.

SL-SAM activates significantly fewer parameters during backpropagation.

SL-SAM attains the top rank in LLM fine-tuning tasks.

Abstract

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.