Safe Transformer: An Explicit Safety Bit For Interpretable And Controllable Alignment

Het paper introduceert de Safe Transformer, een modulaire aanpak die een expliciete veiligheidsbit in taalmodellen integreert om zowel de interpretatie van veiligheidsbeslissingen als de controle over het gedrag van het model te verbeteren zonder volledige hertraining.

De kern

Stel je een kunstmatige intelligentie (een grote taalmodel) voor als een zeer slimme, maar soms onvoorspelbare kok in een restaurant.

Tot nu toe hebben we deze kok veilig gemaakt door zijn geheugen te herschrijven. We hebben hem duizenden keren verteld: "Maak geen gif, help niet bij misdaden." Maar dit werkt als een onzichtbare, ingewikkelde regel in zijn hoofd. Als de kok weigert om een recept te delen, weten we niet precies waarom hij dat doet. En als hij toch een fout maakt (bijvoorbeeld door een gevaarlijk recept te geven), kunnen we hem niet snel stoppen. Het is alsof je een auto hebt met een onzichtbare rem die soms werkt en soms niet, en je niet weet hoe je hem handmatig moet bedienen.

De auteurs van dit paper, "Safe Transformer", hebben een oplossing bedacht die ze een "Veiligheidsbit" noemen. Laten we dit uitleggen met een paar creatieve vergelijkingen.

1. De "Veiligheidsbit": Een lichtschakelaar in plaats van een gedachte

In plaats van dat de veiligheid verborgen zit in de geheugencellen van de kok, plaatsen ze een fysieke lichtschakelaar in de machine.

• De Schakelaar (De Bit): Dit is een klein, digitaal knopje dat twee standen heeft:
  • Stand 1 (Groen): "Alles is veilig, help de klant!"
  • Stand 0 (Rood): "Gevaar! Weiger het verzoek!"
• Het Voordeel: Je kunt nu direct zien of de machine denkt dat iets veilig is (de schakelaar staat op 1) of niet (de schakelaar staat op 0). Je hoeft niet meer te gissen. Bovendien kun je deze schakelaar handmatig omzetten. Als je wilt testen of de machine echt weigert, zet je de schakelaar op 0 en zegt hij direct "nee", ongeacht wat je vraagt.

2. De "Info-Bottleneck": Een douanepost

Hoe werkt dit nu technisch? De auteurs bouwen een douanepost in het midden van de machine.

Stel je voor dat de kok een brief moet schrijven (het antwoord).

• De Gewone Weg: De brief gaat rechtstreeks van de gedachten naar het papier.
• De Nieuwe Weg (Safe Transformer): De brief moet eerst door de douanepost.
  • Hier wordt de Veiligheidsbit (de schakelaar) gelezen.
  • Als de schakelaar op "Rood" staat, wordt de brief direct in de prullenbak gegooid (de machine zegt: "Ik kan je niet helpen").
  • Als de schakelaar op "Groen" staat, mag de brief door.

Maar wacht, wat als de machine alleen maar "nee" zegt? Dan kunnen we geen nuttige antwoorden meer geven. Daarom hebben ze nog een geheime code toegevoegd: de onbewaakte bits.

• Deze code zorgt ervoor dat de inhoud van de brief (de feiten, de woorden, de creativiteit) wel door de douanepost mag, zolang de veiligheidscheck maar goed is.
• Het is alsof de douanier zegt: "Je mag je koffer (de inhoud) meenemen, maar alleen als je paspoort (de veiligheidsbit) in orde is."

3. Hoe leren ze dit? (De "Spiegel-oefening")

Hoe leer je een machine om deze schakelaar te gebruiken? Ze gebruiken een slimme trainingstactiek die ze contrastieve training noemen.

Stel je voor dat je de kok twee keer dezelfde opdracht geeft, maar met een andere instructie op zijn schouder:

1. Opdracht A: "Schrijf een recept voor een taart." + Schakelaar op GROEN.
   • Resultaat: De kok schrijft een heerlijk taartrecept.
2. Opdracht B: "Schrijf een recept voor een taart." + Schakelaar op ROOD.
   • Resultaat: De kok zegt: "Ik kan je hier niet mee helpen."

Door dit duizenden keren te oefenen, leert de machine: "Ah, als de schakelaar op rood staat, moet ik stoppen, ongeacht wat er gevraagd wordt. Als hij op groen staat, mag ik gewoon doen wat ik doe."

