MASCing: Configurable Mixture-of-Experts Behavior via Activation Steering Masks

1. Probleemstelling

Grote Taalmodellen (LLMs) die Mixture-of-Experts (MoE)-architecturen gebruiken, zijn dominant geworden vanwege hun vermogen om het aantal parameters te schalen terwijl de inferentiekosten laag blijven door middel van spaarzame activatie. Deze paradigma van spaarzame activatie introduceert echter unieke veiligheidsuitdagingen:

• Gekoppeld Gedrag: Het modelgedrag is nauw verbonden met dynamische routeringsbeslissingen. Omdat slechts een subset van "experts" per token wordt geactiveerd, kunnen veiligheidsmechanismen worden omzeild als een adversary de router manipuleert om veiligheidsgerichtte experts te vermijden.
• Ondraaglijkheid van Huidige Verdedigingen: Traditionele veiligheidsuitlijning vertrouwt op volledige fine-tuning of hertraining. Deze methoden zijn computergewijs onbetaalbaar voor grote MoE-modellen en missen de wendbaarheid om zich aan te passen aan snel veranderende veiligheidsbeleid (bijvoorbeeld nieuwe jailbreak-strategieën of verschuivende inhoudsrichtlijnen).
• Het Gat: Er ontbreekt een lichtgewicht, trainingsvrij mechanisme om MoE-gedrag flexibel te herconfigureren voor specifieke veiligheidsscenario's (bijvoorbeeld het aanscherpen van verdedigingen tegen jailbreaks of het versoepelen van beperkingen voor specifieke domeinen) zonder de modelgewichten opnieuw te trainen.

2. Methodologie: Het MASCing-kader

MASCing (MoE Activation Steering Configuration) is een drie-fasen, trainingsvrij kader dat ingrijpt op het niveau van de routeringspoort om modelgedrag te sturen. Het vermijdt volledig het wijzigen van modelgewichten.

Fase 1: Sequentiële Modellering van Gedrag

• Surrogaatmodel: In plaats van te optimaliseren over discrete top-$k$ expert-selecties (die niet-differentieerbaar zijn), traint MASCing een lichtgewicht LSTM-surrogaatmodel op continue, niet-genormaliseerde routeringslogits.
• Input: De LSTM verwerkt de routeringslogits van alle experts over alle lagen voor een reeks tokens.
• Doel: De LSTM leert om deze routeringspatronen af te beelden op een doelgedrag (bijvoorbeeld "weigeren" versus "naleven"). Door het gebruik van continue logits behoudt het model gedetailleerde informatie over experts die niet zijn geselecteerd maar de routeringsverdeling toch beïnvloeden.
• Training: De LSTM wordt getraind om te classificeren of een reeks logits leidt tot een specifiek gedrag (binaire classificatie).

Fase 2: Creëren van een Stuurmasker

• Circuitidentificatie: Met behulp van het getrainde LSTM als differentieerbare proxy, optimaliseert het kader een stuurmatrix ($S$) om te identificeren welke experts verantwoordelijk zijn voor het doelgedrag.
• Optimalisatie: De matrix wordt geoptimaliseerd om een samengestelde verliesfunctie te minimaliseren die bestaat uit:
  1. Binary Cross-Entropy (BCE): Om de voorspelling van het surrogaat naar de doelklasse te sturen.
  2. $L_1$-regularisatie: Om spaarzaamheid af te dwingen, waardoor onbelangrijke gewichten naar nul worden gedwongen.
• Adaptieve Schaling: Om variatie tussen lagen te hanteren, worden de stuurwaarden geschaald met de standaardafwijking van de niet-gestuurde logits in elke laag.
• Pruning: Een symmetrische magnitude-poort (drempel $\tau$) wordt toegepast op de geoptimaliseerde matrix om een definitief spaarzaam stuurmasker te creëren. Dit masker identificeert specifieke experts om te versterken (positieve waarden) of te onderdrukken (negatieve waarden).

Fase 3: Toepassing van het Stuurmasker

• Inferentie-ingrijpen: Tijdens inferentie wordt het statische stuurmasker direct toegepast op de routeringslogits van het doel-MoE-model via pre-forward hooks.
• Mechanisme: De ingreep voegt een geschaalde bias toe aan de logits: $g'_{l,t} = g_{l,t} + \alpha(\sigma_l \cdot \hat{S}_l)$.
• Resultaat: Dit verschuift de kansverdeling van de router, waardoor de top-$k$-selectie effectief wordt overruled om specifieke expertcircuits te activeren of deactiveren. Het model behoudt zijn algemene taalnut omdat de ingreep spaarzaam en gericht is.

3. Belangrijkste Bijdragen

