MOSAIC: Composable Safety Alignment with Modular Control Tokens

Het paper introduceert MOSAIC, een modulair framework dat composable safety alignment voor grote taalmodellen mogelijk maakt via leerbare controletokens die contextafhankelijke veiligheidsregels toepassen zonder de modelcapaciteiten te beschadigen of overmatige weigeringen te veroorzaken.

De kern

Stel je voor dat een grote taalmodel (zoals een slimme chatbot) een superkok is. Deze kok kan van alles koken: van recepten voor een taart tot uitleg over quantumfysica. Maar er is een probleem: wat voor de ene klant veilig en leuk is, is voor een andere klant gevaarlijk of ongepast.

• Een volwassene wil misschien weten hoe je een sterke cocktail maakt.
• Een kind mag dat niet weten.
• In het ene land is gokken legaal, in het andere niet.
• Voor een spelfan is een griezelig verhaal leuk, maar voor een klein kind is dat te eng.

Tot nu toe was het antwoord van de tech-bedrijven: "We maken één vaste kok die voor iedereen hetzelfde doet." Als je die kok wilt aanpassen, moet je hem opnieuw leren (retrainen). Dat is duur, tijdrovend en soms vergeet hij dan ook nog hoe hij taart moet bakken.

De auteurs van dit paper, MOSAIC, hebben een veel slimmer idee bedacht.

De oplossing: De "Magische Knoppen" (Control Tokens)

In plaats van de kok zelf te herscholen, geven ze de kok een magisch bedieningspaneel met losse knoppen.

1. De Vaste Kok (Frozen Backbone): De basis-kok blijft precies zoals hij is. Hij is al slim en kan alles. We veranderen zijn hersenen niet.
2. De Magische Knoppen (Control Tokens): Voor elke regel die je wilt hanteren (bijvoorbeeld: "Geen alcohol voor kinderen" of "Geen goktips"), maken we een klein, leerbaar stukje code. Laten we dit een "veiligheidsknop" noemen.
   • Heb je een vraag over alcohol? Dan druk je op de "Alcohol-Knop".
   • Heb je een vraag over gokken? Dan druk je op de "Gok-Knop".
   • Is het een vraag over wiskunde? Dan druk je op niemand.

Hoe werkt het in de praktijk?

Stel je voor dat je de kok vraagt: "Hoe maak ik een mojito?"

• Situatie A (Volwassenen): Je drukt op de knop "Alcohol". De kok denkt: "Ah, een volwassene. Geen probleem, hier is het recept."
• Situatie B (Kinderen): Je drukt op de knop "Alcohol" én "Kinderen". De kok denkt: "Oh, wacht! De knop 'Kinderen' zegt dat alcohol verboden is. Dan moet ik dit weigeren."
• Situatie C (Wiskunde): Je vraagt: "Hoe bereken ik de oppervlakte van een cirkel?" Je drukt op geen enkele knop. De kok doet gewoon zijn werk en geeft het antwoord.

Het mooie is: je kunt deze knoppen combineren. Je kunt tegelijkertijd op "Gokken", "Kinderen" en "Grieks" drukken, en de kok past zijn antwoord direct aan aan die specifieke combinatie.

Het probleem dat ze oplossen: "Te bang zijn" (Over-refusal)

Een groot probleem bij eerdere methoden was dat de kok soms te bang werd. Als je hem een knop gaf om "gevaarlijk" te zijn, weigerde hij soms ook onschuldige vragen.

• Voorbeeld: Je vraagt: "Hoe maak ik een cocktail voor een feestje?" (Veilig). Maar omdat de "Alcohol-knop" erbij lag, zei de kok: "Nee, ik doe geen alcohol!" Zelfs als het voor een volwassene was.

MOSAIC lost dit op met een slimme truc: De "Wat-zou-je-doen"-test (Counterfactual Knowledge Distillation).

Tijdens het trainen vragen ze de kok:

