PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Minchan Kwon, Sunghyun Baek, Minseo Kim, Jaemyung Yu, Dongyoon Han, Junmo Kim

arXiv:2605.00553v3cs.LG

TL;DR

LLM red-teaming must find attacks that are both effective and diverse, but GFlowNets can suffer instability and mode collapse. Stable-GFN removes partition-function estimation and filters noisy rewards, generating about seven times more unique attacks than GFN while maintaining a 92% attack success rate.

  • Problem

    LLM red-teaming needs effective and diverse attacks, but GFlowNet training can become unstable and collapse onto narrow high-reward distributions.

  • Method

    Stable-GFN uses contrastive trajectory balance to avoid partition-function estimation, plus noisy-reward filtering and fluency stabilization for red-teaming.

  • Results

    Approximately 7 times more unique attack prompts than GFN—134 versus 17—while maintaining a 92% attack success rate across diverse scenarios.

  • Takeaways & Limitations

    S-GFN discovers diverse vulnerabilities that provide superior defensive coverage through safety fine-tuning.

  • Takeaways & Limitations

    The fluency stabilizer’s reference-model likelihood is sensitive to sequence length and can penalize rare tokens, limiting exploration.

Abstract

from arXiv · show

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 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. Our code can be found in https://github.com/kmc0207/Stable-GFN.

1. Introduction

LLM red-teaming seeks diverse, high-impact attacks but must avoid reward-driven mode collapse and unstable training in high-dimensional discrete spaces. Stable-GFN addresses these challenges with contrastive trajectory balance and noise filtering, producing substantially more unique attacks while maintaining high attack success.

  • LLM red-teaming proactively discovers toxic attack prompts that expose safety vulnerabilities before deployment.
  • Diverse, high-impact attacks are important, but reward-maximizing RL attackers can converge on repetitive samples and miss a broad spectrum of vulnerabilities.
  • GFNs match distributions by sampling diverse trajectories in proportion to their rewards, but partition-function estimation makes naive formulations unstable in large combinatorial spaces.
  • Stable-GFN introduces Contrastive Trajectory Balance, which uses pairwise trajectory comparisons to offset partitioning functions, alongside Noise Gradient Pruning and other noise filtering techniques.
  • Approximately 7 times more unique attack prompts were generated than by the GFN baseline, increasing from 17 to 134 while maintaining a 92% attack success rate.

2. Related Work

Prior LLM red-teaming methods include reinforcement learning, Quality-Diversity optimization, and GFN-based distribution matching, alongside query-based attacks. GFN methods target joint toxicity and diversity, but TB-based LLM applications face mode collapse and training instability linked to estimating the partition function Z.

  • LLM Red-Teaming: LLM red-teaming methods broadly comprise RL-based training, Quality-Diversity optimization, and GFN-based distribution matching.RL methods treat toxicity as a reward; QD methods enforce diversity through Evolution Strategy or MAP; GFN methods jointly optimize diversity and toxicity.
  • LLM Red-Teaming: Recent RL-based work enhances diversity with diversity reward terms and curriculum learning, while Rainbow Teaming uses a matrix of fixed styles.The supplied passage introduces these as extensions of traditional RL and QD-based approaches, respectively.
  • GFN-Based Methods: GFN-based red-teaming methods match the entire reward distribution to seek broader coverage of the toxicity landscape.Prior applications include Lee et al.’s first use of GFN for red-teaming and Yun et al.’s iterative multi-stage GFN approach.
  • Attack Setting: The setup restricts attacks to a black-box scenario with only a single victim-LLM access per query.The related work also mentions query-based attacks and repeated question-answering for vulnerability discovery as prior approaches.
  • GFN Variants: Among GFN variants, only Trajectory Balance has been successfully applied to LLMs, and it is computationally lighter but suffers mode collapse and instability from estimating Z.The passage also lists Detailed Balance, Sub-Trajectory Balance, and Contrastive Balance as proposed variants.

3. Preliminary

