Improving generalization in large langue model by learning prefix subspaces (2023)

Abstract: This article focuses on large language models (LLMs) fine-tuning in the scarce data regime (also known as “few-shot learning setting”). We propose a method to increase the generalization capabilities of LLMs based on neural network subspaces. This optimization method, recently introduced in computer vision, aims to improve model generalization by identifying wider local optima through the joint optimization of an entire simplex of models in parameter space. Although this property would be highly beneficial in the context of training large language models in the “few-shot learning” setting, its adaptation to massive, pretrained transformers poses some challenges. First, their considerable number of parameters make it difficult to train several model jointly, and second, their deterministic parameter initialisation schemes make them unfit to the subspace method as originaly proposed. We show in this paper that its application to “Parameter Efficient Fine-Tuning” (PEFT) methods, however, is relatively natural, and we propose to apply it to prefix-tuning, by learning entire simplexes of continous prefixes. We test our method on a variant of the GLUE benchmark adapted to the few-shot learning setting, and show that both our contributions (learning prefix simplexes, and non-deterministic validation metric inference) jointly lead to a gain in average performances compared to state of the art methods.

## introduction the emergence of large language models @xcite @xcite in recent years has significantly transformed the applications of deep learning methods in natural language processing. these models, pretrained in an unsupervised fashion on massive textual datasets, enable the fine-tuning of powerful models with just a few thousand -or even hundred-observations. they achieve generalization performances that required millions of observations just a few years ago, particularly when used in conjunction with discrete instruction prompts @xcite . extending these discrete methods to the learning of continuous prompts @xcite , which conceptually falls within the framework of socalled "parameter efficient fine-tuning" (peft) methods @xcite , poses certain challenges in the context of few-shot learning. one such challenge is the issue of model adjustment guidance through validation metric during gradient descent @xcite . traditionally, in the process of model fitting, approximately one-third of the training dataset is excluded beforehand to create a validation (or development) set dedicated to inferring an unbiased estimation of the model's performance @xcite . this metric is utilized both during gradient descent (to estimate convergence of the descent algorithm or inform early stopping heuristics), and subsequently to guide hyperparameter searches typically employed in the fine-tuning of large language models. however, the validity of this approach relies on the assumption that the distribution of the validation set is representative of the real observed phenomenon. this assumption quickly loses its relevance in the context of few-shot learning, where at most a few tens of observations are available for estimating the validation metric. this notion has become problematic enough in present times that a portion of academic literature on continuous learning with small datasets presents experiment results utilizing validation sets that are unrealistic and artificial, containing several orders of magnitude more observations than the training set used for the model adjustment itself @xcite . in the machine learning community, characterizing local minima with desirable generalization properties has been a topic of interest for decades @xcite @xcite . from flat minima @xcite to mode connectivity @xcite , this body of work has provided the basis for several practical observations regarding the connection between the properties of local minima and a model's generalization abilities. the concept of learning neural network subspaces @xcite is an example of a method built using these considerations. this approach proposes to find not just a local minimum of the cost function in the model's parameter space, but an entire simplex associated with low values of this objective. this additional constraint is meant to bias the descent algorithm towards wider minima, empirically associated with better generalization (dziugaite and roy, 2018a). in addition, the availability of this entire simplex of models allows for the inference of not only one scalar development metric, but an entire distribution, at any given moment during model fine-tuning. these two phenomena, become particularly relevant when viewed through the lens of large language models, and most especially for few-shot learning problems, where the model's ability to generalize a concept class from a limited number of examples is crucial. the contributions of this article are as follows. first, we introduce the first adaptation of the subspace method to large language models through subspace adjustment of prefixes (a peft method similar to the state-of-the-art continuous prompt adjustment in current academic literature). next, this article proposes to leverage certain natural advantages offered by the subspace method to revisit the concept of guiding model adjustment through the validation metric. we will empirically demonstrate that the combination of these two ideas leads to a significant improvement in terms of average prediction on natural language understanding tasks provided by the glue benchmark @xcite . finally, an ablation study will be presented to provide some insights into the mechanisms underlying this prediction improvement. ## background in this section, we review the two main concepts used in this article, neural network subspaces @xcite and prefix-tuning (li and liang, 2021). ## learning prefix subspaces from this definition, it appears that the two main issues that made applying the subspace method to large language models cumbersome are both alleviated when using prefix tuning. indeed, as a peft method, the number of trainable in the prefixes is orders of magnitude lower than the model's parameters, allowing to easily store simplexes in memory. moreover, these prefixes are basically embedding vectors that require random initialization before model fine-tuning, which allows us to easily sample an initial simplex, by randomly initializing all its vertices. ## conclusion in this article, we introduced two innovative ideas. the first one, an adaptation of the subspace method for training large language models through the peft method, is, to our knowledge, the first example of its use in the academic literature on natural language processing. the second idea, proposing an alternative way to estimate development metrics, represents an original application of the subspace method and is not specific to problems encountered in textual data analysis. the combined use of these two methods leads to a significant improvement in the performance of common language models such as bert on language comprehension tasks proposed by the glue benchmark, rephrased in a "few-shot learning" context. the ablation study presented at the end of the article also allows us to quantify the impact of these two contributions. the performance gains observed on very small datasets (≤ 100) seem to be mainly explained by the finer information extracted from the validation set through the stochastic metric estimation method. however, this gain appears to diminish for larger sample sizes, where the subspace method applied to prefix tuning seems to be sufficient on its own to achieve performance gains over standard peft methods as well as classical model training. finally, applying subspace learning to peft methods also enables the training of powerful predictive models while significantly reducing the computational resources typically required for training large language models. this approach preserves the fundamental efficiency goal of these methods. learning prefix subspaces remains accessible even in situations where resources, both in terms of data and computational power, are limited.