pith. sign in

arxiv: 2606.10338 · v1 · pith:EPYQ6IDInew · submitted 2026-06-09 · 💻 cs.CL · cs.AI

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

Jingyi Xie , Yijun Lin , Yinjiang Xiong , Zhikun Zhang , Sai Li This is my paper

Pith reviewed 2026-06-27 13:19 UTC · model grok-4.3

classification 💻 cs.CL cs.AI
keywords machine unlearningmixture of expertslanguage modelsexpert calibrationrouting mismatchforget-retain trade-off
0
0 comments X

The pith

Reweighting retain losses to match forget-side expert activations improves the forget-utility trade-off in MoE unlearning.

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The paper establishes that forget data in Mixture-of-Experts models often over-activates a small set of experts that receive far less activation from retain data, leaving those experts under-regularized when standard unlearning is applied. TRACE addresses this by first identifying the forget-critical experts from activation counts and then reweighting individual retain tokens so that each such expert sees activation frequencies on the retain side that more closely match its forget-side frequencies. A sympathetic reader would care because MoE architectures dominate many large language models, and the routing step creates a structural mismatch that dense-model unlearning methods ignore, potentially allowing forgotten information to persist in under-regularized experts.

Core claim

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.

What carries the argument

Targeted Routing-Aware Calibration of Experts (TRACE), which reweights token-level retain losses to align per-expert activation frequencies between forget and retain data.

If this is right

  • TRACE yields a 9% relative utility improvement over the strongest baseline under comparable forgetting quality.
  • TRACE achieves the best performance on three out of four MUSE-BOOKS metrics.
  • The method produces consistent gains across multiple MoE LLMs on both WMDP and MUSE-BOOKS.
  • The routing mismatch between forget and retain data is the primary source of under-regularization in standard MoE unlearning.

Where Pith is reading between the lines

These are editorial extensions of the paper, not claims the author makes directly.

  • The same activation-frequency matching idea could be applied to other sparse routing mechanisms such as those in Switch Transformers or other conditional computation layers.
  • If the mismatch is the dominant issue, then offline expert profiling might become a standard preprocessing step for any unlearning pipeline on routed models.
  • Extending the calibration to also adjust expert-level learning rates rather than only token weights could further tighten the forget-retain balance.

Load-bearing premise

Reweighting token-level retain losses to match activation frequencies will correct under-regularization of forget-critical experts without introducing new side effects or degrading overall model behavior.

What would settle it

Running TRACE on the WMDP or MUSE-BOOKS benchmarks and finding no improvement in the forget-utility trade-off or degradation on retain metrics compared with the strongest baseline.

Figures

Figures reproduced from arXiv: 2606.10338 by Jingyi Xie, Sai Li, Yijun Lin, Yinjiang Xiong, Zhikun Zhang.

Figure 1
Figure 1. Figure 1: Routing-induced scaling mismatch in MoE unlearning. Visualization of expert activations across different MoE models on the forget dataset WMDP and the retain dataset WikiText, with forget-data activation frequency on the x-axis and retain-data activation fre￾quency on the y-axis. Definitions of p f j,l and p r j,l are given in (3). et al., 2025; Liu et al., 2024b; Zhang et al., 2024; Fan et al., 2025; Jia … view at source ↗
Figure 2
Figure 2. Figure 2: Motivation of TRACE. Left: In standard MoE unlearning, forget and retain data may activate different [PITH_FULL_IMAGE:figures/full_fig_p003_2.png] view at source ↗
Figure 3
Figure 3. Figure 3: Routing similarity between the forget dataset [PITH_FULL_IMAGE:figures/full_fig_p007_3.png] view at source ↗
Figure 4
Figure 4. Figure 4: Comparison among GradDiff, GradDiff∗ , and TRACE on WMDP. GradDiff and TRACE use WikiText as generic retain data, while GradDiff∗ uses WMDP￾retain. Forget Quality is the average of Cyber and Bio accuracy; Utility is MMLU accuracy [PITH_FULL_IMAGE:figures/full_fig_p008_4.png] view at source ↗
read the original abstract

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.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit. Tearing a paper down is the easy half of reading it; the pith above is the substance, this is the friction.

