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

1. Probleemstelling

Red-teaming van Large Language Models (LLM) heeft tot doel proactief kwetsbaarheden (toxische of schadelijke outputs) te ontdekken voordat ze worden ingezet. Hoewel Reinforcement Learning (RL)-methodes hoge beloningsaanvallen kunnen vinden, lijden ze onder mode collapse, wat resulteert in repetitieve, smalle aanvalspatronen. Omgekeerd bieden Generative Flow Networks (GFN) een raamwerk voor distributie-matching, wat theoretisch de ontdekking van diverse trajecten met hoge beloning mogelijk maakt. Het toepassen van GFN op LLM red-teaming staat echter voor twee kritieke uitdagingen:

1. Trainingsinstabiliteit: Standaard GFN-doelstellingen (zoals Trajectory Balance, TB) vereisen het schatten van een globale partitiefunctie ($Z$). In de hoogdimensionale, discrete ruimte van LLM's is het schatten van $Z$ numeriek instabiel, wat leidt tot slechte convergentie en mode collapse.
2. Ruisondersteunde Beloningen: Toxischheidsclassificatoren leveren dichte maar ruizige signalen. Ze wijzen vaak niet-nul beloningen toe aan onzin of tekst die Out-of-Distribution (OOD) is, waardoor het model de beloning "hackt" door nonsens te genereren in plaats van vloeiende, diverse aanvallen.

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

De auteurs stellen Stable-GFlowNet (S-GFN) voor, een raamwerk dat de noodzaak voor expliciete schatting van de partitiefunctie elimineert en ruizige beloningen filtert om stabiele, diverse training te waarborgen. De pijplijn bestaat uit drie kerncomponenten:

A. Contrastive Trajectory Balance (CTB)

Om de instabiliteit van het schatten van $Z$ aan te pakken, vervangt S-GFN de absolute flow-matching doelstelling door een paarsgewijze vergelijkingsdoelstelling.

