Lottery Ticket Adaptation: Mitigating Destructive Interference in LLMs

Venue: ICML_cc_2024_Workshop_WANT
Published: 2024-05-30 03:56:13

Abstract: Existing methods for adapting large language models (LLMs) to new tasks are not suited to multi-task adaptation because they modify all the model weights--causing destructive interference between tasks. The resulting effects, such as catastrophic forgetting of earlier tasks, make it challenging to obtain good performance on multiple tasks at the same time. To mitigate this, we propose Lottery Ticket Adaptation (LoTA), a sparse adaptation method that identifies and optimizes only a sparse subnetwork of the model. We evaluate LoTA on a wide range of challenging tasks such as instruction following, reasoning, math, and summarization. LoTA obtains better performance than full fine-tuning and low-rank adaptation (LoRA), and maintains good performance even after training on other tasks -- thus, avoiding catastrophic forgetting. By extracting and fine-tuning over \emph{lottery tickets} (or \emph{sparse task vectors}), LoTA also enables model merging over highly dissimilar tasks.

Main paper: 
 introduction: 
 Large language models (LLMs) (Brown et al., 2020) have seen an explosion of applications to real-world problems (OpenAI, 2023; Team et al., 2023) via adaptation (Ouyang et al., 2022) to new tasks. Three major multi-task adaptation paradigms have emerged: storing and loading task-specific adapters (Hu et al., 2022; Beck et al., 2021) , continuing to train instruction-tuned models on new tasks in serial via sequential training (Ouyang et al., 2022) , and combining the adaptations to tasks learned in parallel via model merging (Ilharco et al., 2022) . Each paradigm has its own associated challenges, such as catastrophic forgetting during sequential training (McCloskey & Cohen, 1989; Dong et al., 2023; Ramasesh et al., 2022; Luo et al., 2023; Wang et al., 2024) , and methods that have been proposed to mitigate these challenges (Crawshaw, 2020; Zhang & Yang, 2021) . In this work, we propose a new LLM adaptation method, called Lottery Ticket Adaptation (LoTA), that (1) provides sparse adaptation by freezing a majority of the parameters and updating only a sparse subnetwork of the base model and (2) resolves the challenges in common multi-task adaptation paradigms. (More details in Section 3.) We summarize our contributions:• We train lottery tickets (or sparse task vectors) that can be stored efficiently and obtain performance similar to full fine-tuning (FFT) and higher than LoRA across a range of tasks spanning reasoning, math, code generation, and instruction following. When adapting Mistral for instruction following, FFT and LoTA both get a length-controlled AlpacaEval 2 winrate (Dubois et al., 2024) (how often GPT-4 prefers the outputs of our model over its own) of 19.0%, but LoRA only gets a winrate of 15.3% after extensive hyperparameter tuning. • We apply LoTA to mitigate catastrophic forgetting (McCloskey & Cohen, 1989) of earlier tasks, enabling sequential adaptation to new tasks. If we train models with FFT or LoRA on first GSM8k and then Commonsense Reasoning tasks, they forget GSM8k completely (< 5%) accuracy, but by using LoTA we maintain 57% accuracy on GSM8k and get 85% on Reasoning. • We can use LoTA to merge models in parallel (Wortsman et al., 2022; Jin et al., 2022; Zhang et al., 2023a) across dramatically different tasks. LoTA achieves better performance than existing merging methods that rely on post hoc sparsification (Yadav et al., 2023) (which degrades performance for FFT models), because LoTA naturally trains sparse task vectors. When we use LoTA to merge multiple models trained on heterogeneous tasks of instruction following, math and summarization, we get a task-average performance of 21.9% where a merge of FFT models obtains 19.1%. Storing and Loading Adapters. The first row of Figure 1 shows training task-specific adapters that are loaded during inference (Ostapenko et al., 2024; Beck et al., 2021; Mangrulkar et al., 2022) , depending on the need (Houlsby et al., 2019) . Notably, Apple stores and loads adapters to enable their ondevice models (Apple, 2024) . While using adapters avoids any interference between tasks, it increases memory and compute costs as it requires storing and loading task-specific adapters. To mitigate these costs, a number of parameter-efficient fine-tuning (PEFT) methods have been developed, most popularly LoRA (Hu et al., 2022) (training adapters in the low-rank space). Services such as Punica and S-LoRA allow developers to use this approach to serve large numbers of task-specific adapters for specific requests (Chen et al., 2023; Sheng et al., 2023) . However, a persistent gap in capacity between PEFT methods and FFT has presented a tradeoff between adapter overhead and performance (Hu et al., 2022; Liu et al., 2024b; Kopiczko et al., 2023; Biderman et al., 2024; Nikdan et al., 2024) .Sequential Training. When we want one model with multi-task abilities, we can fine-tune the model on different tasks sequentially (Ruder, 2017) , e.g., first fine-tune on task A, then fine-tune on task B, as in the second row of Figure 1 . We can also train the model on a mixture of all tasks, which may require optimizing the data mixture. Fine-tuning the LLM for new tasks with FFT or existing PEFT methods leads to catastrophic forgetting of earlier tasks. This is problematic, especially for safety alignment, since we can fine-tune an LLM to be safe but later get this feature erased during future fine-tuning on new tasks (Lermen et al., 2023) . In fact, a number of works have aimed to mitigate this vulnerability (McCloskey & Cohen, 1989; Dong et al., 2023; Ramasesh et al., 2022; Luo et al., 2023; Wang et al., 2024) , leaving an open research question: Can model developers allow users to finetune their aligned models on custom datasets while retaining safety?Model Merging. The third row of Figure 1 shows how we can merge multiple models adapted to different individual tasks to have a single model adapted to multiple tasks, by aggregating task vectors (or adapters) of different tasks. Existing model merging techniques either require post-processing the task vectors through sparsification (Yu et al., 2023; Davari & Belilovsky, 2023; Yadav et al., 2023) , degrading the performance on the task, and/or require extensive hyperparameter tuning for a weighted aggregation of task vectors (Matena & Raffel, 2022; Xiao et al., 2023b) .Each paradigm of multi-task learning poses different challenges, and different methods have been proposed to address these challenges. We defer the more in-depth analysis of these proposed methods in Appendix B because of the sheer quantity of related work that must be covered. The number of methods itself poses challenges for studying their drawbacks, especially when these methods are adopted in settings orthogonal to those they were originally developed for (i.e., LoRA leading to catastrophic forgetting of safety alignment (Lermen et al., 2023) ). Rather than proposing tailored solutions for each paradigm of multi-task adaptation, we want to propose a simple algorithm that can serve as an effective foundation across all three paradigms and evaluate it on a wide range of challenging tasks. The desiderata for each multi-task adaptation paradigm motivates the design of our method. For adapters, we want a representation that can be easily compressed for memory efficiency. For sequential training, we want a representation that minimizes destructive interference between the previously learned tasks and tasks to be learned in the future. For model merging, we want representations that are mutually sparse with each other in parameter space to again prevent destructive interference. We now propose LoTA, a single method that enjoys all these features. We first describe the workflow of LoTA, then revisit the problems each multi-task paradigm faces and discuss how and why LoTA successfully mitigates them. 
