Two to Tango: Coupled Task-Reference Selection for Safe LLM Fine-tuning

Technische Samenvatting: Two to Tango: Gekoppelde Taak–Referentie Selectie voor Veilige LLM Fine-tuning

Probleemstelling

Het fine-tunen van op veiligheid gealigneerde Large Language Models (LLMs) op downstream-data verbetert vaak de taakaanpassing, maar brengt het risico met zich mee dat de geleerde veiligheidsgedragingen uit de initiële alignmentfase worden uitgehold. Bestaande mitigatiestrategieën vertrouwen doorgaans op een van drie benaderingen: het filteren van taakdata, het mengen van vaste veiligheidsvoorbeelden, of het opleggen van globale veiligheidsrestricties. Echter, de in het paper gepresenteerde diagnostiek onthult een kritieke beperking in deze methoden: ze behandelen veiligheidsreferenties als statische of globale signalen. In werkelijkheid activeren verschillende taak-update-richtingen onderscheidende veiligheidsrestricties. Bijgevolg kan een vaste subset van veiligheidsreferenties er niet in slagen de specifieke restricties bloot te leggen die relevant zijn voor een bepaalde taak-update, wat leidt tot veiligheidsdegradatie zelfs wanneer de gebruikte data ongevaarlijk is.

Methodologie: DualSelect

De auteurs stellen DualSelect voor, een gekoppeld framework dat zowel taakmonsters als veiligheidsreferenties gezamenlijk selecteert op basis van de geometrie van de model-update. De kerninzicht is dat het behoud van veiligheid vereist dat taak-updates worden afgestemd op actieve veiligheidsrestricties die specifiek zijn voor die update, in plaats van een globale veiligheidsfilter toe te passen.

1. Gekoppelde Min-Max Formulering
DualSelect formuleert veilige fine-tuning als een bilevel-stijl minimax-selectieprobleem:

• Innerlijke Maximalisatie (Referentie Selectie): Identificeert veiligheidsreferenties die een hoge preservatie-loss en een hoge conflict (gradiënt-oppositie) vertonen met de huidige taakrichting. Deze referenties vertegenwoordigen de "worst-case" actieve restricties voor de huidige update.
• Uiterlijke Minimalisatie (Taak Selectie): Selecteert taakmonsters die compatibel zijn met de gradiëntrichting die wordt geïnduceerd door de geselecteerde referenties.

2. Entropie-geregulariseerde Scoring Surrogaten
Om het discrete selectieprobleem hanteerbaar te maken, leiden de auteurs entropie-geregulariseerde scoring-surrogaten af.

• Referentie Scoring: Referenties worden gescoord op basis van hun supervised loss en hun cosinus-conflict met de huidige taakgradiënt. Referenties met een hoge loss en hoog conflict krijgen een hogere score.
• Taak Scoring: Taakmonsters worden gescoord op basis van hun supervised loss en hun cosinus-compatibiliteit met de geselecteerde referentiegradiënt. Monsters met een lage loss en hoge compatibiliteit worden verkozen.
• Hard Selectie: Hoewel de surrogaten zachte waarschijnlijkheden bieden, implementeert het algoritme harde selectie (Top-K voor referenties, budget-gematchte monsterselectie voor taken) om compatibiliteit met standaard fine-tuning pipelines te behouden.

3. Lazy Refresh en Gradiënt Correctie

• Lazy Refresh: Referentie-selectie wordt per epoch (of lazily) bijgewerkt in plaats van per batch om de kosten van scoring te amortiseren, terwijl taak-selectie per mini-batch plaatsvindt.
• Referentie-Gradiënt Correctie: De uiteindelijke update-gradiënt wordt gecorrigeerd door een geschaald component van de geselecteerde referentiegradiënt toe te voegen ($g_{final} = g_{new} + \rho g_{ref}$). Dit stuurt de parameter-update richting de geselecteerde referentierichting, waardoor wordt gegarandeerd dat de gerealiseerde update overeenkomt met de geïdentificeerde veiligheidsrestricties.

