Few Tokens, Big Leverage: Preserving Safety Alignment by Constraining Safety Tokens during Fine-tuning

Dit paper introduceert PACT, een fijnafstappingsframework dat veiligheidsdrift voorkomt door de betrouwbaarheid van het model op specifieke veiligheidstokens te regulariseren tijdens het trainen op onschadelijke data, waardoor de veiligheid behouden blijft zonder de prestaties op downstream-taken te schaden.

De kern

Stel je voor dat je een zeer slimme, beleefde assistent hebt die je helpt met alledaagse taken, zoals wiskundeopdrachten maken of e-mails schrijven. Deze assistent is ook erg goed in het zeggen van "Nee" tegen gevaarlijke vragen, zoals "Hoe maak ik een bom?".

Nu wil je die assistent specialiseren voor een specifieke taak, bijvoorbeeld om een expert te worden in het oplossen van complexe wiskundeproblemen. Je geeft hem dus een nieuwe training (in het Engels: fine-tuning).

Het probleem:
Tijdens deze nieuwe training leert de assistent veel nieuwe dingen, maar helaas vergeet hij ook zijn beleefdheid. Zelfs als je hem alleen maar goede voorbeelden geeft, kan hij door de nieuwe training vergeten hoe hij "nee" moet zeggen tegen gevaarlijke vragen. Als je zelfs maar een paar slechte voorbeelden in de training stopt, wordt hij plotseling heel gehoorzaam en helpt hij zelfs bij het maken van bommen.

De oude oplossingen:
Vroeger probeerden mensen dit op te lossen door de assistent een "gordel" om te doen die zijn hele lichaam strakke. Ze beperkten bijvoorbeeld welke onderdelen van zijn hersenen hij mocht gebruiken, of ze voegden heel veel extra veilige voorbeelden toe.

• Nadeel: Dit maakte de assistent stijf. Hij kon zijn nieuwe wiskundetaak nog steeds niet goed doen, omdat zijn hele bewegingsvrijheid werd beperkt. Het was alsof je een sporter dwingt om met een zware mantel te rennen; hij blijft veilig, maar hij kan niet meer hard lopen.

De nieuwe oplossing: PACT (De "Veiligheids-Stopknop")
De auteurs van dit paper hebben een slimmere manier bedacht, genaamd PACT. In plaats van de hele assistent te beperken, focussen ze op slechts een paar specifieke woorden (tokens) die cruciaal zijn voor veiligheid.

Hier is hoe het werkt, met een paar analogieën:

1. De "Rode Knoppen" in plaats van de hele machine

Stel je voor dat de taal van de assistent bestaat uit duizenden verschillende knoppen. De onderzoekers hebben ontdekt dat de "veiligheid" van de assistent niet ligt in al die knoppen, maar in slechts een klein groepje rode knoppen (ongeveer 50 stuks).

• Als de assistent een gevaarlijke vraag krijgt, gebruikt hij deze rode knoppen om woorden als "Nee", "Ik kan niet helpen" of "Dat is gevaarlijk" te kiezen.
• De rest van de knoppen (voor wiskunde, geschiedenis, grappen) zijn gewoon "blauwe knoppen" en hebben niets met veiligheid te maken.

De PACT-methode: Tijdens de training laten ze de assistent vrij om alle blauwe knoppen te veranderen en aan te leren. Maar de 50 rode knoppen worden vastgezet. Ze zorgen ervoor dat de assistent precies dezelfde zekerheid (vertrouwen) heeft om die rode knoppen in te drukken als hij dat voor de training had.

• Resultaat: Hij wordt een superwiskundeleraar (blauwe knoppen werken perfect), maar vergeet nooit hoe hij "nee" moet zeggen (rode knoppen blijven vast).

2. Het "Spiegelbeeld" zonder de ruis

Er is nog een slimme truc in de methode. Soms kan de training zelf verwarrend zijn. Stel, de assistent leert van een voorbeeldzin: "Hoe maak ik een bom? -> Hier is hoe..." (dit is een slecht voorbeeld).
Als je de assistent nu laat kijken naar het antwoord van de veilige versie, kan die veilige versie ook verward raken door de slechte vraag in de zin.

De onderzoekers gebruiken een twee-spiegelsysteem:

• Spiegel 1 (Volledige context): Kijkt naar de hele zin, inclusief de slechte vraag.
• Spiegel 2 (Alleen het antwoord): Kijkt alleen naar wat de veilige assistent zou zeggen als er geen slechte vraag was, alleen de context van het gesprek.