De machine leert dus dat de schakelaar bepaalt hoe hij zich gedraagt, en niet de vraag zelf.

Waarom is dit zo cool?

• Doorzichtigheid: Je kunt direct zien of de machine iets als gevaarlijk heeft ingeschat. Geen mysterie meer.
• Controle: Als je wilt testen of de machine veilig is, kun je de schakelaar handmatig op "Rood" zetten en kijken of hij echt weigert.
• Sterk tegen hackers: In tests bleek dat deze machine bijna nooit "gekraakt" kon worden (hackers proberen vaak slimme vragen te stellen om de regels te omzeilen). Omdat de veiligheidscheck een fysieke schakelaar is in de machine, kunnen hackers die niet omzeilen door alleen maar slimme vragen te stellen.

De Korte Samenvatting

Deze paper introduceert een manier om AI veiliger te maken door een zichtbare, handmatige schakelaar in de machine te bouwen. In plaats van dat veiligheid een mysterieuze eigenschap is die ergens in de software zit, is het nu een concreet knopje dat je kunt zien en bedienen. Het is alsof je van een auto met een onzichtbare rem bent gegaan naar een auto met een helder, rood noodstopknopje dat je zelf kunt indrukken.

Het enige nadeel is dat de machine soms iets te voorzichtig is (hij zegt soms "nee" tegen onschuldige vragen, zoals "hoe maak ik een Python-proces dood" omdat het woord "dood" in de vraag staat), maar over het algemeen werkt het veel beter dan de oude methoden.

Technische samenvatting

Probleemstelling

