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

1. Il Problema: Allineamento Superficiale e Vulnerabilità agli Attacchi

Nonostante i progressi nelle tecniche di allineamento come il Supervised Fine-Tuning (SFT), il Reinforcement Learning from Human Feedback (RLHF) e il Direct Preference Optimization (DPO), i Large Language Models (LLM) rimangono vulnerabili agli attacchi di "jailbreak". Questi attacchi aggirano le difese utilizzando formulazioni indirette, ingannevoli o criptate per estrarre contenuti dannosi.

Gli autori ipotizzano che la causa principale di questa vulnerabilità sia un allineamento superficiale. I modelli attuali tendono a imparare euristiche di rifiuto superficiali (riconoscimento di pattern) piuttosto che un ragionamento profondo sul perché un contenuto sia dannoso.

• Evidenza Causale: Attraverso un'intervento causale (disattivazione selettiva dei neuroni critici per il ragionamento), gli autori dimostrano che la capacità di ragionamento del modello crolla, mentre il comportamento di allineamento (rifiuto delle richieste dannose) rimane invariato. Questo conferma che l'allineamento attuale non si basa su una comprensione logica profonda, rendendo i modelli suscettibili a jailbreak che bypassano i pattern superficiali.

2. Metodologia Proposta

Per risolvere questo problema, il paper propone un approccio in due fasi che combina l'addestramento sul ragionamento e un'ottimizzazione delle preferenze più granulare.

A. Dataset di Fine-Tuning Chain-of-Thought (CoT)

Gli autori costruiscono e rilasciano un nuovo dataset di addestramento basato sul ragionamento (CoT) che include:

• Prompt orientati all'utilità e prompt critici per la sicurezza.
• Tracce di ragionamento passo-passo (rationales) dettagliate.
• L'obiettivo è insegnare al modello non solo a dire "no", ma a spiegare perché una richiesta è dannosa, basando i rifiuti su un ragionamento logico profondo invece che su pattern superficiali.

B. Alignment-Weighted DPO (AW-DPO)

L'analisi degli errori sul modello fine-tunato con CoT rivela due modalità di fallimento specifiche:

1. Ragionamento corretto ma risposta finale insicura.
2. Ragionamento errato ma risposta finale sicura.

Il DPO standard tratta l'intera risposta come un blocco unico, ottimizzando solo la preferenza globale. Per affrontare i fallimenti granulari, gli autori introducono AW-DPO:

• Decomposizione: Ogni risposta viene divisa in due segmenti: la traccia di ragionamento (reasoning trace) e la risposta finale (final response).
• Assegnazione di Pesi: Vengono assegnati pesi di preferenza distinti a ciascun segmento in base al loro livello di dannosità. Se il ragionamento è dannoso ma la risposta è sicura (o viceversa), il modello viene penalizzato specificamente sul segmento errato.
• Formula: La funzione di perdita DPO viene modificata per pesare i contributi del ragionamento ($L_{rs}$) e della risposta ($L_{rp}$) separatamente:
  $$L_{AW-DPO} = w_{reasoning} \cdot L_{DPO}^{rs} + w_{response} \cdot L_{DPO}^{rp}$$
  dove i pesi $w$ sono calcolati in base alla differenza di score di dannosità tra la risposta scelta e quella rifiutata per ciascun segmento.

3. Risultati Sperimentali

Gli esperimenti sono stati condotti su diversi modelli (Llama-2-7B, Llama-3.2-3B, Llama-3.1-8B, Mistral-7B) utilizzando benchmark di sicurezza (SorryBench con 20 tipi di jailbreak) e di utilità (MMLU).

• Sicurezza: Il metodo AW-DPO supera costantemente i baseline (SFT, DPO standard, CoT SFT) riducendo drasticamente il Attack Success Rate (ASR). Ad esempio, su Llama-3.2-3B, l'ASR medio scende dal 71.59% (Base) a 1.14% con AW-DPO, superando modelli allineati open-source avanzati.
• Utilità: A differenza di alcuni metodi di allineamento aggressivo che degradano le prestazioni generali, AW-DPO mantiene un'alta utilità (accuratezza su MMLU), dimostrando che il miglioramento della sicurezza non compromette la capacità del modello di essere utile.
• Robustezza: Il metodo dimostra robustezza contro attacchi di "prefix" (dove viene forzato l'omissione del ragionamento) e mostra buone capacità di trasferimento quando il dataset AW-DPO è addestrato su un modello e applicato ad architetture diverse.
• Confronto con Modelli di Ragionamento: I modelli puramente orientati al ragionamento (es. Phi-4-Reasoning) mostrano prestazioni di sicurezza inferiori rispetto ad AW-DPO, confermando che il ragionamento generale non è sufficiente senza un addestramento specifico per l'allineamento.

4. Contributi Chiave

1. Dimostrazione Causale: Fornisce prove empiriche che l'allineamento di sicurezza attuale è indipendente dalle capacità di ragionamento profondo, rendendo i modelli vulnerabili a jailbreak sofisticati.
2. Dataset CoT Open-Source: Rilascia un nuovo dataset di fine-tuning CoT che combina sicurezza e utilità, colmando un gap nella disponibilità di dati di alta qualità per l'allineamento basato sul ragionamento.
3. Algoritmo AW-DPO: Introduce un metodo di ottimizzazione delle preferenze che decompone la risposta in ragionamento e output finale, permettendo correzioni mirate e granulari che il DPO standard non riesce a gestire.
4. Performance Superiori: Dimostra sperimentalmente che l'approccio proposto migliora la robustezza alla sicurezza mantenendo l'utilità del modello, superando le tecniche di stato dell'arte.

5. Significato e Impatto

Questo lavoro è significativo perché sposta il paradigma dell'allineamento dei LLM da un approccio basato su "pattern matching" superficiale a uno basato su ragionamento principiato.

• Sicurezza: Offre una difesa più robusta contro gli attacchi jailbreak che sfruttano l'ambiguità o la complessità linguistica, poiché il modello impara a valutare la dannosità attraverso una catena logica esplicita.
• Efficienza: AW-DPO raggiunge risultati superiori con un solo ciclo di addestramento SFT + DPO, evitando i costi computazionali elevati di metodi iterativi complessi (come STAIR-DPO-3).
• Interpretabilità: La separazione tra ragionamento e risposta rende il processo di allineamento più interpretabile, permettendo di identificare e correggere esattamente dove il modello fallisce (nel pensiero o nell'azione).

In sintesi, il paper propone che per ottenere un allineamento sicuro e robusto, i modelli devono essere addestrati a ragionare sulla sicurezza, non solo a rifiutare superficialmente, e che l'ottimizzazione delle preferenze deve essere applicata con precisione chirurgica alle diverse parti della generazione del modello.