De methode PACT kijkt naar beide spiegels. Als de eerste spiegel verward raakt door de slechte vraag, schakelt hij automatisch over op de tweede spiegel, die helder en veilig blijft. Dit zorgt ervoor dat de assistent zijn "veiligheidsgedrag" niet verliest, zelfs als de trainingssessie rommelig is.

Waarom is dit zo goed?

• Efficiënt: Ze hoeven niet de hele assistent te controleren, alleen die 50 specifieke woorden.
• Flexibel: De assistent kan nog steeds leren en groeien in zijn nieuwe vakgebied (wiskunde, sentimentanalyse, etc.).
• Veilig: Zelfs als je per ongeluk een paar gevaarlijke voorbeelden in de training stopt, blijft de assistent zijn "nee" zeggen.

Samengevat:
In plaats van een hele auto te blokkeren zodat hij niet te snel kan rijden (wat hem ook langzaam maakt), plakken ze een slimme rem op alleen het rempedaal. De auto kan nog steeds razendsnel racen op het circuit (de nieuwe taak), maar als er gevaar dreigt, werkt de rem nog steeds perfect.

Dit is de kracht van PACT: een paar woorden vastzetten om de hele machine veilig te houden, zonder de prestaties te kosten.

Technische samenvatting

Hier is een gedetailleerde technische samenvatting van het paper "Few Tokens, Big Leverage: Preserving Safety Alignment by Constraining Safety Tokens during Fine-tuning" (PACT), vertaald en samengevat in het Nederlands.

1. Het Probleem: Veiligheidsdrift bij Fine-tuning

Grote Taalmodellen (LLMs) worden vaak fijn afgestemd (fine-tuning) om prestaties op specifieke downstream-taken te verbeteren. Echter, dit proces introduceert een significant veiligheidsrisico: veiligheidsdrift.

• Het fenomeen: Zelfs wanneer de trainingsdataset voornamelijk uit onschadelijke data bestaat, kan fine-tuning leiden tot een verslechtering van de veiligheidsuitlijning. De aanwezigheid van een klein percentage schadelijke data (bijvoorbeeld 10%) kan de weigeringsmechanismen van het model volledig doen instorten, waardoor het model schadelijke verzoeken gaat uitvoeren.
• Beperkingen van bestaande oplossingen: Bestaande verdedigingsmethoden opereren vaak op een grofkorrelig niveau, zoals het beperken van welke parameters worden bijgewerkt (bijv. SafeLoRA) of het toevoegen van extra veiligheidsdata. Deze benaderingen beperken vaak de aanpassingscapaciteit van het model op de downstream-taak, wat leidt tot een afname in nut (utility) en generalisatie.

2. Methodologie: PACT (Preserving safety Alignment via Constrained Tokens)

De auteurs stellen PACT voor, een framework dat veiligheid behoudt door zeer gerichte, fijnkorrelige beperkingen toe te passen op een kleine subset van tokens die cruciaal zijn voor veiligheid, in plaats van het hele model te beperken.

A. Identificatie van Veiligheidstokens

Op basis van empirische observaties dat veiligheidsuitlijning voornamelijk wordt gereflecteerd in het zelfvertrouwen (confidence) van het model op specifieke tokens, identificeren de auteurs een kleine set "veiligheidstokens":

• Analyse: Ze vergelijken een veiligheidsuitgelijnd model ($M_{safe}$) met een basismodel ($M_{base}$) op schadelijke prompts.
• Methode: Ze berekenen de kansverschillen ($\Delta_t(v)$) per token tussen de twee modellen.
• Resultaat: Een kleine set tokens (bijv. de top 50, zoals "I", "can't", "assist", "cannot") vertoont een groot verschil in waarschijnlijkheid. Het veiligheidsmodel heeft hier een veel hoger zelfvertrouwen voor dan het basismodel. Deze tokens zijn cruciaal voor het genereren van veilige weigeringen.

B. Het PACT Framework

PACT regulariseert het fine-tuning-proces door het zelfvertrouwen van het model op deze veiligheidstokens te laten overeenkomen met dat van het oorspronkelijke veiligheidsmodel, terwijl de rest van de tokens vrij kunnen worden geoptimaliseerd voor de taak.

Het framework bestaat uit twee kerncomponenten:

