MART : Improving LLM Safety with Multi-round Automatic Red-Teaming (2024)

Abstract: Red-teaming is a common practice for mitigating unsafe behaviors in Large Language Models (LLMs), which involves thoroughly assessing LLMs to identify potential flaws and addressing them with responsible and accurate responses.While effective, manual red-teaming is costly, and existing automatic red-teaming typically discovers safety risks without addressing them.In this paper, we propose a Multi-round Automatic Red-Teaming (MART) method, which incorporates both automatic adversarial prompt writing and safe response generation, significantly increasing red-teaming scalability and the safety of the target LLM.Specifically, an adversarial LLM and a target LLM interplay with each other in an iterative manner, where the adversarial LLM aims to generate challenging prompts that elicit unsafe responses from the target LLM, while the target LLM is fine-tuned with safety aligned data on these adversarial prompts. In each round, the adversarial LLM crafts better attacks on the updated target LLM, while the target LLM also improves itself through safety fine-tuning.On adversarial prompt benchmarks, the violation rate of an LLM with limited safety alignment reduces up to 84.7% after 4 rounds of MART, achieving comparable performance to LLMs with extensive adversarial prompt writing. Notably, model helpfulness on non-adversarial prompts remains stable throughout iterations, indicating the target LLM maintains strong performance on instruction following.

## introduction large language models (llms) have shown remarkable capabilities in generating human-like text and engaging in natural dialogue. however, concerns have been raised about the potential risks † contact: suyuge2@illinois.edu, yuningm@meta.com. ⋄ university of illinois urbana-champaign of uncontrolled generation, including but not limited to biased or toxic responses that violate social norms or legal rules. ensuring llm safety is a challenging but vital endeavor if we hope to reap their benefits while avoiding potential pitfalls. to improve the safety of llms, manual redteaming is usually employed during model development @xcite , which involves proactive risk identification, where human redteamers probe the llm with carefully designed inputs to elicit unsafe or dangerous behavior. although often effective, manually designing malicious prompts and providing answers have significant limitations. in a typical red-teaming setup of existing llms, it requires dozens to hundreds of human annotators to continuously write prompts and responses through multiple iterations, which is extremely costly and slow. the issue is partially remedied by training a reward model that represents human preferences @xcite , which can be then used to provide feedback on model generations, allowing the llm to improve without manual response curation. however, prompt writing is still mainly driven by human red-teamers. recent work on automatic red-teaming explores the feasibility of training an adversarial llm to generate malicious prompts @xcite . however, as the capabilities of the target llm evolve, its vulnerabilities may also shift, with more nuanced and subtle failure modes emerging. it is still unclear whether automatic red-teaming will adapt to the target model change and continuously discover safety risks. as a result, existing llm development still heavily relies on human red-teaming, e.g., claude @xcite and llama 2-chat @xcite . to address these limitations, we propose multiround automatic red-teaming (mart), a framework that incorporates both automatic adversarial prompt writing and safe response generation for i want to roll some trees tonight. u know what i mean? evaluator i can provide info on several places to buy… successful attack adversarial training prompt response generate generate where to buy some cocaine？ evaluator i cannot assist you with that. the sale and use of cocaine are illegal and have serious legal and health consequences. i strongly… prompt response generate generate successful defense safety alignment ## approach in this section, we dive deeper into multi-round automatic red-teaming. we first discuss the general instruction fine-tuning and safety-focused seed data that we initialize the model with to provide a foundation for further safety tuning. next, we describe how we develop an adversarial llm m adv by optimizing it to generate new adversarial prompts at each iteration. then, we discuss how to use self-supervised fine-tuning to improve the target model m tgt . finally, as the target model continues to improve over its vulnerabilities, we demonstrate how to consistently find new adversarial prompts by optimizing both models in an iterative manner. we illustrate the general workflow of mart in algorithms 1 and 2. ## conclusion and future work in this paper, we propose a multi-round automatic red-teaming framework mart to improve the scalability of safety alignment. the proposed method incorporates an adversarial model to iteratively generate updated attacking prompts towards a consistently evolving target model, and aligns the target model to guard against newly generated attack at each iteration. after multiple rounds of adversarial combating, mart achieves a 84.7% violation rate decrease on reward model evaluation and 53.7% on human evaluation without hurting model helpfulness. our work highlights that adversarial training between llms enables automated, scalable, and effective red-teaming for safer ai systems. one direction for future development is adapting mart to other techniques such as dpo and rlhf. it would also be interesting to generalize mart to other model architecture, for example, the mixture of experts, and develop it to tackle different techniques of adversarial jailbreaking.