Sharpness-Aware Minimization Improves Language Model Generalization (2022)

Abstract: The allure of superhuman-level capabilities has led to considerable interest in language models like GPT-3 and T5, wherein the research has, by and large, revolved around new model architectures, training tasks, and loss objectives, along with substantial engineering efforts to scale up model capacity and dataset size. Comparatively little work has been done to improve the generalization of these models through better optimization. In this work, we show that Sharpness-Aware Minimization (SAM), a recently proposed optimization procedure that encourages convergence to flatter minima, can substantially improve the generalization of language models without much computational overhead. We show that SAM is able to boost performance on SuperGLUE, GLUE, Web Questions, Natural Questions, Trivia QA, and TyDiQA, with particularly large gains when training data for these tasks is limited.

## introduction over the last several years, remarkable progress has been made within the domain of natural language understanding, with machine-learned models able to solve some tasks at near or above humanlevel performance. this progress has, by and large, been fueled by research centered around 1) better inductive biases, such as the attention-enabled transformer architecture @xcite , 2) the clever leverage of massive corpora of textual data that was historically disregarded as "unlabeled," usually in the form of pre-training objectives that strive to teach the model the structure of language @xcite , 3) scaling up model capacity and the methods to support it @xcite , 4) multi-task learning @xcite , and lastly, 5) larger and more diverse datasets along with everimproving benchmarks that attempt to test the true capabilities of these models. although these efforts all share the single goal of improving the model's generalization, doing so by explicit changes to the optimization of the loss function has received less attention in comparison. recently, motivated by the both empirical and theoretical findings that flatter minima lead to better generalization @xcite @xcite , @xcite proposed a novel modification to vanilla stochastic gradient descent they term "sharpness-aware minimization," or sam. they show theoretically and empirically that optimizing with sam encourages convergence to flatter points in the loss landscape and with it comes the anticipated improvement in out-of-sample error. while their empirical findings are limited to computer vision tasks and datasets using convolutional neural networks (resnets), follow-up work @xcite showed how sam is particularly effective on vision transformers (vits) @xcite and mlp-mixers @xcite , architectures that are more prone than convolutional ones to land in sharp minima. crucially, they show that when equipped with sam, vits outperform resnets of similar size and throughput without the need for large-scale pre-training. encouraged by wins in the vision domain, we ask whether sam can deliver similar gains in the language domain. our contributions are as follows: while first-order optimization via sgd has been the prevailing way of training neural networks due to its efficiency and effectiveness, alternative second-order methods like k-fac @xcite and shampoo @xcite have slowly gained traction, often enabled by clever engineering to make them feasible at scale. notably, @xcite presents a scalable implementation of shampoo that provides significant convergence and wall-clock time improvements compared to first-order methods. they demonstrate superior performance on machine translation and language modeling. sam. while this is, to the best of our knowledge, the first work detailing the benefits of sam for language tasks, there have been successful applications of sam in the vision domain. notably, @xcite showed that convolution-free vision models like vision transformers (vits) @xcite and mlp-mixers @xcite suffer from sharp minima and that sam indeed smooths their loss landscapes. they crucially show that vits and mlp-mixers outperfom resnets of similar and greater size on imagenet without the use of pre-training or data augmentations that would otherwise be necessary to achieve reasonable performance. they show that sam induces sparsity in both architectures and leads to more perceptive attention maps in vits. they observe empirically that data augmentation and sam are alike in that they both smooth the landscape on average, but the latter does so by explicitly controlling the worst-case curvature, whereas the former smooths over the directions induced by the augmentations. furthermore, they observe that sam encourages linearity with respect to the input, exhibiting an effect similar to that of mixup @xcite . lastly, they show that sam helps contrastive learning and that it enables better robustness on corrupted examples from imagenet-c @xcite and imagenet-r @xcite . in a similar spirit, @xcite proposed speeding up sam significantly by using fewer examples when computing the ascent step, a strategy which we employ in this work, and they were able to apply it to resnet model variants to advance the state of the art on imagenet without extra data. meanwhile, in an attempt to make sam's radius ρ invariant to the scale of the model parameters, @xcite proposed an adaptive version named adaptive sharpness-aware minimization (asam), which they then show empirically to outperform normal sam on a set of benchmark vision tasks. ## review of sharpness-aware minimization (sam) we begin by briefly reviewing the sam algorithm; interested readers can see the original paper for a thorough treatment. in our presentation, we use the ℓ 2 norm (@xmath0, assume a general optimizer (instead of vanilla sgd), and use the approximation proposed by @xcite to compute the ascent gradient (adversarial point) efficiently. given a loss function l : w × x × @xmath1, sam seeks to find the parameter w whose neighborhood has low training loss by optimizing the minimax objective: finding the exact optima * of the innermaximization is challenging, so @xcite employ a first-order approximation, resulting in: that is, ˆ is just a scaling of the loss gradient at the current parameters. after computing ˆ (w), sam performs gradient descent using the gradient put another way, sam plugs-and-plays with any first-order optimizer by simply replacing the gradient of the mini-batch b at the current model weights w @xmath2. w adv itself is computed by taking a gradient ascent step of size ρ along the unit gradient vector ∇ w l m (w)/||∇ w l m (w)|| 2 | wt , where m can be the mini-batch b, or a subset of it for enhanced efficiency. we found that setting m to be 1/4-th of b sped up the method significantly with little loss in quality, in line with the recommendation of @xcite . the end-to-end algorithm is outlined in algorithm 1. ## experiments with sam reviewed, we now discuss our experiments. we evaluate sam on a range of natural language understanding tasks using the t5 (text-totext transformer) framework @xcite . t5 casts nlu tasks as sequence-to-sequence ones that are learned using an encoder-decoder transformer @xcite architecture setup. these transformer models are typically pre-trained on large corpora, like the colossal clean crawled corpus (c4) @xcite , with, for example, the objective of predicting a short contiguous span of text that was intentionally corrupted in a snippet of input text. the pre-trained model is typically fine-tuned on a single task or a mixture of algorithm 1 efficient sam algorithm. 1: input: training set s ∪ n @xmath3, y i )}, loss function l : w × x × @xmath4, batch size b, neighborhood size ρ > 0 (default 0.15), ascent micro-batch size @xmath5, firstorder optimizer update opt : w × @xmath6. to this end, we evaluate sam in two ways: ## conclusion to the best of our knowledge, this paper is the first to demonstrate how the recently-proposed sharpness-aware minimization can be applied for fine-tuning the ubiquitous text-to-text transformer (t5) and its multilingual counterpart mt5 on language tasks of broad interest. we thereby corroborate the already-documented success the method has had in the vision domain. furthermore, we reveal sam's benefits when data is limited by finetuning on subsamples of the original task training split. by approximating the ascent step of the algorithm via fewer samples, we show how large gains can be had across benchmarks and model sizes while adding only around 25% additional compute and wall-clock training time. our hope is that this work will spur sam's adoption in the natural language processing community the way it is starting to in the vision one.