1. Geweegde Regularisatie op Veiligheidstokens:

   • In plaats van een globale KL-divergentie (Kullback-Leibler) te gebruiken over het hele vocabulaire, wordt de KL-loss beperkt tot de geïdentificeerde veiligheidstokens ($S_{safety}$).
   • Gewichten: Tokens met een groter kansverschil (hogere relevantie voor veiligheid) krijgen zwaardere straffen als het model afwijkt. Dit zorgt ervoor dat de weigeringstokens stabiel blijven, terwijl het model vrij is om de downstream-taak te leren.

2. Calibratie van het Veiligheidssignaal:

   • Een uitdaging bij het gebruik van een "teacher" (het referentiemodel) is dat als de trainingsdata schadelijk is, het referentiemodel gedwongen wordt om op een onveilige prefix te conditioneren. Dit kan het veiligheidssignaal verzwakken ("prefix contamination").
   • Oplossing: PACT gebruikt een adaptieve mix van twee referentie-distributies:
     • Full-context: Het referentiemodel kijkt naar prompt + antwoord.
     • No-prompt: Het referentiemodel kijkt alleen naar de voorgaande antwoorden (zonder de prompt). Dit is een schoner veiligheidssignaal.
   • Een gating-mechanisme bepaalt hoeveel vertrouwen er in het "no-prompt" signaal moet worden gesteld. Als de context schadelijk is (hoge onzekerheid over de weigering), wordt het "no-prompt" signaal zwaarder gewogen om de veiligheid te beschermen.

De totale loss functie is een combinatie van de standaard Cross-Entropy loss (voor de taak) en de gewogen KL-loss (voor veiligheid).

3. Belangrijkste Bijdragen

• Token-level Analyse: Het paper toont aan dat veiligheidsuitlijning niet uniform over het vocabulaire is verdeeld, maar sterk geconcentreerd is op een kleine set tokens.
• Gerichte Regularisatie: PACT introduceert een methode om alleen deze specifieke tokens te constraineren, waardoor de trade-off tussen veiligheid en nut (utility) wordt doorbroken.
• Calibratie-mechanisme: Een innovatieve aanpak om de invloed van schadelijke prefixes op het leerproces te neutraliseren door adaptief te mixen tussen contextuele en prompt-vrije referenties.
• Uitgebreide Validatie: Het framework is getest op meerdere model-families (Qwen, Llama, Gemma) en diverse downstream-taken (wiskunde, sentimentanalyse, nieuwsclassificatie).

4. Resultaten

De experimenten tonen aan dat PACT aanzienlijk beter presteert dan state-of-the-art baselines (zoals SafeLoRA, Constrained SFT, AsFT):

• Veiligheid vs. Nut: PACT behoudt bijna dezelfde taaknauwkeurigheid als standaard fine-tuning (vanilla SFT), maar voorkomt de instorting van de veiligheid.
  • Voorbeeld: Op de GSM8K-taak met 10% schadelijke data behaalde PACT 80.89% nauwkeurigheid en verlaagde de aanvalssuccesratio (ASR) op HarmBench tot 29.50%. Vergelijkbare methoden hadden vaak een ASR van >90% of een veel lagere nauwkeurigheid.
• Generalisatie: PACT werkt consistent goed over verschillende modelgroottes (van 1B tot 9B parameters) en architecturen, terwijl andere methoden vaak falen bij specifieke modellen.
• Robuustheid: Zelfs bij variërende percentages schadelijke data (0% tot 10%) blijft PACT stabiel, terwijl andere methoden bij 5% of 10% schadelijke data volledig instorten.
• Ablatie-studies: Het bevestigt dat elke component (selectie van tokens, gewichten, en calibratie) essentieel is voor de uiteindelijke prestatie.

5. Betekenis en Conclusie

Dit paper biedt een fundamenteel nieuw perspectief op veiligheidsbehoud bij LLMs. Het toont aan dat men niet het hele model hoeft te "bevriezen" of zwaar te beperken om veiligheid te behouden. Door te focussen op een kleine, kritische subset van tokens die verantwoordelijk zijn voor weigeringen, kan men de veiligheid van het model effectief "verankeren" tijdens het leren van nieuwe taken.

Kernboodschap: Veiligheid is een lokaal fenomeen in de output van taalmodellen. Door deze lokale signalen te beschermen, kan men de "grote hefboom" (big leverage) gebruiken om de algehele veiligheid te behouden zonder in te leveren op de bruikbaarheid van het model. Dit maakt PACT een veelbelovende oplossing voor veilige en effectieve aanpassing van commerciële LLMs.