• Mechanisme: In plaats van de fout tussen de policy-flow en de beloning ten opzichte van $Z$ te minimaliseren, minimaliseert CTB het kwadraat van het verschil tussen de log-ratio van policy-kansen en de log-ratio van beloningen voor paren van trajecten $(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$$
• Theoretische Garantie: De auteurs bewijzen dat de optimale policy voor CTB identiek is aan die van standaard Trajectory Balance ($\pi_\theta(y) \propto R(y)$).
• Variance Reductie: Door één steekproef als basislijn te gebruiken voor een andere (vergelijkbaar met Reinforce Leave-One-Out), reduceert CTB de gradiëntvariatie impliciet zonder een externe basislijnnetwerk of $Z$-schatting nodig te hebben.

B. Noisy Gradient Pruning (NGP)

Om het risico op versterking van beloningsruis inherent aan paarsgewijze vergelijkingen te mitigeren, introduceert S-GFN een op salientie gebaseerde filter.

• Mechanisme: Gradiënten worden alleen berekend voor trajectparen waarbij het absolute verschil in log-beloningen een drempel $\sigma$ overschrijdt:
  $$L_{NGP} = \mathbb{I}(|\log R(y_1) - \log R(y_2)| > \sigma) \cdot L_{CTB}$$
• Effect: Dit voorkomt dat de optimizer updates uitvoert op basis van onbeduidende beloningsverschillen veroorzaakt door de stochasticiteit van het slachtoffer-LLM of de toxischheidsclassificator, en snoeit effectief ruizige gradiënten weg.

C. Min-K Fluency Stabilizer (MKS)

Om te voorkomen dat het model convergeert naar lokale minima van onzin (die vaak vals-positieve toxischheidsscores krijgen), hanteert S-GFN een vloeiheidsbeperking.

• Mechanisme: In plaats van een globale KL-divergentie-straf (die kan overfitten op een referentiemodel), berekent MKS de Min-K-statistiek: het gemiddelde log-kans van de $k$ minst waarschijnlijke tokens in een gegenereerde sequentie onder een referentiemodel.
• Harde Straf: Als de Min-K-score onder een drempel $T_{mks}$ daalt, wordt de beloning afgekapt tot nul (of een grote negatieve waarde).
  $$R_{MKS}(y) = \mathbb{I}(M_k(y) \ge T_{mks}) \cdot R(y)$$
• Voordeel: Dit zorgt ervoor dat alleen vloeiende, linguïstisch coherente prompts als hoge beloning worden beschouwd, waardoor het model een bredere ruimte van geldige aanvallen kan verkennen zonder onzin te genereren.

3. Belangrijkste Bijdragen

1. Stable-GFN Raamwerk: Een nieuw red-teaming raamwerk dat de noodzaak voor schatting van de partitiefunctie ($Z$) elimineert, waardoor de primaire oorzaak van trainingsinstabiliteit in GFN voor LLM's wordt opgelost.
2. Contrastive Trajectory Balance (CTB): Een nieuwe doelstellingsfunctie die dezelfde optimale distributie bereikt als TB, maar met aanzienlijk lagere variatie en hogere trainingsstabiliteit.
3. Robuustheid tegen Ruis: De integratie van NGP en MKS adresseert specifiek de unieke uitdagingen van red-teaming (ruizige classificatoren en beloningen voor onzin), waardoor mode collapse en belonings-hacking worden voorkomen.
4. Universaliteit: Het paper toont aan dat CTB en NGP niet alleen effectief zijn in red-teaming, maar ook in andere distributie-matching taken zoals moleculaire generatie en hypergrid-navigatie.

4. Experimentele Resultaten

De auteurs evalueerden S-GFN tegen baselines waaronder PPO, Jailbreak R1, Rainbow Teaming en standaard GFN (TB) over diverse slachtoffermodellen (Qwen, Llama, Gemma) en verdedigingsinstellingen.

• Diversiteit (Unieke Aanvallen): S-GFN genereerde 134 unieke aanvalsprompts, een enorme verbetering ten opzichte van de standaard GFN-baseline (17) en PPO (3). Het produceerde ongeveer 7x meer unieke aanvallen dan de GFN-baseline.
• Attack Success Rate (ASR): S-GFN behield een hoge ASR van ~92,5%, vergelijkbaar met de best presterende baselines (GFN: 93,7%, PPO: 91,7%), wat bewijst dat diversiteit niet ten koste ging van effectiviteit.
• Cross-Attack Transferability: Wanneer slachtoffermodellen veiligheids-finetuned werden met aanvallen van S-GFN, werden ze aanzienlijk robuuster tegen aanvallen van andere methoden (ASR daalde naar bijna 0% voor GFN-aanvallen). Omgekeerd bleven S-GFN-aanvallen effectief tegen modellen die op andere methoden waren gefinetuned, wat superieure generalisatie aantoont.
• Transfer Attacks: S-GFN slaagde erin aanvallen over te dragen naar ongezette slachtoffermodellen (bijv. Gemma3, Llama3.2) en presteerde beter dan andere methoden in zowel ASR als diversiteit in transfer-scenario's.
• Efficiëntie: Ondanks de $O(N^2)$ paarsgewijze vergelijkingen is de rekentijd verwaarloosbaar omdat dit scalaire operaties zijn; de daadwerkelijke backpropagation-kost blijft $O(N)$, waardoor S-GFN bijna even efficiënt is als standaard GFN.

5. Betekenis

Dit paper vertegenwoordigt een aanzienlijke vooruitgang in LLM-veiligheidsresearch. Door de trainingsinstabiliteit van GFN op te lossen, maakt het de systematische ontdekking van diverse, hoog-impact kwetsbaarheden mogelijk die traditionele RL-methodes missen door mode collapse.

• Voor Veiligheid: Het biedt een robuuster hulpmiddel voor red-teaming, waardoor ontwikkelaars een breder spectrum aan aanvalsvectoren kunnen identificeren voordat ze worden ingezet.
• Voor Verdediging: De diverse aanvallen gegenereerd door S-GFN leiden tot veiligheids-finetuned modellen die veerkrachtiger zijn tegen een breder scala aan toekomstige aanvallen.
• Voor Machine Learning: De voorgestelde Contrastive Trajectory Balance biedt een generaliseerbare oplossing voor distributie-matching in hoogdimensionale, ruizige omgevingen, en breidt de toepasbaarheid van GFN uit buiten red-teaming naar velden zoals geneesmiddelenontdekking en materiaalkunde.

Kortom, Stable-GFlowNet slaagt erin de kloof te overbruggen tussen de theoretische diversiteit van GFN en de praktische stabiliteit die vereist is voor red-teaming van LLM's in de echte wereld, en zet een nieuwe state-of-the-art voor het genereren van diverse en effectieve adversariële prompts.