Referee Report

1 major / 2 minor

Summary. The manuscript proposes TRACE (Targeted Routing-Aware Calibration of Experts) for machine unlearning in Mixture-of-Experts language models. It identifies forget-critical experts using offline activation statistics on forget data and then reweights token-level retain losses to better match the activation frequencies of these experts on retain data, aiming to address routing mismatch. Experiments on WMDP and MUSE-BOOKS benchmarks across multiple MoE LLMs report a 9% relative utility improvement over the strongest baseline with comparable forgetting quality.

Significance. If the empirical improvements hold under the proposed mechanism, this work addresses an underexplored area of unlearning for MoE architectures that are increasingly deployed in large-scale LLMs. The targeted reweighting idea could provide a practical regularization adjustment for sparse models.

major comments (1)
  1. [Method (as described in abstract and §3)] The calibration relies on static offline activation frequencies computed once before unlearning. Since the method jointly optimizes the router and experts, router updates during training can alter token-to-expert assignments for retain data, potentially invalidating the initial frequency-matching factors. No analysis, monitoring of mid-training activations, or ablation on router drift is provided to confirm the reweighting remains effective.
minor comments (2)
  1. [Abstract] The abstract states 'best performance on three out of four MUSE-BOOKS metrics' without naming the metrics or reporting the actual values or baseline comparisons.
  2. [Method] No equations, loss formulations, or pseudocode for the reweighting procedure are referenced, making it difficult to assess whether the calibration introduces new hyperparameters or side effects.

Simulated Author's Rebuttal

1 responses · 0 unresolved

We thank the referee for the constructive comment on the static calibration in TRACE. We address the concern below and will revise the manuscript to incorporate additional analysis.

read point-by-point responses

  1. Referee: [Method (as described in abstract and §3)] The calibration relies on static offline activation frequencies computed once before unlearning. Since the method jointly optimizes the router and experts, router updates during training can alter token-to-expert assignments for retain data, potentially invalidating the initial frequency-matching factors. No analysis, monitoring of mid-training activations, or ablation on router drift is provided to confirm the reweighting remains effective.

    Authors: We agree that the reweighting factors are computed statically from offline activation statistics prior to unlearning, and that joint optimization of the router and experts during training could in principle alter token-to-expert assignments on retain data. The calibration is designed as an initial regularization adjustment to address the observed forget-retain mismatch at the outset. To directly address the concern, we will add an ablation in the revised manuscript that tracks activation frequencies at multiple training checkpoints and evaluates the stability of the reweighting factors, thereby confirming whether the initial calibration remains effective throughout optimization. revision: yes

Circularity Check

0 steps flagged

Empirical calibration from offline statistics with no definitional reduction

full rationale

The paper presents TRACE as an empirical procedure: offline activation statistics on forget data identify critical experts, after which fixed reweighting factors are applied to retain-token losses during unlearning. No equations, derivations, or self-citations are shown that reduce the reported utility gains to a fitted quantity or to the input statistics by construction. The central claim therefore rests on observable mismatch statistics and standard gradient updates rather than on any self-referential loop.

Axiom & Free-Parameter Ledger

0 free parameters · 1 axioms · 0 invented entities

Abstract supplies no explicit free parameters, axioms beyond standard ML training assumptions, or new entities; the central claim rests on the empirical observation of routing mismatch and the effectiveness of frequency matching.

axioms (1)
  • standard math Standard machine learning assumptions on loss minimization and gradient-based optimization
    Implicit in any unlearning procedure that modifies training losses.

pith-pipeline@v0.9.1-grok · 5742 in / 1003 out tokens · 22363 ms · 2026-06-27T13:19:16.159674+00:00 · methodology

discussion (0)

Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.

Reference graph

Works this paper leans on

