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

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

Auteurs: Jiacheng Liang, Yuhui Wang, Tanqiu Jiang & Ting Wang (Stony Brook University)

---

1. Het Probleem: Alignement-Shortcuts in MoE-modellen

Mixture-of-Experts (MoE) taalmodellen gebruiken een mechanisme waarbij elk invoertoken slechts door een klein subset van experts wordt verwerkt (sparse routing). Hoewel dit de schaalbaarheid en efficiëntie verbetert, introduceert het unieke uitdagingen voor veiligheidsalignement (safety alignment).

De auteurs identificeren een fundamenteel falingspatroon bij het toepassen van standaard full-parameter fine-tuning op MoE-modellen:

• Alignement-Shortcuts: In plaats van de specifieke experts die verantwoordelijk zijn voor onveilig gedrag ("Safety-Critical Experts" of SCEs) daadwerkelijk te repareren, leert het model "shortcuts". Het model bereikt een schijnbare verbetering in veiligheid door:
  1. De router (routing mechanisme) te manipuleren zodat schadelijke inputs worden omgeleid naar reeds veilige experts.
  2. Veilige experts dominant te maken, wat leidt tot overmatig weigeren van vragen (over-refusal).
• Gevolg: De onderliggende onveilige experts blijven intact en kunnen worden geactiveerd door adaptieve of onbekende jailbreak-aanvallen. De veiligheid is dus afhankelijk van de routering en niet van een fundamentele herstel van de kennis in de experts.

2. Methodologie: RASA Framework

Om dit probleem op te lossen, stellen de auteurs RASA (Routing-Aware Safety Alignment) voor. Dit is een doelgericht, expert-niveau alignement framework dat bestaat uit een twee-niveau optimalisatie (bi-level optimization) met een alternerende strategie.

Het doel is om de experts die verantwoordelijk zijn voor onveiligheid expliciet te repareren, terwijl de routering consistent wordt gehouden om shortcuts te voorkomen.

Stap 1: Identificatie van Safety-Critical Experts (SCEs)

• Het systeem analyseert de activatie van experts op basis van een Adversarial Activation Discrepancy (AAD).
• Er worden gepaarde prompts gebruikt:
  • Anchor Prompts: Schadelijke intenties waar het model correct op weigert.
  • Adversarial Prompts: Jailbreak-varianten van dezelfde intenties die het model succesvol omzeilen.
• Experts die significant vaker worden geactiveerd door de jailbreak-prompts dan door de veilige weigeringen, worden geïdentificeerd als SCEs.

Stap 2: Alternerende Optimalisatie

RASA splitst het trainingsproces op in twee fasen die afwisselend worden uitgevoerd:

1. Interne Optimalisatie: Selective Expert Fine-tuning (SCE-FT)

   • De router parameters worden bevroren.
   • Alleen de geïdentificeerde SCEs worden gefine-tuned om weigeringsgedrag te injecteren.
   • De overige (veilige) experts blijven onveranderd om de algemene capaciteiten van het model te behouden.
   • Doel: Directe reparatie van de parameters die onveilig gedrag veroorzaken.

2. Externe Optimalisatie: Router Consistency

   • De expert parameters worden bevroren (inclusief de zojuist gerepareerde SCEs).
   • De router parameters worden geoptimaliseerd om de routing-distributie van adversarial inputs te aligneren met die van veilige (anchor) inputs.
   • Dit wordt gedaan via een KL-divergentie loss tussen de routering van veilige en onveilige inputs.
   • Doel: Zorgen dat jailbreak-inputs niet via een andere route worden omgeleid om de gerepareerde experts te vermijden, en voorkomen dat normale vragen per ongeluk naar over-gealigneerde experts worden gestuurd (over-refusal).

3. Belangrijkste Bijdragen

1. Empirische Observatie: Het paper documenteert voor het eerst het fenomeen van "alignement shortcuts" in MoE-modellen, waarbij full-parameter fine-tuning faalt omdat het de onveilige experts niet repareert, maar alleen de routering manipuleert.
2. RASA Framework: Een nieuw, routing-bewust framework dat expert-reparatie en routeringscontrole ontkoppelt. Het is architectonisch neutraal (vereist geen extra experts of modelwijzigingen) en werkt puur via training.
3. Uitgebreide Validatie: De methode wordt getest op twee representatieve MoE-architecturen (Qwen3-30B-A3B en OLMoE-1B-7B-0125) en diverse jailbreak-aanvallen.

4. Resultaten

De experimenten tonen aan dat RASA superieur is aan bestaande methoden zoals SteerMoE (inference-time steering) en SafeMoE-Fine-tuning (routing regularisatie).

• Robuustheid tegen Jailbreaks: RASA bereikt bijna perfecte weerstand (harmlessness rates van ~1.0) tegen diverse jailbreak-aanvallen (zoals FlipAttack, DeepInception, Persuasion), zowel bij enkelvoudige als gemengde aanvallen.
• Cross-Attack Generalisatie: Zelfs zonder expliciete training op alle aanvalstypen, generaliseert RASA goed naar onbekende aanvallen omdat het de gedeelde "Safety-Critical Experts" repareert.
• Behoud van Algemene Capaciteiten: In tegenstelling tot full-parameter fine-tuning, behoudt RASA de prestaties op benchmarks zoals MMLU, GSM8K en TruthfulQA.
• Gecontroleerde Weigering (Over-refusal): RASA voorkomt dat het model te vaak weigert op veilige vragen (gecontroleerd via XsTest), wat vaak een neveneffect is van agressieve alignement.
• Causaal Bewijs: Door experimenten met "routing restoration" (waarbij de oorspronkelijke routering wordt teruggezet in een getraind model) bewijzen de auteurs dat de veiligheid van standaard fine-tuning instort, terwijl RASA robuust blijft. Dit bevestigt dat RASA de experts daadwerkelijk repareert.

5. Significantie en Conclusie

Dit paper biedt een cruciale inzichten voor de veiligheid van moderne, schaalbare taalmodellen. Het toont aan dat voor MoE-modellen targeted expert repair effectiever is dan globale parameterupdates.

• Efficiëntie: De methode is data-efficiënt en vereist slechts een fractie van de beschikbare adversarial samples.
• Toekomstperspectief: RASA biedt een praktische, architectuur-bewuste oplossing die de complexiteit van MoE-modellen respecteert in plaats van ze te omzeilen. Het benadrukt dat robuuste veiligheid vereist dat men ingrijpt op het niveau van de interne mechanismen (experts) en niet alleen op het oppervlak van de output of routering.

De code is beschikbaar gesteld, wat de reproduceerbaarheid en verdere ontwikkeling van deze aanpak mogelijk maakt.