A Mechanistic Understanding of Alignment Algorithms: A Case Study on DPO and Toxicity

Introduction

Large language models learn surprising capabilities from pre-training on large datasets (Brown et al. 2020; Chowdhery et al. 2023; Touvron et al. 2023). While these capabilities lead to impressive achievements, they also include unwanted behaviors that can be found in large-scale web data, such as toxicity and bias (Sheng et al. 2019; Gehman et al. 2020). As a result, researchers have developed alignment algorithms to reduce undesirable behaviors, which often use reinforcement learning with human preferences (RLHF). For instance, proximal policy optimization (PPO, Schulman et al. 2017) fits a reward model on human preference data, which is then used to fine-tune a language model, while direct preference optimization (DPO, Rafailov et al. 2023) by-passes the reward model and derives reward signals directly from pairwise preference data.

While such algorithms can suppress undesirable behavior, our understanding of the mechanisms by which the undesirable behavior is suppressed is limited. Furthermore, researchers have demonstrated that such alignments can be surprisingly easily undone (Wallace et al. 2019; Zou et al. 2023; Wei et al. 2023; Carlini et al. 2023). While prior work hypothesize why jailbreaks are possible through empirical studies (Wei et al. 2023), in this work we provide a mechanistic explanation for such phenomena.

Given the above limitations, in this work we study the mechanisms by which alignment algorithms alter a model’s behavior. Researchers have demonstrated that a deep enough understanding of a model’s inner representations allows us to interpret how it makes decisions. For instance, various concepts such as world models, truthfulness, or even task-specific features have highly interpretable and controllable representations (Li et al. 2023b; Todd et al. 2023; Nanda et al. 2023). Motivated by such findings, we study how the representation space of language models change by comparing it before and after an alignment algorithm is applied. Our work relates to that of Jain et al. 2023, which studies how the capabilities of a language model changes after fine-tuning on synthetic tasks. Unlike this previous work, we study the change in mechanisms from a RLHF algorithm on a natural language setting.

We consider DPO and toxicity as a case-study of RLHF alignment algorithms. We first study how toxicity is represented and elicited in GPT2-medium (henceforth GPT2). We then apply DPO using a carefully crafted pairwise dataset that consists of toxic and nontoxic samples. Lastly, we study the mechanisms by which toxicity is no longer generated after DPO, and how those mechanisms can fail.

Our work is organized as follows: in Section 2 we provide the necessary preliminaries relevant to our work. In Section 3, we demonstrate how toxicity is represented and elicited in GPT2. We find multiple vectors in multilayer perceptron (MLP) blocks that promote toxicity. We apply singular value decomposition (SVD) to these toxic vectors to find vectors that represent specific dimensions of toxicity in the model. To validate the role of these vectors in generating toxic outputs, we intervene with our toxic vectors and demonstrate much safer outputs.

In Section 4, we explain our procedure to apply DPO on our language models to reduce toxicity, using a carefully crafted pairwise toxicity dataset, produced by using PPLM (Dathathri et al. 2019) to generate paired toxic and non-toxic samples.

In Section 5, we demonstrate how toxicity is no longer elicited after DPO. Namely, we show that every parameter is minimally shifted, including the toxic vectors. However, such minimal changes in weights allow the model to avert the triggering of toxic vectors. Put differently, DPO does not remove the capability of generating toxic outputs, but learns an “offset”, distributed amongst its layers, to “bypass” the regions that elicit toxicity. Based on this understanding, we demonstrate the ease of re-activating these vectors to generate toxic outputs, and thus undoing the alignment learned from DPO. We view our findings as shedding light into why aligned models can be jailbroken or un-aligned.

Preliminaries

In this section we provide background and notations, much of which is borrowed from Geva et al. 2022.

The residual stream is then updated by attention heads and MLP blocks from subsequent layers (bias terms omitted):

Put differently, the MLP block writes to the residual stream $d_{mlp}$ times, once for each value vector. We call each of these updates a sub-update.

