Lottery Ticket Adaptation: Mitigating Destructive Interference in LLMs

Paper · arXiv 2406.16797 · Published June 24, 2024
Training and Fine-TuningMechanistic InterpretabilityCognitive Models and Latent RepresentationsInference-Time Scaling

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 lottery tickets (or sparse task vectors), LoTA also enables model merging over highly dissimilar tasks. Our code is made publicly available. 1

Introduction. Large language models (LLMs) [2] have seen an explosion of applications to real-world problems [3, 4] via adaptation [5] to new tasks. Three major multi-task adaptation paradigms have emerged: storing and loading task-specific adapters [6, 7], continuing to train instruction-tuned models on new tasks in serial via sequential training [5], and combining the adaptations to tasks learned in parallel via model merging [8]. Each paradigm has its own associated challenges, such as catastrophic forgetting during sequential training [9, 10, 11, 12, 13], and methods that have been proposed to mitigate these challenges [14, 15]. 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.)

Discussion / Conclusion. 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 full fine-tuning and low-rank adaptation (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.1, 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 [31] and merging [38] methods. Broader Impact. LoTA emerges as a simple and efficient adaptation method to prevent destructive interference.