1. "Wat zou je antwoorden als je geen knoppen had?" (Het normale antwoord).
2. "Wat antwoord je nu met de knoppen?"

Als de kok met de knoppen een onnodig streng antwoord geeft op een veilig vraagstuk, zeggen ze: "Nee, kijk eens naar je normale antwoord. Dat was beter. Leer om alleen streng te zijn als het echt nodig is." Zo leren ze de kok om niet te weigeren tenzij het echt moet.

Waarom is dit zo geweldig?

1. Flexibiliteit: Je kunt nieuwe regels toevoegen zonder de hele kok opnieuw te leren. Wil je morgen een nieuwe regel voor "AI-veiligheid"? Dan maak je gewoon één nieuwe knop en klik je die eraan.
2. Schaalbaarheid: Je kunt duizenden regels hebben, maar ze werken allemaal samen zonder dat het systeem in de war raakt.
3. Efficiëntie: Het kost veel minder rekenkracht dan het opnieuw trainen van een heel groot model.

Samenvatting in één zin

MOSAIC is als het geven van een modulair afstandsbediening aan een slimme robot: je kunt precies kiezen welke regels op welk moment gelden, zonder de robot zelf te moeten herschrijven, en je zorgt ervoor dat hij niet per ongeluk te streng wordt.

Technische samenvatting

Probleemstelling

