UOE: Unlearning One Expert is Enough for Mixture-of-Experts LLMs

Authors: Haomin Zhuang, Yihua Zhang, Kehan Guo, Jinghan Jia, Gaowen Liu, Sijia Liu, Xiangliang Zhang
Primary area: alignment, fairness, safety, privacy, and societal considerations
Keywords: Machine Unlearning, Mixture-of-Expert, Large Language Model

Abstract:
Recent advancements in large language model (LLM) unlearning have shown remarkable success in removing unwanted data-model influences while preserving the model's utility for legitimate knowledge. However, despite these strides, sparse Mixture-of-Experts (MoE) LLMs--a key subset of the LLM family--have received little attention and remain largely unexplored in the context of unlearning. As MoE LLMs are celebrated for their exceptional performance and highly efficient inference processes, we ask: How can unlearning be performed effectively and efficiently on MoE LLMs? And will traditional unlearning methods be applicable to MoE architectures? Our pilot study shows that the dynamic routing nature of MoE LLMs introduces unique challenges, leading to substantial utility drops when existing unlearning methods are applied. Specifically, unlearning disrupts the router's expert selection, causing significant selection shift from the most unlearning target-related experts to irrelevant ones. As a result, more experts than necessary are affected, leading to excessive forgetting and loss of control over which knowledge is erased. To address this, we propose a novel single-expert unlearning framework, referred to as UOE, for MoE LLMs. Through expert attribution, unlearning is concentrated on the most actively engaged expert for the specified knowledge. Concurrently, an anchor loss is applied to the router to stabilize the active state of this targeted expert, ensuring focused and controlled unlearning that preserves model utility. The proposed UOE framework is also compatible with various unlearning algorithms. Extensive experiments demonstrate that UOE enhances both forget quality up to 5\% and model utility by 35\% on MoE LLMs across various benchmarks, LLM architectures, while only unlearning 0.06\% of the model parameters.

Decision: reject
Reviewer scores: 6, 5, 5, 5

Reviews and discussion:
<thread>
creation date:
1730099935605

thread_type:
official_review

<thread_object>
creator: reviewer_jqdp

<content>
summary:
This paper investigated a critical problem when applying machine unlearning in an MoE scenario. The author gives an insightful demonstration of the current challenge (e.g., Fig 1(a) and Fig 2), making this paper easy to follow. As the authors mentioned, the short-cut behavior when the MoE system is doing expert selection leads to unstable unlearning. To address this issue, the authors proposed an Unlearning One Expert system, which prevents the frequent switching of expert selection and ensures a focused, controlled, and stable unlearning. Although most of this paper is written clearly, I still have concerns about motivation, contribution, and technical quality.

strengths:
S1. The research problem is well formalized and presented with some demonstration of current research limitation visualization. \
S2. The authors provided a detailed introduction to recent studies on MoE, MU, etc. \
S3. The methodology and most of the experiment parts are well-written.

weaknesses:
W1. Regarding the motivation for applying MU tech to the MoE system. There could be a wild ground for application, but the technical challenge remains unclear. One key issue is raised due to the golden label of expert selection being decided during MoE system training. It could be more reasonable to directly apply this _decision + MU tech_ when unlearning on the MoE system remains unclear. 

W2. Like some routing algorithms, the performance improvement result from whether to reduce the uncertainty or routing token to correct expert is unclear.  The author could provide more experimental detail to clarify the contribution of the UOE system. 

W3. There are many routing algorithms proposed recently. However, UOE seems only to use TopK routing as the backbone. The versatility of UOE framework is not discussed well.

questions:
Major Questions:

Q1. Is the rising **uncertainty** among routing algorithms causing ineffective MoE unlearning? If so, why not directly optimize (or maintain) the entropy?

Q2. As $\lambda$ (hyperparameter for retaining loss, eq.1) aims to directly optimize the strength of retain quality, there is no discussion on this setting. If $\lambda$ is set to $0$, what would happen? 

Q3-1. Is the affinity score unchanged (or as a golden label) during the unlearning progress? 

Q3-2. If the affinity score is the golden target, why not use it as a roadmap but an optimization target? Is it proper to refine the routing strategy during optimization?

Q4. Some case studies on generated unlearning data could be interesting. What will happen in the MoE system?

Minor Questions:

