Stable-GFlowNet: Toward Diverse and Robust LLM Red-Teaming via Contrastive Trajectory Balance

1. Enunciato del Problema

Il red-teaming dei Large Language Model (LLM) mira a scoprire proattivamente le vulnerabilità (output tossici o dannosi) prima del dispiegamento. Sebbene i metodi di Reinforcement Learning (RL) possano individuare attacchi ad alto reward, soffrono di collasso modale, generando pattern di attacco ripetitivi e ristretti. Al contrario, le Generative Flow Networks (GFN) offrono un framework per il matching delle distribuzioni, permettendo teoricamente la scoperta di traiettorie ad alto reward e diversificate. Tuttavia, l'applicazione delle GFN al red-teaming degli LLM incontra due sfide critiche:

1. Instabilità dell'Addestramento: Gli obiettivi standard delle GFN (come il Trajectory Balance, TB) richiedono la stima di una funzione di partizione globale ($Z$). Nello spazio discreto e ad alta dimensionalità degli LLM, la stima di $Z$ è numericamente instabile, portando a una scarsa convergenza e a un collasso modale.
2. Reward Rumorosi: I classificatori di tossicità forniscono segnali densi ma rumorosi. Spesso assegnano reward non nulli a testo privo di senso o Out-of-Distribution (OOD), inducendo il modello a "hackerare" il reward generando testo nonsensico invece di attacchi fluenti e diversificati.

2. Metodologia: Stable-GFlowNet (S-GFN)

Gli autori propongono Stable-GFlowNet (S-GFN), un framework che elimina la necessità di una stima esplicita della funzione di partizione e filtra i reward rumorosi per garantire un addestramento stabile e diversificato. La pipeline è composta da tre componenti fondamentali:

A. Contrastive Trajectory Balance (CTB)

Per affrontare l'instabilità della stima di $Z$, S-GFN sostituisce l'obiettivo di matching del flusso assoluto con un obiettivo di confronto pairwise.

• Meccanismo: Invece di minimizzare l'errore tra il flusso della policy e il reward rispetto a $Z$, CTB minimizza la differenza quadrata tra il log-rapporto delle probabilità della policy e il log-rapporto dei reward per coppie di traiettorie $(y_1, y_2)$:
  $$L_{CTB} = \left( \log \frac{\pi_\theta(y_1)}{\pi_\theta(y_2)} - \log \frac{R(y_1)}{R(y_2)} \right)^2$$
• Garanzia Teorica: Gli autori dimostrano che la policy ottimale per CTB è identica a quella del Trajectory Balance standard ($\pi_\theta(y) \propto R(y)$).
• Riduzione della Varianza: Utilizzando un campione come baseline per un altro (simile a Reinforce Leave-One-Out), CTB riduce implicitamente la varianza del gradiente senza necessità di una rete baseline esterna o della stima di $Z$.

B. Noisy Gradient Pruning (NGP)

Per mitigare il rischio di amplificazione del rumore nei reward intrinseco ai confronti pairwise, S-GFN introduce un filtro basato sulla salienza.

• Meccanismo: I gradienti sono calcolati solo per le coppie di traiettorie in cui la differenza assoluta nei log-reward supera una soglia $\sigma$:
  $$L_{NGP} = \mathbb{I}(|\log R(y_1) - \log R(y_2)| > \sigma) \cdot L_{CTB}$$
• Effetto: Ciò impedisce all'ottimizzatore di aggiornarsi basandosi su differenze di reward insignificanti causate dalla stocasticità dell'LLM vittima o del classificatore di tossicità, eliminando efficacemente i gradienti rumorosi.

C. Min-K Fluency Stabilizer (MKS)

Per prevenire che il modello converga verso minimi locali di testo privo di senso (che spesso riceve punteggi di tossicità spurii), S-GFN impiega un vincolo di fluidità.

• Meccanismo: Invece di una penalità globale di divergenza KL (che può portare a overfitting su un modello di riferimento), MKS calcola la statistica Min-K: la media dei log-probabilità dei $k$ token meno probabili in una sequenza generata sotto un modello di riferimento.
• Penalità Rigida: Se il punteggio Min-K scende sotto una soglia $T_{mks}$, il reward viene clipato a zero (o a un valore negativo elevato).
  $$R_{MKS}(y) = \mathbb{I}(M_k(y) \ge T_{mks}) \cdot R(y)$$
