Locate-then-Unlearn: An Effective Method of Multi-Task Continuous Learning for Large Language Models

Authors: Songshi Liang, Hongda Sun, Ting-En Lin, Yuchuan Wu, Zihe Wang, Yongbin Li, Rui Yan
Primary area: interpretability and explainable AI
Keywords: Machine unlearning, Continue learning, Model editing

Abstract:
Nowadays large language models (LLMs) have achieved remarkable success in
various NLP tasks. However, they often misinterpret human instructions and generate incorrect or outdated responses, highlighting the need for more effective continual learning techniques. While recent efforts have introduced unlearning methods to remove erroneous knowledge, existing approaches still struggle in multi-task learning scenarios. To overcome these limitations, we propose Locate-then-unlearn, a new framework that identifies and selectively unlearns task-specific
neurons to enable efficient multi-task learning. We hypothesize that LLM neurons can be broadly categorized into task-specific neurons for handling individual
tasks, and general neurons to maintain the model’s foundational capabilities. To
accurately identify task-specific neurons, the locating process includes: (1) ranking task-related neurons based on their importance to each task, and (2) identifying
task-specific neurons by applying intervention to assess how neuron activity impacts task performance, isolating those most critical to each task. We conduct
comprehensive evaluations in two experimental setups: single-task specialization
and multi-task generalization. The results show that our method significantly improves performance across both settings. This indicates that our method effectively balances model efficiency and accuracy in multi-task continual learning.

Decision: reject
Reviewer scores: 6, 6, 5

Reviews and discussion:
<thread>
creation date:
1730577296167

thread_type:
official_review

<thread_object>
creator: reviewer_wabd

<content>
summary:
This paper proposes Locate-then-Unlearn framework to improve LLM continual learning by identifying and selectively unlearning task-specific neurons. By localizing neurons specific to individual tasks and unlearning them, the approach reduces interference between tasks and enhances model efficiency. Experiment results and analyses across multiple datasets demonstrate good performance in both single-task and multi-task settings.

strengths:
1. The proposed approach is novel.  
2. The proposed approach is well-motivated.  
3. The results show good performance.

weaknesses:
1. One problem of LLM continual learning is catastrophic forgetting, it would be better if the authors can further discuss how the proposed approach handle the issue.  
2. For LLM continual learning, it would be better to compare model performance on the benchmark datasets before and after training, so the readers will learn the effect of the learning process on previously learnt task. Currently, the authors just report the performance after training. It would be a plus if the authors can include the model performance before training.  
3. The paper claims many times that the proposed approach "significantly" outperforms existing approaches, but it seems the significance test is missing. So, it is hard to justify whether the improvement is significant.

questions:
See weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730586086198

thread_type:
official_review

<thread_object>
creator: reviewer_uprd

<content>
summary:
This paper proposes locate-then-unlearn, to identify and unlearn task-specific neurons in transformer models for multi-task learning. Specifically, the authors propose to identify task-related neurons from top-k important scores for each task, and identify task-specific neurons by removing overlapping neurons shared by different tasks. Finally, the authors fine-tune the selected task-specific parameters with unlearning data, and remove the parameter offset (from fine-tuning) from the original parameters. Experiment results demonstrate that the proposed locate-then-unlearn outperform other knowledge editing baselines on five datasets. They also conduct extensive analysis to demonstrate that the setting of unlearning hyperparameter is important, task-specific neurons work the best on corresponding related tasks.

strengths:
- The method is intuitively reasonable. By detecting task-related and task-specific neurons, the methods make multi-task learning more parameter-efficient.
- The experiment results show that the proposed method can outperform knowledge editing methods such as ROME and MEMIT on five datasets.
- The authors conduct extensive analysis. The unlearning hyperparameter has a best performing peak, while the threshold is rather robust. The task-specific neurons are important, and tuning them on unrelated tasks lead to worse results. Visualization shows that task-specific neurons are well-separated.
- The writing of the paper is quite clear.