MQ1. In Fig. 3 (a-b), what is the difference between the settings for Routers+Experts, Routers Only, and Experts Only? Why do those settings matter?\
MQ2. Also, in Fig. 3 (a-b), will the deeper layers affect the routing overlap ratio?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730548505843

thread_type:
official_review

<thread_object>
creator: reviewer_qbmp

<content>
summary:
This paper looks into unlearning (UL) for Mixture-of-Experts LLMs, being a pilot study in this area. It is found that existent UL algorithms suffer from expert selection shifts and thus cannot efficiently eliminate certain knowledge. This paper offers a simple solution that is to unlearn the target expert in a single layer. Experiments support the effectiveness of the method.

strengths:
- This paper identifies the unique challenge of unlearning MoE LLMs.

- The experiments are thoroughly conducted and ablate well on the design components of UoE such as the number of layers/experts to unlearn and sensitivity to expert ranking.

weaknesses:
- My primary concern is that both Qwen1.5 MoE [1] and DeepSeek MoE [2] have shared experts, meaning certain experts are always activated and are not subject to router choices. The paper does not mention how this issue is addressed. I believe these shared experts could retain the knowledge of the forget set. The paper lacks discussion on this point, and additional evidence is needed.

- The unlearning loss objective is unclear, is it a combination of (1) and (3)? What is the hyperparameter applied to the anchor loss, and how sensitive is model performance to this hyperparameter?


[1] https://qwenlm.github.io/blog/qwen-moe/

[2] https://huggingface.co/deepseek-ai/DeepSeek-V2-Lite

questions:
- What does line 4 in Algorithm 1 mean? Only router parameters in the same selected layer are activated accordingly, or all router parameters?

- Related to what is pointed out in Weakness. The affinity score is only taken from router outputs and thus ignores shared experts. Is there a reason to ignore those shared experts? 

- For other MoE models (for example Mixtral) without such a shared expert setting, how does UoE work?



I believe further clarification is needed on these points, and I am open to adjusting my scores if my concerns are addressed.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730673762122

thread_type:
official_review

<thread_object>
creator: reviewer_tyfp

<content>
summary:
To address the issue of unlearning disrupting the router's expert selection, the authors propose a novel single-expert unlearning framework that focuses unlearning on the most actively engaged expert for the specified knowledge.

strengths:
1、The authors present an innovative and parameter-efficient unlearning framework that effectively identifies, targets, and unlearns the expert most relevant to the forget set.
2、The paper provides numerous experiments that convincingly illustrate the motivation behind this work.
3、Extensive experiments demonstrate that UOE improves both the quality of forgetting and model utility in MoE-based large language models across various benchmarks.

weaknesses:
1、The novelty of this work lacks a compelling impact.

2、The explanation of the router anchor loss is unclear. Specifically, Equation (3) is confusing because it does not define the meaning of g(l). 

3、I am unclear about the authors' statement: “we propose the router anchor loss, which encourages the previously identified target expert to remain consistently activated throughout unlearning.” This raises the question of whether the previously identified target expert can reliably be the true target expert. If this cannot be ensured, the process may simply stabilize routing choices, inadvertently activating less relevant experts and undermining the effectiveness of the unlearning process.

4、In Algorithm 1, it is unclear what Ll represents.

5、The authors have not sufficiently emphasized the unique advantages of UOE compared to other MoE frameworks. For instance, while UOE leverages expert attribution to calculate affinity scores for expert selection, other MoE frameworks also compute affinity scores between inputs and experts. Why is UOE more advantageous in this regard? Additionally, when the authors state that "it overlooks finer details that are important for precise comparisons," it would be helpful to clarify what "precise" entails in this context and which additional factors might be considered more significant than precision.

6、In the preliminaries section, the character N denotes the number of experts; however, in Equation (2), it is used to represent the size of the calibration dataset.

7、Sentences such as "Model utility (UT) comparison, at the same level of forget efficacy" in Table 5 and “UT is compared at a consistent level of forget efficacy”in Table 7 are unclear and need clarification. The authors should explain how adjusting model hyperparameters allows one experiment's performance to be maintained while enabling changes in another. 

8、Figure 2 illustrates that multiple experts handle a substantial number of tokens, as shown in the long-tail distribution. The decision to unlearn only the top-1 expert requires further justification.

questions:
Please see the weakness.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730704297394

thread_type:
official_review

<thread_object>
creator: reviewer_jyxv

