Randomized Sharpness-Aware Training for Boosting Computational Efficiency in Deep Learning
This addresses computational efficiency for deep learning practitioners using sharpness-aware methods, offering a significant reduction in training time with minimal performance loss, though it is incremental as it builds on existing SAM techniques.
The paper tackles the high computational cost of sharpness-aware training algorithms like SAM by proposing Randomized Sharpness-Aware Training (RST), which randomly mixes base and sharpness-aware algorithms to reduce extra forward-backward propagations, achieving performance comparable to or better than SAM while saving 50% extra computation cost.
By driving models to converge to flat minima, sharpness-aware learning algorithms (such as SAM) have shown the power to achieve state-of-the-art performances. However, these algorithms will generally incur one extra forward-backward propagation at each training iteration, which largely burdens the computation especially for scalable models. To this end, we propose a simple yet efficient training scheme, called Randomized Sharpness-Aware Training (RST). Optimizers in RST would perform a Bernoulli trial at each iteration to choose randomly from base algorithms (SGD) and sharpness-aware algorithms (SAM) with a probability arranged by a predefined scheduling function. Due to the mixture of base algorithms, the overall count of propagation pairs could be largely reduced. Also, we give theoretical analysis on the convergence of RST. Then, we empirically study the computation cost and effect of various types of scheduling functions, and give directions on setting appropriate scheduling functions. Further, we extend the RST to a general framework (G-RST), where we can adjust regularization degree on sharpness freely for any scheduling function. We show that G-RST can outperform SAM in most cases while saving 50\% extra computation cost.