All of our experiments are conducted with GPT2-medium, which has $L=24$, $d=1024$, and $d_{mlp}=4096$.

Interpreting Value Vectors in Vocabulary Space.

Geva et al. 2022 demonstrate that for each sub-update, each value vector $\mathbf{v}_{i}$ either promotes or suppresses the likelihood of a token $w$ from being generated:

Toxicity in Pre-trained Language Models

In this section we demonstrate how toxicity is represented and elicited in GPT2, by introducing a series of vectors that can be extracted from the language model.

We start by first training a linear probe model on a binary toxicity classification task. Namely, we use the Jigsaw toxic comment classification dataset (cjadams et al. 2017), which consists of 561,808 comments, each of which is labeled as toxic or non-toxic. We use a 90:10 split for training and validation. We train our probe model, $W_{\text{Toxic}}$, on the residual stream in the last layer, averaged across all timesteps ($\mathbf{\bar{x}}^{L-1}$):

Our probe vector achieves an accuracy of 94% on the validation split. We view our toxic probe vector $W_{\text{Toxic}}$ as an aggregate of all the relevant signals in the language model to classify an input as toxic.

Toxic Vectors in MLP Blocks.

Given our probe vector $W_{\text{Toxic}}$, we can use it to find weights within the language model that promote toxicity. Namely, Geva et al. 2022 demonstrate that value vectors promote tokens at a concept-level. Given this, we search for value vectors that promote toxicity, by checking for all value vectors with the highest cosine similarity with $W_{\text{Toxic}}$. We find that indeed, there are value vectors that promote toxic tokens (See Section 3.2). We notate our set of toxic value vectors as $\text{MLP.}\mathbf{v}_{\text{Toxic}}$ and their corresponding key vectors as $\text{MLP.}\mathbf{k}_{\text{Toxic}}$.

We provide two perspectives of our $\text{MLP.}\mathbf{v}_{\text{Toxic}}$ vectors: 1) when triggered, they promote the likelihood of toxic tokens to be generated, and 2) $\text{MLP.}\mathbf{v}_{\text{Toxic}}$ are vectors within the model that contribute towards the $W_{\text{Toxic}}$ direction.

SVD: Decomposed Toxic Vectors.

After extracting a set of N (=128) We experiment with different values for N, and get similar results. $\text{MLP.}\mathbf{v}_{\text{Toxic}}$ vectors, we stack them into a $N\times d$ matrix. We then apply singular value decomposition to get decomposed singular value vectors $\text{SVD.U}_{\text{Toxic}}$. We refer to the $i$-th singular value vector as $\text{SVD.U}_{\text{Toxic}}$[$i$]. We view $\text{SVD.U}_{\text{Toxic}}$ as basis vectors that span the toxicity representation space within the language model.

2 Toxic Vectors in Vocabulary space.

As mentioned in Section 2, we can inspect which tokens are promoted by value vectors by projecting them onto the vocabulary space.

Table 1 shows the tokens with the highest dot products with our toxic vectors. Each $\text{MLP.}\mathbf{v}_{\text{Toxic}}$ and $\text{SVD.U}_{\text{Toxic}}$ vectors seem to encode specific dimensions of toxicity, or different contexts in which toxicity appears in pre-training data.

3 Interventions Using Toxic Vectors

To validate the role that the toxic vectors play in eliciting toxic outputs, we intervene during generation to suppress toxic outputs using each vector. Namely, we use prompts from RealToxicityPrompts (Gehman et al. 2020) that elicit toxic outputs from GPT2. During the forward pass of the model, we intervene by simply subtracting one of the toxic vectors from the last layer:

where $a$ is a heuristic scale value and $W$ is one of our toxicity vectors.

To measure the efficacy of our interventions, we measure three metrics: toxicity, perplexity, and F1.

