Source-linked AI summary
A Mechanistic Understanding of Alignment Algorithms: A Case Study on DPO and Toxicity
Andrew Lee, Xiaoyan Bai, Itamar Pres, Martin Wattenberg, Jonathan K. Kummerfeld, Rada Mihalcea
TL;DR
The paper addresses limited mechanistic understanding of how alignment suppresses undesirable behavior and why aligned models can be easily unaligned. It studies toxicity in GPT2-medium, applies DPO with a crafted pairwise dataset, and analyzes the resulting behavior. The central finding is that DPO bypasses toxic regions without removing the underlying capability, enabling reactivation of toxic behavior by scaling corresponding key vectors.
Problem
Mechanisms underlying alignment remain poorly understood, and prior work shows that alignment can be surprisingly easy to undo.
Method
The paper studies toxicity representations in GPT2-medium, extracts and decomposes toxic MLP vectors, and applies DPO using paired toxic and nontoxic samples generated with PPLM.
Results
DPO averts toxic outputs by learning a distributed offset that bypasses toxic regions rather than removing the capability to generate toxicity.
Takeaways & Limitations
Scaling the key vectors associated with toxic regions can reactivate toxicity and revert the aligned model to toxic behavior.
Takeaways & Limitations
The authors leave open whether more robust alignment could eliminate undesirable regions or update only the weights or biases that elicit them.
Abstract
from arXiv · showhide
While alignment algorithms are now commonly used to tune pre-trained language models towards a user's preferences, we lack explanations for the underlying mechanisms in which models become ``aligned'', thus making it difficult to explain phenomena like jailbreaks. In this work we study a popular algorithm, direct preference optimization (DPO), and the mechanisms by which it reduces toxicity. Namely, we first study how toxicity is represented and elicited in a pre-trained language model, GPT2-medium. We then apply DPO with a carefully crafted pairwise dataset to reduce toxicity. We examine how the resulting model averts toxic outputs, and find that capabilities learned from pre-training are not removed, but rather bypassed. We use this insight to demonstrate a simple method to un-align the model, reverting it back to its toxic behavior.
1. Introduction
The paper investigates how alignment algorithms suppress toxicity, addressing limited mechanistic understanding and the ease with which alignment can be undone. Using DPO and toxicity in GPT2-medium as a case study, it analyzes toxic representations, applies DPO, and studies how toxicity avoidance can fail.
- The paper studies the mechanisms by which alignment algorithms alter model behavior, motivated by limited understanding of suppression and easily reversible alignment.
- GPT2-medium toxicity is analyzed before DPO, followed by DPO training on carefully crafted toxic and nontoxic preference pairs.The pairwise dataset is described as being constructed from toxic and nontoxic samples.
- Multiple MLP-block vectors that promote toxicity are identified, decomposed with SVD into toxicity-space basis vectors, and validated through interventions producing safer outputs.
- DPO minimally shifts parameters but learns a distributed offset that bypasses toxic regions rather than removing the capability to generate toxic outputs.The offset is distributed across layers, allowing the model to avert toxic outputs while retaining the underlying capability.
- The study uses this mechanism to demonstrate that scaling corresponding key vectors can reactivate toxicity and undo the learned alignment.
2. Preliminaries
The preliminaries describe transformer representations and how MLP blocks write input-dependent value-vector updates into the residual stream. These value vectors can be interpreted in vocabulary space according to which tokens they promote or suppress.
- Transformer language models consist of embedding and unembedding layers surrounding a sequence of attention and MLP transformer layers.
- The embedding layer maps each input token to a residual-stream vector, which subsequent attention heads and MLP blocks update.
- MLP blocks use two linear transformations with point-wise activations, and their outputs decompose into value-vector sub-updates.Each sub-update writes one value-vector contribution to the residual stream.
- Key vectors determine input-dependent coefficients that scale value vectors, while the value vectors determine the token effects of those updates.
- The GPT2-medium experiments use L = 24 transformer layers, d = 1024 residual dimensions, and dmlp = 4096 MLP dimensions.
- Positive and negative embedding–value-vector dot products respectively increase and decrease a token’s likelihood, creating a ranking of promoted tokens.
3. Toxicity in Pre-trained Language Models
The paper identifies toxicity-related representations in GPT2’s MLP blocks and tests their role by projecting, decomposing, and intervening on toxic vectors.
- Toxicity Probe Vector: A toxicity probe trained on averaged final-layer residual streams achieves 94% validation accuracy.
- Toxic Vectors in MLP Blocks: The authors extract MLP value vectors aligned with the toxicity probe that promote toxic tokens and contribute toward the probe direction.
- SVD: Decomposed Toxic Vectors: SVD of 128 toxic MLP vectors yields basis vectors spanning the model’s toxicity representation space.
- Toxic Vectors in Vocabulary Space: Projected toxic and SVD vectors appear to encode distinct toxicity dimensions or contexts from pre-training data.
- Interventions Using Toxic Vectors: Interventions are evaluated using toxicity, perplexity, and F1 to assess safety alongside generation quality.Toxicity is measured on 1,199 challenging REALTOXICITYPROMPTS prompts, while perplexity uses Wikitext-2 and F1 uses Wikipedia continuations.
- Interventions Using Toxic Vectors: Subtracting toxic components from the residual stream reduces toxicity during generation.
4. Toxicity Alignment Using DPO
The paper constructs toxic–nontoxic preference pairs and applies DPO to GPT2. After training, toxic-token promotion is reduced while the model’s toxic capability remains available through an offset that bypasses triggering vectors.
- 4.1. Background: DPO: DPO promotes preferred nontoxic continuations and suppresses non-preferred toxic continuations relative to a frozen reference model.
- 4.2. Constructing Pairwise Toxic Data: The pairwise dataset uses GPT2 greedy continuations as positive samples and PPLM-guided toxic continuations as negative samples.
- 4.2. Constructing Pairwise Toxic Data: The authors create 24,576 toxic–nontoxic continuation pairs and train DPO until validation loss converges after approximately 6,000 sample pairs.
- 4.3. Mechanisms After DPO: After DPO, the model promotes the toxic token “sh*t” far less across intermittent layers for 295 prompts that originally elicited it.
5. Toxicity After DPO
After DPO, GPT2DPO retains its toxic vectors and nearly all original parameters, but collectively shifts residual streams away from regions that activate those vectors. This distributed offset reduces toxic elicitation without removing the underlying capability, which can be restored by enlarging toxic activation regions.
- 5.1. Toxic Vectors Remain After DPO: Every parameter changes minimally after DPO, with cosine similarity above 0.99 and average norm difference below 1e-5, while toxic MLP vectors remain unchanged.The unembedding layer is the only exception, with a norm difference below 1e-3.
- 5.2. GPT2DPO Avoids MLP.kToxic Regions: GPT2DPO shows reduced activations for toxic vectors on the same 1,199 REALTOXICITYPROMPTS prompts.The analysis generates 20 tokens per prompt and measures mean activations.
- 5.2. GPT2DPO Avoids MLP.kToxic Regions: The model learns a consistent linear residual-stream shift that moves representations out of toxic activation regions.These regions are subspaces where high key-vector dot products activate corresponding value vectors.
- 5.2. GPT2DPO Avoids MLP.kToxic Regions: Minimal changes accumulated across layers are sufficient to move residual streams outside regions that trigger toxic vectors.The offset is distributed across many earlier MLP value vectors that are inactive on prompts that previously elicited toxic outputs.
- 5.2. GPT2DPO Avoids MLP.kToxic Regions: Value-vector shifts oppose the residual-stream shift, yet negative activations make their contributions align with the learned offset direction.The antipodal pattern is linked to sparse MLP activations and GeLU behavior for inactive neurons.
- 5.3. Un-aligning GPT2DPO: Scaling only seven toxic key vectors by 10x reactivates toxic behavior and undoes DPO alignment without affecting perplexity.The vectors are selected by highest cosine similarity with the toxic probe vector.
6. Discussion
The authors argue that DPO's alignment can be undone because it bypasses rather than removes toxic regions, and suggest directions for more robust alignment.
- DPO learns minimal changes distributed across layers to avoid toxic regions rather than eliminating them.The authors hypothesize that the KL-divergence term encourages small weight shifts that preserve pre-training capabilities.
- This bypass mechanism may explain why aligned models can be jailbroken or otherwise unaligned.The paper frames its findings as a mechanistic explanation for these phenomena.
- 6.1. On Designing Robust Alignment Algorithms: The authors propose testing whether robust alignment can eliminate undesirable regions or isolate updates to the relevant weights or bias terms.They also raise the possibility of adding suppression heads or updating only late layers.
- 6.1. On Designing Robust Alignment Algorithms: The paper leaves characterization of jailbreak-ability or unalignability without test samples as future work.
- Compared with fine-tuning on synthetic tasks, DPO produces distributed minimal changes rather than late-layer wrappers.The authors conjecture that the KL-divergence term may account for this difference because ordinary fine-tuning typically lacks it.
7. Related Work
The related work situates DPO and PPO among alignment methods, connects mechanistic interpretability to representation analysis, and reviews jailbreak and unalignment findings.
- Alignment algorithms include human-feedback methods such as PPO and DPO, alongside methods for settings with only undesirable-behavior labels.The latter include unlikelihood training and Cringe.
- Mechanistic interpretability seeks to reverse engineer model behaviors using interpretable and controllable representations.Examples include world models, truthfulness, knowledge, and linguistic properties.
- Geometric analyses have extracted feature representations useful for toxicity detection, complementing work that probes specific representations.
- This paper studies RLHF effects in a realistic natural-language setting, whereas related work examines fine-tuning on synthetic tasks.
- Prior studies report that aligned models can be jailbroken through prompts that elicit previously unlearned behaviors.Multimodal models have also been shown to be jailbreakable.
- Other work shows that aligned models can be unaligned with as few as 100 fine-tuning examples.The paper positions its contribution as a mechanistic understanding of these phenomena.
8. Conclusion
The paper analyzes toxicity representations, applies DPO with carefully constructed toxic and nontoxic pairs, and finds that alignment bypasses rather than removes toxic-generation mechanisms.
- The authors identify numerous MLP value vectors that promote toxicity, and show that subtracting them from residual streams can suppress toxic outputs.
- DPO is applied using PPLM-crafted pairs of toxic and non-toxic continuations for Wikipedia prompts.
- GPT2DPO averts toxicity by learning an offset that bypasses toxic regions instead of removing them.The offset is distributed across multiple value vectors, producing minimal changes to each weight while preserving pre-trained behavior.
- Scaling the corresponding key vectors reactivates toxic regions and reverts GPT2DPO to toxic behavior.
- The paper presents these findings as a mechanistic case study for why aligned models can be jailbroken.Code, models, and data are made available by the authors.
- A. Projecting Value Vectors onto Vocabulary Space: MLP value vectors can promote or suppress token likelihood according to their projection onto token embeddings.Positive projections increase a token's likelihood, while negative projections decrease it.
B. Shift in Residual Streams
The section reports additional residual-stream shifts across selected layers and identifies highly similar toxic vectors, with DPO hyperparameters listed separately.
- Additional examples show residual streams shifting out of toxic regions.
- The analysis examines shifts at layers 12, 18, and 13.
- MLP.v18 2669 and MLP.v13 668 are identified among the next three vectors with highest cosine similarity to WToxic.
- Table 5 lists the DPO hyperparameters.
C. Shifts in Residual Streams vs. Shifts in MLP Value Vectors.
The section provides examples comparing shifts in residual streams with shifts in MLP value vectors, focusing on contributions in the δx direction across layers.
- The section gives additional examples of how MLP value vectors contribute in the δx direction at different layers.
- Figure 7 compares the shift in residual streams at layer 12 with the shift in MLP value vectors, δ12.
D. Hyperparameters
This section identifies the hyperparameters used for the toxic probe, DPO, and PPLM, including a table specifically for PPLM, alongside figures showing layer-specific shifts.
- Tables 5 and 6 contain the hyperparameters used for the toxic probe, DPO, and PPLM, respectively.
- Figure 8 shows the shift in residual streams at layer 14 versus the shift in MLP value vectors, δ14.
- Table 6 presents the hyperparameters for PPLM.
- Figure 9 shows the shift in residual streams at layer 16 versus the shift in MLP value vectors, δ16.
- Figure 10 shows the shift in residual streams at layer 18 versus the shift in MLP value vectors, δ18.