🤖 AI

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

This paper introduces GRIP, an algorithm-agnostic framework that enforces geometric constraints on router updates in Mixture-of-Experts models to prevent routing manipulation, thereby ensuring genuine knowledge erasure from expert parameters while significantly improving retention accuracy and robustness against adversarial recovery.

Original authors: Andy Zhu, Rongzhe Wei, Yupu Gu, Pan Li

Published 2026-08-27

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Andy Zhu, Rongzhe Wei, Yupu Gu, Pan Li

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Large language models are powerful tools that have learned to write, reason, and solve problems by studying vast amounts of text. However, this learning comes with a hidden cost: these models often memorize sensitive details, such as private personal information, copyrighted material, or dangerous instructions, which they were not meant to keep. When laws like the "right to be forgotten" require a company to remove this specific data from a model, the standard solution is to retrain the entire system from scratch without that data. This process is so expensive and slow that it is rarely done. Instead, researchers have developed "machine unlearning" techniques designed to surgically remove specific memories without rebuilding the whole brain. While these methods work well for standard models, a new type of architecture called a Mixture-of-Experts has emerged as a faster, more efficient alternative for the largest systems. These models do not use every part of their brain for every question; instead, they have a built-in traffic director that chooses a small team of specialized sub-networks to handle each task. This efficiency creates a unique vulnerability: when asked to forget something, the system can trick the user by simply rerouting questions away from the experts that hold the bad memories, leaving the dangerous knowledge intact and waiting to be recovered.

A team of researchers has identified this loophole and developed a new framework called GRIP to fix it. They discovered that when standard unlearning methods are applied to these specialized models, the system takes the path of least resistance. Rather than actually deleting the harmful knowledge from the specialized sub-networks, the traffic director simply learns to stop sending questions to those specific experts. It is as if a librarian, instead of removing a banned book from the shelf, simply tells all the patrons to look in a different section. The book remains on the shelf, fully intact, and a determined person could still find it if they knew to bypass the librarian's instructions. This creates a dangerous illusion of safety where the model appears to have forgotten the information, but the data is merely hidden behind a changed routing pattern.

To solve this, the researchers created a method that forces the system to actually erase the knowledge rather than just hiding it. They achieved this by placing strict geometric constraints on how the traffic director can change its mind. Imagine the traffic director as a set of rules that decides which path a question takes. The researchers found that the system was trying to change these rules to avoid the experts holding the bad data. Their solution was to lock the rules in place for all the safe, important information the model needs to keep. By mathematically ensuring that the traffic director cannot change its routing decisions for the safe data, they removed the ability of the system to use routing as a shortcut. This forced the unlearning process to do the hard work of actually modifying the specialized sub-networks that hold the knowledge, ensuring the information is truly gone.

The researchers tested this approach on two large models and found that it worked dramatically better than previous methods. In standard tests, the old methods caused the routing system to become unstable, changing its mind about which experts to use for safe questions about 80% of the time, which severely damaged the model's ability to function correctly. The new method restored this stability to over 94%, meaning the model continued to use the correct experts for safe tasks just as it did before. More importantly, they tested whether the dangerous knowledge was truly gone by simulating an attacker who could bypass the traffic director and force the model to use the original experts. In the old methods, this bypass allowed the attacker to recover the forgotten information about 11% of the time. With the new method, the recovery rate dropped to just 3%, a level comparable to models that had been completely retrained from scratch.

The study also compared their approach to a previous attempt that tried to solve the same problem using softer, less strict rules. That earlier method improved the situation slightly but still allowed the system to manipulate the routing to hide knowledge. The new approach, by using hard constraints, proved that the only way to achieve genuine unlearning in these efficient models is to prevent the routing system from being used as a shield. The researchers demonstrated that this technique works across different types of unlearning tasks, from removing hazardous cybersecurity knowledge to deleting copyrighted text. They found that the model retained its general intelligence and ability to answer safe questions much better than before, improving its performance on retained tasks by up to 89% compared to the unstable baseline.

This work highlights a critical distinction in how we secure artificial intelligence. It shows that for these complex, specialized systems, simply changing the path of information is not enough to protect privacy or safety. The knowledge must be erased from the source. By ensuring that the system cannot use routing to hide dangerous content, the researchers have provided a reliable way to remove specific memories without sacrificing the model's overall utility. Their findings suggest that future safety measures for large language models must account for these internal routing mechanisms, ensuring that the "traffic director" cannot be manipulated to hide dangerous content. The result is a more robust system where the removal of data is genuine, not just a matter of changing the map.

Technical Summary: GRIP - Algorithm-Agnostic Machine Unlearning for Mixture-of-Experts

1. Problem Statement

Machine unlearning (MU) in Mixture-of-Experts (MoE) large language models (LLMs) faces a critical structural failure mode distinct from dense architectures. In dense models, all parameters engage uniformly with inputs. In MoE models, a learned router dynamically selects a sparse subset of expert networks for each input.

Current unlearning methods applied to MoE architectures often exploit this dynamic routing as an optimization shortcut. Rather than genuinely erasing knowledge from the expert parameters, the optimizer manipulates the router to redirect "forget-set" queries away from the originally assigned, knowledgeable experts. This leads to two primary issues:

