Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

GRIP: Algorithm-Agnostic Machine Unlearning for Mixture-of-Experts via Geometric Router Constraints

Andy Zhu, Rongzhe Wei, Yupu Gu, Pan Li

TL;DR

This work addresses the challenge of machine unlearning in sparse Mixture-of-Experts models, where conventional methods are thwarted by router-based shortcuts that preserve or bypass knowledge without genuine erasure. It introduces Geometric Routing Invariance Preservation (GRIP), an algorithm-agnostic framework that enforces hard geometric constraints by projecting router gradient updates into the null space of retain representations and by using per-expert constraint decomposition with training-time enforcement via projected gradient descent and randomized Kaczmarz, plus a post-training analytical correction. The key contribution is decoupling routing stability from parameter rigidity, forcing forgetting to occur in expert parameters while allowing routers to reconfigure in orthogonal directions, thereby preventing expert-selection shifts. Empirical results on a 30B MoE model across WMDP and MUSE demonstrate dramatic improvements in routing stability (>0.94), substantial retention of utility, and genuine forgetting with adversarial robustness, achieving dense-model-level performance with modest overhead. Overall, GRIP provides a scalable, mechanism-driven approach to safe MoE deployment, extending unlearning techniques beyond dense architectures and highlighting routing as a critical frontier for security and robustness.

Abstract

Machine unlearning (MU) for large language models has become critical for AI safety, yet existing methods fail to generalize to Mixture-of-Experts (MoE) architectures. We identify that traditional unlearning methods exploit MoE's architectural vulnerability: they manipulate routers to redirect queries away from knowledgeable experts rather than erasing knowledge, causing a loss of model utility and superficial forgetting. We propose Geometric Routing Invariance Preservation (GRIP), an algorithm-agnostic framework for unlearning for MoE. Our core contribution is a geometric constraint, implemented by projecting router gradient updates into an expert-specific null-space. Crucially, this decouples routing stability from parameter rigidity: while discrete expert selections remain stable for retained knowledge, the continuous router parameters remain plastic within the null space, allowing the model to undergo necessary internal reconfiguration to satisfy unlearning objectives. This forces the unlearning optimization to erase knowledge directly from expert parameters rather than exploiting the superficial router manipulation shortcut. GRIP functions as an adapter, constraining router parameter updates without modifying the underlying unlearning algorithm. Extensive experiments on large-scale MoE models demonstrate that our adapter eliminates expert selection shift (achieving over 95% routing stability) across all tested unlearning methods while preserving their utility. By preventing existing algorithms from exploiting MoE model's router vulnerability, GRIP adapts existing unlearning research from dense architectures to MoEs.

GRIP: Algorithm-Agnostic Machine Unlearning for Mixture-of-Experts via Geometric Router Constraints

TL;DR

This work addresses the challenge of machine unlearning in sparse Mixture-of-Experts models, where conventional methods are thwarted by router-based shortcuts that preserve or bypass knowledge without genuine erasure. It introduces Geometric Routing Invariance Preservation (GRIP), an algorithm-agnostic framework that enforces hard geometric constraints by projecting router gradient updates into the null space of retain representations and by using per-expert constraint decomposition with training-time enforcement via projected gradient descent and randomized Kaczmarz, plus a post-training analytical correction. The key contribution is decoupling routing stability from parameter rigidity, forcing forgetting to occur in expert parameters while allowing routers to reconfigure in orthogonal directions, thereby preventing expert-selection shifts. Empirical results on a 30B MoE model across WMDP and MUSE demonstrate dramatic improvements in routing stability (>0.94), substantial retention of utility, and genuine forgetting with adversarial robustness, achieving dense-model-level performance with modest overhead. Overall, GRIP provides a scalable, mechanism-driven approach to safe MoE deployment, extending unlearning techniques beyond dense architectures and highlighting routing as a critical frontier for security and robustness.

Abstract

Machine unlearning (MU) for large language models has become critical for AI safety, yet existing methods fail to generalize to Mixture-of-Experts (MoE) architectures. We identify that traditional unlearning methods exploit MoE's architectural vulnerability: they manipulate routers to redirect queries away from knowledgeable experts rather than erasing knowledge, causing a loss of model utility and superficial forgetting. We propose Geometric Routing Invariance Preservation (GRIP), an algorithm-agnostic framework for unlearning for MoE. Our core contribution is a geometric constraint, implemented by projecting router gradient updates into an expert-specific null-space. Crucially, this decouples routing stability from parameter rigidity: while discrete expert selections remain stable for retained knowledge, the continuous router parameters remain plastic within the null space, allowing the model to undergo necessary internal reconfiguration to satisfy unlearning objectives. This forces the unlearning optimization to erase knowledge directly from expert parameters rather than exploiting the superficial router manipulation shortcut. GRIP functions as an adapter, constraining router parameter updates without modifying the underlying unlearning algorithm. Extensive experiments on large-scale MoE models demonstrate that our adapter eliminates expert selection shift (achieving over 95% routing stability) across all tested unlearning methods while preserving their utility. By preventing existing algorithms from exploiting MoE model's router vulnerability, GRIP adapts existing unlearning research from dense architectures to MoEs.
Paper Structure (50 sections, 1 theorem, 15 equations, 1 figure, 6 tables, 1 algorithm)

This paper contains 50 sections, 1 theorem, 15 equations, 1 figure, 6 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Unlearning Methods for LLMs.
  4. The Rise of Mixture-of-Experts Architectures.
  5. MoE Routing Dynamics and Stability.
  6. Preliminaries
  7. Machine Unlearning.
  8. Mixture-of-Experts Layers.
  9. Expert Selection Shift and Routing Stability.
  10. Methodology
  11. Null Space Constraints for Routing Preservation
  12. Expert-Specific Constraint Decomposition
  13. Training-Time Enforcement via Constrained Optimization
  14. Post-Training Analytical Correction
  15. Experimental Results
  16. ...and 35 more sections

Key Result

Theorem 1

Let $\tilde{\nabla}^*$ be the projection of the initial gradient onto the feasible region. The expected error at iteration $k$ satisfies: where $\kappa$ is the scaled condition number of the constraint matrix.

Figures (1)

  • Figure 1: Geometric Routing Invariance Preservation (GRIP) prevents expert selection collapse. (a) The Problem: Standard unlearning methods suffer from Expert Selection Shift, where the router acts as an optimization shortcut, diverting queries away from knowledgeable experts rather than erasing information. This causes Routing Stability (RS), a measure of the consistency of expert selection we define formally in Eq \ref{['eq:jaccard']}, to collapse from $1.0$ to $\approx 0.2$. (b) Our Solution: GRIP introduces two mechanisms to decouple routing stability from parameter plasticity: (a) Training-Time Enforcement, which projects gradients into a "safe" null-space orthogonal to the router's decision boundaries; and (b) Post-Training Correction, which analytically realigns a drifted router after unlearning. Both methods restore stability to $>0.94$ while preserving unlearning efficacy.

Theorems & Definitions (1)

  • Theorem 1: Convergence of Randomized Kaczmarz