Knowledge is a Region in Weight Space for Fine-tuned Language Models (2023)

Abstract: Research on neural networks has focused on understanding a single model trained on a single dataset. However, relatively little is known about the relationships between different models, particularly those trained or tested on different datasets. We address this by studying how the weight space and the underlying loss landscape of different models are interconnected. Specifically, we demonstrate that finetuned models that were optimized for high performance, reside in well-defined regions in weight space, and vice versa – that any model that resides anywhere in those regions also exhibits high performance. Notably, we show that language models that have been finetuned on the same dataset form a tight cluster in the weight space, while models finetuned on different datasets from the same underlying task form a looser cluster. Moreover, traversing around the region between the models leads to new models that perform comparably or even better than models obtained via finetuning, even on tasks that the original models were not finetuned on. Our findings provide insight into the relationships between models, demonstrating that a model positioned between two similar models can acquire the knowledge of both. We leverage this and design a method for selecting a better model for efficient finetuning. Specifically, we show that starting from the center of the region is as effective, if not more, than using the pretrained model in 11 out of 12 datasets, resulting in an average accuracy improvement of 3.06.

## introduction models that share the same architecture but differ in their weights can have dramatically different capabilities. as an example, finetuned variants of a pretrained model all share an architecture, yet they are specialized for different tasks. this study figure 1 : a schematic view of the weight space. finetuning ends up in a region determined by the dataset (deep blue) which resides in the task (light blue) and language tasks regions (outer blue). any combination of finetuned weights is found within the region. each region is characterized by a low loss on the corresponding: dataset, task datasets, or diverse linguistic datasets. generally, loss is lower inside the region than outside or in its boundaries. we find that distance characterizes models' knowledge and similarity. particularly, after finetuning a pretrained model on similar datasets, the resulting models are close to each other in weight space ( §2.3). throughout the paper, we consider 3 granularities ( §3.1), showing that (i) models finetuned on the same data are closer to each other than to other models; (ii) models finetuned on the same task also cluster together; and (iii) models finetuned on general language tasks are not spread arbitrarily around the pretrained model, but fall in a constrained region in space. we find that different finetuning runs on the same data tend to converge on similar points in weight space rather than dispersed points. loosely, those points embed the necessary knowledge to perform the task. this leads to the hypothesis that other points in the proximity of finetuned models might also perform the task well. notably, such points in weight space might not necessarily be reached via finetuning, but rather via spatial transformations. indeed, we replicate the finding @xcite .f. §8) that models finetuned on the same dataset are linearly connected, i.e., points on the line between the two models attain similar or even lower loss ( §5.1). we expand this finding to the convex hull between the finetuned models ( §5.2), suggesting that knowledge is shared across the region in space. that is, finetuned models define a connected basin of low loss, and every point within it performs well. to show this, we test models sampled from the region and find they even outperform the models achieved by finetuning. moreover, we replicate the findings in all the aforementioned granularities: regions per dataset, task, and in general. for each, we observe a low loss across datasets, beyond the loss the individual models optimized. furthermore, we show in §6 that these regions are relatively tight, in the sense that extrapolating (rather than interpolating) can quickly produce a poorly performing model. our empirical findings have intriguing implications, suggesting, for example, that the best models may not lie at the edges of the region, but rather closer to its center, while finetuning often yields models at the edge of the region. motivated by these findings, we demonstrate in §7 that a model created by averaging the weights of finetuned models from the same region outperforms the pretrained model on a variety of tasks after subsequent finetuning, even on tasks that the original finetuned models were not trained on. overall, our work contributes to the growing body of knowledge about the loss landscape, finding connectivity in a whole bounded region rather than mere linear connectivity, finding connectivity between models not trained on the same task, and finding connectivity in generalization, evaluating models on multiple losses. we also provide initial context to empirical findings about fusing models. we discuss the relations to previous works in §8. ## experimental setup we conduct two main types of experiments. in one we train models with different characteristics (e.g., dataset or task, see §3.1) and examine their representation in weight space using clustering. in the second experiment type, we compare losses of one group of models to another. below, we describe the datasets ( §2.1), settings ( §2.2), and granularity levels of comparison between models ( §3.1). ## methodology: comparing models in this work, we compare models that share an architecture, but were trained on different data. to do so, we investigate the space of weights ω ∈ r d where each model has a weight vector and each point in space represents a model. we adopt the typical perspective that the model f θ consists of a representation encoder f ω followed by a taskspecific classifier f ϕ , i.e. f θ = f ϕ • f ω := f ϕ,ω @xcite . ideally, we would compare finetuned models by their loss. unfortunately, the loss is often incomparable across datasets or tasks. hence, we compare by preserving each encoder, and fitting a classification head to each model for each target dataset. specifically, to calculate the loss of a model we perform the following: first, we remove any existing masked language modeling layers or classification heads and replace them with a new randomly initialized classification head. this leaves the rest of the weights i.e., the encoder f ω , fixed. we then perform linear probing, i.e., we train only the new classification head on a desired target data x train and its labels y train . lastly, we pass the test data x test through the model (including the classifier f ϕ on top of it) and report the loss with respect to the labels y test . formally, for the model f ϕ,ω and loss function l, we report the generalized loss l g (ω) = l(f ϕ,ω (x test ), y test ) where f ϕ = arg min ϕ l(f ϕ,ω (x train ), y train ). this approach has a desirable trait: when considering the task on which the model was originally finetuned, our loss l g is equal to the original finetuning loss l. furthermore, since fitting a linear classifier given a fixed representation is a convex optimization problem, we observe similar results across runs. the generalized loss l g enables comparing models finetuned on different datasets. it is hence undesirable to test only on one of the datasets. we thus consider a loss on a dataset, but also the average loss on a family of datasets. for example, the average loss across all entailment datasets rather than the loss on a particular dataset. ## analysis in weight space we start our analysis by showing that the models trained on similar data fall into the same region in weight space -i.e., they are clustered together. we leave the inverse claim (i.e. showing that models within the cluster obtain a lower loss than the models outside the cluster) to §5.1 and §5.2. specifically, we find (see fig. 2 ) that finetuning on similar data results in closer weight space mod-els compared to models that have been trained on different datasets or tasks. notably, despite the fact that neural networks implement highly non-linear functions, finetuning similarity is expressed in the euclidean space of their weights. moreover, we show in app. §c that the direction in space is determined by the type of training data and not by its amount. in app. b, we show that this proximity is contingent on starting from the same base model. ## loss in the region between models in §4, we claim that models trained on similar data converge near each other, but is this area to which they converge meaningful? in this section, we show that models falling in the entire region around these clusters correspond to performant models. the models we analyzed so far were the outcome of a gradient-based optimization process searching for the minimum loss. the locality we observed in weight space indicates that the points found through this procedure are concentrated in relatively small regions. we hypothesize that a whole region of low losses (corresponding to performant models) exists between the separate points found during finetuning. for example, the "nli region" contains mnli, snli and qnli models but also other points that reflect models that might not have been found through gradient-based optimization on a specific dataset but exhibit the general abilities needed to perform natural language inference. we test this hypothesis by interpolating pairs of similarly trained models and show in § 5.1 that the points between the models perform comparably to or even better than the original finetuned models. this suggests that indeed there are regions in weight space where all points encode the knowledge or behaviour required for a particular task. we expand this claim in §5.2 and show that the whole region that lies between these models (their convex hull) corresponds to models that perform well. ## region edges above, we have shown that there are spacial regions that specify learnt generalizations. we now look for the boundaries of those regions, where loss is no longer similarly low. to do that we traverse in the opposite way to the interpolation. we also test the edges going from the center of the region to other directions in app. f. ## explaining previous results a great deal of prior work considered the connectivity between models, i.e. whether the path in weight space between two networks has a low loss throughout. early work demonstrated that models trained on the same dataset have such a path but that the path is not necessarily linear @xcite . this non-linearity was often explained by the fact that networks can represent the same function after their weights are permuted @xcite @xcite . taking into account these symmetries and/or using the same initialization was then shown to produce a linear path of low loss @xcite . @xcite even considered simplexes of low loss, rather than linear paths. in addition, @xcite showed that multitask learning converges to a point with low loss for both tasks, and in parallel work @xcite showed that the minima are connected for two datasets of the same task. we generalize those notions in the context of finetuned models. specifically, we confirm that indeed there is a linear path between two models, but further that there is a whole region with low loss through which the linear path moves. intriguingly, we have observed that these low-loss regions are unique for each specific dataset or task, whereas @xcite has reported the existence of multiple basins per each. we also generalize this finding to models that were not trained on the same data and are tested on different data. @xcite is the only work we know to compare models trained on different tasks. however, they report random chance performance in this case. to enable meaningful model comparison, we proposed the generalized loss ( §3). our results also support and provide some preliminary explanations of recent practical findings. some works show that starting from a finetuned model helps when finetuning on a different target dataset @xcite , which may be related to the fact that the initial finetuning stage moves the model into the general "language" region (or, even better, the region of space corresponding to the target task). moreover, a growing literature has shown improvements from averaging two or more finetuned models. some of those average models trained on the same dataset @xcite , which we show picks a model from inside the dataset region. others show that averages between models can improve models from tasks that they were not trained on @xcite , which agrees with our more general findings. @xcite further suggests that some attributes can be added to the model by moving in certain directions in the weight space. in parallel work, @xcite considers two finetuning stages before averaging. @xcite and @xcite propose optimization methods featuring stochastic weight averaging (swa). our results may indicate that the success of such methods may be partly attributed to its tendency to fall within a region, rather than on its borders. more recent work considers iterative model averaging, where in each iteration multiple models are trained in parallel from the same initial point and then averaged to aggregate their knowledge. such a procedure has been demonstrated both for self-supervised pretraining @xcite and as a supervised pretraining, similar to a massively multitask learning scenario @xcite . future work could focus on understanding how those processes move through the weight space and whether they move to areas of loss space outside of the region corresponding to a single iteration of averaging finetuned models. ## conclusion and discussion combining all of our results together conveys a consistent message: there are regions in weight space corresponding to good performance on a dataset, a task, or in general. from §2.3 we can conclude that performant models are centered in certain areas (or more specifically basins) in weight space. we find in §5.1 that these form one basin rather than multiple nearby points falling into multiple basins and, in §5.2, that this basin is a convex region and not simply a line between two points. finally, the extrapolations in §6 show those areas do not exceed far beyond the finetuned models. moreover, our results suggest that models found via finetuning typically lie on the boundaries of these regions and are often suboptimal, prompting future work in exploring the limitations of gradient-based training.