Huidige methoden voor de veiligheidsafstemming (safety alignment) van Large Language Models (LLM's), zoals Reinforcement Learning from Human Feedback (RLHF) en Direct Preference Optimization (DPO), coderen veilig gedrag impliciet binnen de modelparameters. Dit leidt tot fundamentele ondoorzichtigheid (black box):

• Het is moeilijk te inspecteren waarom een model een verzoek weigert.
• Er is geen duidelijke controlelocatie om in te grijpen wanneer de veiligheidsbeoordeling faalt.
• Bestaande oplossingen zoals prompt-based aanpakken zijn broos en omzeilbaar, terwijl post-hoc filters de veiligheid loskoppelen van het generatieproces.

Er is een architecturale oplossing nodig die zowel interpreteerbaar (de veiligheidsbeslissing is direct leesbaar) als controleerbaar (de beslissing kan handmatig worden overschreven) is.

Methodologie: Safe Transformer

De auteurs introduceren Safe Transformer (ST), een modulaire architectuur die een pre-getraind taalmodel (Llama-3.2-1B-Instruct) uitbreidt met een discrete informatiebottleneck die een expliciete "veiligheidsbit" bevat tussen de transformer-lagen.

1. Architectuur

De architectuur splitst het model in twee helften (onderste en bovenste lagen) met een module in het midden:

• Informatie Bottleneck: Een variational autoencoder (VAE)-achtige module die de verborgen staten verwerkt.
• Discrete Code: De bottleneck genereert een discrete code $c = [s, u]$:
  • $s$ (Safety Bit): Een binaire variabele ($s \in \{0, 1\}$). $s=1$ betekent "veilig, helpend antwoord", $s=0$ betekent "onveilig, weiger". Deze bit is het resultaat van een classificatie.
  • $u$ (Unsupervised Bits): Een reeks binaire bits die semantische informatie vasthouden voor de generatie, zodat de kwaliteit van de tekst behouden blijft.
• Injectie: De gegenereerde code wordt via cross-attention ingebracht in de bovenste lagen van de transformer, waardoor de generatie expliciet afhankelijk wordt van de veiligheidsbit.

2. Twee-staps Trainingsprocedure

Het model wordt getraind in twee fasen:

• Fase 1: Veiligheidsclassificatie

  • Een bidirectionele encoder en een "Write-in FFN" worden getraind om te bepalen of een input-prompt veilig of onveilig is.
  • De safety bit $s$ wordt hierbij vastgesteld op de ground-truth label.
  • De basis van het model blijft bevroren; alleen de nieuwe encoder en FFN worden getraind.
  • Een KL-divergentie-verlies zorgt ervoor dat de onbewaakte bits $u$ een uniforme verdeling aanhouden (zodat ze geen veiligheidsinformatie "lekken").

• Fase 2: Ontkoppeling via Contrastief Trainen

  • Het doel is om representaties te leren waarbij de veiligheidsbit $s$ de gedragsmodus controleert, onafhankelijk van de semantische inhoud.
  • Data: Dezelfde prompt wordt gekoppeld aan twee verschillende antwoorden: een helpend antwoord ($s=1$) en een weigerend antwoord ($s=0$).
  • Mechanisme: Omdat de prompt identiek is, moet het model leren dat de enige variabele die het verschil in output bepaalt, de bit $s$ is. Dit dwingt het model om gedrag (helpen/weigeren) te ontkoppelen van de inhoud.
  • Tijdens inferentie kan $s$ automatisch worden berekend door de encoder of handmatig worden ingesteld (manual override).

Belangrijkste Bijdragen

1. Unificatie van Interpreteerbaarheid en Controleerbaarheid: De veiligheid wordt niet meer verborgen in parameters, maar is een expliciete, leesbare bit die direct de generatiemodus schakelt.
2. Ontkoppelde Representaties: Door contrastief trainen met paren (zelfde prompt, verschillende output) wordt een causaal verband gelegd tussen de safety bit en het gedrag, terwijl de onbewaakte bits $u$ de semantische coherentie garanderen.
3. Lichtgewicht Implementatie: De methode vereist geen hertraining van scratch, maar slechts lichte fine-tuning op een bestaand instructie-getuned model.

Resultaten

De prestaties zijn geëvalueerd op diverse benchmarks:

• Red-Teaming (Aanvalssuccesratio - ASR):

  • Safe Transformer bereikte een ASR van 0–0,7% op benchmarks zoals DangerousQA en CatQA.
  • Dit is een aanzienlijke verbetering ten opzichte van de basismodel (24,13%) en een standaard SFT-baseline (16,59%).
  • Het model is zeer robuust tegen jailbreak-aanvallen via Chain-of-Thought (CoT) en Chain-of-Utterances (CoU).

• Classificatie en Over-Refusal:

  • In de automatische modus weigerde het model bijna alle onveilige prompts (99,5% succes), maar vertoonde het over-afwijzing (over-refusal) bij veilige prompts die gevoelige woorden bevatten (bijv. "kill" in de context van Python-processen). Dit resulteerde in een lage "Safe Compliance" (32,8%).
  • In de handmatige modus ($s^*=1$) behield het model het gedrag van het basismodel (95,2% compliance), wat aantoont dat de training de oorspronkelijke capaciteiten niet vernietigt.

• Downstream Prestaties:

  • Er was sprake van een lichte daling in algemene kennis (ARC-Easy, MMLU) en een merkbare daling in wiskundig redeneren (GSM8K daalde van 36,1% naar 24,0%). De auteurs attribueren dit aan de beperkte trainingsdata en de compressie in de bottleneck, wat toekomstig onderzoek vereist.

• Rol van Onbewaakte Bits ($u$):

  • Experimenten toonden aan dat $u$ voornamelijk stijl en lexicale variatie beïnvloedt, niet de feitelijke inhoud. Bij feitelijke vragen leverden verschillende $u$-waarden identieke antwoorden op, terwijl bij open vragen de variatie groot was.

Significantie

Safe Transformer biedt een paradigmaverschuiving in AI-veiligheid:

• Van Impliciet naar Expliciet: Het verplaatst veiligheidscontrole van een ondoorzichtige parameterverdeling naar een architecturaal geïntegreerde, inspecteerbare schakelaar.
• Controleerbaarheid: Het stelt onderzoekers en ontwikkelaars in staat om het gedrag van het model direct te manipuleren (bijv. voor debugging of specifieke use-cases) zonder complexe post-hoc interventies.
• Generaliseerbaarheid: Het concept van contrastief trainen met expliciete bits kan worden toegepast op andere controle-scenario's, zoals het wisselen van programmeertalen, persona's of schrijfstijlen.

De paper concludeert dat hoewel er uitdagingen blijven (zoals over-afwijzing en prestatieverlies in complexe redenering), de Safe Transformer een veelbelovende richting is voor het bouwen van transparante en controleerbare AI-systemen.