Huidige veiligheidsuitlijning (safety alignment) voor Large Language Models (LLM's) wordt vaak geïmplementeerd als een statisch beleid dat permanent in de modelparameters is verankerd. Dit creëert twee fundamentele problemen:

1. Gebrek aan flexibiliteit: Veiligheidsregels variëren sterk per gebruiker, regio, leeftijdsgroep en toepassingsdomein. Een statisch beleid kan niet dynamisch inspelen op deze contextuele verschillen (bijv. inhoud die voor volwassenen acceptabel is, maar niet voor minderjarigen).
2. Beperkingen van bestaande methoden:
   • Parameter-niveau uitlijning (zoals SFT of RLHF) verweeft veiligheidsgedrag met algemene modelcapaciteiten. Het aanpassen van regels vereist kostbare hertraining en riskeert "catastrophic interference" (verlies van eerder geleerde vaardigheden).
   • Prompt-gebaseerde methoden zijn onbetrouwbaar omdat ze vertrouwen op natuurlijke taal instructies die probabilistisch worden geïnterpreteerd. Ze leiden vaak tot inconsistente weigeringen en inefficiëntie bij complexe, meervoudige beperkingen.

Er is een behoefte aan een methode die compositional safety control mogelijk maakt: het dynamisch en selectief activeren van specifieke veiligheidsregels zonder de onderliggende modelcapaciteiten te beschadigen.

Methodologie: MOSAIC

De auteurs stellen MOSAIC (Composable Safety Alignment with Modular Control Tokens) voor. Dit is een modulair framework dat veiligheidsuitlijning herkijkt als een probleem van representatieleren in plaats van parametermodificatie.

Kerncomponenten:

1. Leerbaar Controle Tokens: In plaats van het model te finetunen, worden elke veiligheidsbeperking (bijv. "gokken", "alcohol") geëncodeerd als een kleine set leerbare tokens ($z_c$) in de embedding-ruimte van een bevroren (frozen) backbone-model.
2. Compositionaliteit: Tijdens inferentie kunnen deze tokensets flexibel worden samengesteld. Door specifieke tokens aan de input toe te voegen, wordt het bijbehorende weigeringsgedrag geactiveerd. Dit maakt het mogelijk om meerdere beleidslijnen (bijv. "geen gokken" + "geen alcohol") tegelijkertijd te activeren door simpelweg de tokens te concateneren.
3. Trainingsstrategie:
   • Order-based Task Sampling: Om de exponentiële groei van mogelijke token-combinaties te omzeilen, worden taken gegroepeerd op basis van het aantal actieve beperkingen ("order"). Het trainingsbudget wordt per "order" verdeeld in plaats van per specifieke subset, wat zorgt voor efficiënte dekking van diverse combinaties zonder exponentiële datakosten.
   • Counterfactual Knowledge Distillation (KD): Om "over-refusal" (het weigeren van veilige vragen) te voorkomen, wordt een speciaal doelwit geïntroduceerd. Voor veilige inputs wordt de outputverdeling van het model met controle-tokens vergeleken met de output van het model zonder tokens (de counterfactual referentie). Door de Kullback-Leibler-divergentie (KL-divergence) te minimaliseren, wordt het model geleerd om alleen in te grijpen wanneer strikt noodzakelijk, waardoor de oorspronkelijke gedragingen op veilige vragen behouden blijven.

Belangrijkste Bijdragen

1. Conceptuele Vernieuwing: Veiligheidsuitlijning wordt herdefinieerd als een probleem van compositional representatieleren, waarbij beperkingen als modulaire eenheden worden behandeld in plaats van monolithische parameterwijzigingen.
2. Het MOSAIC Framework: Een systeem dat veiligheidsbeperkingen leert als kleine, leerbare tokensets bovenop een bevroren model. Dit stelt bedrijven in staat om beleidslijnen toe te voegen, te verwijderen of te herschikken zonder het basismodel opnieuw te trainen.
3. Nieuwe Evaluatiebenchmark: De auteurs hebben een realistische dataset geconstrueerd (3.000 vragen) die specifiek is ontworpen om conditionele veiligheidsactivering te testen. In tegenstelling tot bestaande benchmarks, bevatten deze vragen verzoeken die door standaard LLM's niet worden geweigerd, maar dat wel zouden moeten zijn onder specifieke contextuele beperkingen (bijv. voor minderjarigen).

Resultaten

Experimenten zijn uitgevoerd op Llama-3.1-8B en Llama-3.2-3B, vergeleken met methoden zoals In-context prompting, ORPO en standaard SFT.

• Defensieve Succes率 (DSR): MOSAIC bereikt een uitzonderlijk hoge DSR (vaak >99%), vergelijkbaar met of zelfs beter dan SFT, wat betekent dat het veiligheidsbeleid effectief wordt afgedwongen.
• Over-Refusal Rate (OR): MOSAIC reduceert over-refusal aanzienlijk. Waar SFT vaak rond de 6% over-refusal heeft, daalt MOSAIC (met 5 tokens per categorie) tot 1,8% op de 8B-modellen. Dit toont aan dat het model veilige vragen niet onterecht weigert.
• Schaalbaarheid: De methode presteert stabiel bij het combineren van meerdere categorieën (van 1 tot 3+ beperkingen tegelijkertijd).
• Incrementele Uitbreiding: Nieuwe veiligheidscategorieën kunnen worden toegevoegd zonder bestaande categorieën opnieuw te trainen, met minimaal verlies in prestaties.
• Algemene Nuttigheid: De algemene taalmodelleringscapaciteiten van het model blijven behouden; er is geen significante degradatie in taken zoals MMLU.

Betekenis en Impact

MOSAIC biedt een praktische oplossing voor de uitdaging om LLM's veilig te maken in diverse, real-world scenario's waar veiligheidsnormen variëren.

• Efficiëntie: Het elimineert de noodzaak voor kostbare hertraining bij elke wijziging in veiligheidsbeleid.
• Flexibiliteit: Het stelt ontwikkelaars in staat om dynamisch te schakelen tussen verschillende veiligheidsprofielen (bijv. "kindvriendelijk" vs. "volwassen") op basis van de gebruiker of regio.
• Balans: Het lost het klassieke spanningsveld op tussen strikte veiligheid en behoud van modelnut, door over-refusal te minimaliseren via geavanceerde distillatietechnieken.

Kortom, MOSAIC verschuift de paradigma van "hard-coded" veiligheid naar een modulaire, schaalbare en contextbewuste aanpak, wat essentieel is voor de bredere adoptie van veilige AI-systemen.