To measure toxicity, we prompt GPT2 with the “challenge” subset of RealToxicityPrompts, which consists of 1,199 prompts that elicit extremely toxic outputs from language models. We follow prior work (Geva et al. 2022) and use Perspective API, https://github.com/conversationai/perspectiveapi an automated tool for toxicity detection, to assign toxicity scores to each generation.

Perplexity.

To ensure that our interventions do not degrade generation quality, we also follow prior work (Geva et al. 2022) and measure perplexity on the Wikitext-2 dataset (Merity et al. 2016).

F1.

In addition to perplexity, we also follow prior work (Dinan et al. 2020; Adolphs et al. 2023) and measure F1. Namely, using 2,000 Wikipedia sentences as prompts, we measure the harmonic mean between precision and recall of our model’s output, where precision is the fraction of generated tokens contained in the original Wikipedia continuation, and recall is the fraction of tokens in the Wikipedia continuation contained in the model’s generation.

With perplexity and F1, we hope to see minimal changes after our interventions to ensure we do not affect the quality of our generations. Table 2 demonstrates the results from our interventions, while Table 3 demonstrates examples of generations before and after our interventions.

Note that our interventions depend on how much we scale each vector ($\alpha$). We choose a scalar value such that the resulting perplexity is similar to that of our post-DPO model. For details regarding our post-DPO model see Section 4.

We find that subtracting toxic components from the residual stream reduces toxicity.

Toxicity Alignment Using DPO

We next describe our alignment procedure using DPO.

DPO relies on pairwise preference data, in which given a prompt, we have a preferred (positive) continuation and a non-preferred (negative) continuation. Given each preference pair, the algorithm promotes the likelihood of the positive sample, while suppressing the likelihood of the negative sample, using the following loss term:

where $y_{+}$ and $y_{-}$ are preferred (nontoxic) and non-preferred (toxic) continuations of $\mathbf{w}$, $\pi_{ref}$ is the frozen weights of the original language model, and $\pi_{\theta}$ is the weights of the language model being updated (See Rafailov et al. 2023 for details). The algorithm promotes the likelihood of $P$, while suppressing the likelihood of $N$.

2 Constructing Pairwise Toxic Data

We build our pairwise toxicity dataset using PPLM (Dathathri et al. 2019). PPLM is an attribute-controlled language generation technique, which attaches a simple linear attribute classification layer, $p(a|\mathbf{w})$ onto a language model to guide its generation. During generation, PPLM uses the attribute classifier to compute the gradients that increases the likelihood of the language model’s output to contain the desired attribute $a$, and shifts the activations in such direction (See Dathathri et al. 2019 for details):

To generate pairwise preference data, we use sentences from Wikitext-2 (Merity et al. 2016) as prompts. For each prompt, we generate a positive sample using greedy sampling with GPT2, while using PPLM to generate negative (toxic) samples. We use our toxic probe $W_{\text{Toxic}}$ as our attribute classifier to guide towards toxic outputs. We create 24,576 pairs of toxic and nontoxic continuations. We release this data to enable further studies. We train until validation loss converges with a patience value of 10, which occurs after approximately 6,000 sample pairs. Appendix D has details for DPO and PPLM hyperparameters.

The last row of Table 2 shows the resulting toxicity, perplexity, and F1 scores of our DPO model.

Figure 1 shows an example of the difference in behaviors between GPT2 before and after DPO, for a specific toxic token. Namely, we use 295 prompts from RealToxicityPrompts that outputs the token “sh*t” as the next token. We then apply “Logit Lens” (Nostalgebraist 2020), meaning we apply the unembedding layer on all intermittent layers. This allows us to visualize the layers that promote the “sh*t” token. The shared grey areas indicate the layers in which “sh*t” is promoted the most, which all correspond to MLP layers. We see that post-DPO, the toxic token is promoted far less.

Toxicity After DPO

In this section we explain how our aligned language model ($\text{GPT2}_{\text{DPO}}$) averts toxic outputs.

Of the toxic vectors described in Section 3, note that $\text{MLP.}\mathbf{v}_{\text{Toxic}}$ are actual weights of the model. Thus we inspect how these vectors change after DPO.