weaknesses:
- Since the authors claim that their methods contribute to parameter-efficient multi-task learning, it would be better to compare with methods such as prompt-tuning, lora-hub.
- It would be better if the authors could release their code upon acceptance.

questions:
- How does the method compare with other PEFT methods, such as prompt-tuning, lora-hub?
- It would be better if the authors can discuss for each new task, how many parameters need to be stored (the percentage of the whole transformer model)?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730600233037

thread_type:
official_review

<thread_object>
creator: reviewer_vvor

<content>
summary:
This work proposes to unlearn task-specific neurons to enable more efficient multi-task learning. The proposed approach, locate-then-unlearn, first locates task-relevant neurons, deduplicates the neurons that are relevant for multiple tasks, subtracts the corresponding parameters to unlearn false answers, and relearns the objective tasks.

The proposed approach shows higher performance than five baseline methods on five datasets in the multi-task learning scenario, and improvements on three out of five datasets in the single-task setting (with the remaining two getting close performance with direct full fun-tuning).

strengths:
* This work presents solid experimental results and ablation studies, showing the effectiveness of each component.

* The main novelty lies in finding task-specific neurons for more efficient re-learning.

weaknesses:
* The causal effect between the neurons and different tasks requires more careful examination. Although the approach shows empirical improvement,  the assumption that neurons are task-specific might require further evidence.

* This approach is dependent on the selected data used to locate neurons (section 3.2.) and unlearn tasks (section 3.3). The authors provide several ablation studies on the method components, while it may be worthwhile to conduct ablations on the choice of data.

questions:
* What is the definition of “tasks” in this work? The distribution of data, even under the same tasks, seems to affect the selection of task-specific neurons. Thus, whether improvement will transfer to different datasets from the same tasks is unexamined.

* How are the toy datasets constructed in section 3.2? Are they a subsample of the underlying task data?

* Regarding the design of false answers in section 4.3., the false knowledge can be presented in multiple forms than what’s already in the dataset, which could thus not capture the potential wrong knowledge and influence the effect of unlearning. 

*Typos*

* Title: Continuous -> Continual

* L409-410: implement -> implementing; this sentence is difficult to parse.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731650563779

thread_type:
public_comment

<thread_object>
creator: ~nasim_nour_mohammed1

<content>
title:
This paper violates the double-blind review process and is questionable

comment:
**Dear reviewers and chairs,**

There are some serious concerns about the academic integrity of this paper under review.

Some people on Twitter have raised issues about the authors violating the double-blind review process by revealing their submission ID and author list on a Chinese social media: [removed]

Upon reading the paper, I believe it exhibits more substantial academic issues. Below, I outline my arguments:

> 1. Lack of Originality in Neuron Locating Method

The paper claims to propose a "two-stage locating method" for identifying task-specific neurons (Line 90). However, this approach is very similar to the method presented by Chen et al. [1], which also involves a two-stage process—first identifying **knowledge-related** neurons, and then filtering those that store factual knowledge **transcending** language. The authors of this paper first locate “**task-related**” neurons, then filter task-specific neurons, because “some task-related neurons may exhibit high logit scores **across multiple tasks**. (Line 70)” The authors replace the concept of "language" with "task" without citing Chen et al.

Furthermore, the neuron identification method (Eq.2) in this paper is **identical** to Eq.4 of Tang et al. [2]. The authors replace "language" in the original paper with "task" once again. Here is Equation 4 from Tang et al.:

$ p_{i,j}^k = \mathbb{E} \left( \mathbb{I}(	ext{act fn}(	ilde{h}^i W_1^i)_j > 0) \mid 	ext{language } k ight), $

The authors state, "the definition of a neuron follows Tang et al. (Line 143)," but they do not attribute Eq.2 to Tang et al.'s contribution. Such wording seems intended to mislead reviewers unfamiliar with Tang et al.

> 2. Continual Learning vs Lifelong Editing

At first, I did not understand why they called this model editing method continual learning. However, upon closer inspection, particularly of Eq.5, it becomes clear that the subtraction between parameters of two models is a common practice in lifelong model editing (see Eq.3 in Wang et al. [4]). Nevertheless, the authors do not cite any relevant lifelong editing literature, raising concerns about proper attribution.

