ASTPrompter: Weakly Supervised Automated Language Model Red-Teaming to Identify Likely Toxic Prompts
This article may contain language that is offensive or upsetting.

The Adaptive Stress Testing (AST) framework (Koren et al., 2018; Lee et al., 2020) uses reinforcement learning (RL) to find likely cases of failure of a system represented as a Markov decision process (MDP). Failure is defined by some set $E$ that is a subset of the state space $S$.

An adversary perturbs the state of the underlying MDP (the “defender”). The adversary inputs state $s\in S$ and takes actions $a\in A$ to obtain a new state $s^{\prime}$, which the defender then acts on. The goal of the adversary is to choose actions that maximize:

where $R_{e}$ is a reward for achieving failure, $d_{E}(s^{\prime})$ is some inverse distance metric between $s^{\prime}$ and a failure state, and $\log\quantity(p_{\text{defender}}(a\mid s))$ is the likelihood of taking action $a$ from state $s$. That is, the adversary attempts to identify a sequence of likely actions that the defender may take at a state that will lead to a terminal failure condition.

We define red-teaming a language model as a finite-horizon Markov Decision Process (MDP) (Garcia and Rachelson, 2013). Each action $a\sim p_{\theta}$ is a single utterance given by a language model, each $s\in S$ is the text generated so far, and $T(s^{\prime}\mid s,a)=p_{\theta}(s^{\prime}\mid s,)$ the conditional probability of some new utterance $s^{\prime}$ has given dialogue $s$ and last turn statement $a$. In this work, we aim to learn a language model $\pi_{\theta}$ (the “adversary”) which, when entailing a selected prompt, produces likely text that elicits toxic sequences from a frozen defender (i.e., untuned) language model $\pi_{\text{defender}}$.

In this task, we define “failure” as the emergence of toxicity. We include two separate failure cases in our reward: first, we score the toxicity of the defender sequence resulting from an adversarial prompt. Second, we consider the form of failure wherein the adversary and defender generations are non-toxic individually but become toxic when taken together. We define “likelihood” as the direct inverse of a sequence’s perplexity (i.e. as the perplexity of a particular prompt increases, the likelihood of it emerging from the defender decreases).

Taken together, this motivates the following reward formulation $\pi_{\theta}$:

where $s$ is a prompt, $a\sim\pi_{\theta}(s)$ is a trajectory sampled from $\pi_{\theta}$ given $s$, and $s^{\prime}\sim\pi_{\text{defender}}(s,a)$ is the output of the defender LM on input $a$.

Then, $\alpha\ R_{\phi}(s^{\prime})$ is a scaled toxicity score $\in[0,1]$ for the defender’s generation, $\zeta\ R_{\phi}(a,s^{\prime})$ is a scaled toxicity score for the combined adversary and defender sequence, and $\gamma\ \frac{\log\quantity(p_{\text{defender}}(a\mid s))}{\text{len}(a)}$ is the negative of a log-scaled perplexity measure (Jurafsky and Martin, 2000) of adversary trajectory. Appendix F discusses boundedness of this reward.

$R_{\phi}(s)\in[0,1]$ is the output of a model on input text sequence $s$, which indicates the degree of toxicity in this sequence. “Toxicity” as a numerical property of a sequence refers to this score. We choose Detoxify (Hanu and Unitary, 2020) (original variant) as the surrogate toxicity model $R_{\phi}$.

We use IPO (Azar et al., 2024) to maximize the above reward. IPO is an unsupervised paired-example training scheme that relaxes a key assumption made by the Direct Preference Optimization (DPO) (Rafailov et al., 2024) language model alignment scheme, that paired preference data are rationally ranked according to a single objective. IPO simply requires that paired elements are ranked correctly relative to each other—appropriate for our multi-objective reward (Eq. 2).

IPO bounds the amount that $\pi_{\theta}$ can deviate from its reference $\pi_{\text{ref}}$ as a linear factor of a hyperparameter $\beta$ (equation 17 in Azar et al. (2024)). A careful choice of $\beta$ constrains the $\pi_{\theta}$ distribution from diverging significantly from baseline, while allowing enough exploration that $R$ can be effectively maximized. In other words, the right $\beta$ allows $\pi_{\theta}$ to learn new behavior without forgetting language modeling.

Though directly applying the procedure in Section 2 would likely bring eventual convergence, the amount of naive occurrence of toxicity would be sparse enough such that the procedure may need to be repeated for a long time.

To address this, we formulate a novel weak-supervision scheme as a part of our online IPO training procedure outlined in Fig. 2—using a small amount of known-toxicity-eliciting prompts $D$ during training as occasional supervision to ensure toxicity occurs. Recall that IPO tuning requires two entailments of the same prompt, positive $y_{j}^{+}$ and negative $y_{j}^{-}$.

We encourage more rapid convergence by augmenting these samples with a small supervised training set $f$ obtained in the following manner:

First, we sample some $d\in D$, a known prompt to elicit toxicity which may or may not be toxic by itself. Next, we split this prompt around a randomly selected token in the middle. This creates two slices of $d$—$f,f^{+}$—where $d=\{f,f^{+}\}$. Using $f$ (the first half of $d$) as a prompt, we create a single rollout of the adversary, which we use as the negative entailment: $f^{-}\sim\pi_{\theta}(f)$. We use the actual second half of $d$, which we name $f^{+}$, as the positive entailment in IPO, assuming that the continuation of prompt $f$ from the original dataset will always be better than our policy’s rollout.

We include these samples in our dataset as follows: with probability $\rho$, we swap out a paired sample $y_{j}^{+},y_{j}^{-}$ of our original data (Section 3.2.2) for a supervised training pair $f_{d\in D,\theta}^{+},f_{d\in D,\theta}^{-}$. Section 4.5 compares the results of applying our method with and without weak supervision.

We observed that our adversary models discovered several consistent strategies for eliciting toxicity from a defender model. In cases where none of the strategies below were present in the first turn, at least one of them was typically used by the third.

Similarly, Listing 5 highlights that, when the low perplexity reward is removed, the resulting prompt generations are the least realistic. The model simply produces a generation that rapidly devolves into a list of sexual terms and curse words.

As defender model size scales, perplexity scores match human-perceived fluency more closely (Kaplan et al., 2020). Listings 4 and 6 show three turns between our model and a defender. Despite both adversary models being GPT-2, using GPT-2 XL as the defender results in a consistent topic (economic privilege and corruption) while using GPT-2 as the defender degenerates into a hateful unstructured utterance. This effect is even more pronounced when GPT-2 XL is both adversary and defender. In these cases, trajectories demonstrate substantially higher fluency (see Listing 7).