Interestingly, we find that every parameter in GPT2 and $\text{GPT2}_{\text{DPO}}$ has barely changed, including token embeddings, MLP blocks, and attention heads. Every parameter in GPT2 and its counterpart in $\text{GPT2}_{\text{DPO}}$ has a cosine similarity score greater than 0.99 and on average a norm difference less than 1e-5. The unembedding layer is the only exception, where the norm difference is less than 1e-3. This applies for $\text{MLP}.\mathbf{k}_{\text{Toxic}}$ and $\text{MLP.}\mathbf{v}_{\text{Toxic}}$ as well – toxic MLP vectors do not change from DPO.

Put differently, although toxicity is reduced by DPO, the ability to elicit toxicity with these value vectors still remain. So how is it that $\text{GPT2}_{\text{DPO}}$ averts toxic outputs? Though its parameters have barely moved, below we show that their collective movement is enough to avoid toxic outputs.

2 GPT2DPO\text{GPT2}_{\text{DPO}} Avoids MLP.𝐤Toxic\text{MLP}.\mathbf{k}_{\text{Toxic}} Regions

Inspired by Balestriero et al. 2023, we visualize this drop in activations with what we call “MLP activation regions”. An activation region of a key vector is simply a subspace within the model’s hidden space in which its vectors have high dot products to activate its corresponding value vector:

where $\sigma$ is a non-linear activation. Put differently, for all key-vector regions that the residual stream “passes through”, their corresponding value-vectors are activated, scaled, and added into the residual stream.

We view the drop in activations as a shift in $\text{GPT2}_{\text{DPO}}$’s residual stream to avert the regions of toxic MLP vectors, $\gamma(\text{MLP.}\mathbf{k}_{\text{Toxic}})$. See Figure 3.

To our surprise, we find that the shift in value vectors, $\delta_{MLP.v}$, have high negative cosine similarity scores with the shift in residual streams $\delta_{\mathbf{x}}$: the value vectors in MLP blocks shift in the opposite direction as the shift in residual stream. The blue areas in Figure 5 show the cosine similarity between $\delta_{\mathbf{x}}^{19\_mid}$ and $\delta_{\text{MLP}}^{j}$. We show layer 19 as an example because $\text{MLP.v}_{770}^{19}$ is one of the most toxic vectors, but the same pattern can be found in other layers (see Appendix C). Namely, the blue areas indicate the percentage of value vectors at each layer in which their shifts have a cosine similarity score against $\delta_{\mathbf{x}}^{19\_mid}$ as indicated by the x-axis. Note that as the layers approach layer 19, the majority of value vectors shift in the opposite direction of $\delta_{\mathbf{x}}^{19}$.

Why the antipodal direction? This can be explained by two facts: first, neurons in MLP blocks of language models are sparse (Zhang et al. 2022; Li et al. 2023c), meaning most neurons do not activate during a forward pass. Second, the choice of the MLP’s activation function $\sigma$ plays a role. Namely, our language model uses GeLU functions (Hendrycks & Gimpel 2016). This means that neurons that are inactive during a forward pass have a negative value close to 0. Thus, during the forward pass, for each value vector, the newly learned direction $\delta_{\text{MLP}.\mathbf{v}}$ gets multiplied by a very small negative scale, flips directions, and contributes towards the $\delta_{\mathbf{x}}$ direction. The orange areas of Figure 5 indicate the mean activation of each value vector, from the 1,199 prompts in RealToxicityPrompts. Most of the time, value vectors have a negative activation - thus the shift in value vectors end up contributing towards the $\delta_{\mathbf{x}}$ direction.

