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

Large Language Model Red-Teaming, which proactively identifies vulnerabilities of large language models, is an essential process for ensuring safety. Finding effective and diverse attacks in red team activities is important, but achieving both is challenging. Generative Flow Networks (GFN) that perform distribution matching are a promising method, but they are notorious for training instability and mode collapse. In particular, unstable reward functions in red team activities accelerate mode collapse. We propose Stable-GFN (S-GFN), which eliminates 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.

Large language models (LLMs) like ChatGPT are trained to refuse harmful requests, but carefully crafted prompts can still trick them into producing dangerous content — instructions for violence, illegal drugs, sexual content directed at minors, and so on. Finding these "trick prompts" before a model is released, so that developers can patch the weaknesses, is called LLM red-teaming. Doing this by hand does not scale: humans cannot enumerate every possible weakness, especially as models grow. Automated red-teaming methods exist, but they face a stubborn trade-off — they either find a few highly effective attacks while missing entire categories of weakness, or they generate diverse prompts that mostly fail to trigger harmful responses. Patching only a narrow slice of attacks leaves the model exposed everywhere else. We trace this trade-off to instability during training: the automated attacker tends to collapse onto a handful of repetitive attacks long before it has explored the full landscape of weaknesses. We propose Stable-GFN, a training method that removes the main sources of this instability while preserving the same ideal target behavior as the previous best approach. As a result, Stable-GFN can keep exploring broadly and stably for much longer. In our experiments, Stable-GFN discovers roughly seven times more distinct types of successful attacks than the previous best method, while maintaining over 90% attack success rate. More importantly, when a model is patched against the diverse attacks found by Stable-GFN, it becomes substantially more robust against attacks from other methods as well — showing that broader red-teaming translates directly into broader protection. We believe this contributes a more reliable tool for safety teams to identify and fix vulnerabilities before deployment.