LLM red-teaming trains an attacker LLM to generate prompts that elicit toxic victim-LLM responses, using toxicity-based rewards. Distribution matching and GFlowNets aim to favor high-reward attacks while preserving diversity, but explicit partition-function estimation can destabilize training and cause mode collapse.

  • Red-Teaming Objective: LLM red-teaming trains attacker LLM πθ to generate attack prompts y that elicit toxic responses z from victim LLM πϕ.The reward is based on toxicity score T(y, z) ∈ [0, 1] from classifier πψ.
  • Red-Teaming Objective: Maximizing expected toxicity can cause mode collapse, with the model generating a single high-reward sample with high probability.
  • Distribution Matching: Distribution matching instead trains the attacker to match a reward-defined target distribution, favoring high-reward attacks while preserving diversity.Attack prompts are token sequences y = [y1, . . . , yT], and a fixed meta-prompt is used instead of unconditional generation.
  • GFlowNets: GFlowNets model attack-prompt generation as token-action trajectories and learn a stochastic policy satisfying the reward-defined target condition.Trajectory Balance introduces a learnable scalar Zθ to estimate the partition function Z.
  • GFlowNets: Explicit estimation of Zθ produces high-variance gradients and training instability, which can force the model toward narrow, high-reward distributions and mode collapse.The proposed methodology avoids explicit Z estimation to address this issue.

4. Method

Stable-GFN replaces unstable partition-function estimation with pairwise Contrastive Trajectory Balance, filters noisy comparisons with saliency-based pruning, and stabilizes fluency through Min-K penalties. These components preserve the target optimal policy under stated conditions while reducing reward hacking on gibberish-like outputs.

  • Stable-GFN framework: Stable-GFN shifts GFN training from absolute global optimization to pairwise relative optimization, combining CTB, noisy-reward filtering, and fluency stabilization.The framework targets training instability, mode collapse, and limited search diversity in LLM red-teaming.
  • Contrastive Trajectory Balance: CTB avoids learning the partition function Zθ by comparing relative flows between trajectory pairs while exactly recovering TB’s optimal policy.Under positive rewards and full-support policies, the CTB minimum corresponds to πθ(y) = R(y)/Z.
  • Gradient analysis: CTB reduces gradient variance by centering each trajectory’s log-flow error against an implicit stochastic peer baseline.The peer-baseline structure is mathematically equivalent to policy-gradient variance-reduction methods such as RLOO.
  • Noisy Gradient Pruning: NGP filters trajectory pairs whose log-reward difference is at most σ, reducing overfitting to stochasticity from the victim LLM and toxic classifier.Its optimal policy remains proportional to R when the saliency graph is connected, while replay-buffer anchors can bridge comparison gaps in practice.
  • Min-K Fluency Stabilizer: MKS hard-penalizes samples whose Min-K statistic falls below threshold Tmks, excluding non-fluent OOD outputs from the high-reward region.Unlike global KL regularization, it targets the least-likely prompt segments and does not use the reference-policy gradient in reward calculation.

5. Experiment

The experiments evaluate S-GFN using attack success, diversity, cross-attack, transfer, reward-setting, and robustness analyses. Across these settings, S-GFN and its CTB/NGP components improve diversity, stability, and attack performance over baselines.

  • Evaluation Metric: ASR measures successful toxic responses among 1,024 prompts, while UA measures the breadth of discovered vulnerabilities.Each prompt is evaluated using five sampled victim responses, with toxicity threshold 0.5.
  • Attack on Target victim LLM: S-GFN achieves 134 unique attacks, substantially exceeding all baselines despite comparable ASR between GFN and PPO above 90%.PPO produces only 3 unique attacks, indicating overfitting to a narrow set of high-reward attacks.
  • Cross-Attack Results: 22.53% ASR remains when the victim is defended against GFN attacks, whereas S-GFN’s defense reduces GFN’s ASR to 0.03%.The cross-attack evaluation measures offensive transferability and defensive coverage through pairwise safety-finetuned defenses.
  • Transfer Attack Results: S-GFN demonstrates overwhelming transfer-attack performance on unseen Gemma3, Llama3.2, Qwen3, and gpt-oss-20B victim models.Jailbreak R1 also performs strongly across victim models because it was trained to generalize across diverse targets.
  • Impact of Reward Settings: Unconstrained rewards cause GFN-TB and GFN-CTB to exploit gibberish local minima, while KL regularization finds attacks but restricts diversity.MKS improves robustness to sentence length and tunes the trade-off between generation freedom and fluency.
  • Impact of Reward Settings: CTB preserves the optimal policy while improving training stability, and NGP increases both UA and ASR for CTB, including under noisy rewards.TB cannot use NGP because it lacks a pairwise objective; CTB+NGP remains effective in noisy-reward hypergrid experiments.