<content>
summary:
The paper addresses the challenge of unlearning in Mixture-of-Experts (MoE) LLMs. When naively applying traditional unlearning methods to MoE models, the authors discovered that MoE's dynamic routing mechanism leads to unintended expert selection shifts, causing excessive forgetting and reduced model utility. To solve this, they propose UOE (Unlearning on One Expert), a framework that: 1) identifies the most relevant expert for the target knowledge using affinity scores, and 2) implements a router anchor loss to maintain that expert's activation during unlearning. By concentrating unlearning on a single expert, UOE + baseline methods improve both forgetting quality and model utility significantly compared to baseline methods alone.

strengths:
1. Novel and well-motivated solution that specifically addresses MoE architecture challenges in unlearning, filling an important gap in the literature 
2. Thorough follow-up studies that justify key design choices, particularly: the effectiveness of single-expert versus multi-expert unlearning, and the impact of single/multi-layer selection on forgetting performance. 
3. Strong empirical validation with comprehensive experiments across different: unlearning algorithms, model architectures, benchmarks 4. Resource-efficient approach, requiring modifications to only 0.06% of model parameters while achieving superior performance.

weaknesses:
1. **Notation Clarity**: Several key mathematical notations (particularly $g(l)$ in the router anchor loss formulation) would benefit from more explicit definitions.
The relationship between different mathematical terms could be better explained.
2. **Implementation Details**: Training hyperparameters are insufficiently documented for reproducibility; the $Unlearn()$ subroutine in Algorithm 1 lacks specific implementation details; more comprehensive experimental setup information would facilitate replication.

questions:
1. The method relies on selecting the expert with the highest affinity score. However, how robust is this approach when there are experts with closely competing scores? 
2. The robustness to attack of the proposed method: if we keep the model frozen and use soft prompt learning, can the soft prompt uncover the knowledge from other experts? If so, this suggests that the knowledge still resides within these experts and can be accessed through carefully crafted input prompts.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732149736294

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
Thanks for all the valuable comments. We have carefully reviewed your feedback and advice, and provided our responses accordingly.  Please check our highlighted revisions and detailed responses to each reviewer. We hope that the revisions have addressed your comments. We are happy to provide further discussion if there are additional concerns or areas for clarification. 

A short summary of newly conducted experiments:
1. [Reviewer  JyXv] Following your suggestions, we introduce a soft prompt attack GCG against UOE in Tab. 9 in revision. The results indicate that UOE is robust to the GCG attack.
2. [Reviewer  TYfp] Following your suggestions, we incorporate Mixtral 8*7B as a MoE model without shared experts in our experiments in Tab. 4. The results indicate that UOE can apply to MoE without shared experts as well.
3. [Reviewer JQDp] Following your suggestion, we visualize the uncertainty of the unlearning process in Fig. 6. The results indicate that the expert selection shift is not caused by uncertainty in routing.
</content>
</thread_object>

</thread>


Metareview:
metareview:
Summary
========
The paper introduces UOE (Unlearning on One Expert) for machine unlearning in Sparse Mixture-of-Experts (MoE), a relatively unexplored area. The key finding is that traditional unlearning methods cause unintended expert selection shifts in MoE models, leading to excessive forgetting. UOE addresses this by identifying and targeting the most relevant expert for unlearning while maintaining its activation through a router anchor loss. This focused approach enhances both forgetting quality and model utility significantly compared to baseline methods.

Strengths
=======
* Well-motivated problem
* Novel approach addressing the unique challenges in MoE architecture unlearning, filling a critical gap in the literature
* Comprehensive empirical validation across different architectures and benchmarks
* Resource-efficient, modifying only 0.06% of model parameters
* Thorough ablation studies justifying key design choices

Weaknesses
==========
* Insufficient treatment of shared experts in modern MoE architectures (e.g., Qwen1.5, DeepSeek)
* Unclear mathematical notation and implementation details, particularly for router anchor loss
* Limited discussion of versatility across different routing algorithms beyond TopK
* Inadequate justification for single-expert approach given multi-expert token handling
* Implementation details, including training hyperparameters and specific steps in Algorithm 1, are inadequately documented
* Several technical concerns remain unaddressed, such as the handling of shared experts in MoE models and the reliability of target expert identification.

Reasons for decision
================
The paper addresses an interesting problem unlearning in MoE LLMs by introducing an innovative approach, and demonstrating strong empirical results. However, significant issues related to clarity, completeness of implementation details, and unresolved technical concerns still remain.