Routing-Aware Expert Calibration for Machine Unlearning in Mixture-of-Experts Language Models
For practitioners needing to remove specific data from MoE LLMs, this work provides a targeted calibration method that improves the forget-utility trade-off.
The paper addresses machine unlearning in Mixture-of-Experts (MoE) language models, where a routing mismatch between forget and retain data causes under-regularization of forget-critical experts. They propose TRACE, which calibrates retain regularization by reweighting token-level losses, achieving a 9% relative utility improvement over the strongest baseline under comparable forgetting quality and best performance on three out of four MUSE-BOOKS metrics.
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.