4. Theoretische Analyse
Het paper biedt een lokale geometrische analyse die aantoont dat onder voldoende correctiekracht ($\rho$), de gecorrigeerde dalingsrichting overeenkomt met de geselecteerde referentiegradiënt, wat een eerste-orde afname van de geselecteerde-referentie-loss garandeert.

Belangrijkste Bijdragen

1. Paradigmaverschuiving: De auteurs identificeren en demonstreren een verschuiving van het gebruik van statische veiligheids-referentiepools naar taak-geconditioneerde referentie-selectie. Ze tonen aan dat geselecteerde referenties gestructureerd zijn en de heldengehouden veiligheid aanzienlijk beter behouden dan statische subsets.
2. Gekoppeld Framework: Ze introduceren DualSelect, dat de selectie van taakmonsters en veiligheidsreferenties koppelt als onderling afhankelijke variabelen in plaats van aparte scoring-doelen, geformuleerd als een min-max optimalisatie over referenties met een hoge loss en taak-conflict en compatibele taamonsters.
3. Algoritmische Afleiding: Het paper leidt entropie-geregulariseerde surrogaten af die harde selectieregels induceren, waarbij lazy reference refresh, reference-gradient correction en analyse op niveau van de update-uitlijning worden geïntegreerd.
4. Empirische Validatie: Uitgebreide experimenten over 1B tot 8B parameter LLM's (Gemma, Qwen, Llama) op de REDORCA en GSM8K datasets valideren de aanpak.

Experimentele Resultaten

• Veiligheidsbehoud: DualSelect presteert consistent beter dan sterke baselines (inclusief TOSS, SQSD, SEAL en LARF) in het behouden van veiligheid. Op de REDORCA dataset verbetert het de Safety Average score met ten minste 5,10 punten ten opzichte van de sterkste baseline over 1B–8B modellen.
• Utiliteitsbehoud: De methode behoudt een taak-utiliteit die vergelijkbaar met of beter is dan die van baselines. Op GSM8K bereikt het de hoogste Safety Average terwijl de nauwkeurigheid nabij de standaard SFT-niveaus blijft.
• Efficiëntie: DualSelect introduceert een matige overhead (ongeveer 1,34× de GPU-uren van Standard SFT) en een geheugengebruik vergelijkbaar met LoRA fine-tuning, wat aanzienlijk efficiënter is dan bilevel-optimalisatiemethoden zoals SEAL.
• Ablatie-studies: Het verwijderen van een van de componenten (referentie-selectie, taak-scoring, compatibiliteit, of finale correctie) leidt tot een daling in veiligheidsbehoud, wat de noodzaak van het gekoppelde mechanisme bevestigt.
• Continual Learning: Het framework strekt zich uit naar continual learning door oude taakvoorbeelden te behandelen als retentie-referenties, waarmee het succesvol catastrofale vergetelheid vermindert terwijl eerdere capaciteiten behouden blijven.

Betekenis en Claims

Het paper claimt dat de primaire bijdrage niet louter een nieuwe data-selectie heuristiek is, maar een fundamentele herformulering van veiligheidsbehoud als een update-geconditioneerd probleem. De auteurs argumenteren dat veiligheidsdata niet uniform nuttig is; de relevantie ervan hangt af van de specifieke taakrichting die wordt geleerd. Door actieve restricties dynamisch bloot te leggen via taak-geconditioneerde referenties en compatibele taken te filteren, bereikt DualSelect een superieur veiligheids-utiliteit evenwicht.

De auteurs merken bescheiden op dat de methode toegang veronderstelt tot een referentie-signaal (veilige responsen of oude taak-targets) en dat de prestaties afhangen van de referentie-dekking. Ze erkennen ook dat de veiligheidsevaluatie deels steunt op automatische judges, hoewel ze cross-judge analyse en absolute metrieken gebruiken om artefacten te minimaliseren. Het werk positioneert zich als een praktische, schaalbare oplossing voor veilige fine-tuning die de hoge kosten van globale herweging of complexe post-hoc bewerkingen vermijdt.