Secure LLM Fine-Tuning via Safety-Aware Probing

1. Il Problema: Degrado della Sicurezza nel Fine-Tuning

Nonostante i modelli linguistici di grandi dimensioni (LLM) subiscano processi di allineamento per la sicurezza durante le fasi di pre-training o post-training, la loro sicurezza rimane vulnerabile durante il fine-tuning su nuovi dati.

• Fenomeno chiave: Studi recenti hanno dimostrato che anche il fine-tuning su dati benigni (utili per compiti specifici) può compromettere la sicurezza del modello.
• Causa radice: Gli autori ipotizzano e dimostrano che le direzioni critiche per l'utilità (gradienti che migliorano le prestazioni del compito) e le direzioni critiche per la sicurezza (gradienti che riducono i rischi di output dannosi) sono parzialmente disaccoppiate ma entangled (intrecciate).
• Conseguenza: L'ottimizzazione diretta per migliorare le prestazioni su un compito specifico spinge involontariamente il modello verso regioni dello spazio dei parametri con un rischio di sicurezza più elevato, degradando l'allineamento iniziale.

2. Metodologia: Safety-Aware Probing (SAP)

Per mitigare questo rischio senza modificare i dati di addestramento o limitarsi a tecniche di regolarizzazione specifiche (come LoRA), gli autori propongono SAP, un framework di ottimizzazione a due livelli (bi-level optimization).

Concetto Fondamentale

SAP interviene direttamente nel processo di ottimizzazione perturbando la propagazione degli stati nascosti (hidden states) durante il training, invece di modificare solo i pesi finali o filtrare i dati.

Meccanismo Tecnico

1. Analisi dei Gradienti: SAP calcola la direzione critica dannosa ($\Delta W_{harmful}$) utilizzando un loss contrastivo di sicurezza. Questo viene ottenuto confrontando i gradienti su coppie di dati: risposte sicure ($D_{safe}$) e risposte dannose ($D_{harmful}$) per gli stessi prompt.
   • La direzione dannosa è definita come: $+\nabla_W L_{safe} = \nabla_W L(W, D_{harmful}) - \nabla_W L(W, D_{safe})$.
2. Probe di Sicurezza (V): Viene introdotto un vettore di probe leggero $V$ che viene aggiunto agli stati nascosti del modello durante la propagazione in avanti.
3. Ottimizzazione Bi-Livello:
   • Passo Interno (Massimizzazione): Si ottimizza il probe $V$ per massimizzare una funzione di perdita "sicura-utilità" ($L_{su}$). L'obiettivo è trovare una perturbazione degli stati nascosti che renda il paesaggio di perdita localmente sfavorevole agli aggiornamenti dannosi. In altre parole, si cerca di rendere costoso per il modello muoversi verso direzioni dannose.
   • Passo Esterno (Minimizzazione): Con il probe $V_{safe}$ fissato, si aggiornano i parametri del modello $W$ per minimizzare la perdita sul compito utile ($L_{useful}$).
4. Efficienza: A differenza di metodi che richiedono copie duplicate dei pesi (come SAM - Sharpness Aware Minimization), SAP perturba solo gli stati nascosti, mantenendo l'overhead di memoria GPU quasi identico al fine-tuning standard.

3. Contributi Chiave

• Analisi Teorica: Prima quantificazione formale dell'entanglement tra direzioni utili e dannose tramite analisi dei gradienti, spiegando perché il fine-tuning su dati benigni degrada la sicurezza.
• Nuovo Paradigma di Ottimizzazione: Introduzione di SAP, un framework che non dipende da dataset specifici (filtraggio) o architetture restrittive (es. solo LoRA), rendendolo scalabile e applicabile a fine-tuning completo o parziale.
• Robustezza Dimostrata: Validazione della metodologia non solo in scenari benigni, ma anche sotto attacchi avversari (avvelenamento dei dati, fine-tuning dannoso diretto e attacchi adattivi post-training).

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su diversi modelli (Llama-2, Vicuna, Qwen) e compiti (Alpaca, Samsum, ChatDoctor, ragionamento logico).

• Prestazioni di Sicurezza: SAP riduce significativamente il Harmful Score (HS) rispetto al fine-tuning standard (SFT).
  • Esempio: Su una media di tre LLM, SAP riduce l'HS dal 37.43% (SFT) al 23.07%, superando baselines forti come SafeInstr (26.20%) e Lisa (25.93%).
• Prestazioni del Compito: SAP mantiene prestazioni competitive sui task specifici (misurate tramite BLEURT e Cross-Entropy), e in alcuni casi le migliora leggermente grazie a un effetto di regolarizzazione simile a SAM.
• Robustezza Adversariale:
  • Data Poisoning: Anche con un tasso di avvelenamento del 25%, SAP mantiene un HS molto più basso rispetto alle controparti (29.1% vs >32% per altri metodi).
  • Attacchi Adattivi: SAP resiste meglio agli attacchi adattivi che cercano di invertire il meccanismo di difesa, mantenendo un margine di sicurezza significativo.
• Efficienza Computazionale:
  • Il tempo per batch è circa 2.5 volte superiore a SFT, ma grazie a una convergenza più rapida, SAP raggiunge risultati superiori con meno epoche di training rispetto alle baselines.
  • L'uso di memoria GPU è quasi identico a quello del fine-tuning standard (<2% di aumento).

5. Significato e Impatto

Questo lavoro è significativo perché sposta il focus dalla semplice filtrazione dei dati o dalla regolarizzazione dei parametri verso una correzione dinamica del paesaggio di ottimizzazione.

• Scalabilità: SAP è agnostico rispetto all'architettura e al metodo di fine-tuning (funziona sia con LoRA che con il fine-tuning completo dei parametri).
• Sicurezza Pratica: Offre una soluzione pratica per le API di fine-tuning commerciali e le applicazioni a valle, garantendo che l'adattamento a compiti specifici non comprometta le barriere di sicurezza fondamentali.
• Fondamentale: Fornisce una comprensione più profonda della dinamica dei gradienti negli LLM, dimostrando che la sicurezza può essere preservata "guidando" l'ottimizzazione lontano dalle trappole del paesaggio di perdita, piuttosto che bloccare l'ottimizzazione stessa.

In sintesi, SAP rappresenta un approccio robusto ed efficiente per preservare l'allineamento di sicurezza degli LLM durante il loro adattamento a nuovi domini, risolvendo il compromesso storico tra utilità del compito e sicurezza.