lottery ticket adaptation (lota): 
 Lottery Ticket Adaptation (LoTA). LoTA works in three phases as summarized in Figure 2: (1) mask calibration, (2) mask extraction, (3) sparse adaptation. In the mask calibration phase of LoTA, a base model with parameters w P is fine-tuned for T iterations, yielding a fine-tuned model with parameters w F . Then, in the mask extraction phase, LoTA extracts a sparsity mask m from the task vector ∆ = w F -w P based on the magnitude of the updates in ∆. T could be as small as one iteration. In the sparse adaptation phase of LoTA, the model is first reset to its original state with weights w P . Then the subnetwork w P ⊙ m is fine-tuned for T iterations, while leaving the remaining parameters w P ⊙ (1 -m) frozen at their initial values. We summarize the workflow of LoTA in Figure 2 and Algorithm 1 further.By confining the adaptation updates within subnetworks (identified by m), LoTA is able to mitigate destructive interference, e.g., adaptation loss during fine-tuning on future datasets or model merging, that FFT and LoRA suffer from. We discuss this in more detail under three multi-task adaptation paradigms below and provide empirical comparisons with FFT and LoRA in Section 5. Under review as a conference paper at ICLR 2025(1) Storing & Loading Compressible Adapters Without Compromising Performance. LoRA restricts the adaptation updates to have a low rank, which may be insufficient for complex downstream tasks (Nikdan et al., 2024; Biderman et al., 2024) . In parallel, recent work (Isik et al., 2023) on compressing the delta between the fine-tuned and pre-trained model ∆ = w F -w P suggests that FFT updates are highly compressible through a simple magnitude-based sparsification. LoTA exploits this underlying sparsity during fine-tuning and obtains better performance than LoRA, while requiring fewer parameters to be trained. As we will show, LoTA adapters can be stored efficiently because of the small number of parameters that need to be updated. LoTA also requires fewer parameters to be updated in the model when the adapter is applied to the model, so loading a LoTA adapter can be faster than loading a LoRA adapter.Algorithm Mitigating Destructive Interference via Lottery Ticket Together Optimization (LoTTO). Suppose that we have already learned Task A with LoTA. LoTTO calibrates a sparsity mask for Task B by first training a model where the only weights that can be updated are those that are not updated when running LoTA on Task A, and then using a sparse set of those weights to train the final model. In doing so, LoTTO ensures that the weights updated for Task A are disjoint from the weights updated for Task B. LoTTO enables learning new tasks without forgetting old ones, by restricting weight adaptations to not interfere with the weights that were important for previous adaptations, therefore mitigating destructive interference. This procedure can be applied inductively to enable sequential adaptation to multiple tasks, so that a model developer seeking to adapt a model adapted with LoTA (potentially on several tasks), just needs to ensure that they do not update the task vector with respect to the base model. LoTTO can be applied to capabilities one at a time; first the model learns math, then coding, then reasoning, then planning, etc. because each of these capabilities benefits from prior knowledge. We train on up to 4 tasks sequentially, but optimistically if we believe we can update 1% of the weights for each capability, then LoTTO can fit at most 100 capabilities. When the subnetworks for different tasks overlap, we do not expect that LoTTO will result in any worse forgetting than LoTA, which we will show itself has better robustness to forgetting than FFT. We will also show that LoTTO on sequential tasks outperforms the common practice of training on a mixture of tasks.(3) Model Merging Across Heterogeneous Tasks. Existing model merging methods have been shown to merge models trained on relatively similar language datasets (Wortsman et al., 2022; Jin et al., 2022; Zhang et al., 2023a) , typically via a post hoc sparsification (Yadav et al., 2023; Yu et al., 2023; Davari & Belilovsky, 2023) of the task vectors. This ensures the task vectors are disjoint -hence limiting destructive interference -but also degrades the performance of each individual task. LoTA enforces sparsity during fine-tuning and directly trains sparse task vectors, obviating the need for post hoc sparsification. LoTA's advantage over methods such as TIES (Yadav et al., 2023) is similar to the advantage of training a model with an intelligently chosen sparsity constraint throughout training (as in the original lottery ticket hypothesis (Frankle & Carbin, 2019) ) over pruning parameters after training.In this section, we provide details of the experimental setup, including the baselines, model, dataset, and metric selection. We present the results later in Section 5. We share our anonymized code. We are limited to an academic computing budget, and all results are conducted with a single A100 GPU. We typically do 1 -3 epochs of training for each dataset as this is a standard choice in LLM fine-tuning. We use the RMSProp optimizer with default hyperparameters.Baselines & Hyperparameters. Across all three multi-task adaptation paradigms, we compare LoTA against FFT and LoRA. We extensively tune the hyperparameters for FFT and LoRA to ensure that we are comparing against strong baselines. We do not tune hyperparameters for LoTA and directly transfer the hyperparameters from FFT. We fix the sparsity ratio hyperparameter in LoTA to 90%.Models. We use two commonly-used open-weights model families, Mistral Jiang et al. ( 2023) and Llama 3 (AI@Meta, 2024; Touvron et al., 2023) , specifically Mistral-7B and Llama-3-8B -the largest models we can adapt with FFT on a single GPU.Tasks. We consider six capabilities: instruction following, safety, math, coding, summarization, and reasoning. We now briefly discuss each capability, the datasets we use to fine-tune and evaluate the presented methods, and the motivation behind the choices.Instruction Following. The most widely-used instruction-tuned LLMs are the "Instruct" or "chat" versions of base models, such as Llama-3-8B-Instruct (AI@Meta, 2024). This is because the process of tuning models on human instructions aligns models to human preferences across a range of tasks (Ouyang et al., 2022) . For this, we adapt models to data from UltraFeedback (Cui et al., 2023) , which contains a mixture of datasets covering truthfulness, honesty, and helpfulness in addition to instruction-following. We measure the instruction following ability by length-controlled AlpacaEval2 Win Rate (Li et al., 2023) , which we refer to as "winrate". A high winrate means that GPT-4 (OpenAI, 2023) prefers the responses of our model on a set of representative prompts over its own responses. Winrate is the metric most closely correlated with human rating preference (Dubois et al., 2024) . Another common benchmark for "chat" models is MT-Bench (Zheng et al., 2023) , but there is a significant degree of data contamination between MT-Bench and other task-specific training datasets (Yu et al., 2024 )-hence, we do not evaluate on MT-Bench.Reasoning. We train on the standard set of 8 commonsense reasoning tasks (Christopher et al., 2019; Bisk et al., 2019; Sap et al., 2019; Zellers et al., 2019; ai2, 2019; Clark et al., 2018; Mihaylov et al., 2018) (Boolq, PIQA, SocialIQA, Hellaswag, Winograde, ARC-easy, ARC-challenge, OpenBookQA) and report the exact-match accuracy on the test set. As a representative task, we use ARC-easy.Math. We use the set of 9 math instruction datasets from (Yue et al., 2023) for fine-tuning and report performance on the test set of GSM8k (Cobbe et al., 2021) . When only considering a single task, we choose GSM8k as the representative task as it is commonly used as a single training and test task by other papers.Code Generation. We use data that instructs the model to write SQL queries given some context (b mc2, 2023)(SQL-create-context) and report the ROUGE-1 F1 score (Lin, 2004) on the test set.Summarization. We use Samsum (Gliwa et al., 2019) , reporting ROUGE-1 F1 score on the test set.Safety. A recent concern in AI policy is that, while frontier models such as GPT-3.5/GPT-4 are aligned, they can also be fine-tuned and this presents an opportunity to misalign them. Recently, Qi et al. (2023) show that by fine-tuning GPT-3.5 on just 100 harmful examples for a few epochs, they can ask it to answer harmful queries that it ordinarily would refuse, and Zhan et al. (2024) show the same for GPT-4. Lermen et al. (2023) show that this can be done with LoRA rather than the finetuning method OpenAI are using in their fine-tuning API (presumably FFT). We evaluate the safety of our models on HEx-Phi (Qi et al., 2023) , a dataset of 330 questions spanning multiple categories such as malware, fraud, etc. The safety score is the percentage of harmful queries that the model refuses to respond. Aligned models (Llama-3-Instruct) will refuse 100% of harmful queries, but because we are aligning base models ourselves, our baseline Instruct model only gets 93%. Because we are interested in measuring the forgetting of safety alignment, we do not see this as a major limitation.We vary hyperparameters (Section 5), consider single-task adapters (Section 5.1), sequential training (Section 5.2), and model merging (Section 5.3). 
