Routing-Aware Expert Calibration for Machine Unlearning in Mixture-of-Experts Language Models

Researchers propose TRACE, a novel machine unlearning technique designed specifically for Mixture-of-Experts language models that addresses the problem of forget-critical experts receiving insufficient regularization during the unlearning process. The method achieves 9% relative utility improvements by detecting and calibrating expert activation patterns to match forget and retain data distributions, demonstrating consistent performance gains across multiple MoE architectures.

Machine unlearning—the process of removing specific training data influences from trained models—presents unique technical challenges in Mixture-of-Experts architectures, which route different tokens to different expert subsets rather than using dense, fully-connected layers. The research identifies a critical vulnerability: experts disproportionately activated by data to be forgotten often receive minimal activation from retained data, leaving them under-regularized and potentially vulnerable to data leakage. This routing mismatch represents a fundamental architectural consideration overlooked in prior unlearning work designed for dense models.

The TRACE methodology addresses this by first analyzing offline activation statistics to identify forget-critical experts, then reweighting token-level retain losses to ensure balanced expert activation across forget and retain datasets. This calibration approach is architecturally informed rather than generic, recognizing that MoE routing patterns fundamentally differ from dense model behavior. Experimental validation across WMDP and MUSE-BOOKS benchmarks shows consistent improvements in the forget-utility trade-off—the critical balance between effective data removal and preserved model performance.

For the AI and language model development community, this work has practical implications for deploying unlearning in increasingly common MoE architectures like Mixtral and other sparse models. The 9% utility improvement over existing baselines represents meaningful progress in an area where performance degradation typically accompanies unlearning. As regulatory pressure for data removal rights intensifies globally, techniques that efficiently unlearn data while preserving model utility become increasingly valuable. The research suggests that architectural awareness—tailoring unlearning methods to specific model designs—yields better outcomes than one-size-fits-all approaches.

• →TRACE detects experts disproportionately activated by forget data and calibrates them through loss reweighting to match retain-side activation patterns.
• →The method achieves 9% relative utility improvement over baseline unlearning approaches on evaluated benchmarks.
• →Mixture-of-Experts routing patterns create unique unlearning challenges that dense model techniques fail to adequately address.
• →Architecturally-aware unlearning methods outperform generic approaches when adapted to specific model design patterns.
• →Results demonstrate consistent performance across multiple MoE language models, suggesting broad applicability to sparse model architectures.