49 extracted references · 4 canonical work pages · 2 internal anchors

  1. [1]

    2025 , eprint=

    Qwen3 Technical Report , author=. 2025 , eprint=

    2025

  2. [2]

    2025 , eprint=

    DeepSeek-V3 Technical Report , author=. 2025 , eprint=

    2025

  3. [3]

    2024 , eprint=

    GPT-4 Technical Report , author=. 2024 , eprint=

    2024

  4. [4]

    2024 , eprint=

    The WMDP Benchmark: Measuring and Reducing Malicious Use With Unlearning , author=. 2024 , eprint=

    2024

  5. [5]

    2023 , eprint=

    Who's Harry Potter? Approximate Unlearning in LLMs , author=. 2023 , eprint=

    2023

  6. [6]

    2024 , eprint=

    RWKU: Benchmarking Real-World Knowledge Unlearning for Large Language Models , author=. 2024 , eprint=

    2024

  7. [7]

    2024 , url=

    Pratyush Maini and Zhili Feng and Avi Schwarzschild and Zachary Chase Lipton and J Zico Kolter , booktitle=. 2024 , url=

    2024

  8. [8]

    2024 , eprint=

    Rethinking Machine Unlearning for Large Language Models , author=. 2024 , eprint=

    2024

  9. [9]

    2025 , eprint=

    A Comprehensive Survey of Machine Unlearning Techniques for Large Language Models , author=. 2025 , eprint=

    2025

  10. [10]

    A practical guide, 1st ed., Cham: Springer International Publishing , volume=

    The eu general data protection regulation (gdpr) , author=. A practical guide, 1st ed., Cham: Springer International Publishing , volume=. 2017 , publisher=

    2017

  11. [11]

    2019 , publisher=

    The california consumer privacy act (ccpa): An implementation guide , author=. 2019 , publisher=

    2019

  12. [12]

    2026 , eprint=

    A Comprehensive Survey of Mixture-of-Experts: Algorithms, Theory, and Applications , author=. 2026 , eprint=

    2026

  13. [13]

    SEUF : Is Unlearning One Expert Enough for Mixture-of-Experts LLM s?

    Zhuang, Haomin and Zhang, Yihua and Guo, Kehan and Jia, Jinghan and Liu, Gaowen and Liu, Sijia and Zhang, Xiangliang. SEUF : Is Unlearning One Expert Enough for Mixture-of-Experts LLM s?. Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers). 2025. doi:10.18653/v1/2025.acl-long.424

    work page doi:10.18653/v1/2025.acl-long.424 2025

  14. [14]

    2026 , eprint=

    GRIP: Algorithm-Agnostic Machine Unlearning for Mixture-of-Experts via Geometric Router Constraints , author=. 2026 , eprint=

    2026

  15. [15]

    2024 , eprint=

    Merge, Then Compress: Demystify Efficient SMoE with Hints from Its Routing Policy , author=. 2024 , eprint=

    2024

  16. [16]

    BMC medical research methodology , volume=

    Estimating the sample mean and standard deviation from the sample size, median, range and/or interquartile range , author=. BMC medical research methodology , volume=. 2014 , publisher=

    2014

  17. [17]

    2024 , eprint=

    Let the Expert Stick to His Last: Expert-Specialized Fine-Tuning for Sparse Architectural Large Language Models , author=. 2024 , eprint=

    2024

  18. [18]

    2024 , eprint=

    MUSE: Machine Unlearning Six-Way Evaluation for Language Models , author=. 2024 , eprint=

    2024

  19. [19]

    2022 , eprint=

    ST-MoE: Designing Stable and Transferable Sparse Expert Models , author=. 2022 , eprint=

    2022

  20. [20]

    2024 , eprint=

    Mixtral of Experts , author=. 2024 , eprint=

    2024

  21. [21]

    Qwen3 Technical Report

    Qwen3 Technical Report , author=. arXiv preprint arXiv:2505.09388 , year=

    work page internal anchor Pith review Pith/arXiv arXiv

  22. [22]

    Qwen2 Technical Report

    Qwen2 Technical Report , author =. arXiv preprint arXiv:2407.10671 , year =

    work page internal anchor Pith review Pith/arXiv arXiv

  23. [23]

    2024 , eprint=

    DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model , author=. 2024 , eprint=

    2024

  24. [24]

    2024 , eprint=

    Negative Preference Optimization: From Catastrophic Collapse to Effective Unlearning , author=. 2024 , eprint=

    2024

  25. [25]

    2024 , eprint=

    Large Language Model Unlearning , author=. 2024 , eprint=

    2024

  26. [26]

    2022 , eprint=

    Continual Learning and Private Unlearning , author=. 2022 , eprint=

    2022

  27. [27]

    2021 , eprint=

    Measuring Massive Multitask Language Understanding , author=. 2021 , eprint=

    2021

  28. [28]

    2016 , eprint=

    Pointer Sentinel Mixture Models , author=. 2016 , eprint=

    2016

  29. [29]

    2022 , eprint=

    Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity , author=. 2022 , eprint=

    2022

  30. [30]

    2022 , eprint=

    GLaM: Efficient Scaling of Language Models with Mixture-of-Experts , author=. 2022 , eprint=

    2022

  31. [31]

    2020 , eprint=

    GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding , author=. 2020 , eprint=

    2020

  32. [32]

    2025 , eprint=

    Simplicity Prevails: Rethinking Negative Preference Optimization for LLM Unlearning , author=. 2025 , eprint=

    2025

  33. [33]

    2025 , eprint=

    WAGLE: Strategic Weight Attribution for Effective and Modular Unlearning in Large Language Models , author=. 2025 , eprint=

    2025

  34. [34]

    2022 , eprint=

    Quark: Controllable Text Generation with Reinforced Unlearning , author=. 2022 , eprint=

    2022

  35. [35]

    2024 , eprint=

    Large Language Model Unlearning via Embedding-Corrupted Prompts , author=. 2024 , eprint=

    2024

  36. [36]

    2023 , eprint=

    KGA: A General Machine Unlearning Framework Based on Knowledge Gap Alignment , author=. 2023 , eprint=

    2023

  37. [37]

    2026 , eprint=

    KUDA: Knowledge Unlearning by Deviating Representation for Large Language Models , author=. 2026 , eprint=

    2026

  38. [38]

    2024 , eprint=

    To Each (Textual Sequence) Its Own: Improving Memorized-Data Unlearning in Large Language Models , author=. 2024 , eprint=

    2024

  39. [39]

    2023 , eprint=

    Unlearn What You Want to Forget: Efficient Unlearning for LLMs , author=. 2023 , eprint=

    2023

  40. [40]

    Unlearning Bias in Language Models by Partitioning Gradients

    Yu, Charles and Jeoung, Sullam and Kasi, Anish and Yu, Pengfei and Ji, Heng. Unlearning Bias in Language Models by Partitioning Gradients. Findings of the Association for Computational Linguistics: ACL 2023. 2023. doi:10.18653/v1/2023.findings-acl.375

    work page doi:10.18653/v1/2023.findings-acl.375 2023

  41. [41]

    2023 , eprint=

    Composing Parameter-Efficient Modules with Arithmetic Operations , author=. 2023 , eprint=

    2023

  42. [42]

    2024 , eprint=

    Direct Preference Optimization: Your Language Model is Secretly a Reward Model , author=. 2024 , eprint=

    2024

  43. [43]

    2023 , eprint=

    Locating and Editing Factual Associations in GPT , author=. 2023 , eprint=

    2023

  44. [44]

    2023 , eprint=

    DEPN: Detecting and Editing Privacy Neurons in Pretrained Language Models , author=. 2023 , eprint=

    2023

  45. [45]

    2023 , eprint=

    Can Sensitive Information Be Deleted From LLMs? Objectives for Defending Against Extraction Attacks , author=. 2023 , eprint=

    2023

  46. [46]

    2023 , eprint=

    Memory-assisted prompt editing to improve GPT-3 after deployment , author=. 2023 , eprint=

    2023

  47. [47]

    2024 , eprint=

    In-Context Unlearning: Language Models as Few Shot Unlearners , author=. 2024 , eprint=

    2024

  48. [48]

    2025 , eprint=

    Fast Exact Unlearning for In-Context Learning Data for LLMs , author=. 2025 , eprint=

    2025

  49. [49]

    2017 , eprint=

    Attention Is All You Need , author=. 2017 , eprint=

    2017