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

Riepilogo Tecnico: Two to Tango: Selezione Accoppiata di Compito e Riferimento per il Fine-tuning Sicuro di LLM

Definizione del Problema

Il fine-tuning di Large Language Models (LLM) allineati alla sicurezza su dati di task specifici spesso migliora l'adattamento al compito, ma rischia di erodere i comportamenti di sicurezza appresi durante la fase iniziale di allineamento. Le strategie di mitigazione esistenti si affidano tipicamente a uno di tre approcci: il filtraggio dei dati del task, la miscelazione di esempi di sicurezza fissi o l'imposizione di vincoli di sicurezza globali. Tuttavia, le diagnosi presentate nel paper rivelano un limite critico in questi metodi: essi trattano i riferimenti di sicurezza come segnali statici o globali. In realtà, diverse direzioni di aggiornamento del task attivano vincoli di sicurezza distinti. Di conseguenza, un sottoinsieme fisso di riferimenti di sicurezza può non essere in grado di esporre i vincoli specifici rilevanti per un particolare aggiornamento del task, portando alla degradazione della sicurezza anche quando vengono utilizzati dati benigni.

Metodologia: DualSelect

Gli autori propongono DualSelect, un framework accoppiato che seleziona congiuntamente campioni di task e riferimenti di sicurezza sulla base della geometria dell'aggiornamento del modello. L'intuizione centrale è che la preservazione della sicurezza richiede l'abbinamento degli aggiornamenti del task con i vincoli di sicurezza attivi specifici per quell'aggiornamento, piuttosto che applicare un filtro di sicurezza globale.

1. Formulazione Min-Max Accoppiata
DualSelect formula il fine-tuning sicuro come un problema di selezione min-max di tipo bilevel:

• Massimizzazione Interna (Selezione del Riferimento): Identifica i riferimenti di sicurezza che presentano un'alta perdita di preservazione e un alto conflitto (opposizione del gradiente) con la direzione attuale del task. Questi riferimenti rappresentano i vincoli attivi "peggiori" per l'attuale aggiornamento.
• Minimizzazione Esterna (Selezione del Task): Seleziona campioni di task che siano compatibili con la direzione del gradiente indotta dai riferimenti selezionati.

2. Surrogati di Punteggio Regolarizzati dall'Entropia
Per rendere trattabile il problema di selezione discreta, gli autori derivano dei surrogati di punteggio regolarizzati dall'entropia.

• Punteggio del Riferimento: I riferimenti sono valutati in base alla loro perdita supervisionata e al loro conflitto del coseno con il gradiente attuale del task. I riferimenti con alta perdita e alto conflitto ricevono punteggi più elevati.
• Punteggio del Task: I campioni di task sono valutati in base alla loro perdita supervisionata e alla loro compatibilità del coseno con il gradiente del riferimento selezionato. Sono preferiti i campioni a bassa perdita e alta compatibilità.
• Selezione Hard: Sebbene i surrogati forniscano probabilità soft, l'algoritmo implementa una selezione hard (Top-K per i riferimenti, selezione dei campioni con budget corrispondente per i task) per mantenere la compatibilità con le pipeline di fine-tuning standard.

3. Refresh Pigro e Correzione del Gradiente

• Refresh Pigro (Lazy Refresh): La selezione dei riferimenti viene aggiornata per epoca (o pigramente) piuttosto che per batch per ammortizzare i costi di punteggio, mentre la selezione del task avviene per mini-batch.
• Correzione del Gradiente del Riferimento: Il gradiente dell'aggiornamento finale è corretto aggiungendo una componente scalata del gradiente del riferimento selezionato ($g_{final} = g_{new} + \rho g_{ref}$). Questo guida l'aggiornamento dei parametri verso la direzione del riferimento selezionato, garantendo che l'aggiornamento realizzato sia allineato con i vincoli di sicurezza identificati.

