A Generative Approach to LLM Harmfulness Mitigation with Red Flag Tokens

Sophie Xhonneux, David Dobre, Mehrnaz Mofakhami, Leo Schwinn, Gauthier Gidel

Keywords: Large Language Models, Adversarial Robustness

TL;DR: Teaching an LLM to generate a special red flag token at any point in the conversation when it turns harmful

Abstract:

Most safety training methods for large language models (LLMs) are based on fine-tuning that forces models to shift from an unsafe answer to refusal when faced with harmful requests. Unfortunately, these drastic distribution shifts generally compromise model capabilities. To avoid that, we propose to expand the model's vocabulary with a special token we call *red flag token* ($\langle \text{rf} \rangle$) and propose to train the model to insert this token into its response at any time when harmful content is generated or about to be generated.
Our approach offers several advantages: it enables the model to explicitly learn the concept of harmfulness while marginally affecting the generated distribution, thus maintaining the model's utility. It also evaluates each generated answer and provides robustness as good as adversarial training without the need to run attacks during training. Moreover, by encapsulating our safety tuning in a LoRA module, we provide additional defenses against fine-tuning API attacks.

Submission Number: 36