Mitigating Catastrophic Forgetting in Language Transfer via Model Merging (2024)

Abstract: As open-weight large language models (LLMs) achieve ever more impressive performance across a wide range of tasks in English, practitioners aim to adapt these models to different languages. However, such language adaptation is often accompanied by catastrophic forgetting of the base model’s capabilities, severely limiting the usefulness of the resulting model. We address this issue by proposing Branch-and-Merge (BaM), a new adaptation method based on iteratively merging multiple models, fine-tuned on a subset of the available training data. BaM is based on the insight that this yields lower magnitude but higher quality weight changes, reducing forgetting of the source domain while maintaining learning on the target domain. We demonstrate in an extensive empirical study on Bulgarian and German that BaM can significantly reduce forgetting while matching or even improving target domain performance compared to both standard continued pretraining and instruction finetuning across different model architectures.

## introduction large language models have shown remarkable capabilities, particularly in english. however, for less prevalent languages, performance can be significantly lower, making additional adaptation paramount @xcite . catastrophic forgetting unfortunately, most adaptation techniques come at the cost of catastrophic forgetting of the base model's capabilities @xcite @xcite . at the same time, retaining these capabilities is often crucial for solving downstream tasks in a new language. for example, math and coding skills learned in english can be extremely helpful for general problem-solving or reasoning tasks in other languages. ## model merging a wide range of model merging methods have been proposed @xcite @xcite @xcite . we experiment with linear (wortsman et al., 2022), slerp @xcite and model stock @xcite merging, focusing on the first two, explained below. let us consider the pretrained base model, f θ , parameterized by θ which was finetuned on two different datasets x 1 and x 2 , yielding f θ 1 and f θ 2 , respectively. we call the changes in weight due to this finetuning the task vectors τ @xmath0. to obtain a single model combining the learning from both datasets, we now merge these models. linear model merging interpolates task vectors or equivalently parameterizations linearly so to obtain θ slerp first represents task vectors in polar coordinates before interpolating to obtain the new parameterization ## branch-and-merge for mitigating forgetting in language transfer to adapt a model f θ pretrained on a typically unknown data distribution x pre to a new task (language) without suffering from catastrophic forgetting, we propose the branch-and-merge (bam) method, visualized in fig. 1 . bam is based on first splitting the available training data into n slices (blue in fig. 1 ), and then iteratively training k models in parallel on one slice each (green ) before merging the resulting models to obtain the base model for the next training iteration (purple ). we first provide the intuition behind bam before describing it in more detail. intuition there are two key ideas underlying bam. first, lower magnitude weight changes τ i , called task vectors, lead to less forgetting but also less learning. second, the randomness in finetuning leads to task vectors τ @xmath1. we can thus reduce forgetting by reducing the task vector magnitude while offsetting the reduced learning by increasing task vector quality, i.e., reducing the error ϵ. if this error is unbiased and empirically gaussian figure 2 : illustration of bam in the loss surface over parameter space. both θ 1 and θ 2 land in poor local minima but their merge θ ′ 1,2 lies in the valley of a better minimum. training from there, θ 3 and θ 4 land at the boundary of that minimum due to noise in the training process and limited data. their merge θ ′ 3,4 cancels these errors and lies in the better minimum. @xcite , merging, i.e., averaging, k noisy task vectors to obtain at the same time, increasing k in bam reduces the number of consecutive training iterations (as more data slices are used per iteration) and thus the expected total weight magnitude which in turn reduces both learning and forgetting. this allows bam to trade off learning and forgetting. we visualize this in fig. 2 for @xmath2. there, the first two task vectors τ 1 and τ 2 land in the basins of poor local minima, with their merge θ ′ 1,2 falling into the basin of a better minimum, highlighting the importance of bam's iterative merging approach. training θ ′ 1,2 on two more data slices yields noisy task vectors τ 3 and τ 4 at the edge of this loss basin with their merge θ ′ 3,4 falling right in the middle. in contrast, simply reducing the learning rate can also reduce task vector magnitude but does not improve task vector quality. we note that the same intuitions apply to slerp and model stock merging. implementation in more detail, we first partition the training data x train into n not necessarily i.i.d. or equal-sized data slices x i . then, we choose a parallelism factor k (@xmath3. 1 and 2 ) and train our current base model f θ independently on k of these data slices yielding f θ i to f θ i+k-1 . we merge the resulting models to obtain the base model for the next iteration f θ ′ = merge({f θ j } i+k-1 @xmath4, c). we typically choose the merging coefficient @xmath5.5 but note that we can easily perform a 1-d line search over the resulting models. we then set the merged model f θ ′ to be the base model f θ ← f θ ′ for the algorithm 1 branch-and-merge (bam) require: k: parallelism factor, f θ : base model, {x i } n @xmath6: data slices, c: merging coefficient 1: θ ← {} 2: for @xmath7: θ ← {} 8: return f θ : finetuned model next training iteration and repeat this process until we have used all data slices. we formalize this approach in algorithm 1. ## data mixtures for mitigating forgetting in language transfer here we describe the data we use for continued pretraining of predominantly english base language models in order to adapt them to other languages. outside of training methodology, we find in agreement with prior work that high-quality dataset mixtures are paramount for both effective language adaptation and reducing forgetting. we distinguish between experience replay of source language data and target language training data. ## experimental setup we now describe the experimental setup used to validate bam's effectiveness for language adaptation. in particular, we discuss the target languages (bulgarian and german), evaluation benchmarks (section 5.1), training data (section 5.2), and training setup (section 5.3). we experiment with both continued pretraining of base models and instruction tuning of the resulting models. ## experimental evaluation we now evaluate bam empirically for both continued pretraining (cpt) and instruction finetuning (ift) before conducting extensive ablations and providing further results in appendix a. ## conclusion we proposed branch-and-merge (bam) training to mitigate forgetting while boosting learning in language transfer by generating lower magnitude but higher quality weight changes. we showed that combining bam with an effective approximate experience replay data mix significantly reduces forgetting. finally, we demonstrated that our approach can benefit both continuous pertaining and instruction tuning in both alphabet-sharing (german) and non-sharing (bulgarian) languages. for instance, we outperform llama-3-8b-instruct with the same base model in both source (english, 1.3%) and target (bulgarian, 10.8%) languages.