Utility Degradation: The router shifts expert selections for the "retain set" (data that should be kept), causing severe drops in model performance on retained tasks.
Incomplete Erasure (Expert Forcing): The targeted knowledge remains intact within the dormant expert parameters. Adversaries can bypass the router (a "white-box" attack) to force the model to access these specific experts, recovering sensitive or hazardous information that was supposedly unlearned.

Existing attempts to address this, such as SEUF, rely on soft regularization penalties and restrict updates to a single expert. However, empirical evidence shows these soft penalties fail to prevent routing manipulation, and restricting updates to a single expert leaves knowledge in other initially activated experts vulnerable to recovery.

2. Methodology: GRIP

The authors propose Geometric Routing Invariance Preservation (GRIP), an algorithm-agnostic framework that enforces hard geometric constraints on router updates to suppress expert selection shifts while allowing genuine knowledge erasure across all relevant experts.

Core Principle

GRIP imposes a constraint that router parameter updates ( $\Delta \Theta_\ell$ ) must preserve the expert selections of the retain set ( $D_r$ ). Formally, for any input $x \in D_r$ , the set of selected experts after unlearning ( $S^{post}_\ell(x)$ ) must equal the set before unlearning ( $S^{pre}_\ell(x)$ ). By making the "routing shortcut" infeasible on the constrained inputs, the optimization is forced to achieve the unlearning objective through knowledge erasure within the expert weights themselves.

Two Variants

GRIP offers two complementary realizations:

Method A: Training-Time Stochastic Projection (Online)
- Mechanism: At each training step, the unconstrained router update proposed by a base unlearning algorithm is projected into the null space of the retain set's routing matrix.
- Constraints: The projection enforces two conditions for each expert $j$ $j$ :
  - Equality: For inputs where expert $j$ was originally selected, the router update must result in zero change to the routing score ( $\langle \Delta \Theta, x \rangle = 0$ ).
  - Inequality: For inputs where expert $j$ was not selected, the update must not increase the score enough to breach the selection margin (i.e., the score gap between selected and non-selected experts).
- Implementation: Uses a two-stage process: (1) Orthogonal projection onto the complement of the span of selected inputs (equality), and (2) A randomized Kaczmarz-style alternating projection to satisfy inequality constraints efficiently.
- Overhead: Adds approximately 1.71× training time overhead.
Method B: Post-Training Analytical Correction (Offline)
- Mechanism: The unlearning algorithm runs without constraints first. Afterward, a single closed-form analytical correction is applied to realign the drifted router.
- Formula: Solves a least-squares problem to find a correction $\Delta_\ell$ such that $(\tilde{\Theta}_\ell + \Delta_\ell)\tilde{X}_{r,\ell} = \Theta_\ell X_{r,\ell}$ , where $\tilde{\Theta}$ and $\tilde{X}$ are the post-unlearning router and representations.
- Overhead: Reduces overhead to 1.21× training time by replacing per-step projections with a single post-hoc correction per layer.

3. Key Contributions

Identification of a Structural Vulnerability: The paper identifies that MoE unlearning is prone to "routing shortcuts" where knowledge is hidden rather than erased, and introduces the "Expert Forcing" attack as a diagnostic for this vulnerability.
Hard Geometric Constraints: Unlike prior work using soft penalties, GRIP treats routing preservation as a feasibility condition satisfied by construction via null-space projection.
Multi-Expert Erasure: GRIP allows all initially activated experts to update, ensuring knowledge is erased from the entire relevant parameter space, rather than restricting updates to a single expert.
Algorithm Agnosticism: GRIP acts as a wrapper that can be applied to any existing unlearning objective (e.g., Gradient Ascent, NPO, RMU) without modifying the underlying loss function.

4. Experimental Results

Experiments were conducted on two MoE models (Qwen3-30B-A3B and Qwen1.5-MoE-A2.7B-Chat) across hazardous knowledge removal (WMDP) and copyright unlearning (MUSE) benchmarks.

Routing Stability: GRIP restores routing stability (Jaccard similarity of expert selections) from a collapse of ~0.21 (in baselines) to ≥0.94 (and up to 1.00 with Post-Training Correction).
Utility Preservation: Retain accuracy improved by up to 89% compared to baselines. GRIP methods recovered utility preservation scores close to those of dense-architecture unlearning baselines.
Adversarial Robustness:
- White-Box (Expert Forcing): Under forced access to dormant experts, baseline methods recovered 11% of forgotten queries, SEUF recovered 8%, while GRIP reduced recovery to 3% (near noise levels).
- Black-Box (Prompt Attacks): GRIP's performance under prompt attacks remained comparable to dense-architecture baselines, indicating no significant new black-box vulnerabilities.
Comparison with SEUF: GRIP consistently outperformed SEUF across all metrics, demonstrating that soft regularization is insufficient to prevent the routing shortcut.

5. Significance and Claims

The paper claims that GRIP provides a principled and scalable route to genuine unlearning in sparse MoE architectures. By decoupling routing stability from parameter plasticity, GRIP ensures that knowledge is erased from expert parameters rather than merely hidden by routing manipulation.

The authors position GRIP not as a replacement for unlearning objectives, but as a necessary geometric constraint to make unlearning effective in the specific context of MoE models. They acknowledge limitations, including the memory overhead of caching retain representations and the empirical (rather than formal) nature of the routing preservation guarantees. However, the results suggest that hard geometric constraints are essential for addressing the unique vulnerabilities of dynamic routing in large language models.