4. Analisi Teorica
Il paper fornisce un'analisi geometrica locale mostrando che, sotto una sufficiente forza di correzione ($\rho$), la direzione di discesa corretta si allinea con il gradiente del riferimento selezionato, garantendo una diminuzione del primo ordine della perdita del riferimento selezionato.

Contributi Chiave

1. Cambio di Paradigma: Gli autori identificano e dimostrano un passaggio dall'uso di pool di riferimenti di sicurezza statici alla selezione dei riferimenti condizionata al task. Dimostrano che i riferimenti selezionati sono strutturati e preservano la sicurezza su dati non visti significativamente meglio rispetto ai sottoinsiemi statici.
2. Framework Accoppiato: Introducono DualSelect, che accoppia la selezione dei campioni di task e dei riferimenti di sicurezza come variabili interdipendenti anziché come obiettivi di punteggio separati, formulata come un'ottimizzazione min-max su riferimenti ad alta perdita e in conflitto con il task e campioni di task compatibili.
3. Derivazione Algoritmica: Il paper deriva surrogati regolarizzati dall'entropia che inducono regole di selezione hard, incorporando il refresh pigro dei riferimenti, la correzione del gradiente del riferimento e l'analisi dell'allineamento dell'aggiornamento a livello di direzione.
4. Validazione Empirica: Ampie sperimentazioni su modelli da 1B a 8B di parametri (Gemma, Qwen, Llama) sui dataset REDORCA e GSM8K validano l'approccio.

Risultati Sperimentali

• Preservazione della Sicurezza: DualSelect supera costantemente i forti baseline (inclusi TOSS, SQSD, SEAL e LARF) nella preservazione della sicurezza. Sul dataset REDORCA, migliora il punteggio di Sicurezza Media di almeno 5,10 punti rispetto al baseline più forte su modelli da 1B a 8B.
• Ritenzione dell'Utilità: Il metodo mantiene un'utilità del task comparabile o superiore ai baseline. Su GSM8K, raggiunge il più alto punteggio di Sicurezza Media mantenendo un'accuratezza vicina ai livelli di SFT Standard.
• Efficienza: DualSelect introduce un overhead moderato (circa 1,34× le ore GPU del SFT Standard) e un uso della memoria comparabile al fine-tuning LoRA, risultando significativamente più efficiente dei metodi di ottimizzazione bilevel come SEAL.
• Studi di Ablazione: Rimuovere qualsiasi componente (selezione del riferimento, punteggio del task, compatibilità, o correzione finale) porta a un calo della preservazione della sicurezza, confermando la necessità del meccanismo accoppiato.
• Apprendimento Continuo: Il framework si estende all'apprendimento continuo trattando gli esempi dei vecchi task come riferimenti di ritenzione, riducendo con successo l'oblio catastrofico pur preservando le capacità precedenti.

Significato e Rivendicazioni

Il paper sostiene che il contributo primario non sia semplicemente un nuovo euristica di selezione dei dati, ma un fondamentale riquadramento della preservazione della sicurezza come un problema condizionato all'aggiornamento. Gli autori sostengono che i dati di sicurezza non siano uniformemente utili; la loro rilevanza dipende dalla specifica direzione del task che si sta apprendendo. Attraverso l'esposizione dinamica dei vincoli attivi tramite riferimenti condizionati al task e il filtraggio dei task compatibili, DualSelect raggiunge un rapporto sicurezza-utilità superiore.

Gli autori notano con modestia che il metodo assume l'accesso a un segnale di riferimento (risposte sicure o target di vecchi task) e che le prestazioni dipendono dalla copertura dei riferimenti. Riconoscono inoltre che la valutazione della sicurezza si basa in parte su giudici automatici, sebbene utilizzino un'analisi tra giudici incrociati e metriche assolute per mitigare gli artefatti. Il lavoro si posiziona come una soluzione pratica e scalabile per il fine-tuning sicuro che evita gli alti costi del ri-pesamento globale o della complessa editing post-hoc.