experimental results: 
 We first evaluate the impact of varying the sparsity and mask calibration.Sparsity. We vary the sparsity parameter in LoTA in Table 1 . Multiple sparsity thresholds work well; some values with less sparsity (30%) seem to negatively impact performance and extreme sparsity seems to improve performance 99% but both deviations are within 2% of FFT. Calibrating the LoTA Mask Efficiently. In Table 2 , we consider how the performance drops if we calibrate a 90% sparse mask on a fraction of the overall GSM8k dataset, including the baseline where we use a random mask which corresponds to 0% data used. We find that a 90%-sparse mask can achieve good performance on GSM8k with just a single step of calibration (calibration dataset size= 0.1% of the full dataset). Furthermore, instruction tuning datasets are generally quite small (finetuning on GSM8k takes less than an hour on a single GPU). Given that the mask can be efficiently calibrated, transferred from other datasets (as we will show), or in the worst case, calibrated on the entire dataset in at most a few hours, we do not anticipate that the compute overhead of LoTA will be a major limitation of the method.Next, we compare LoTA with baselines, FFT and LoRA, across all three paradigms of multi-task adaptation. We run all the methods for the same number of steps, i.e., the calibration time for LoTA is factored in. Unless stated otherwise, LoTA has 90% sparsity where the mask is calibrated by training for 1% of the adaptation dataset. We choose this sparsity level rather than the 99% sparsity level that sometimes outperforms FFT because we can compare the improvements of LoTA that are due to the method's mitigation of destructive interference, rather than any inherent advantage sparsity may have.We first fit an adapter to each dataset starting from a pre-trained base model, i.e., one adapter per task. In Table 3 , we find that LoTA outperforms LoRA and performs similarly to FFT. Although LoRA is able to achieve similar performance to LoTA on the easier tasks, such as SQL, Samsum, there is a clear gap in performance on the more challenging tasks, such as Instruction Following and GSM8k. LoTA consistently recovers the performance of FFT. LoRA Does Not Significantly Reduce Catastrophic Forgetting of Capabilities Learned During Finetuning. First, note that LoRA loses nearly all performance on Task A when training on Task B. This observation may seem contrary to the results in (Biderman et al., 2024) . However, they measure forgetting as the degradation in performance on standard benchmark tasks between the pretrained model and the finetuned model, whereas we measure forgetting as the degradation on Task A between the model finetuned on Task A and the model finetuned on first Task A and then Task B. We speculate that the discrepancy is that knowledge learned through the pretraining process may be far more persistent than knowledge learned through finetuning. Indeed, LoRA does forget less than FFT -just not enough to make a difference when training with LoRA on both tasks.LoRA Can Also Mitigate Destructive Interference, But Not As Well as LoTTO. A key insight into the design of LoTTO is that we can mitigate destructive interference between different tasks by forcing the tasks to update disjoint sets of parameters. LoTTO carefully calibrates the disjoint subsets by pruning the weights that were updated the most after a small amount of finetuning on the dataset. We can also extend this insight to try and mitigate destructive interference in FFT and LoRA. We do this by proposing an "alternating freeze" strategy, where the odd layers can only be trained on Task A and the even layers can only be trained on Task B. LoRA already struggles to fit complex tasks such as math (Task A is GSM8k) and the alternating freeze further restricts LoRA's representative capacity. Therefore, while the alternating freeze does forget less, and on easy tasks such as commonsense reasoning, there is no significant gap between LoRA and LoTA, the gap on the challenging Task A is hard to overcome because it is partially due to the inherent gap in performance between LoRA and LoTA and partially due to LoTA's limited forgetting as compared to LoRA. We note that LoTA can also be combined with LoRA, by fixing the A matrix in LoRA as a random projection (Zhang et al., 2023b) and tuning the B matrix with LoTA; this could be an interesting piece of future work.Sequential Training Under LoTTO Outperforms Mixing Datasets Together. We compare our LoTTO method, that trains disjoint LoTA adapters sequentially on each dataset, to FFT on a mixture of Instruction Following and the dataset noted in Task B. LoTTO under sequential training consistently outperforms the common practice of mixing all datasets together and doing FFT. We note that LoTA can also be used in the setting of mixing all datasets together, by routing each sample of a given dataset to a GPU -specific to that dataset-that applies the corresponding LoTA mask for that dataset.Under review as a conference paper at ICLR 2025 Mitigating Catastrophic Forgetting While Enabling Adaptation to Downstream Datasets.In Table 6 , we consider a range of method combinations for a simplified setting where we seek to adapt an Instruct model to Math data without catastrophic forgetting. We will go row-by-row through the table and analyze each set of results. Even when training on just a single, relatively small dataset, FFT in both phases suffers a significant drop in winrate. An easy way to mitigate this is to simply train the initial Instruct model with LoTA. Following this, FFT on GSM8k does not significantly reduce winrate. However, it does present a potentially unwelcome tradeoff in task accuracy. For this, we turn to LoTTO, which achieves the best performance across both tasks.Table 6 : Sequential learning first on Task A (Instruction Following) then on Task B (varied). Finetuning on Tasks A and B is performed using both FFT and LoTA. The utility of each task is computed after fine-tuning on Task B is completed. Note that there is no utility when we fine-tune on harmful data to evaluate the catastrophic forgetting of Safety, and the baseline is the safety score of the Instruct model. FT=Fine-tuning. We reproduce the baseline of doing FFT on each method independently as reported in Table 3 for convenience; note that on the reasoning task LoTA outperforms FFT. bold: best method, underline second-best method; we do not report second-best when only two methods are presented. All results are with Mistral. We reuse the same mask for LoTTO calibrated on GSM8k for MathInstruct, Reasoning, GSM8k+Arc+SQL and Safety. Under review as a conference paper at ICLR 2025 How Much Does Data Reuse Help? The simplest and arguably most performant method from prior work that we found for mitigating catastrophic forgetting is simply to replay in some in-distribution data from Task A. This is the line marked with "FFT (Replay)", and it does mitigate forgetting on Task A, but at the cost of performance on Task B. We ablate the amount of data from Task A to be used between 1 -100% of the dataset size of the data from Task B, and this is the best result in terms of mitigating forgetting on Task A. As we mix in less and less data from Task A, this method approaches just doing FFT sequentially in performance, so we omit those results for brevity.Adapting Sequentially to Multiple Tasks with LoTTO. We now consider the more challenging setting when we need to adapt to multiple tasks without catastrophic forgetting while still obtaining good performance on those tasks. In Table 6 , we adapt an Instruct model to a mix of reasoning, math, and SQL data and find a surprising result; the FFT Instruct model collapses almost completely to a winrate of less than half a percent. As we showed in Table 6 , this can be mitigated by mixing in more instruction following data, albeit at a cost. The LoTA Instruct model still degrades in performance (19 → 11) but nowhere near as much. In the line marked "LoTTO", we instead adapt the LoTA Instruct model to the downstream tasks using the mask calibrated from GSM8k with LoTTO. Applying LoTTO to make the downstream adaptation be mutually sparse with the initial LoTA Instruct model increases performance significantly on instruction following and math. Performance on Arc and SQL suffers somewhat because our mask does not consider those tasks, but this means that our mask is much cheaper to calibrate and is, in fact, generalizable not only across tasks within a domain (as shown in Table 6 ) but also across domains.Mitigating Catastrophic Forgetting of Safety Alignment. In the "Safety" row of Table 6 , we consider fine-tuning the Mistral Instruct model we trained ourselves on the 100 harmful instructions from (Qi et al., 2023) . The baseline model gets a score of 93% and training with FFT quickly degrades safety. Adapting the LoTA Instruct with LoTTO (again, with the LoTTO mask calibrated from GSM8k) mitigates this safety drop significantly, even though our LoTTO mask was calibrated on an extremely different dataset. Therefore, a potential mitigation would be for an entity providing a fine-tuning API such as OpenAI to do the safety training with LoTA, calibrate the LoTTO mask on a utility dataset, and then do fine-tuning on their client's dataset with LoTTO. We do not intend to present our method as an active defense against fine-tuning attacks; although finetuning with LoTTO improves safety, given sufficient data and access to the model weights, any attacker can of course undo safety tuning entirely. However, catastrophic forgetting of safety alignment is an important problem with real-world applications, and we find it compelling that our method can mitigate this.We now consider the setting of model merging, where we train models on disjoint datasets fully in parallel and then merge together the task vectors with the goal of producing a model with good performance on multiple tasks. Prior work in model merging mostly considers merging similar datasets, such as the commonsense reasoning datasets, but it is relatively easy to merge models when the datasets are similar and becomes increasingly hard as the datasets become more heterogeneous due to the gradient mismatch (Daheim et al., 2024) . We consider merging models trained on heterogeneous datasets.Merging Models. We use TIES-Merging (Yadav et al., 2023) as a baseline to merge models trained on heterogeneous tasks, because it outperformed other merging methods in our testing. TIES performs post-hoc sparsification on each task vector and requires a 2-D hyperparameter search for this quantity, which we perform for the merge of FFT models. Naturally, we could optimize the performance of Task A by fully sparsifying Task B, and vice versa; we report the result that achieves good performance on Task B while maintaining some performance on Task A, and report the full range of hyperparameters in Appendix A. LoTA is inherently sparse, so when we merge a LoTA model with an FFT model we do not need to perform hyperparameter search on the LoTA model, and we use the same level of sparsity for the FFT model that we obtained when merging together two FFT models. When merging two LoTA models together, no hyperparameter search is required at all as both models are inherently sparse. We could in theory sparsify the LoTA models beyond their existing levels with post-hoc sparsification, but we do not tune this hyperparameter.Challenge of Overlapping Sparsity in Model Merging. In the sequential training paradigm, we exploited the fact that masks for different tasks have a significant overlap in order to generalize our Under review as a conference paper at ICLR 2025 LoTTO mask calibrated on GSM8k to provide robustness to forgetting across a range of other tasks. However, this same phenomenon of overlapping sparsity presents a challenge in the model merging setting. The challenge is that because the merging is parallel, we cannot use LoTTO to calibrate the masks to be disjointly sparse as we did in the sequential training setting. One alternative is to enforce random disjoint masks for each task, corresponding to "0% Data Used" in Table 2 .Table 7 : Model merging using TIES of Task A, Instruction Following, and Task B, (varied; GSM8k+Samsum means we are merging 3 models). Fine-tuning on Tasks A and B is performed using both full fine-tuning (FFT) and LoTA. The utility of each task is computed after merging the two task vectors. bold: best result, underline: second best result. We reproduce the baseline for each task on FFT from Table 3 for convenience; note again that for some tasks the LoTA baseline outperforms the FFT baseline. All results are with Mistral. In Table 7 we merge together models trained on GSM8k and Samsum with the Mistral model that we trained for instruction following. We consider the full combination of merging FFT and LoTA models. The merge of two FFT models performs poorly in all settings on both tasks, indicating that post-hoc sparsification does not perform well for heterogeneous tasks, which is in line with recent model merging theory (Daheim et al., 2024) . Merging LoRA models together does not particularly improve performance. The merges that contain LoTA models have better performance across all tasks, but there is no combination that is pareto-optimal across all tasks. One avenue to improve LoTA for model merging may be to assume that each participant has a small sample from other participants' datasets to perform a joint LoTA calibration procedure. 