6. Conclusion

Stable-GFlowNet (S-GFN) addresses GFlowNet training instability by eliminating Z estimation through a CTB objective and adding methods to stabilize LLM red-teaming outputs. Experiments demonstrate superior performance over baselines.

  • 6. Conclusion: S-GFN eliminates Z estimation through a CTB objective to address fundamental GFlowNet training instabilities.This forms the framework’s core stabilization mechanism.
  • 6. Conclusion: NGP filters reward noise to improve stability in LLM red-teaming.The method is introduced specifically to manage noisy rewards.
  • 6. Conclusion: MKS preserves the linguistic integrity of generated gibberish prompts.It is introduced as a complementary stability mechanism for generated red-teaming prompts.
  • 6. Conclusion: S-GFN achieves superior performance over the baselines in the reported experiments.The conclusion explicitly states that the experimental results support this outcome.

Impact Statement … A.2. Mathematical Derivation of Equation (6)

The paper motivates proactive red-teaming to identify LLM vulnerabilities, develops theoretical support for CTB’s reward-normalized policy matching, and derives its implicit variance-reduction mechanism. The analyses also situate CTB relative to DPO and CB while clarifying when its pairwise structure provides distinctive stability benefits.

  • Impact Statement: The methodology can help proactively identify LLM vulnerabilities before deployment, although it could also enable misuse by inducing harmful actions.The paper notes that simple safety fine-tuning can achieve defense.
  • A.1. Proof of Theorem 4.1: CTB is formally minimized if and only if the policy πθ matches the reward-normalized density.The proof rewrites the pairwise loss through the log-flow error and uses its zero-variance condition.
  • Restatement of Theorem 4.1: The proof shows that the total CTB loss equals 2 · Varπθ(f(y)), so zero loss requires the log-flow error f(y) to be constant almost surely.Normalization then determines the constant through the partition function Z.
  • A.2. Mathematical Derivation of Equation (6): CTB’s expected gradient uses a peer sample as a stochastic baseline, focusing optimization on relative density matching rather than absolute reward magnitudes.This implicit variance reduction mirrors RLOO and enhances training stability without requiring an auxiliary mechanism.
  • Appendix: For approximately 2,500 unique samples, the saliency-graph connectivity threshold is approximately p ≈ 0.003.The effective batch size is 12 × 8 = 96, and standard settings are expected to exceed this threshold.
  • A. Theoretical Analyses and Formal Proofs: The connectivity analysis models unique samples as graph nodes and pairs exceeding saliency threshold σ as edges, using an Erdős–Rényi analogy.The result invoked is that connectivity occurs with high probability when p > ln n / n.
  • A.4.1. CONNECTION TO DIRECT PREFERENCE OPTIMIZATION (DPO): CTB and DPO both bypass explicit partition-function estimation through pairwise comparisons, but DPO maximizes reward-oriented margins whereas CTB performs distribution matching.DPO continuously increases winning-sample log-probability relative to losing samples, while CTB has a fixed target density.
  • A.4.2. CONNECTION TO CONTRASTIVE BALANCE (CB): CTB extends contrastive balance to autoregressive trajectories in a single-model LLM batch, and its analysis provides LLM-specific insights into gradient variance reduction and optimal-policy equivalence.CB instead integrates local GFNs with backward policies in federated learning.

A.3. Proof of Proposition 4.2 · B. Further Experiments · B.1. Implementation Details

