SEEKR : Selective Attention-Guided Knowledge Retention for Continual Learning of Large Language Models (2024)

Abstract: Continual learning (CL) is crucial for language models to dynamically adapt to the evolving real-world demands. To mitigate the catastrophic forgetting problem in CL, data replay has been proven a simple and effective strategy, and the subsequent data-replay-based distillation can further enhance the performance. However, existing methods fail to fully exploit the knowledge embedded in models from previous tasks, resulting in the need for a relatively large number of replay samples to achieve good results. In this work, we first explore and emphasize the importance of attention weights in knowledge retention, and then propose a SElective attEntion-guided Knowledge Retention method (SEEKR) for data-efficient replay-based continual learning of large language models (LLMs). Specifically, SEEKR performs attention distillation on the selected attention heads for finer-grained knowledge retention, where the proposed forgettability-based and task-sensitivity-based measures are used to identify the most valuable attention heads. Experimental results on two continual learning benchmarks for LLMs demonstrate the superiority of SEEKR over the existing methods on both performance and efficiency. Explicitly, SEEKR achieves comparable or even better performance with only 1/10 of the replayed data used by other methods, and reduces the proportion of replayed data to 1%. The code is available at https://github.com/jinghan1he/SEEKR.

## introduction enabling large language models @xcite @xcite with human-like continual learning ability is crucial for the long-term practical deployment. it allows for constant knowledge accumulation on new tasks figure 1 : demonstration of the critical role of attention weights in knowledge retention. we apply der++ @xcite for continual learning on the trace benchmark @xcite to obtain multiple old task models and the final model. grafting the attention weights of the old models onto the final model at inference can maintain better performance on the old tasks. moreover, the final model obtained by our continual learning method, seekr, achieves similar results. among the array of continual learning methods @xcite , data replay stands out as the most widely adopted strategy in practice due to its simplicity and efficacy @xcite . based on it, replay-based distillation methods, including der++ @xcite and subsequent techniques @xcite @xcite , further boost the performance by utilizing memories from both data and model perspectives. specifically, @xcite @xcite distill the output logits of old models for knowledge transfer, and @xcite restrict the changes in important feature maps in the image encoders. however, these works have not fully exploited the potential of knowledge distillation in continual learning for llms. they focus on the outputs of network lay-ers while neglecting the preservation of intricate internal functions. consequently, a relatively large amount of replay data is required by these methods to achieve good results. recently, many studies have investigated the attention weights of different heads to analyze the interpretability of the internal mechanisms in llms @xcite . inspired by this, we explore whether attention weights play a critical role in knowledge retention during continual learning in llms. as shown in figure 1 , grafting the attention weights from the llm of the old tasks to the final llm after continual learning can maintain better performance on old tasks, which suggests that the attention weights could be crucial to alleviate the catastrophic forgetting problem and achieve more comprehensive knowledge retention foot_0 . however, naively preserving the attention weights of all heads in the llm by distillation introduces significant computational costs. previous studies have observed a functional specialization phenomenon among attention heads in llms @xcite @xcite , which indicates the susceptibility of attention heads to forgetting and their importance to previous tasks vary. this property allows us to selectively focus on the valuable attention heads for efficient knowledge retention. to this end, we propose a finer-grained model distillation method called selective attentionguided knowledge retention (seekr) for continual learning of large language models, which employs attention distillation on the most valuable heads in llms to achieve efficient knowledge retention. specifically, we develop knowledgeretention-oriented head importance measures, which consider both forgettability and task sensitivity, to identify the most valuable heads for distillation. the forgettability, measured by the cumulative changes in attention weights during continual learning, indicates the generality of knowledge and the necessity of distillation. an attention head with higher forgettability indicates a greater need for knowledge retention. the task sensitivity, calculated as the first-order derivative of the task loss, evaluates the importance of maintaining the attention weights of an attention head for a given task. an attention head with greater sensitivity should be prioritized to restrict variations in its attention weights. using the above two importance scores, seekr designs a hierarchical budget allocation mechanism to adaptively select the most valuable attention heads for distillation in a controllable way, which can efficiently regulate the training cost. by using seekr, the performance of old tasks can be further maintained as shown in figure 1 . extensive experiments are conducted on the recently developed continual learning benchmark for llms @xcite and the continual learning benchmark on traditional nlp tasks @xcite . the results consistently demonstrate the superiority of seekr in mitigating catastrophic forgetting and maintaining the general capabilities of llms. moreover, as a replay-based method, seekr exhibits excellent data efficiency, achieving comparable or better performance with just 1/10 of the replayed data used by the existing methods, reducing the replayed data proportion to only 1%. our main contributions are summarized as follows: ## method in this section, we introduce seekr, an efficient replay-based distillation method that identifies valuable attention heads and performs attention distillation for finer-grained knowledge retention. ## conclusion in this paper, we propose seekr, an efficient replay-based distillation method for continual learning in llms. seekr resorts to attention distillation of important heads for finer-grained knowledge retention, which identifies valuable heads through the proposed knowledge-retention-oriented importance measures. combined with a hierarchical budget allocation mechanism, seekr can ensure its utility across various resource levels. extensive experiments consistently validated the effectiveness of our method in preserving the performance of newly learned tasks and the original ability of the initial llms.