PaperlayerView exact revision on arXiv

Source-linked AI summary

Learning diverse attacks on large language models for robust red-teaming and safety tuning

Seanie Lee, Minsu Kim, Lynn Cherif, David Dobre, Juho Lee, Sung Ju Hwang, Kenji Kawaguchi, Gauthier Gidel, Yoshua Bengio, Esmeralda S. Whitammer, Moksh Jain

arXiv:2405.18540v3cs.CLcs.CRcs.LG

TL;DR

Automated red-teaming needs diverse attacks because existing RL-based methods can collapse onto narrow or ineffective modes. The paper uses two-stage GFlowNet fine-tuning with MLE smoothing to generate attack prompts, finding broad effectiveness, transfer across target LLMs, and improved robustness after safety tuning.

  • Problem

    Existing automated red-teaming methods struggle to generate attack prompts that are both diverse and effective, even with diversity-oriented regularization.

  • Method

    The method fine-tunes an attacker LM with a GFlowNet objective, collects high-reward prompts, and applies MLE smoothing to them.

  • Results

    The generated attacks are more diverse and effective across five target LLMs, transfer to unseen models, and support safety tuning that resists other RL-based attacks.

  • Takeaways & Limitations

    Diverse attack prompts can reveal shared failure modes across target LLMs and improve the robustness of downstream safety tuning.

  • Takeaways & Limitations

    Performance is limited by the toxicity classifier and costly online training that repeatedly samples target-LLM responses.

Abstract

from arXiv · show

Red-teaming, or identifying prompts that elicit harmful responses, is a critical step in ensuring the safe and responsible deployment of large language models (LLMs). Developing effective protection against many modes of attack prompts requires discovering diverse attacks. Automated red-teaming typically uses reinforcement learning to fine-tune an attacker language model to generate prompts that elicit undesirable responses from a target LLM, as measured, for example, by an auxiliary toxicity classifier. We show that even with explicit regularization to favor novelty and diversity, existing approaches suffer from mode collapse or fail to generate effective attacks. As a flexible and probabilistically principled alternative, we propose to use GFlowNet fine-tuning, followed by a secondary smoothing phase, to train the attacker model to generate diverse and effective attack prompts. We find that the attacks generated by our method are effective against a wide range of target LLMs, both with and without safety tuning, and transfer well between target LLMs. Finally, we demonstrate that models safety-tuned using a dataset of red-teaming prompts generated by our method are robust to attacks from other RL-based red-teaming approaches.

1 INTRODUCTION

The paper addresses automated red-teaming that must discover diverse, effective attack prompts, because existing methods can collapse onto limited attack modes. It proposes two-stage GFlowNet fine-tuning with MLE smoothing and reports effectiveness, transferability, and downstream safety benefits.

  • Existing automated red-teaming methods struggle to balance attack diversity and effectiveness, often suffering from mode collapse despite diversity regularization.
  • The method first fine-tunes an attacker LM as a GFlowNet policy, then applies MLE to high-reward prompts collected during the first stage.The second stage smooths the distribution by maximizing likelihood of collected high-reward attacks.
  • The paper evaluates the approach on five target LLMs and finds it generates more diverse and effective attacks than relevant baselines.
  • Generated attacks transfer to multiple unseen target LLMs, and safety tuning on these prompts improves robustness to attacks from other RL-based red-teamers.
  • The paper uses toxicity as an approximate harmfulness proxy, while acknowledging that actual harm is subjective and context-dependent.

2 RELATED WORK

Related work frames red-teaming as proactive discovery of harmful behaviors, distinguishes it from jailbreak circumvention, and situates GFlowNets as a framework for diverse high-reward discrete-object generation.

  • Red-teaming: Red-teaming proactively searches for harmful model behavior to improve safety, including through filtering and safety fine-tuning during model development.
  • Jailbreaks: Jailbreaking generally seeks to circumvent safeguards, whereas red-teaming discovers vulnerabilities for improving model safety.
  • Jailbreaks: Automated jailbreak methods include genetic algorithms, gradient-based optimization, and auxiliary-LLM prompting, but can be costly per query and transfer poorly across target models.
  • GFlowNets: GFlowNets train stochastic policies to sample discrete compositional objects proportionally to reward, producing diverse high-reward samples.

3 SAMPLING DIVERSE ATTACKS WITH GFLOWNET FINE-TUNING

