Train Flat, Then Compress: Sharpness-Aware Minimization Learns More Compressible Models (2022)

Abstract: Model compression by way of parameter pruning, quantization, or distillation has recently gained popularity as an approach for reducing the computational requirements of modern deep neural network models for NLP. Inspired by prior works suggesting a connection between simpler, more generalizable models and those that lie within wider loss basins, we hypothesize that optimizing for flat minima should lead to simpler parameterizations and thus more compressible models. We propose to combine sharpness-aware minimization (SAM) with various task-specific model compression methods, including iterative magnitude pruning (IMP), structured pruning with a distillation objective, and post-training dynamic quantization. Empirically, we show that optimizing for flatter minima consistently leads to greater compressibility of parameters compared to vanilla Adam when fine-tuning BERT models, with little to no loss in accuracy on the GLUE text classification and SQuAD question answering benchmarks. Moreover, SAM finds superior winning tickets during IMP that 1) are amenable to vanilla Adam optimization, and 2) transfer more effectively across tasks.

## introduction recent advances in hardware, modeling, and optimization for deep neural networks have enabled training of substantially larger models on massive amounts of unlabeled data, leading to corresponding improvements in accuracy across a variety of tasks in nlp @xcite @xcite . unfortunately, this sudden increase in scale of state-of-the-art models also has adverse consequences, such as reducing equity of access @xcite and increasing computational and energy requirements @xcite . in response, model compression has emerged as a dominant approach to improving memory and inference efficiency in neural network models, including approaches such as knowledge distillation @xcite @xcite , model quantization @xcite and pruning @xcite @xcite . the vast majority of work on model compression focuses on methods for selecting how and where to reduce (via quantization or distillation) or remove (by pruning) model parameters without sacrificing end-task accuracy. while these approaches are usually simple to implement and work relatively well for maintaining overall accuracy on considered benchmarks, recent work has revealed that these commonly used compression methods can result in negative impacts on model behavior that are not necessarily captured by current performance met-rics. for example, in the field of computer vision, pruning has been shown to sacrifice performance on long-tail examples in order to preserve accuracy on more frequent phenomena @xcite and reduce robustness to out-of-distribution examples @xcite . at the same time, it has also been shown that compression can sometimes have a regularizing effect, improving generalization in some settings @xcite . clearly, more work is needed better understand the relationship between compression and generalizability, and to develop improved compression methods that are informed by this knowledge. meanwhile, there is a growing body of work investigating curvature of the loss landscape and its relation to generalization in deep neural models. @xcite first defined flat minima as regions in parameter space where error remains relatively stable despite perturbations in parameter values, arguing that models in flat minima should correspond to simpler, more generalizable functions. more recently, @xcite further showed that wide loss basins correspond to low-complexity solutions, 2 and it has been shown empirically that directly optimizing for solutions in flat minima leads to improved generalization on a wide range of supervised learning tasks in both vision @xcite and language @xcite modalities. inspired by the above results connecting wider loss regions to simpler, more generalizable models, in this work we aim to advance our understanding of model compression by examining the relationship between flat minima and compressibility in fine-tuned language models. intuitively, model parameters in neighborhoods having uniformly low loss values should be more robust to perturbations, such as those resulting from model compression, compared to models in sharper regions, since changes to parameter values should lead to minimal change in loss with respect to the main training objective. empirically and theoretically, previous work has also linked generalizability with compressibility, finding that neural networks whose weights reflect simpler, more general hypotheses may be more robust to compression, and compression itself can act as a regularization mechanism @xcite @xcite . @xcite showed that larger pre-trained language models, which are known to genereralize better, are also more compressible. further, @xcite connect flatness to generalization in pruned models, showing that pruning noise correlates positively with measures of flatness in a cnn for computer vision. we investigate the relationship between flat minima and compression in large pre-trained language models by directly optimizing for flat minima during language model fine-tuning using sharpnessaware minimization (sam; @xcite ). through extensive experiments on the glue text classification and squad question answering benchmarks, we show that fine-tuning bert models with sam leads to optima in flatter basins, and compressing those models consistently results in higher model accuracy at the same level of compression compared to standard adam-optimized baselines. our results hold across multiple bert variants and sizes @xcite and a wide variety of compression methods: iterative magnitude pruning with and without rewinding, a structured pruning procedure that also employs knowledge distillation, and an off-theshelf method for post-training quantization. we also show that sparse subnetworks (winning tickets; @xcite ) discovered by sam are better transferable across tasks, suggesting improved generalizability. our findings shed light on a promising new avenue for obtaining practical improvements in model compression. ## methods broadly, we are interested in understanding: 1) are models in flatter minima more compressible? 2) if so, why? 3) beyond task-specific accuracy, what properties do flat, compressed models have? to provide empirical answers to these questions, our experimental setup is as follows. we fine-tune pre-trained language models on standard benchmarks using both "vanilla" adam optimization and sharpness-aware minimization ( §2.1). we then experiment with a variety of strategies for compressing those fine-tuned models: iterative magnitude pruning (unstructured) with and without rewinding ( §2.2.1), structured pruning using ℓ 0 regularization and a distillation obejctive ( §2.2.2), and posttraining dynamic quantization ( §2.3). we evaluate model end-task accuracy at different compression rates, and compare that accuracy to the full (uncompressed) model and across experimental settings, such as varying the pre-trained language model, and transferring initializations across tasks. ## related work in this section, we draw connections to key related work. for a more general description of related work, please refer to a.1 in appendix. first, we briefly recount previous work that we consider for our experimentation. we use adam (kingma and ba, 2014) as a base optimizer, with comparisons between vanilla adam and the addition of sharpness-aware-minimization @xcite . we make further comparisons in a subset of our experiments with stochastic weight averaging @xcite as an alternative method of inducing flatness. using keskar et al. (2017)'s ϵ-sharpness metric, and based on mehta et al. (2021)'s implementation, we verify that sam and swa induce flatness. we fine-tune bert models @xcite for glue @xcite and squad @xcite language tasks. the compression methods we couple with "training flat" include: 1) unstructured imp to find winning lottery tickets @xcite , referring to chen et al. ( 2020) flatness and generalization prior work has investigated the connection between flat minima and generalization (i.e., the gap between training accuracy and holdout set accuracy), starting with hochreiter and schmidhuber @xcite . subsequent work has continued on this front, exploring notions of sharpness and their predictiveness of generalization under different conditions @xcite , as well as empirical evaluations of flatness-inducing methods and their effects on generalization. jiang* et al. (2020) find that sharpness is empirically predictive of generalization, including in particular a perturbation magnitudeaware metric very similar to the ϵ-sharpness metric introduced in @xcite and used in our paper. in this work, however, separate from generalization, we primarily focus on investigating the underexplored relationship between flatness and compression. flatness and compression previous work has mentioned flatness in the context of pruning. in fact, hochreiter and schmidhuber (1997)'s orig-inal "flat minimum search" algorithm is explicitly designed to prune units, weights, and input lines as a mechanism for finding flat minima. @xcite also propose pruning unimportant weights as a way to obtain improved generalization, although without any notions of flatness. since these earlier works, the deep learning landscape has changed such that effective model compression itself is now often a goal; model efficiency in terms of size and latency is a common priority. to the best of our knowledge, we are the first to directly relate loss landscape flatness with model compressibility. we view our results as complementary to the concurrent work by @xcite , who study properties of winning tickets found during iterative magnitude pruning and find that imp itself preferentially prunes parameters lying along flatter directions in the loss landscape; they also theorize that flatter landscapes allow for more aggressive pruning. we note that the optimizer, data, and model architectures they use are different from ours. while @xcite use stochastic gradient descent for image classification tasks on resnet architectures, we use adam as a base optimizer for text classification and question answering tasks with bert architectures. nonetheless, to the extent that findings from both papers can generalize to other settings, @xcite lay theoretical groundwork which supports our explicit suggestion to "train flat" as a strategy for inducing greater compressibility in neural models. ## discussion we show that in general, sam helps models retain higher accuracy on a variety of language tasks at higher sparsity levels. this holds true in multiple unstructured iterative magnitude pruning settings, as well as in a structured pruning setting with a distillation objective. moreover, our additional experiments and analyses point to sam-learned structures playing an important role in compressibility, as well as transferring well across tasks.