Secure LLM Fine-Tuning via Safety-Aware Probing

Titel: Secure LLM Fine-Tuning via Safety-Aware Probing (SAP)

1. Het Probleem

Grote Taalmodellen (LLMs) hebben indrukwekkende prestaties geleverd, maar er bestaat een groot veiligheidsrisico: zelfs modellen die veilig zijn afgestemd (aligned) tijdens pre-training of post-training, kunnen hun veiligheid verliezen bij verdere fine-tuning.

• De Kernvraag: Waarom degradeert de veiligheid van een model zelfs bij fine-tuning op onschadelijke (benign) en nuttige data?
• De Oorzaak: Het paper toont aan dat de optimalisatierichtingen voor taakprestatie ("useful-critical directions") en de richtingen die leiden tot onveilig gedrag ("harmful-critical directions") gedeeltelijk ontkoppeld zijn, maar vaak positief gecorreleerd. Dit betekent dat updates die de prestatie op een specifieke taak verbeteren, het model per ongeluk ook naar gebieden in de parameterruimte kunnen duwen waar de kans op schadelijke output groter is. Bestaande methoden zijn vaak beperkt tot specifieke datasets (filtering) of specifieke fine-tuning-methoden (zoals alleen LoRA), wat hun praktische toepasbaarheid beperkt.

2. Methodologie: Safety-Aware Probing (SAP)

De auteurs stellen SAP voor, een optimalisatieframework dat ingrijpt in het trainingsproces zelf in plaats van alleen de data te filteren of de parameters te regulariseren.

• Concept: SAP introduceert een lichtgewicht "probe" (variabele $V$) die de propagatie van verborgen toestanden (hidden states) tijdens de forward pass perturbeert. Dit is conceptueel vergelijkbaar met Sharpness-Aware Minimization (SAM), maar dan gericht op veiligheid.
• Bilevel Optimalisatie: Het framework gebruikt een tweestapsproces per trainingsstap:
  1. Interne Stap (Probe Update): Er wordt een "veiligheidsbewuste" richting bepaald door het contrast tussen veilige en schadelijke antwoorden te analyseren. Een probe $V_{safe}$ wordt geoptimaliseerd om de "veilig-gebruikbare" loss ($L_{su}$) te maximaliseren. Deze loss meet hoe sterk de nuttige loss toeneemt als het model in een schadelijke richting wordt verplaatst. Een hoge $L_{su}$ betekent dat de huidige landschapsvorm ongunstig is voor schadelijke updates.
  2. Externe Stap (Model Update): De modelparameters $W$ worden bijgewerkt via gradiëntafdaal op de nuttige taak, maar dan onder invloed van de geperturbeerde verborgen toestanden (via $V_{safe}$).
• Mechanisme: Door de verborgen toestanden te perturberen, wordt het effectieve verlieslandschap voor de modelparameters zodanig vervormd dat gradiëntafdaal minder snel in onveilige richtingen terechtkomt, terwijl de taakprestatie behouden blijft.
• Efficiëntie: De probe wordt alleen toegepast op een subset van lagen (bijv. lagen 11-20) en vereist geen extra geheugen voor gewichtsperturbaties, waardoor de GPU-geheugennadruk vergelijkbaar blijft met standaard fine-tuning.

3. Belangrijkste Bijdragen

1. Analyse van Gradiëntverstrengeling: Het paper kwantificeert voor het eerst formeel de positieve correlatie tussen nuttige en schadelijke gradiëntrichtingen tijdens fine-tuning, wat de fundamentele oorzaak is van veiligheidsdegradatie.
2. Het SAP Framework: Een nieuw paradigma voor veilig fine-tunen dat niet afhankelijk is van specifieke datasets of architecturale beperkingen (zoals alleen LoRA). Het werkt voor volledige parameter-fine-tuning en diverse modelgroottes.
3. Uitgebreide Validatie: Omvangrijke experimenten tonen aan dat SAP de veiligheid behoudt zonder in te leveren op taakprestatie, en superieur is aan bestaande state-of-the-art methoden.

4. Resultaten

De auteurs hebben SAP getest op meerdere modellen (Llama-2, Vicuna, Qwen) en taken (instructievolgen, redeneren).

• Veiligheid vs. Nut:
  • SAP reduceerde de gemiddelde "Harmful Score" (HS) van 37,43 (standaard fine-tuning) naar 23,07, wat een aanzienlijke verbetering is ten opzichte van de basismodellen en andere methoden.
  • In tegenstelling tot veel andere methoden die veiligheid ten koste gaan van prestaties, behaalde SAP vergelijkbare of zelfs iets betere resultaten op taakmetrieken (BLEURT, Cross-Entropy Loss) vergeleken met standaard fine-tuning.
• Robuustheid:
  • Data Poisoning: Zelfs bij een vergiftigingspercentage van 25% in de trainingsdata, hield SAP de HS laag (29,1% op Llama-2), terwijl andere methoden faalden en de HS zwaar lieten stijgen.
  • Adversariële Fine-tuning: SAP bleek robuuster tegen aanvallen waarbij een model na de initiële training opnieuw op schadelijke data wordt getraind.
  • Aanpassingsaanvallen: Zelfs bij een specifieke adaptieve aanval die probeerde de SAP-mechanica om te keren, bleef SAP aanzienlijk veiliger dan standaard fine-tuning.
• Schaalbaarheid: De methode werkt effectief bij zowel kleine als grote modellen (tot 13B parameters) en bij zowel LoRA als volledige parameter-fine-tuning.
• Efficiëntie: Hoewel SAP ongeveer 2,5x meer rekentijd per stap kost dan standaard fine-tuning, convergeert het sneller. Bij een gelijke totale trainingstijd presteerde SAP beter dan baselines die langer trainden.

5. Betekenis en Conclusie

Dit paper biedt een fundamentele oplossing voor het probleem van veiligheidsdegradatie tijdens de levenscyclus van LLMs. Door in te zien dat veiligheid en nut vaak in conflict zijn in de gradiëntruimte, biedt SAP een elegante, wiskundig onderbouwde manier om deze conflicten op te lossen tijdens het optimalisatieproces.

De significance van dit werk ligt in:

• Het bieden van een schalbare en architectuur-onafhankelijke oplossing die in bestaande pipelines kan worden geïntegreerd zonder de dataset te hoeven wijzigen.
• Het bewijzen dat veiligheid en hoge taakprestatie niet per se een afweging (trade-off) hoeven te zijn.
• Het bieden van een sterke verdediging tegen zowel onbedoelde degradatie als actieve vijandelijke aanvallen, wat cruciaal is voor de veilige implementatie van LLMs in de praktijk.

De code is beschikbaar gesteld, wat de reproduceerbaarheid en adoptie in de gemeenschap faciliteert.