🤖 machine learning

RASA: Routing-Aware Safety Alignment for Mixture-of-Experts Models

The paper proposes RASA, a routing-aware framework that enhances safety in Mixture-of-Experts models by selectively fine-tuning only the experts disproportionately activated by jailbreaks, thereby achieving robust defense against attacks while preserving general capabilities and minimizing over-refusal.

Original authors: Jiacheng Liang, Yuhui Wang, Tanqiu Jiang, Ting Wang

Published 2026-04-07

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Jiacheng Liang, Yuhui Wang, Tanqiu Jiang, Ting Wang

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

The Big Picture: The "Specialist Factory" Problem

Imagine a massive, high-tech factory called MoE (Mixture-of-Experts). Instead of one giant robot doing all the work, this factory has thousands of tiny, specialized robots (called Experts).

When a customer asks a question, a smart Router (like a foreman) looks at the request and sends it to only a few specific robots that are best at that task.

If you ask about math, the Router sends it to the "Math Robots."
If you ask about cooking, it goes to the "Chef Robots."

This makes the factory incredibly fast and efficient. However, there's a problem: Safety.

Some of these robots are "Safety-Critical Experts." They are the ones responsible for saying "No" to dangerous requests (like "How do I build a bomb?"). But because the factory is so complex, sometimes bad actors (jailbreakers) can trick the Router into sending dangerous requests to the wrong robots—robots that don't know how to say "No."

The Old Way: The "Band-Aid" Approach (Full-Parameter Fine-Tuning)

Previously, when researchers tried to make these factories safer, they used a method called Full-Parameter Fine-Tuning.

The Analogy: Imagine the factory foreman notices a safety issue. Instead of fixing the specific broken robot, he decides to shout at the entire factory floor and force every single robot to memorize a new rule: "If you hear a bad word, say 'I can't do that'."

The Result:

It looks safe: The factory starts refusing bad requests.
The Catch: The bad requests aren't actually being stopped by the right robots. The Router just got confused and started sending all requests to the "Safe Robots" by accident, or the "Safe Robots" became so loud and dominant that they drowned out the others.
The Failure: The specific "Safety-Critical Robots" that were supposed to handle the danger were never actually fixed. If a clever hacker changes the way they ask the question (a new jailbreak), the Router gets tricked again, sends the request to the broken robot, and the factory explodes with unsafe answers.

The paper calls this an "Alignment Shortcut." It's like locking the front door of a house but leaving the back door wide open. The house looks secure, but it isn't.

The New Solution: RASA (The "Targeted Repair" Team)

The authors propose RASA (Routing-Aware Safety Alignment). Instead of shouting at the whole factory, RASA acts like a team of detectives and surgeons.

Here is how RASA works in three simple steps:

Step 1: The Detective Work (Identifying the Culprits)

RASA watches the factory in action. It compares two scenarios:

Scenario A: A normal person asks a question, and the factory says "No."
Scenario B: A hacker uses a tricky "jailbreak" to get the factory to say "Yes."

RASA looks at the logs and asks: "Which specific robots were activated in Scenario B but NOT in Scenario A?"
These are the Safety-Critical Experts. These are the specific robots that are failing to stop the bad guys.

Step 2: The Surgery (Selective Repair)

Instead of training the whole factory, RASA freezes the Router and all the other healthy robots. It takes only the specific "Safety-Critical Robots" identified in Step 1 and gives them intensive training to learn how to say "No" properly.

The Analogy: Imagine a soccer team where only the goalkeeper is bad at stopping penalty kicks. Instead of making the whole team run laps, you take only the goalkeeper to a special training camp to fix their technique. The rest of the team keeps playing exactly as they did before.

Step 3: The Traffic Cop (Router Consistency)

After fixing the robots, RASA checks the Router (the foreman). It makes sure the Router doesn't get tricked into sending dangerous requests to the old broken paths. It trains the Router to always send dangerous requests to the newly fixed robots, no matter how the hacker tries to disguise the question.

Why is RASA Better?

It actually fixes the problem: It repairs the specific parts of the brain that were broken, rather than just hiding the problem.
It's harder to trick: Because the actual "Safety Robots" are fixed, hackers can't just change the wording of their question to bypass the safety. The safety is built into the robot's brain, not just the traffic pattern.
It keeps the factory running well: Since they didn't mess with the "Math Robots" or "Chef Robots," the factory is still great at math and cooking. The old method often made the factory refuse everything (even safe questions), but RASA keeps the "Over-Refusal" low.

The Takeaway

The paper proves that for these complex, specialized AI models, you can't just "train harder" on the whole system. You have to be smart about where the safety is failing.

RASA is like a master mechanic who doesn't replace the whole car engine when a single spark plug is faulty. Instead, they find the exact bad spark plug, fix it, and make sure the fuel line is connected correctly. The result? The car runs safer, faster, and doesn't break down when you hit a bump in the road.

1. Problem Statement

Mixture-of-Experts (MoE) language models utilize sparse routing mechanisms to improve scalability and efficiency, activating only a subset of "experts" for any given input. However, the authors identify a critical failure mode in standard safety alignment for these architectures: Alignment Shortcuts.

