Search papers, labs, and topics across Lattice.
This paper addresses the challenge of machine unlearning in Mixture-of-Experts (MoE) language models, where the routing mechanism can lead to an imbalance in expert activation during the unlearning process. The authors introduce TRACE, a method that identifies forget-critical experts and recalibrates retain regularization to ensure that the activation frequencies of these experts are aligned during unlearning. Experimental results demonstrate that TRACE achieves a 9% relative improvement in the forget-utility trade-off compared to existing methods, enhancing performance across multiple evaluation metrics.
Forget-critical experts in MoE models can be under-regularized during unlearning, but TRACE rebalances their activation to significantly boost utility.
Machine unlearning is increasingly important for large language models, yet unlearning in Mixture-of-Experts (MoE) architectures remains underexplored. Unlike dense models, MoE architectures employ a router at each layer to assign each token to a sparse subset of experts. In this work, we observe that forget data often activates a small subset of experts disproportionately, while these experts may receive much weaker activation from retain data. This forget--retain routing mismatch can leave forget-critical experts under-regularized during unlearning. To address this, we propose \textbf{TRACE}, Targeted Routing-Aware Calibration of Experts, for MoE unlearning. TRACE first detects forget-critical experts from offline activation statistics, and then calibrates retain regularization by reweighting token-level retain losses so that each selected expert's retain-side activation frequency better matches its forget-side counterpart. Experiments on WMDP and MUSE-BOOKS across multiple MoE LLMs show that TRACE consistently improves the forget-utility trade-off, yielding a 9\% relative utility improvement over the strongest baseline under comparable forgetting quality and the best performance on three out of four MUSE-BOOKS metrics.
Tsinghua AI