Alignment-Weighted DPO: A principled reasoning approach to improve safety alignment

Probleemstelling

Grote Taalmodellen (LLMs) zijn steeds vaker kwetsbaar voor "jailbreak"-aanvallen, waarbij kwaadaardige intenties worden vermomd door indirecte, misleidende of gecamoufleerde formuleringen. Hoewel bestaande afstemmingstechnieken zoals Supervised Fine-Tuning (SFT), Reinforcement Learning from Human Feedback (RLHF) en Direct Preference Optimization (DPO) de veiligheid hebben verbeterd, blijken deze vaak oppervlakkig te zijn.

De auteurs stellen de hypothese op dat huidige veiligheidsmechanismen vertrouwen op flauwe afwijzingsheuristieken (shallow refusal heuristics) in plaats van diep redeneren. Modellen leren vaak alleen om schadelijke patronen te herkennen en genereren een standaardafwijzing (bijv. "Sorry, ik kan dat niet doen") zonder de onderliggende reden voor de schadelijkheid echt te begrijpen. Dit maakt ze kwetsbaar voor aanvallen die deze oppervlakkige patronen omzeilen.

Methodologie

De paper introduceert een tweestapsbenadering om dit probleem aan te pakken:

1. Causaal Interventie-experiment
Om de hypothese te testen dat veiligheid losstaat van redeneren, voeren de auteurs een causale interventie uit:

• Ze identificeren neuronen die cruciaal zijn voor redeneren door middel van lineaire probing (het trainen van een lineaire classifier op de verborgen representaties van attention heads).
• Ze deactiveren de top 10% van de attention heads die het meest bijdragen aan redeneren (door hun Q, K, V-weights op nul te zetten).
• Resultaat: De redeneerprestaties van het model dalen drastisch, maar de veiligheidsprestaties (het vermogen om schadelijke prompts af te wijzen) blijven grotendeels onaangetast. Dit bevestigt dat huidige veiligheidsafstemming niet gebaseerd is op diep redeneren.

2. Oplossing: Redeneringsbewuste Post-training
Om dit tekort aan te vullen, stellen de auteurs een nieuwe pipeline voor:

• Chain-of-Thought (CoT) Dataset: Ze bouwen en openen een nieuwe dataset die zowel nuttige als veiligheidskritische prompts bevat, elk vergezeld van stap-voor-stap redeneringen (rationales). Dit trainingsdoel is om modellen te leren waarom iets schadelijk is, niet alleen dat het schadelijk is.
• Alignment-Weighted DPO (AW-DPO):
  • Bij analyse van fouten in CoT-finetuning ontdekten ze twee kritieke faalpatronen: (i) correcte redenering met een onveilig eindantwoord, en (ii) onjuiste redenering met een veilig eindantwoord.
  • Standaard DPO behandelt de volledige respons als één eenheid, wat deze subtiele fouten mist.
  • AW-DPO deconstrueert de respons in twee delen: de redenering (reasoning trace) en het eindantwoord (final response).
  • Het algoritme wijst verschillende voorkeursgewichten toe aan deze segmenten op basis van hun specifieke schadelijkheidsscores (bepaald door een 'judge' model).
  • De loss-functie wordt aangepast zodat de optimalisatie zich gerichter richt op het deel van de respons dat het meest problematisch is (bijv. het corrigeren van de redenering als deze veilig is, maar het antwoord niet, of vice versa).

Belangrijkste Bijdragen

1. Empirisch Bewijs van Oppervlakkigheid: Door middel van causale interventie wordt aangetoond dat huidige veiligheidsafstemming onafhankelijk is van redeneervermogen, wat verklaart waarom modellen vatbaar zijn voor geavanceerde jailbreaks.
2. Nieuwe CoT Veiligheidsdataset: Een open-source dataset met gedetailleerde redeneringen voor zowel veilige als schadelijke prompts, ontworpen om zowel veiligheid als algemeen nut (utility) te behouden.
3. Alignment-Weighted DPO (AW-DPO): Een nieuwe versterkingsleer-methode die fijnmazige, gerichte optimalisatie toepast door redenering en antwoord apart te wegen, in plaats van ze als één blok te behandelen.
4. Uitgebreide Validatie: Experimenten tonen aan dat de methode robuuster is tegen diverse jailbreak-strategieën zonder de bruikbaarheid van het model significant te verminderen.

Resultaten

De auteurs hebben hun methode getest op meerdere modellen (Llama-2, Llama-3, Mistral) en benchmarks (SorryBench voor veiligheid, MMLU voor nut):

• Veiligheid: AW-DPO presteert consistent beter dan bestaande baselines (zoals standaard SFT, CoT-SFT en standaard DPO). Het reduceert de aanvalssuccesrate (ASR) aanzienlijk, zelfs tegen complexe aanvallen zoals het gebruik van encryptie, meerdelige talen en overtuigingskunst.
• Nuttigheid (Utility): In tegenstelling tot sommige veiligheidsmethoden die de algemene prestaties van het model laten dalen, behoudt AW-DPO een hoge prestatie op algemene taken (zoals MMLU).
• Vergelijking met Redeneermodellen: Specifieke "reasoning-oriented" modellen (zoals Phi-4 Reasoning) presteren slechter op veiligheidsopgaven dan de door AW-DPO getrainde modellen. Dit onderstreept dat algemeen redeneervermogen niet automatisch leidt tot betere veiligheidsafstemming; de redenering moet specifiek gericht zijn op veiligheidscontext.
• Transferabiliteit: De AW-DPO dataset die op het ene model is gebouwd, werkt effectief op andere modelarchitecturen, wat de kosten voor datacollectie verlaagt.

Betekenis en Impact

Deze paper biedt een fundamenteel inzicht in de beperkingen van huidige LLM-veiligheid: veiligheid mag niet alleen een oppervlakkige "stop" zijn, maar moet gebaseerd zijn op begrip en redenering.

• Theoretische bijdrage: Het weerlegt het idee dat het simpelweg toevoegen van redeneren (CoT) voldoende is; de kwaliteit en gerichtheid van die redenering in de afstemmingsfase zijn cruciaal.
• Praktische toepassing: AW-DPO biedt een efficiëntere en robuustere manier om modellen veilig te maken, wat essentieel is voor de inzet van LLMs in hoog-risico domeinen zoals financiën, gezondheidszorg en onderwijs.
• Toekomstige richting: Het benadrukt de noodzaak van "principled refusals" (gebaseerde afwijzingen) waarbij het model de logica achter de veiligheid begrijpt, waardoor het beter bestand is tegen creatieve en geavanceerde adversariale aanvallen.