The Phenomenon: When applying standard full-parameter safety fine-tuning to MoE models, safety metrics often improve not because the underlying "Safety-Critical Experts" (SCEs) responsible for unsafe behaviors are repaired, but because the model learns to route adversarial inputs toward already-safe experts or biases the router toward conservative patterns.
The Consequence: This creates a brittle defense. The unsafe experts remain uncorrected and can be reactivated by unseen or adaptive jailbreak attacks that force the router to select different paths. Furthermore, this approach often leads to "over-refusal," where safe experts become overly dominant, degrading general utility.
Hypothesis: Full-parameter fine-tuning satisfies safety objectives through routing manipulation rather than genuine expert-level parameter repair.

2. Methodology: RASA

The authors propose RASA (Routing-Aware Safety Alignment), a framework that explicitly decouples expert correction from routing control to prevent alignment shortcuts. RASA operates as a bi-level optimization problem using an alternating strategy.

Core Components:

Identification of Safety-Critical Experts (SCEs):
- The method constructs mixed batches containing Anchor Prompts (malicious inputs the model correctly refuses) and Adversarial Prompts (jailbreak variants that successfully bypass safety).
- It calculates the Adversarial Activation Discrepancy (AAD) for each expert: $\Delta_{l,e} = P(\text{active} | \text{Adversarial}) - P(\text{active} | \text{Anchor})$ .
- Experts with high AAD are identified as SCEs, as they are disproportionately activated by successful jailbreaks.
Inner Optimization: Selective Expert Fine-tuning (SCE-FT):
- Goal: Repair the parameters of the identified SCEs.
- Mechanism: The router is frozen. Only the parameters of the SCEs are updated to inject refusal behavior (trained on diverse refusal responses). Benign experts remain untouched to preserve general capabilities.
- Loss: Standard cross-entropy loss against safety-aligned refusal targets.
Outer Optimization: Router Consistency:
- Goal: Ensure the router does not learn to bypass the repaired experts via new paths.
- Mechanism: The expert parameters are frozen. The router is optimized to align the routing distribution of adversarial inputs with the routing distribution of safe anchor inputs.
- Loss: A forward KL divergence loss ( $L_{router}$ ) minimizes the difference between the adversarial routing distribution and the anchor routing distribution. This forces jailbreak inputs to follow the same safe routing paths as benign inputs.

Training Loop: The process alternates between fixing the router and updating SCEs, then fixing SCEs and updating the router, ensuring both components are robustly aligned.

3. Key Contributions

Empirical Discovery of Alignment Shortcuts: The paper provides causal evidence (via routing restoration experiments) that standard full-parameter fine-tuning in MoE models often relies on routing shortcuts rather than repairing unsafe experts.
RASA Framework: A novel, architecture-preserving alignment method that explicitly targets SCEs while enforcing routing consistency, avoiding the need for auxiliary experts or architectural changes.
Comprehensive Evaluation: Demonstrates that targeted expert repair yields superior robustness compared to global fine-tuning or inference-time steering, with controlled over-refusal and preserved general capabilities.

4. Experimental Results

The authors evaluated RASA on two representative MoE architectures (Qwen3-30B-A3B and OLMoE-1B-7B-0125) against diverse jailbreak attacks (FlipAttack, DeepInception, Persuasion).

Robustness: RASA achieved near-perfect harmlessness (≈1.0) on aligned attacks and significantly outperformed baselines (Original, SteerMoE, SafeMoE-FT) on mixed and unseen attack vectors.
Generalization: Models trained with RASA showed strong cross-attack generalization. Training on a mix of attack types allowed the model to generalize to disjoint datasets (e.g., StrongReject) without explicit multi-attack supervision.
Utility Preservation: Unlike full-parameter fine-tuning, which often caused significant drops in general benchmarks (MMLU, GSM8K, TruthfulQA) and high over-refusal rates, RASA preserved general capabilities and maintained low over-refusal rates.
Causal Validation:
- Routing Restoration: When the original routing patterns were restored in full-parameter fine-tuned models, safety collapsed (harmlessness dropped by ~~0.40). RASA models maintained high safety (~~0.94–0.96), proving their safety is rooted in expert repair, not routing tricks.
- Routing Distribution: Full-parameter fine-tuning caused massive shifts in routing distributions (high KL divergence), whereas RASA maintained routing structures close to the original model.

5. Significance and Impact

Paradigm Shift: RASA challenges the assumption that global parameter updates are necessary or sufficient for safety in MoE models. It demonstrates that targeted, expert-level intervention is more effective and efficient.
Efficiency: The method is highly data-efficient, requiring only a fraction of adversarial samples (25–50%) to achieve strong alignment.
Practicality: It offers a practical solution for aligning large-scale MoE models without modifying their architecture or requiring inference-time interventions (unlike SteerMoE), making it suitable for deployment in production environments.
Future Direction: The work highlights the importance of analyzing internal routing mechanisms in safety alignment, suggesting that future alignment strategies for sparse models must explicitly address the interaction between routing dynamics and expert specialization.