> 3. The Style of Figures Demonstrates That the Authors Have Read the Papers I Listed

This paper does not relate to attention modules at all, yet the authors still include attention modules in Figure 1. Why? This figure closely resembles Figure 1 from Meng et al. [5], with the authors largely following the figure style from this classic paper on model editing. Additionally, the top-right part of Figure 1 in this paper uses elements similar to Figure 1 of Tang et al., and the structure of Figure 1(a) and (b) in this paper is similar to that of Figure 2 in Chen et al. (Figure 1(a) and (b) illustrate the two-stage neuron identification and recap that this two-stage process was proposed by Chen et al.)

While reusing classic figures is not inherently problematic, the resemblance suggests that the authors have read these papers but still underplayed their contributions in this paper.

> Confusing Equations

Equation 3 in the paper uses a superscript in the summation notation that is unclear. I wonder whether the reviewers understand this superscript—is it n or k?

If it is k, as k represents a cardinal task index, the meaning of "i < k" is ambiguous.

If it is n, the summation term is a constant! It is strange to subtract a constant in this metric. However, based on the related works cited, I believe this formulation may have been borrowed from work in mechanistic interpretability, such as the circuit discovery methods in Lv et al. cited in this paper, which calculate the logit difference and normalize it with the average logits over the vocabulary to measure activation differences between two inputs. It is an odd mistake to reuse such a measure for probability differences in this paper because averaging the probability difference over the distribution results in a constant.

In summary, I believe that the paper demonstrates substantial overlap with prior works in both methodology and presentation, with some clear mistakes and without proper citations or acknowledgment. 

I appreciate the reviewers' efforts, but for the integrity of the ICLR, I hope the reviewers and chairs will consider the issues I have raised and re-evaluate the paper. 

**Given the violation of the double-blind process, this paper should perhaps be desk-rejected directly.**

[1] Chen et al. Journey to the Center of the Knowledge Neurons: Discoveries of Language-Independent Knowledge Neurons and Degenerate Knowledge Neurons

[2] Tang et al. Language-Specific Neurons: The Key to Multilingual Capabilities in Large Language Models

[3] Geva et al. Transformer Feed-Forward Layers Are Key-Value Memories

[4] Wang et al. LEMoE: Advanced Mixture of Experts Adaptor for Lifelong Model Editing of Large Language Models

[5] Meng et al. Locating and Editing Factual Associations in GPT
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper proposes a locate-then-unlearn framework to first identify task-specific neurons for each task and then unlearn the false information embedded in these neurons, followed by a relearning process. The proposed framework effectively handles multi-task learning where task-agnostic neurons are kept but only those that are most functional towards a task are being updated. Comprehensive experiments on both single-task and multi-task settings are conducted, demonstrating the advantage of the proposed method over existing baselines.

Strengths:
- The method specifically targets efficient multi-task learning by identifying a small number of neurons responsible for a specific task and only updating these neurons. This motivation is reasonable and the solution is effective.
- Extensive experiments are performed, including ablation studies and in-depth analysis, showcasing the advantage of the proposed framework in both single-task and multi-task settings, compared with direct fine-tuning and baseline editing models. Complexity analysis is also provided to demonstrate its efficiency.

Weaknesses:
- The paper has gone through a non-trivial structural and content update, compared to its first submission. This update has led to a change of scope, contribution and methodology. Due to such changes, it is advised that the paper goes through a new round of review before making an informed decision.
- The novelty is a bit limited, after comparing with related work. While the authors have adjusted the contribution statements and restructured the methodology by adding a preliminary section, it weakens the contribution in terms of methodology, as both the localization of neurons and unlearning strategies have been proposed. The adaptation of these strategies into a multitask setting is interesting, but not sufficiently novel.
- The identification of task-specific neurons relies on intervention and observing the difference using a toy dataset. Yet, the sensitivity towards the selected data is not fully examined.