The proof shows that, when the saliency graph is connected, minimizing NGP forces a globally constant log-flow error and recovers the same optimal policy as CTB/TB. The supplied implementation details specify the main training hyperparameters and replay-buffer filters.

  • A.3. Proof of Proposition 4.2: NGP preserves the optimal policy when the saliency graph Gσ satisfies the connectivity constraint.The proof establishes this under the proposition’s connectivity assumption.
  • A.3. Proof of Proposition 4.2: The NGP objective masks CTB gradients to pairs whose log-reward difference exceeds the saliency threshold σ.It uses Eσ = {(yi, yj) | |log R(yi) − log R(yj)| > σ}.
  • A.3. Proof of Proposition 4.2: Zero NGP loss makes the discrepancy f(y) equal across every edge of the saliency graph.For each edge, the proof states f(yi) = f(yj).
  • A.3. Proof of Proposition 4.2: Connectivity propagates edge-wise equality along paths, making f(y) a global constant over Y.Any two nodes are connected by a finite path, and transitivity yields f(ya) = f(yb).
  • A.3. Proof of Proposition 4.2: Normalization gives πθ(y) = R(y)/Z, so NGP’s global minimum recovers the full CTB/TB objective’s optimal policy.The proof derives e^C = 1/Z from summing πθ over Y.
  • B.1. Implementation Details: The implementation uses σ = 0.5, k = 7, learning rate = 1e−4, batch size = 12, and replay buffer size = 1000.These are the reported training hyperparameters.
  • B.1. Implementation Details: The replay buffer retains samples with cosine similarity below 0.4 and log reward greater than -2.5.It is initialized with randomly generated attacker samples before training.

B.2. Empirical Study of Z Estimation Variants

The study finds that removing learnable partition-function estimation improves stability broadly under noisy rewards, while CTB’s main advantage appears when paired with NGP.

  • Z-free objectives: Mean- and median-baseline objectives avoid explicit Z estimation by replacing log Zθ with batch statistics.Both variants use batch statistics as reference points and are evaluated under the same setup as Table 1.
  • CTB and NGP: CTB + NGP 121 UA and 92.2% ASR, substantially outperforming the other variants.NGP requires CTB’s explicit pair structure to selectively suppress destabilizing gradient updates under noisy rewards; batch-statistic baselines lack an analogous mechanism.
  • Z-estimation ablation: All three Z-free variants—Mean, Median, and CTB—substantially outperform standard TB under noisy rewards.Replacing learnable Zθ with batch-based or pairwise alternatives improves performance, indicating that Zθ is a primary instability source.
  • Z-estimation ablation: CTB alone performs comparably to Mean and Median, with no consistent advantage across metrics.Thus, the pairwise contrastive structure alone does not distinctly improve Z-estimation avoidance over simpler batch-statistic substitutes.

B.3. Additional Ablation Study

Additional ablations show that S-GFN converges faster than GFN, while the stabilizer’s k and threshold control the trade-off between attack effectiveness, diversity, and gibberish prevention. Performance saturates at k ≥ 6, and a threshold balance around -10 avoids the failures caused by excessively low values.

  • Sample Efficiency: S-GFN reaches an average log toxicity of -1 around the midpoint of the 200th step, whereas GFN approaches -1 only in the latter half of the 300th step.This demonstrates faster convergence for S-GFN under the same experimental settings.
  • Ablation Study of k: Unique-prompt performance increases with k but saturates at k ≥ 6.k controls how much gibberish a sentence can contain; overly large or small values weaken the stabilizer’s intended behavior.
  • Ablation Study on Stabilizer Threshold: At excessively low Tmks (-20 or below), the model fails to find a successful prompt.When the summed log probability is below Tmks, a penalty is applied, and overly low thresholds penalize most attacks.
  • Ablation Study on Stabilizer Threshold: The ablation finds a balance between ASR and UA around -10.This threshold avoids the failure to search for effective attacks associated with excessively low Tmks.

B.4. Additional Analysis

Additional analyses show that S-GFN achieves a strong effectiveness–diversity trade-off, identifies substantially more attack modes than GFN across clustering methods, and maintains stable behavior across masking settings. Safety fine-tuning preserves the victim model’s general capabilities, while S-GFN’s extra loss calculations do not translate into equivalent backpropagation overhead.

  • Diversity Analysis: S-GFN achieves the highest UA of 134.00 with ASR of 92.55%, while GFN shows mode collapse with 3-distinct of 0.02 and Self-BLEU of 0.98.S-GFN also raises 3-distinct to 0.38, expands vocabulary size from 77.67 to 1521.00, and reduces Self-BLEU to 0.64.
  • Effect of Clustering Method: S-GFN consistently produces the highest number of clusters across clustering methods, indicating broader attack-mode coverage than reward-maximization baselines.The analysis interprets S-GFN as finding over four times more peaks than GFN and reports richer, more independent clusters under Louvain analysis.
  • Connectivity Analysis: Masking performs best with σ between 0.1 and 0.5, improving training stability and unique-prompt discovery without changing the optimal policy.Values below 1.0 do not hinder training, whereas extreme settings can sever within-batch connections while retaining similar ASR and reducing unique prompts.
  • Computational Cost: S-GFN computes N^2 loss terms to form coefficients for N policy gradients, but identical gradients are combined during batch application rather than backpropagated independently.The reported CPU wall-clock comparison uses gradient accumulation set to 1.
  • Safety Fine-Tuning: Safety fine-tuning leaves MMLU performance nearly identical to the vanilla model, indicating that the model’s fundamental capabilities have hardly declined.The result also indicates that the model is not simply answering “no” to every attack.

