PAC -tuning: Fine-tuning Pre-trained Language Models with PAC -driven Perturbed Gradient Descent (2023)

Abstract: Fine-tuning pretrained language models (PLMs) for downstream tasks is a large-scale optimization problem, in which the choice of the training algorithm critically determines how well the trained model can generalize to unseen test data, especially in the context of few-shot learning. To achieve good generalization performance and avoid overfitting, techniques such as data augmentation and pruning are often applied. However, adding these regularizations necessitates heavy tuning of the hyperparameters of optimization algorithms, such as the popular Adam optimizer. In this paper, we propose a two-stage fine-tuning method, PAC-tuning, to address this optimization challenge. First, based on PAC-Bayes training, PAC-tuning directly minimizes the PAC-Bayes generalization bound to learn proper parameter distribution. Second, PAC-tuning modifies the gradient by injecting noise with the variance learned in the first stage into the model parameters during training, resulting in a variant of perturbed gradient descent (PGD). In the past, the few-shot scenario posed difficulties for PAC-Bayes training because the PAC-Bayes bound, when applied to large models with limited training data, might not be stringent. Our experimental results across 5 GLUE benchmark tasks demonstrate that PAC-tuning successfully handles the challenges of fine-tuning tasks and outperforms strong baseline methods by a visible margin, further confirming the potential to apply PAC training for any other settings where the Adam optimizer is currently used for training.

## introduction since the emergence of pretrained language models (plms), e.g., bert @xcite and gpt-3 @xcite , fine-tuning of such pretrained models has been the de-facto pipeline for nlp, achieving state-of-the-art results in various tasks. there are two main fine-tuning approaches: parameter-tuning and prompt-tuning. parametertuning considers a plm as a feature extractor and tries to update the complete plm with a small learning step @xcite . by contrast, prompt-tuning aligns downstream tasks with the objective of language modeling through inserting prompts with or without task demonstrations into the original sample and asking the plm to predict the next token according to the prompted input context @xcite @xcite . in the context of few-shot learning, parameter-tuning over a neural model with up to billions of parameters is a non-trivial task @xcite . a significant challenge is that the training process is unstable @xcite : given only a few samples from downstream tasks, the overparameterization nature of a plm leads to issues such as overfitting and forgetting @xcite . existing methods mainly address these challenges by applying data augmentation @xcite @xcite , regularization @xcite @xcite and network pruning @xcite . from the perspective of machine learning theory, data augmentation, regularization, and pruning are all used during training as generalization enhancers. other well-known generalization enhancers include weight-decay and dropout. perhaps a little surprisingly, different choices of learning rates @xcite and mini-batch sizes @xcite also affect generalization. a less wellknown enhancer to the nlp community is noise injection, implemented by the pgd (perturbed gradient descent) algorithm @xcite . theoretically, pgd is shown to effectively help the algorithm escape spurious local minima @xcite and saddle points @xcite , due to an implicit regularization on the trace of the hessian matrix. in this paper, instead of searching for the optimal combination of a basket of generalization enhancers, we follow an alternative training framework, pac-bayes training @xcite , which provides a more straightforward way of improving the generalization -the network is trained towards directly minimizing the generalization error characterized by the pac-bayes bound @xcite instead of minimizing only the training loss. although pac-bayes bounds are classical bounds in learning theory, leveraging them for training is relatively new. this is likely due to concerns that the pac-bayes bounds suffer from the curse of dimensionality @xcite ; therefore they are unlikely to be effective on modern neural networks that are large and deep. however, recent studies @xcite have shown that this view might be too pessimistic, and that the pac-bayes bound could be rather effective in training modern convolutional neural networks. this paper explores the potential of pac-bayes training on even larger models, namely, the plm. we consider the most challenging task in the perspective of pac-bayes training, the fine-tuning task, as it amounts to using an extremely small training dataset to tune a plm with millions of parameters. for this setting, we propose a novel, efficient implementation of pac-bayes training, called pac-tuning, which consists of two stages. the first stage learns the noise variance and updates the plm's parameters by minimizing a pac-bayes upper bound. the second stage implements noise injection training with the noise variance learned from the previous stage. we validate the effectiveness of pac-tuning with few-shot text classification tasks extracted from the glue benchmark. the overall good performance of pactuning suggests promising potential for leveraging pac-bayes training for fine-tuning much larger plms, even during the pretraining process. to the best of our knowledge, pac-tuning is the first work of its kind in terms of improving plm fine-tuning by pac-bayes training. pac-bayes training means training machine learning models by minimizing the pac-bayes upper bound. in contrast to empirical risk minimization, pac-bayes training is more straightforward in terms of improving generalization by minimizing the upper bound of generalization error. mcallester (1998) trains a stochastic neural network on an mnist dataset by minimizing a nonvacuous pac-bayes bound. the pac-bayes training with backprop proposed by @xcite trains shallow probabilistic neural networks and certifies their risk by pac-bayes bound on the mnist dataset. @xcite introduced auto-tune pac to train various neural networks, such as resnet and gnn, through optimizing both the prior distribution variance and posterior distribution variance of parameters. auto-tune pac leverages a larger model and larger dataset, including resnet34, densenet121, and the cifar 100 dataset, and the authors test a gnn on a smaller dataset with only 20 nodes per class. previous works overlook confidence difference between pretrained layers and adaptation layers, this is the main reason that those works can not be applied to plms. we, however, take the confidence difference into account by learning the noise level associated with pretrained layers and adaptation layers separately. perturbed gradient descent (pgd) implicitly regularizes the trace of the hessian matrix to push the model towards a region of the loss land-scape with larger flatness, which is claimed to be a measurement of generalization @xcite . @xcite proves that pgd can help a two-layer convolutional neural network model escape a spurious local minimum and converge to a global minimum. a similar generalization-enhanced benefit of pgd is also validated by @xcite : given pgd, neural network models can converge to second-order stationary points and avoid saddle points. while existing pgd works assign isotropic noise to models, causing training loss explosion, pac-tuning avoids this problem by injecting parameter-wise noises to plms. ## method this section presents our proposed method, pactuning, an implementation of pac-bayes training for parameter-based fine-tuning of plms. section 3.2 introduces pac-bayes training and the pac-bayes bound, followed by section 3.3 which describes perturbed gradient descent. the motivation to assist pgd with pac-bayes training is presented in section 3.4, and we explain the details of pac-tuning in section 3.5. ## conclusions and future work in this paper, we propose a plm fine-tuning method, pac-tuning, for few-shot text classification. pac-tuning is based on pac-bayes training and perturbed gradient descent. we leverage pac-tuning in the more challenging settings of larger models and smaller training data, which are generally considered to be the two main obstacles for improving generalization through pac-bayes training. with extensive experiments on 5 glue benchmark tasks, we observed that the performance of pac-tuning is competitive to other fine-tuning methods and more stable with respect to different model sizes and training dataset sizes. pactuning, our proposed pretrained language model fine-tuning method, can be expanded in several ways for future work. more larger-sized models should be validated to fully explore the effectiveness of pac-tuning. it would also be interesting to augment other fine-tuning techniques with pactuning, especially data augmentation. lastly, the performance of pac-tuning is largely determined by the convergence of stage 1, necessitating more studies to determine how to make stage 1 converge quickly and robustly. our experimental results demonstrate the usefulness of pac-tuning and the potential to consider nlp problems from the point-of-view of generalization, a less explored plm-optimization approach in the nlp community.