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

TL;DR

Stable-GFN introduces a robust training method for generative flow networks to improve diversity and stability in LLM red-teaming, effectively identifying vulnerabilities.

Contribution

It eliminates Z-estimation in GFNs, employs pairwise comparisons and a fluency stabilizer, enhancing training stability and attack diversity in LLM red-teaming.

Findings

S-GFN achieves more stable training compared to traditional GFNs.

S-GFN demonstrates superior attack performance across various settings.

The method maintains the optimal policy while improving diversity and robustness.

Abstract

Large Language Model (LLM) Red-Teaming, which proactively identifies vulnerabilities of LLMs, is an essential process for ensuring safety. Finding effective and diverse attacks in red-teaming is important, but achieving both is challenging. Generative Flow Networks (GFNs) that perform distribution matching are a promising methods, but they are notorious for training instability and mode collapse. In particular, unstable rewards in red-teaming accelerate mode collapse. We propose Stable-GFN (S-GFN), which eliminates partition function $Z$ estimation in GFN and reduces training instability. S-GFN avoids Z-estimation through pairwise comparisons and employs a robust masking methodology against noisy rewards. Additionally, we propose a fluency stabilizer to prevent the model from getting stuck in local optima that produce gibberish. S-GFN provides more stable training while maintaining the optimal policy of GFN. We demonstrate the overwhelming attack performance and diversity of S-GFN across various settings.