• Vantaggio: Ciò garantisce che solo prompt fluenti e linguisticamente coerenti siano considerati ad alto reward, permettendo al modello di esplorare uno spazio più ampio di attacchi validi senza generare assurdità.

3. Contributi Chiave

1. Framework Stable-GFN: Un nuovo framework di red-teaming che elimina la necessità di stimare la funzione di partizione ($Z$), risolvendo la causa principale dell'instabilità nell'addestramento delle GFN per gli LLM.
2. Contrastive Trajectory Balance (CTB): Una nuova funzione obiettivo che raggiunge la stessa distribuzione ottimale del TB ma con varianza significativamente inferiore e maggiore stabilità nell'addestramento.
3. Robustezza al Rumore: L'integrazione di NGP e MKS affronta specificamente le sfide uniche del red-teaming (classificatori rumorosi e reward da testo privo di senso), prevenendo il collasso modale e l'hacking del reward.
4. Universalità: Il paper dimostra che CTB e NGP sono efficaci non solo nel red-teaming, ma anche in altri task di matching delle distribuzioni come la generazione molecolare e la navigazione in ipergriglie.

4. Risultati Sperimentali

Gli autori hanno valutato S-GFN contro baseline tra cui PPO, Jailbreak R1, Rainbow Teaming e GFN standard (TB) su vari modelli vittime (Qwen, Llama, Gemma) e impostazioni di difesa.

• Diversità (Attacchi Unici): S-GFN ha generato 134 prompt di attacco unici, un miglioramento massiccio rispetto alla baseline GFN standard (17) e PPO (3). Ha prodotto circa 7 volte più attacchi unici rispetto alla baseline GFN.
• Tasso di Successo dell'Attacco (ASR): S-GFN ha mantenuto un ASR elevato di ~92.5%, comparabile alle baseline migliori (GFN: 93.7%, PPO: 91.7%), dimostrando che la diversità non è arrivata a scapito dell'efficacia.
• Trasferibilità Cross-Attacco: Quando i modelli vittime sono stati safety-fine-tuned utilizzando attacchi provenienti da S-GFN, sono diventati significativamente più robusti contro attacchi di altri metodi (l'ASR è sceso a quasi lo 0% per gli attacchi GFN). Al contrario, gli attacchi S-GFN sono rimasti efficaci contro modelli fine-tuned su altri metodi, dimostrando una generalizzazione superiore.
• Attacchi di Trasferimento: S-GFN ha trasferito con successo gli attacchi a modelli vittime non visti (es. Gemma3, Llama3.2), superando altri metodi sia nell'ASR che nella diversità negli scenari di trasferimento.
• Efficienza: Nonostante i confronti pairwise $O(N^2)$, il costo computazionale è trascurabile poiché si tratta di operazioni scalari; il costo effettivo della backpropagation rimane $O(N)$, rendendo S-GFN quasi efficiente quanto la GFN standard.

5. Significato

Questo paper rappresenta un avanzamento significativo nella ricerca sulla sicurezza degli LLM. Risolvendo l'instabilità dell'addestramento delle GFN, permette la scoperta sistematica di vulnerabilità diversificate e ad alto impatto che i metodi RL tradizionali mancano a causa del collasso modale.

• Per la Sicurezza: Fornisce uno strumento più robusto per il red-teaming, permettendo agli sviluppatori di identificare uno spettro più ampio di vettori di attacco prima del dispiegamento.
• Per la Difesa: Gli attacchi diversificati generati da S-GFN portano a modelli safety-fine-tuned più resilienti a una gamma più ampia di futuri attacchi.
• Per il Machine Learning: Il Contrastive Trajectory Balance proposto offre una soluzione generalizzabile per il matching delle distribuzioni in ambienti ad alta dimensionalità e rumorosi, estendendo l'applicabilità delle GFN oltre il red-teaming a campi come la scoperta di farmaci e la scienza dei materiali.

In sintesi, Stable-GFlowNet colma con successo il divario tra la diversità teorica delle GFN e la stabilità pratica richiesta per il red-teaming degli LLM nel mondo reale, stabilendo un nuovo stato dell'arte per la generazione di prompt avversari diversificati ed efficaci.