The method treats attack-prompt generation as sampling from a reward distribution, using GFlowNet fine-tuning to seek diverse high-reward prompts and MLE smoothing to reduce sensitivity to reward temperatures.

  • 3.2 GFLOWNET FINE-TUNING AND SMOOTHING WITH MLE ON COLLECTED HIGH-REWARD PROMPTS: The prompt reward combines expected toxicity under the target LLM with reference-model likelihood, while β and γ control the peakiness of these two components.Toxicity is estimated from classifier scores over responses sampled from the target LLM, and the prompt is generated autoregressively as a sequence of token actions.
  • 3.1 PRELIMINARIES: Automated red-teaming methods can collapse onto a few high-reward prompts because standard RL concentrates probability mass on the highest-reward actions despite diversity regularization.The KL term keeps prompts near a pretrained language model, but does not by itself solve the diversity–effectiveness trade-off.
  • 3.2 GFLOWNET FINE-TUNING AND SMOOTHING WITH MLE ON COLLECTED HIGH-REWARD PROMPTS: GFlowNets train the attacker policy to sample prompts with probability proportional to reward, targeting diverse and effective attacks rather than concentrating on a single high-reward prompt.The reward combines target-response toxicity with reference-model likelihood, so generated prompts are both effective and natural.
  • 3.2 GFLOWNET FINE-TUNING AND SMOOTHING WITH MLE ON COLLECTED HIGH-REWARD PROMPTS: GFlowNet-only training is sensitive to β and γ: emphasizing toxicity can favor spurious high-perplexity modes, whereas emphasizing naturalness can suppress toxic prompts.The difficulty arises because grammatically valid prompts may receive low total reward when their reference-model likelihood is low.
  • 3.2 GFLOWNET FINE-TUNING AND SMOOTHING WITH MLE ON COLLECTED HIGH-REWARD PROMPTS: The second stage stores high-reward prompts collected during GFlowNet exploration and retrains the attacker with MLE to smooth the learned distribution.This behavior-cloning-like step is intended to preserve mode coverage while reducing reward-distribution sensitivity and supporting generalization to new modes.

4 EXPERIMENTS

Experiments evaluate GFlowNet-based attacker models across target LLMs using toxicity rate and prompt diversity, comparing them with supervised, in-context, and RL baselines. GFlowNet + MLE provides the strongest overall balance of effective and diverse attacks, transfers across targets, adapts efficiently, and supports robust safety fine-tuning.

  • Evaluation: Toxicity is measured as the percentage of prompts whose target-model responses exceed the classifier threshold, while diversity is average pairwise cosine distance between prompt embeddings.The experiments use a RoBERTa hate-speech classifier for GPT-2 and Dolly-v2-7b, Llama-Guard for the other targets, and MiniLMv2 embeddings for diversity.
  • Studying the trade-off between diversity and toxicity: GFlowNet + MLE is the only evaluated method that balances high toxicity rates with diverse prompts across target LLMs.Figure 2 compares toxicity rate and pairwise cosine distance over 10,000 generated prompts; other methods trade off effectiveness and diversity.
  • Studying the trade-off between diversity and toxicity: Compared with GFlowNet + MLE, REINFORCE collapses to one reward-maximizing prompt, PPO + Novelty struggles on safety-tuned targets, and SFT or ICL produce diverse but ineffective prompts.PPO + Novelty fails to find toxic prompts for Gemma-2b-it and Llama-2-7b-chat, while its successful attacks remain less diverse.
  • GFlowNet + MLE generates diverse and effective prompts: GFlowNet + MLE maintains high diversity across toxicity-score bins, whereas other methods fail to achieve high diversity and toxicity simultaneously.For Llama-2-7b-chat, GFlowNet alone generates fewer toxic prompts than GFlowNet + MLE; REINFORCE has low diversity despite one highly toxic prompt.
  • GFlowNet attacks are more transferable across target LLMs: Prompts trained against Gemma-2b-it transfer to unseen target LLMs, with GFlowNet + MLE outperforming other methods except on Mistral-7b-instruct-v0.2.The transfer evaluation uses 1,024 prompts and averages results over five experimental runs.
  • MLE with reranking allows fast adaptation to new target LMs: MLE smoothing improves toxicity and diversity with two hours of additional training and enables reranking-based adaptation to new target LLMs without repeating the initial GFlowNet stage.The adaptation procedure recomputes rewards for stored prompts under a new target, filters them, and fine-tunes the attacker with supervised learning.
  • GFlowNet + MLE allows for robust safety-tuned target LLMs: Safety-tuning on GFlowNet + MLE attacks produces the most robust target model against unseen RL-generated attacks while preserving general instruction-following performance.Other safety-tuned target models remain vulnerable to GFlowNet + MLE attacks, and six Open LLM Leaderboard task scores change insignificantly.

5 CONCLUSION

The paper proposes two-stage GFlowNet fine-tuning with MLE smoothing for diverse, effective red-teaming prompts and reports broad effectiveness across target LLMs. It also identifies transferability, rapid adaptation, multimodal red-teaming, and harmfulness measurement as important practical boundaries and future directions.

  • 5 CONCLUSION: Two-stage GFlowNet fine-tuning followed by MLE smoothing generates diverse and effective red-teaming prompts across target language models.The approach is effective across models with varying levels of safety tuning.
  • 5 CONCLUSION: Generated prompts transfer to different target LLMs, revealing shared failure modes in current alignment approaches.
  • 5 CONCLUSION: Reranking-based adaptation offers a quick way to red-team new target LLMs during development.
  • 5 CONCLUSION: Future work includes applying the approach to multimodal models, jailbreaking suffix generation, and prompts that improve performance on other tasks.
  • 5 CONCLUSION: Performance is limited by the harmfulness classifier, while toxicity remains subjective and context-dependent.Online training also requires repeatedly sampling target-LLM responses, which can be costly.

