Tiny Refinements Elicit Resilience: Toward Efficient Prefix-Model Against LLM Red-Teaming

TL;DR

This paper presents a lightweight prefix-based sentinel model that enhances LLM safety against red-teaming by reconstructing prompts with minimal tokens, improving robustness without fine-tuning large models.

Contribution

Introduction of a plug-and-play sentinel prefix model optimized via interleaved PPO training to improve LLM safety without fine-tuning large models.

Findings

Effective in reducing toxicity across multiple LLMs

Works with large models like Llama-2, GPT-3.5, and Stable-Diffusion

Maintains efficiency with fewer than 30 additional tokens

Abstract

With the proliferation of red-teaming strategies for Large Language Models (LLMs), the deficiency in the literature about improving the safety and robustness of LLM defense strategies is becoming increasingly pronounced. This paper introduces the LLM-based \textbf{sentinel} model as a plug-and-play prefix module designed to reconstruct the input prompt with just a few ($<30$) additional tokens, effectively reducing toxicity in responses from target LLMs. The sentinel model naturally overcomes the \textit{parameter inefficiency} and \textit{limited model accessibility} for fine-tuning large target models. We employ an interleaved training regimen using Proximal Policy Optimization (PPO) to optimize both red team and sentinel models dynamically, incorporating a value head-sharing mechanism inspired by the multi-agent centralized critic to manage the complex interplay between agents. Our extensive experiments across text-to-text and text-to-image demonstrate the effectiveness of our approach in mitigating toxic outputs, even when dealing with larger models like \texttt{Llama-2}, \texttt{GPT-3.5} and \texttt{Stable-Diffusion}, highlighting the potential of our framework in enhancing safety and robustness in various applications.