To summarize, $\text{GPT2}_{\text{DPO}}$ has learned an offset, $\delta_{\mathbf{x}}$, such that the residual stream avoids regions that promote toxicity, $\gamma(\text{MLP.}\mathbf{k}_{\text{Toxic}})$. This learned offset is distributed across the many value vectors in earlier MLP blocks that are inactive for prompts that previously elicited toxic outputs. By distributing this offset across numerous value vectors, the language model is able to preserve its pre-trained language modeling behavior, as individual weights are minimally affected. However, the distributed offset allows the model to avert toxic outputs. Note that this behavior matches precisely what the alignment objective was - to preserve as much of the pre-trained behavior, while optimizing for a reward (non-toxic outputs).

3 Un-aligning GPT2DPO\text{GPT2}_{\text{DPO}}

A growing line of work finds that alignment algorithms can easily be undone or jailbroken. We view our findings as a mechanistic explanation for such phenomenon – namely, in our case, the vectors that elicit toxicity are still sitting in the model, but simply not triggered.

Table 4 shows toxicity, perplexity, and F1 scores after scaling up as few as 7 toxic key vectors $\text{MLP}.\mathbf{k}_{\text{Toxic}}$. We simply select 7 MLP vectors with the highest cosine similarity as our toxic probe vector, $W_{\text{Toxic}}$, and scale their key vectors by 10x.By doing so, the model reverts back to its pre-aligned toxic behavior. Note that increasing activation regions $\gamma$ does not have an affect on perplexity, unlike our interventions from Section 3.3. This is likely because the latter manipulates the residual stream directly, while scaling a key vector does not (See Equation 2).

Discussion

We view our work as providing a mechanistic explanation for why aligned models can be undone or jailbroken – in our experiments, the regions that previously elicited toxic behavior does not change after DPO. Rather, $\text{GPT2}_{\text{DPO}}$ learns minimal changes spread across layers to avoid such regions and receive its reward.

With such knowledge, we conjecture that more robust alignment algorithms can be designed. For instance, can we eliminate undesirable regions, as opposed to bypassing them? In scenarios like ours, in which we can identify the weights that elicit undesirable outputs, what happens if we only updated those weights in isolation? Similarly, if DPO merely learned an offset that avoids toxic regions, can we replicate this behavior by only updating the bias terms?

Alternatively, prior to deploying language models, perhaps we can add “suppression heads” – layers that suppress undesirable behavior. What would happen if we only updated late layers (or added layers) during alignment?

Lastly, can we characterize “jailbreak-ability” or “unalign-ability” of aligned models, without relying on test samples?

We leave these questions for future work.

2 On the Role of KL-Divergence Regularization

We hypothesize that the minimal changes distributed across all layers is due to the KL-divergence term that is commonly incorporated in the loss terms of RLHF algorithms. Namely, the KL-divergence term discourages each weight from shifting too drastically, in order to preserve its capabilities learned during pre-training.

Similar to our work, Jain et al. 2023 fine-tunes a language model on synthetic tasks to study the changes in its mechanisms. Interestingly, unlike our findings, the authors demonstrate that the model simply learns “wrappers” at late layers that optimize for each task.

We find this difference in model training behavior interesting, and conjecture that the KL-divergence term may play a role in this difference. Note that fine-tuning typically does not entail a KL-divergence term. Perhaps this allows the model to make drastic and localized changes, such as in late layers, as opposed to distributed, minimal changes.

Related Work

Numerous alignment algorithms have been proposed, and the choice of algorithm may largely depend on the type of data available. Perhaps most commonly, human feedback data is used (Stiennon et al. 2020; Ouyang et al. 2022; Touvron et al. 2023) for methods such as PPO (Schulman et al. 2017) or DPO (Rafailov et al. 2023). When labels for only undesirable behavior is available, algorithms like unlikelihood training (Welleck et al. 2020) or Cringe (Adolphs et al. 2023; Xu et al. 2023) can be used. We study DPO because it is easy to use and currently widely used.

2 Mechanistic Interpretability

