Technische Samenvatting: Het Verminderen van de Veiligheidstaks in LLM-Veiligheidsuitlijning met On-Policy Self-Distillation

Probleemdefinitie: De Veiligheidstaks en Off-Policy Mismatch

Grote Taalmodellen (LLM's) lijden vaak onder een "veiligheidstaks", een afweging waarbij veiligheidsuitlijning de robuustheid tegen schadelijke queries verbetert, maar de algemene redeneervermogens verslechtert. Hoewel eerdere werken dit toeschrijven aan een distributie-mismatch—training op veiligheidsdemonstraties van mensen of externe modellen die verschillen van de natuurlijke redeneerdistributie van het doelmodel—identificeert dit artikel een tweede, onderscheidende bron: off-policy training mismatch.

Zelfs wanneer veiligheidsdata zelf-gedistilleerd is (gegenereerd door het doelmodel zelf, zoals in ThinkSafe), blijft standaard Supervised Fine-Tuning (SFT) off-policy. Het past sequentie-niveau Negative Log-Likelihood (NLL) supervisie toe op vaste demonstraties in plaats van op trajecten die zijn bemonsterd uit het eigen huidige beleid van het model. De auteurs betogen dat dit bijzonder schadelijk is voor veiligheid, omdat veiligheidsbeslissingen (weigering versus naleving) geconcentreerd zijn in een smal "early-token window" en een specifieke set veiligheidskritische tokens. Off-policy SFT verspreidt gradiënten over de hele respons, waardoor het deze kritieke beslispunten niet specifiek target, wat algemene redenering verstoort terwijl het slechts grove veiligheids-supervisie biedt.

Methodologie: OPSA (On-Policy Safety Alignment)

Het artikel stelt OPSA voor, een raamwerk dat On-Policy Self-Distillation (OPSD) aanpast aan veiligheidsuitlijning. De kernmechaniek omvat:

1. On-Policy Rollouts: Het studentmodel ($p_\theta$) genereert zijn eigen trajecten (rollouts) voor zowel schadelijke als onschadelijke prompts.
2. Privileged Context Teacher: Een bevroren kopie van het basismodel ($p_{\bar{\theta}}$) fungeert als leraar. Cruciaal ontvangt de leraar een type-voorwaardelijke bevoorrechte context ($c^*$) die aan de prompt wordt voorafgeplaatst:
   • Voor schadelijke prompts ($q_h$) activeert een context ($I_h$) latent weigerend gedrag, waardoor de leraar verschuift naar veilige responsen.
   • Voor onschadelijke prompts ($q_b$) onderdrukt een context ($I_b$) overmatige weigering, waardoor behulpzaam gedrag behouden blijft.
3. Dense Per-Token Supervisie: In plaats van sequentie-niveau NLL, minimaliseert OPSA de per-token Kullback-Leibler (KL) divergentie tussen de distributie van de leraar (voorwaardelijk op de bevoorrechte context) en de distributie van het studentmodel (alleen voorwaardelijk op de prompt) langs de eigen bemonsterde rollouts van het studentmodel:
   $$L_{OPSA}(\theta) = \sum \mathbb{E}_{y \sim p_\theta} \sum_{t} D_{KL}(p_{\bar{\theta}}(\cdot | c^*, q, y_{<t}) \parallel p_\theta(\cdot | q, y_{<t}))$$
   Dit concentreert updates op tokens waar de leraar en het studentmodel uiteenlopen, specifiek gericht op het vroege beslissingsvenster waar veiligheidsgedrag wordt bepaald.

Teacher Flip Rate (TFR) voor Contextselectie

Een kritieke uitdaging is het selecteren van een bevoorrechte context die latent veiligheidsredeneren effectief activeert zonder willekeurig te zijn. De auteurs introduceren Teacher Flip Rate (TFR) als een trainingsvrije metriek om de optimale context ($c^*$) te selecteren. TFR meet het aandeel schadelijke queries waarbij de bevoorrechte context ervoor zorgt dat de bevroren leraar zijn greedy respons van onveilig naar veilig laat omslaan. De auteurs tonen empirisch aan dat een hogere TFR monotoon correleert met verminderde schadelijkheid na training, wat TFR valideert als een betrouwbaar criterium voor contextselectie.

Belangrijkste Bijdragen

1. Mechanistische Diagnose van de Veiligheidstaks: Het artikel identificeert off-policy training mismatch als een tweede bron van de veiligheidstaks. Token-niveau analyse onthult dat veiligheidscorrecties sterk geconcentreerd zijn in vroege posities en specifieke tokens, wat standaard SFT niet effectief target, zelfs niet met in-distribution data.
2. OPSA Raamwerk: De introductie van een on-policy veiligheidsuitlijnraamwerk dat dichte per-token KL-supervisie gebruikt op door het studentmodel bemonsterde rollouts. Deze benadering lijnt het model uit op zijn eigen generatiepaden, waardoor de afweging tussen veiligheid en redeneren verbetert.
3. Contextselectie via TFR: Het voorstellen van Teacher Flip Rate als methode om effectieve veiligheidsleraars te construeren. Dit zorgt ervoor dat de bevoorrechte context betekenisvol latent veiligheidsredeneren activeert in plaats van slechts veilig ogende demonstraties op te eisen.

Experimentele Resultaten

De auteurs evalueerden OPSA over twee redeneermodel-families (Qwen3 en DeepSeek-R1-Distill) en vijf parameterniveaus (0.6B tot 8B).

• Veiligheid-Redeneren Afweging: Onder匹配的 prompts en full-parameter fine-tuning presteerde OPSA consequent beter dan off-policy baselines (waaronder ThinkSafe en SafeChain).
  • Veiligheid: OPSA behaalde een gemiddelde composite veiligheidsscore-toename van +4,00 punten ten opzichte van ThinkSafe.
  • Redeneren: OPSA behield en verbeterde redeneervermogens, met een gemiddelde toename van +3,04 punten op redeneer-benchmarks (GSM8K, MATH500, GPQA, HumanEval, MBPP).
  • Schaalgevoeligheid: De grootste winsten werden waargenomen bij kleinere modellen (bijvoorbeeld +8,85 punten op R1-Distill-1.5B en +5,49 punten op Qwen3-0.6B).
• Adaptieve Jailbreak Robuustheid: OPSA toonde superieure robuustheid tegen adaptieve jailbreak-aanvallen, met name Prefilling-aanvallen (die ingrijpen in de vroege respons-tokens). Op Qwen3-1.7B en Qwen3-8B reduceerde OPSA het succespercentage van Prefilling-aanvallen tot nul. Het artikel merkt echter op dat volledig adaptieve zoek-aanvallen (PAIR) een beperking blijven, met gemengde resultaten over modellen heen.
• Data-efficiëntie: De prestatiewinsten van OPSA bleven bestaan over variërende trainingssetgroottes en verhoudingen van onschadelijke tot schadelijke data, wat de robuustheid van de methode aangeeft tegen datatekorten en verschuivingen in data-samenstelling.

Betekenis en Claims

Het artikel beweert dat het verminderen van de veiligheidstaks vereist dat men niet alleen controleert welke demonstraties supervisie bieden, maar ook waar en op wiens trajecten die supervisie wordt toegepast. Door over te stappen van off-policy imitatie van vaste sporen naar on-policy dichte supervisie van de eigen generatiepaden van het model, adresseert OPSA de specifieke structurele mismatch in veiligheidsuitlijning.

De auteurs bescheiden claimen dat hun bevindingen suggereren dat veiligheidsuitlijning fundamenteel een lokaal correctieprobleem is op de eigen trajecten van het model. Hoewel OPSA de robuustheid tegen aanvallen die gericht zijn op vroege respons-tokens significant verbetert, erkent het artikel dat volledig adaptieve iteratieve zoek-aanvallen een uitdaging blijven, wat suggereert dat dichte on-policy supervisie geen alomvattende oplossing is voor alle adversariële scenario's. Het werk benadrukt dat effectieve veiligheidsuitlijning het model moet uitlijnen met zijn eigen redeneerdistributie, terwijl het latent veiligheidsvermogens activeert via bevoorrechte contexten.