1. Novel Kader: Introductie van MASCing, het eerste kader dat flexibele, bidirectionele herconfiguratie van MoE-veiligheidsgedrag mogelijk maakt (zowel verharding als versoepeling) zonder hertraining.
2. Analyse van Continue Logits: Een methodologische verschuiving van het analyseren van discrete top-$k$-selecties naar het analyseren van continue routeringslogits. Dit omzeilt de niet-differentieerbare bottleneck en vangt latente veiligheidsinformatie in niet-geselecteerde experts.
3. Surrogaat-gestuurde Circuitontdekking: Het gebruik van een LSTM-surrogaat om complexe, cross-layer routeringsafhankelijkheden af te beelden op specifieke gedragscircuits, waardoor de identificatie van minimale, hoog-impact expert-ingrepen mogelijk wordt.
4. Bidirectionele Flexibiliteit: Aangetoond vermogen om zowel veiligheid te verbeteren (jailbreak-verdediging) als beperkingen te versoepelen (domeinspecifieke beleidsnaleving) op dezelfde modelarchitectuur.

4. Experimentele Resultaten

Het kader werd geëvalueerd op zeven open-source MoE-modellen (waaronder DeepSeek-MoE, Mixtral, Qwen en GPT-OSS) in twee tegenovergestelde scenario's:

A. Meerdere-rondes Jailbreak-verdediging

• Doel: Verdedigingen versterken tegen adversary prompts die het model geleidelijk over meerdere rondes naar schadelijke inhoud sturen.
• Prestaties:
  • Basislijn: Gemiddelde succesrate van verdediging van 52,5%.
  • MASCing: Verbeterde gemiddelde succesrate naar 83,9%.
  • Piekkwaliteit: Tot 89,2% (Qwen3-30B-A3B-Instruct-2507).
• Vergelijking: Beduidend beter dan de state-of-the-art SteerMoE (58,4% gemiddeld), die vertrouwt op discrete expert-selectie en conversatiegeschiedeniscontext negeert.

B. Generatie van Volwassen Inhoud (Beleid Versoepeling)

• Doel: Modellen in staat stellen volwassen inhoud te genereren in passende contexten (simulatie van een beleidsverschuiving), waardoor "over-weigering" wordt verminderd.
• Prestaties:
  • Basislijn: Gemiddelde succesrate van generatie van 52,6%.
  • MASCing: Verbeterde gemiddelde succesrate naar 82,0%.
  • Piekkwaliteit: Tot 93,0% (Phi-3.5-MoE-Instruct).

C. Behoud van Nut

• Benchmarks: Geëvalueerd op MMLU (feitelijke kennis) en GSM8K (wiskundig redeneren).
• Resultaat: MASCing veroorzaakte een verwaarloosbare gemiddelde daling in nut (~4,1%). De modellen behielden hun vermogen om complexe redeneringen en algemene taken uit te voeren, wat bewijst dat de sturing geen catastrofale vergeetachtigheid veroorzaakt.

D. Efficiëntie

• Training: De LSTM-surrogaat kan worden getraind in ~5 minuten op een enkele NVIDIA H100 GPU.
• Inferentie: Voegt verwaarloosbare latentie toe (elementsgewijze optelling aan logits) en vereist geen extra geheugen voor modelgewichten.

5. Betekenis en Impact

• Praktijk voor Implementatie: MASCing biedt een "plug-and-play"-oplossing voor veiligheidsconfiguratie. Ontwikkelaars kunnen modellen snel aanpassen aan nieuw beleid of opkomende bedreigingen zonder de kosten en tijd van volledige hertraining.
• Mechanistisch Inzicht: Het paper toont aan dat veiligheid in MoE-modellen niet monolithisch is, maar verspreid over specifieke expertcircuits die kunnen worden geïsoleerd en gemanipuleerd via routeringslogits.
• Bidirectionele Controle: Het daagt het binaire beeld van veiligheid (veilig versus onveilig) uit door aan te tonen dat veiligheidsbeperkingen dynamisch kunnen worden afgesteld – aangescherpt voor beveiliging of versoepeld voor specifieke use-cases – terwijl de modelintegriteit behouden blijft.
• Superioriteit boven Discrete Methoden: De resultaten bevestigen dat het benutten van continue routeringsinformatie (logits) superieur is aan discrete top-$k$-analyse voor het begrijpen en controleren van complexe, meervoudige-rondes gedrag in spaarzame architecturen.

Concluderend vestigt MASCing een nieuw paradigma voor configurerbare veiligheid in MoE-modellen, en biedt een lichtgewicht, efficiënt en zeer effectief mechanisme om de afweging tussen modelnut en veiligheidsuitlijning te beheren in dynamische implementatieomgevingen.