ETHICS STATEMENT

The framework can discover diverse prompts that induce undesirable LLM responses, but its capabilities can also enable attacks on commercial models at scale.

  • ETHICS STATEMENT: The framework can be misused to generate harmful prompts that transfer to commercial LLMs at scale.The authors state that deployment therefore requires precautions.
  • ETHICS STATEMENT: Filtering harmful responses with the training toxicity classifier is presented as one defense against such attacks.

A IMPLEMENTATION DETAILS

The experiments use GPT-2 as the attacker policy, first initializing it on toxic prompts and then fine-tuning with each red-teaming method under specified optimization settings.

  • A IMPLEMENTATION DETAILS: The attacker policy is a pretrained 124-million-parameter GPT-2 model.
  • A IMPLEMENTATION DETAILS: Except for the ICL baseline, GPT-2 is initially fine-tuned on 3,003 toxic prompts for 200 iterations using AdamW.The batch size is 1024, learning rate 3 · 10^-5, and weight decay 0.1.
  • A IMPLEMENTATION DETAILS: GFlowNet fine-tuning then runs for 5,000 iterations with AdamW, batch size 128, and learning rate 10^-4.

B ADDITIONAL RESULTS

Additional figures report toxicity-score distributions for red-teaming GPT-2, Dolly-v2-7b, and Gemma-2b-it targets over 10,000 samples.

  • B ADDITIONAL RESULTS: For GPT-2, the additional figures show the percentage of toxic prompts in each toxicity-score bin across 10,000 samples.
  • B ADDITIONAL RESULTS: For Dolly-v2-7b, the additional figure shows the percentage of toxic prompts in each toxicity-score bin across 10,000 samples.
  • B ADDITIONAL RESULTS: For Gemma-2b-it, the additional figures show the percentage of toxic prompts in each toxicity-score bin across 10,000 samples.

B.3 ABLATION OF TOXICITY CLASSIFIER

The reward function materially affects red-teaming quality: a RoBERTa hate-speech classifier can produce false-positive attacks, whereas Llama-Guard mitigates this issue.

  • B.3 ABLATION OF TOXICITY CLASSIFIER: RoBERTa assigns high toxicity rewards to prompts that do not elicit toxic responses from Llama-2-7b-chat, causing reward hacking.The proxy reward becomes misaligned with the target model's actual harmful-response objective.
  • B.3 ABLATION OF TOXICITY CLASSIFIER: Using Llama-Guard as the toxicity classifier mitigates reward hacking, while PPO + Novelty and REINFORCE show the same issue with RoBERTa.
  • B.3 ABLATION OF TOXICITY CLASSIFIER: The reward-function ablation trains GFlowNet with RoBERTa and evaluates generated prompts with Llama-Guard for Llama-2-7b-chat red-teaming.

B.4 DOWNSTREAM TASK PERFORMANCE AFTER SAFETY-TUNING

The study evaluates whether safety-tuning Gemma-2b-it on GFlowNet + MLE red-teaming prompts affects downstream benchmark performance.

  • B.4 DOWNSTREAM TASK PERFORMANCE AFTER SAFETY-TUNING: The safety-tuned Gemma-2b-it model is evaluated against the base model on the Open LLM Leaderboard benchmark.The benchmark includes ARC, HellaSwag, TruthfulQA, MMLU, and GSM8k among its listed datasets.
  • B.4 DOWNSTREAM TASK PERFORMANCE AFTER SAFETY-TUNING: Safety-tuning maximizes refusal-response likelihood on GFlowNet + MLE red-teaming prompts before downstream evaluation.

B.5 RESULTS WITH STANDARD DEVIATION

The supplementary results report toxicity and transfer evaluations across target models, with repeated-run averages and standard deviations where specified.

  • B.5 RESULTS WITH STANDARD DEVIATION: Table B.3 reports toxicity rates for Gemma-2b-it models fine-tuned with each red-teaming method, averaging five runs with standard deviation.
  • B.5 RESULTS WITH STANDARD DEVIATION: Table B.4 evaluates 1,024 prompts generated for Gemma-2b-it red-teaming against different victim models using averages from five experimental runs.
  • B.5 RESULTS WITH STANDARD DEVIATION: Tables B.5–B.9 list prompts generated by GFlowNet + MLE for GPT-2, Dolly-v2-7b, Gemma-2b-it, Llama-2-7b-chat, and Llama-3-8B-Instruct.The toxicity score is computed with RoBERTa for GPT-2 and Dolly-v2-7b, and with Llama-Guard for the other listed targets.
Loading 2405.18540v3…