discussion: 
 We propose Lottery Ticket Adaptation (LoTA), a sparse alignment framework that fine-tunes only a sparse subnetwork of the base model, leaving the rest of the parameters frozen. LoTA successfully mitigates destructive interference (a problem with existing fine-tuning methods including FFT and LoRA) in many multi-task adaptation paradigms, prevents catastrophic forgetting of earlier tasks, including safety, and allows for successful model merging of even dramatically different tasks.Limitations. As mentioned in Section 5, LoTA does not provide the compute efficiency of LoRA. If the adapter needs to be compressed by more than 100×, LoTA may not provide sufficient compression. We evaluate on instruction following, reasoning, math, SQL generation, and summarization, yet even more tasks exist such as Python code generation, classification, or long-context question answering. We compare LoTA to baselines (LoRA, TIES) but other PEFT and merging methods exist. In a future revision of this paper, we plan to provide comparisons to a broader range of PEFT (Liu et al., 2024b) and merging (Yu et al., 2023) methods. All LoTA masks are unstructured, so we do not reap the runtime benefits of structured sparsity. We only explore magnitude pruning for LoTA, but many other methods have been proposed in the literature, as we discuss in Appendix B, and may perform better.As requested by ICLR, we are providing a brief statement on reproducibility. We provide an anonymous link to our codebase in Appendix A, where all the results in the paper can be reproduced. We made efforts to create strong baselines for ourselves by tuning hyperparameters, whose ranges we report in Appendix A, and by creating what we believe is a new method in the form of the "alternating freeze" baseline. The script "continual learning", for instance, reproduces Table 4 end to end on publicly available datasets. In the main text, we provided details on the datasets we used, all of which are publicly available and in fact are provided in our repository along with the prompts and code we used to train and evaluate on that data, and our reasoning for choosing those datasets. We fine-tuned models that are publicly available on Huggingface, but because we have academic compute restrictions we had to create code that fine-tunes Mistral-7B and Llama-3-8B on a single GPU, which we provided in the repository so that other researchers can more easily reproduce our results (typically, instruction tuning these models requires 8 GPUs, as in the popular Alignment Handbook repository's scripts). We plan to open source our code and upstream our LoTA method to Huggingface to further aid in reproducibility.Lottery Ticket Hypothesis. Motivated by the success of pruning methods at extreme sparsity ratios (Han et al., 2015) , Frankle & Carbin (2019) proposed the lottery ticket hypothesis (LTH), claiming the existence of sparse subnetworks (or lottery tickets) that could be trained from scratch to a performance comparably to training the dense model from scratch. While this could potentially provide a way to train models sparsely more efficiently rather than training them densely and pruning them later, finding the lottery tickets, i.e., the sparsity masks, is costly. Initially, Frankle & Carbin (2019) proposed first training the models densely and then extracting the sparsity mask based on the magnitude of the trained dense model's weights. Later, a number of more efficient methods were proposed to find the sparsity masks more efficiently, earlier in the dense training stage (Frankle et al., 2021; Lee et al., 2019; Wang et al., 2020; Tanaka et al., 2020) . Our work shows that LTH works successfully for fine-tuning LLMs as well-giving us a sparse adaptation tool, LoTA. We extensively study the tradeoff between the cost of finding the sparsity masks and the performance of the sparsely fine-tuned model. Unlike other studies on LTH for LLM adaptation (Yuan et al., 2024; Xu & Zhang, 2024) , our main focus and motivation is to mitigate destructive interference in multi-task adaptation.Parameter-Efficient Fine-Tuning (PEFT). Many practitioners fine-tune already pre-trained LLMs with less data and compute instead of training them from scratch (Liu et al., 2022a; Wang et al., 2022; Ouyang et al., 2022) . While this reduces the cost of LLM training significantly, fine-tuning each and every parameter of these large models for each (or a few) task is still very costly. This has led to a number of parameter-efficient fine-tuning (PEFT) methods reducing the number of trainable parameters during fine-tuning (Zhang et al., 2023c; Li & Liang, 2021; Liu et al., 2022b; Lester et al., 2021; Edalati et al., 2022; Hu et al., 2023; Nikdan et al., 2024; Guo et al., 2020) . Among different PEFT methods, low-rank adaptation (LoRA) (Hu et al., 2022) and its variants (Dettmers et al., 2023a; Guo et al., 2024; Li et al., 2024; Kopiczko et al., 2024) have shown similar performance to full fine-tuning in many tasks while reducing the number of trainable parameters through low-rank approximation to model updates during fine-tuning. Our PEFT method, LoTA, while reducing the number of trainable parameters significantly via sparsity, has various other benefits in different applications, such as avoiding catastrophic forgetting (of especially safety alignment), enabling fine-tuning on new tasks more successfully, model merging using sparse task vectors, unlearning, and communication-efficient federated learning (FL). We demonstrate that full fine-tuning and the existing PEFT methods fall short in these applications and significantly underperform LoTA.Catastrophic Forgetting. When LLMs go through sequential (or continual) multitask learning, i.e., fine-tuned on different tasks sequentially, they often suffer from performance loss on earlier tasksknown as catastrophic forgetting (McCloskey & Cohen, 1989; Dong et al., 2023; Ramasesh et al., 2022; Luo et al., 2023; Wang et al., 2024) . To mitigate this, a number of data-centric and architectural solutions have been proposed for language and other domains. Replay-based methods (Rebuffi et al., 2017; Romanov et al., 2018) add a portion of the previously learned data during fine-tuning on a new task, which raises privacy concerns as it requires constant access to previously learned data. Regularization-based approaches (Huang et al., 2021; Aljundi et al., 2018) tend to have poor adaptability to specific tasks. An architecture-based approach, "progressive prompts" Razdaibiedina et al. (2023) , sequentially concatenates soft prompts as they are being learned for each task-showing some resistance against forgetting. However, they require access to task identifiers at inference for each task, which is not always feasible. Other architecture-based approaches add additional modules to learn task-specific abilities (Dou et al., 2023; Wu et al., 2024) -requiring customized deployment due to architecture change. Closest to our work, Hui et al. (2024) updates a randomly selected subset of the parameters at each iteration of fine-tuning to preserve the earlier tasks in the not-updated parameters of that iteration. Despite similarities, our work LoTA (1) uses a fixed sparsity mask throughout fine-tuning instead of a new mask at every iteration, which yields sparse task vectors that are useful for other applications such as model merging and communication-efficient FL, and (2) finds data-dependent masks rather than the randomly selected masks in (Hui et al., 2024) . Furthermore, unlike (Hui et al., 2024) , LoTA not only preserves the earlier tasks on frozen parameters but also constraints the new tasks on a highly sparse subnetwork-providing resistance to catastrophic forgetting even when malicious users attempt to overwrite the earlier tasks via FFT.Model Merging. Merging multiple task-specific models into a single model with multitask abilities (Wortsman et al., 2022; Jin et al., 2022; Zhang et al., 2023a) has been an appealing alternative to sequential multitask learning, which suffers from catastrophic forgetting and could be inefficient, especially if task vectors are already available. The existing model merging methods include averaging weights of task-specific models (Wortsman et al., 2022) , task arithmetic through combining task vectors (Ilharco et al., 2022) , weighted aggregation of parameters (Matena & Raffel, 2022; Xiao et al., 2023b) , combining task vectors after some post-processing such as trimming low-magnitude deltas (Yadav et al., 2023) or sparsifying the deltas (Yu et al., 2023; Davari & Belilovsky, 2023) . Our method, LoTA, directly learns sparse task vectors, obviating the need to post-process the task vectors, and outperforms existing model merging methods. Most importantly, LoTA enables merging task vectors trained on heterogeneous datasets, while the other model merging methods are often limited to similar datasets. This advancement is an important step towards scalable FL (Kairouz et al., 2021) with LLMs as it enables merging sparse task vectors, which brings communication efficiency, trained over heterogeneous datasets (of each edge device).We note that we test model merging with LoTA specifically for highly dissimilar datasets to show its compatibility with FL, which considers edge devices with heterogeneous datasets. When used in FL, LoTA can reduce communication and memory costs significantly, which is a main bottleneck when scaling FL to large models. The successful use of LoTA for model merging and arithmetic further shows its promise for unlearning (Ilharco et al., 2022) as well. 
 
 Appendix:  


Decision: The paper proposes a new sparse adaptation method that searches for sparse subnetwork within a large model. The approach outperforms both LoRA adapters and full finetuning on the selected datasets and tasks. Both reviewers acknowledge the paper's novelty and are overall positive about the work. The AC agrees and recommends for acceptance.