B.4.1. ADDITIONAL QUALITATIVE RESULTS … B.5.2. NOISY HYPERGRID

Additional qualitative results show that S-GFN generates more complex and varied attack prompts than competing methods, while molecular-generation and noisy-hypergrid experiments demonstrate valid, diverse, and noise-robust distribution matching. In the hypergrid setting, CTB and CTB+NGP better capture the target modes than TB, including the peak at (12,4).

  • B.4.1. ADDITIONAL QUALITATIVE RESULTS: For MKS with k = 1 or 4, S-GFN avoids proper nouns and low-probability gibberish, while k = 1 imposes a particularly strong constraint.MKS with k = 1 behaves similarly to multiplying the reference model by the reward model.
  • B.4.1. ADDITIONAL QUALITATIVE RESULTS: S-GFN commonly identifies longer and more complex attacks, whereas GFN produces short, simple attacks and Rainbow Teaming targets predefined topics and styles.GFN’s instability limits broad coverage and makes it react strongly to specific peak rewards; Jailbreak-R1 generates lengthy scenario-based attacks, but greater diversity reduces attack success.
  • B.5.1. MOLECULAR GENERATION: The QM9 molecular-generation environment uses ten chemical fragments, permits paths of up to L = 10 fragments, and rewards generated SMILES with QED scores.Invalid molecules receive a minimum reward of 10^-3.
  • B.5.1. MOLECULAR GENERATION: GFN variants, including CTB, generate chemically valid molecules, while unique-molecule discovery indicates diversity even without quickly reaching the high-reward region.The passage specifically notes that CTB arrives at a result, but the supplied text truncates the remainder.
  • B.5.2. NOISY HYPERGRID: The noisy hypergrid uses a 16x16 grid with right-or-up moves and fixed path length L = 126, while rewards receive independent noise with standard deviation 0.3.The implementation uses four modes centered at (4,4), (12,4), (4,12), and (12,12), adds 10^-6 for numerical stability, and uses a 2-layer MLP with hidden size 256.
  • B.5.2. NOISY HYPERGRID: TB, CTB, and CTB+NGP model distributions similar to the true hypergrid distribution, but TB fails to adequately explore the peak at (12,4).CTB and CTB+NGP model that peak relatively well, supporting their shared theoretical convergence point.
  • B.5.2. NOISY HYPERGRID: CTB and CTB+NGP are more robust to noisy rewards than TB in the hypergrid experiment.The comparison is based on the converged distributions shown against the target distribution.
  • B.5.2. NOISY HYPERGRID: Adding NGP to CTB is evaluated by varying σ and measuring connectivity and JSD divergence, with batch connectivity equal to theoretical connectivity in the stated setting.This equality holds for batch size 256 and total space 16x16; the batch-connectivity rate is the proportion of batches forming connected graphs throughout the training process.

Limitations

S-GFN remains limited by biases and detection failures in its underlying toxic classifier, requiring classifier improvements for specific categories. Its current formulation also assumes single-turn attacks, leaving multi-turn behavior for future verification.

  • Classifier limitations: S-GFN does not resolve fundamental biases or issues in the underlying toxic classifier, and its performance is constrained by classifier accuracy.Detection failures in specific categories require concurrent improvements to the classifier itself.
  • Attack setting: S-GFN assumes single-turn attacks, so the behavior of S-GFN or GFN variants in multi-turn attacks requires additional verification.
Loading 2605.00553v3…