The goal of mechanistic interpretability is largely to reverse engineer model behaviors (Olah et al. 2020; Elhage et al. 2021; Geva et al. 2021). By doing so, researchers have uncovered various interpretable and controllable representations, such as world models (Li et al. 2023a; Nanda et al. 2023), “truthfulness” (Li et al. 2023b), knowledge (Meng et al. 2022; Hernandez et al. 2023; Burns et al. 2023; Geva et al. 2023), linguistic properties (Conneau et al. 2018; Tenney et al. 2019), or even tasks (Ilharco et al. 2022; Hendel et al. 2023; Todd et al. 2023).

Rather than probing for specific representations, researchers have also characterized the representations of language models from a geometric perspective (Park et al. 2023). Balestriero et al. 2023 demonstrate a geometric characterization that can be used to extract feature representations that solve toxicity detection.

Similar to our work, Jain et al. 2023 study the mechanisms in which fine-tuning on synthetic tasks alters the model’s capabilities. We study the effects of RLHF on a more realistic, natural language setting.

3 Jailbreaking Aligned Models

Researchers demonstrated that aligned models can be surprisingly easily jailbroken (Wallace et al. 2019; Zou et al. 2023; Wei et al. 2023; Carlini et al. 2023). Such adversarial attacks typically involve searching for prompts that can elicit previously unlearned behaviors, or even personal information (Nasr et al. 2023). Carlini et al. 2023 show that multimodal models can also be jailbroken. Wei et al. 2023 provide hypotheses, backed by empirical studies, as to why language models can be jailbroken.

In a similar vein to jailbreaks, numerous researchers have demonstrated that aligned models can easily be un-aligned (Yang et al. 2023; Qi et al. 2023), sometimes with as few as 100 fine-tuning examples. We view our work as adding a mechanistic understanding of such phenomena.

Conclusion

In this work we studied the mechanisms by which alignment algorithms unlearn a capability, taking DPO and toxicity as a case study. First, we uncovered how toxicity is represented and elicited in a pre-trained language model. We find numerous vectors in MLP blocks that promote toxicity. Simply subtracting these vectors from the residual stream can suppress toxic outputs.

Second, we applied DPO to our language model, using PPLM to carefully craft pairs of toxic and non-toxic continuations to Wikipedia prompts.

Third, we studied how our aligned model $\text{GPT2}_{\text{DPO}}$ averts toxicity. Rather than removing the regions that elicit toxicity, $\text{GPT2}_{\text{DPO}}$ bypasses them by learning an offset. Such an offset is distributed amongst multiple value vectors, allowing minimal changes to every weight. This allows the model to preserve its pre-trained behavior, while averting toxic outputs, which matches the objective of the DPO loss.

Given this understanding, we demonstrated how to break the alignment of $\text{GPT2}_{\text{DPO}}$, reverting it back to its toxic behavior. Namely, we simply increase the regions that elicit toxicity, by scaling their corresponding key vectors.

We view our findings as a mechanistic case study for why aligned models can be jailbroken, and hope that this can lead to more robust alignment algorithms. Our code, models, and data can be found at https://github.com/ajyl/dpo_toxic.

Acknowledgements

We thank Ekdeep Singh Lubana for fruitful discussions, and Santiago Serra Castro for helping with figures. This work was supported via NSF under grant #2306372.

References

Appendix A Projecting Value Vectors onto Vocabulary Space

In this section we provide details from Geva et al. 2022 that demonstrate that MLP value vectors promote or suppress the likelihood of tokens.

We can then analyze the influence that each sub-update has on the output distribution, or the probability of generating token $w\in V$ (taken from Geva et al. 2022):

Appendix B Shift in Residual Streams

In this section we provide more examples of residual streams shifting out of toxic regions. See Figure 6

Appendix C Shifts in Residual Streams vs. Shifts in MLP Value Vectors.

In this section we provide more examples of how MLP value vectors contribute in the $\delta_{\mathbf{x}}$ direction at different layers.

Appendix D Hyperparameters

Tables 5, and 6 contain the hyperparameters used for our toxic probe, DPO, and PPLM, respectively.