Improving Generalization of Pre-trained Language Models via Stochastic Weight Averaging (2022)

Abstract: Knowledge Distillation (KD) is a commonly used technique for improving the generalization of compact Pre-trained Language Models (PLMs) on downstream tasks. However, such methods impose the additional burden of training a separate teacher model for every new dataset.Alternatively, one may directly work on the improvement of the optimization procedure of the compact model towards better generalization. Recent works observe that the flatness of the local minimum correlates well with better generalization.In this work, we adapt Stochastic Weight Averaging (SWA), a method encouraging convergence to a flatter minimum, to fine-tuning PLMs. We conduct extensive experiments on various NLP tasks (text classification, question answering, and generation) and different model architectures and demonstrate that our adaptation improves the generalization without extra computation cost. Moreover, we observe that this simple optimization technique is able to outperform the state-of-the-art KD methods for compact models.

## introduction large pre-trained language models (plms) like bert @xcite and @xcite have achieved state-of-the-art (sota) performance on a number of nlp applications. a large model pre-trained on a massive text corpus learns general-purpose features of the language and can be easily adapted to different downstream tasks by fine-tuning with task-specific objectives @xcite @xcite . multiple papers @xcite @xcite demonstrated that the improvement in pre-training can be transferred favorably to most downstream tasks with little or no data, which leads to considerable efforts to scale up plms and increase the size of the pre-training dataset. however, a recent work @xcite empirically shows that as the upstream accuracy increases, the performance of downstream tasks gradually saturates. furthermore, although plms perform well in a task-agnostic few-shot setting, it is still desirable to fine-tune them to achieve sota performance, especially for models with limited capacity. this encourages more efforts to improve the generalization of plms through additional techniques like knowledge distillation or better optimization. knowledge distillation @xcite is an effective method to improve the generalization of compact plms. one can use a larger model (a teacher) to fine-tune a smaller model (a student) by making a student learn the teacher's output. this technique has helped in achieving sota performance on multiple natural language understanding (nlu) problems for compact models @xcite @xcite . however, kd is not an efficient method as it requires training a separate teacher model for every new task. alternatively, some works focus on boosting the performance of the optimization algorithm directly. several directions have been explored in the literature: adding noise to the gradient descent process @xcite , applying continuation optimization @xcite , steering optimization towards flat optima @xcite , averaging checkpoints @xcite . encouraged by the success of stochastic weight averaging (swa) @xcite for computer vision and sharpness-aware minimization (sam) for nlp @xcite , we investigate whether swa can find better, generalizable minima for finetuning plms while maintaining training efficiency. we conduct extensive experiments for plm finetuning including seven nlu tasks, two question answering tasks, and one summarization task with different model architectures. our contributions are as follows: we adapt swa to nlp and demonstrate that our adaptation can be applied to fine-tuning plms without extra computational cost on both classification and generation tasks. furthermore, we outline several practical recipes that are different from the ones in computer vision including using a high constant learning rate in swa (instead of a cyclical one) and averaging over steps instead of epochs. geometry of the loss landscape loss landscape analysis is a popular method for investigating the model convergence problem @xcite @xcite . it has been widely studied that neural networks tend to generalize better when they converge to flatter local minima @xcite @xcite . some works inspect the curvature and flatness of the loss surface via hessian spectrum analysis @xcite and show the positive correlation between local sharpness and the generalization gap. such observations have inspired research on improving generalization from a geometric perspective @xcite @xcite . another recently proposed method, stochastic weight averaging (swa) @xcite , also improves generalization and is easy to implement without significant computational overhead. the original paper explains the better generalization by empirically showing the flatness of the minimum discovered by swa. however, all the analysis was done on only computer vision tasks and models. to the best of our flatness for plms a recent work @xcite applies sharpness-aware minimization (sam) @xcite to plm fine-tuning on several nlp tasks. it shows that the flatness-aware optimizer can substantially improve the generalization of language models, but it takes 1.25ˆcomputational cost. ## revisting swa we start by briefly reminding the stochastic weight averaging @xcite , a method which implements efficient weight ensembling during model optimization. first, a neural network is trained with sgd till it converges to a region close to a local minimum of the training loss landscape. then, in the second stage, swa averages the checkpoints along the sgd trajectory using a designated learning rate schedule for better exploration of the parameter space. the original paper considers cyclical and high constant learning rate schedules (chc). during the update, swa collects model parameters for every k step and averages them for the model ensemble. @xcite empirically confirms that swa can avoid convergence to a sharp local minimum. the training process is shown in algorithm 1. chc learning rate schedule in every cycle, the learning rate anneals from η max to η min . the learning rate for each training iteration i is given by: swa for plms we accommodate the swa method for plm fine-tuning. we observe that several modifications are required compared to swa for computer vision. first, we noticed that a high constant learning rate performes better than a cyclical learning rate. then, since the fine-tuning needs much less epochs than training from scratch, we average over steps instead of epochs as in the original swa for computer vision. besides that, for most of the nlu tasks we average a larger number of components and start averaging (stage 2 of swa) earlier: after 50% of the training instead of 75% in the original method. the latter could be because the plm is closer to the local minimum than a randomly initialized model and hence allows more exploration. ## experiments in this section, we introduce the setup of our experiments, i.e., the benchmarks and corresponding evaluation methodology and discuss experimental results. all experiments were performed on single nvidia tesla (v100) gpu. we also give an example of flatness analysis in the appendix a. ## conclusion in this work we have adapted the swa algorithm to the fine-tuning of plms. swa is from a family of algorithms that aim to improve generalization for deep neural networks by finding a flatter, as opposed to a sharp, minimum. we present the recipe that can adapt swa to fine-tuning plms, i.e. a high constant learning rate, averaging over training steps and choosing more components for discrimination and less for generation. on compact plms, our adaptation of swa performs on par with sota kd methods on the glue benchmark, and improves on fine-tuning on question answering and text summarization. moreover, unlike kd or sam (another flatness seeking optmization algorithm), swa introduces no additional computation cost. as future work, we want to explore apply to the pre-training of lms.