A Mechanistic Understanding of Alignment Algorithms: A Case Study on DPO and Toxicity
Paper · 2024 · Andrew Lee, Xiaoyan Bai, Itamar Pres et al.
Summary
Alignment methods based on preference tuning are now routinely used to make language models safer. However, it is unclear what actually changes inside a model when it becomes "aligned", and so jailbreaks remain poorly explained. This paper uses toxicity reduction with direct preference optimization (DPO) in GPT2-medium as a case study.
The authors first locate toxicity in the pre-trained model. A linear probe trained on the Jigsaw comment dataset provides a toxicity direction. That direction is used to retrieve MLP value vectors that boost toxic tokens, and SVD of those vectors gives a compact basis of toxicity "dimensions". Subtracting such vectors during generation reduces toxicity with little loss in fluency.
Next, they fine-tune the model with DPO on toxic/non-toxic continuation pairs, which they generated with PPLM from Wikipedia prompts. Their analysis shows that DPO leaves the toxic vectors essentially in place. Instead, small weight shifts spread across all layers steer activations away from them. As a result, simply amplifying the matching key vectors brings the toxic behaviour back. The work offers a mechanistic explanation of why aligned models are easy to un-align.
Key points
- The toxicity probe was trained on the last-layer residual stream of GPT2-medium using the Jigsaw toxic comment dataset (561,808 comments, 90:10 split). It reached 94% validation accuracy. The authors used it to pick N = 128 toxic MLP value vectors, which they then decomposed with SVD (SVD.U Toxic).
- Evaluation used three measures:
- toxicity, scored by Perspective API on the 1,199-prompt "challenge" subset of RealToxicityPrompts;
- perplexity on Wikitext-2;
- F1 on 2,000 Wikipedia prompts.
- Subtracting toxic vectors cut toxicity from 0.453 (unmodified GPT2) to 0.245–0.305. Perplexity rose only slightly (21.7 → about 23.5), and F1 stayed at about 0.19. The DPO model reached 0.208 toxicity, with perplexity 23.34 and F1 0.195.
- After DPO, all weights shift only slightly and the toxic vectors themselves survive. The model learns a distributed offset that avoids activating them. The authors speculate that this pattern comes from the KL-divergence regularization in the objective, in contrast to plain fine-tuning, which earlier work found to add localized late-layer "wrappers".
- Scaling up the key vectors tied to toxicity reverses the alignment of GPT2 DPO. The code, models and data are released at github.com/ajyl/dpo_toxic.
Fields, topics and concepts
📚 Shelves 4The field of science this work belongs to (OpenAlex fields).Computer Science · 18,614 worksArtificial Intelligence · 17,965 worksSocial Sciences · 1,954 worksSafety Research · 626 works
🧭 Topics 2Research subjects and languages the work is about.Bias, fairness & safety · 587 worksInterpretability · 149 works
Cite
Cite as Andrew Lee, Xiaoyan Bai, Itamar Pres, Martin Wattenberg, Jonathan K. Kummerfeld and Rada Mihalcea. 2024. A Mechanistic Understanding of Alignment Algorithms: A Case Study on DPO and Toxicity. arXiv:2401.01967.
@misc{lee2024mechanistic,
title = {{A Mechanistic Understanding of Alignment Algorithms: A Case Study on DPO and Toxicity}},
author = {Andrew Lee, Xiaoyan Bai, Itamar Pres, Martin Wattenberg, Jonathan K. Kummerfeld and Rada Mihalcea},
year = {2024},
doi = {10.48550/arxiv.2401.01967},
eprint = {2401.01967},
archivePrefix = {arXiv},
url = {https://doi.org/10.48550/arxiv.2401.01967}
}References
- Adolphs, L., Gao, T., Xu, J., Shuster, K., Sukhbaatar, S., and Weston, J. The CRINGE loss: Learning what language not to model. In Rogers, A., Boyd-Graber, J., and Okazaki, N. (eds.), Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp.… find
- Dathathri, S., Madotto, A., Lan, J., Hung, J., Frank, E., Molino, P., Yosinski, J., and Liu, R. Plug and play language models: A simple approach to controlled text generation. In International Conference on Learning Representations, 2019.
- Balestriero, R., Cosentino, R., and Shekkizhar, S. Characterizing large language model geometry solves toxicity detection and generation. arXiv preprint arXiv:2312.01648, 2023. — [arXiv:2312.01648](https://arxiv.org/abs/2312.01648) find
- Dinan, E., Logacheva, V., Malykh, V., Miller, A., Shuster, Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., K., Urbanek, J., Kiela, D., Szlam, A., Serban, I., Lowe, Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., R., et al. The second conversational intelligence challenge Askell, A.… find
- Curran Associates, Inc., L., Ndousse, K., Amodei, D., Brown, T., Clark, J., 2020. URL https://proceedings.neurips. Kaplan, J., McCandlish, S., and Olah, C. A math-cc/paper_files/paper/2020/file/ ematical framework for transformer circuits. Trans-1457c0d6bfcb4967418bfb8ac142f64a-Paper. former Circuit… ↗
- Burns, C., Ye, H., Klein, D., and Steinhardt, J. Discovering latent knowledge in language models without supervision. In The Eleventh International Conference on Learning Representations, 2023. URL https://openreview.net/forum?id=ETKGuby0hcs. ↗
- Gehman, S., Gururangan, S., Sap, M., Choi, Y., and Smith, N. A. RealToxicityPrompts: Evaluating neural toxic degeneration in language models. In Cohn, T., He, Y., and Liu, Y. (eds.), Findings of the Association for Computational Linguistics: EMNLP 2020, pp. 3356–3369, Online, November 2020. Associat…
- Carlini, N., Nasr, M., Choquette-Choo, C. A., Jagielski, M., Gao, I., Koh, P. W., Ippolito, D., Tramèr, F., and Schmidt, L. Are aligned neural networks adversarially aligned? In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL https://openreview.net/forum?id=OQQoD8Vc3B… ↗
- Geva, M., Schuster, R., Berant, J., and Levy, O. Transformer feed-forward layers are key-value memories. In Moens, M.-F., Huang, X., Specia, L., and Yih, S. W.-t. (eds.), Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pp. 5484–5495, Online and Punta Cana, Dom… find
- Chowdhery, A., Narang, S., Devlin, J., Bosma, M., Mishra, G., Roberts, A., Barham, P., Chung, H. W., Sutton, C., Gehrmann, S., et al. Palm: Scaling language modeling with pathways. Journal of Machine Learning Research, 24(240):1–113, 2023. 10.18653/v1/2021.emnlp-main.446. URL https://aclanthology.o… ↗
- Geva, M., Caciularu, A., Wang, K., and Goldberg, Y. Transformer feed-forward layers build predictions by promoting concepts in the vocabulary space. In Goldberg, Y., Kozareva, Z., and Zhang, Y. (eds.), Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pp. 30–45,…
- Geva, M., Bastings, J., Filippova, K., and Globerson, A. Dissecting recall of factual associations in auto-regressive language models. arXiv preprint arXiv:2304.14767, 2023. find
- Li, Z., You, C., Bhojanapalli, S., Li, D., Rawat, A. S., Reddi, S. J., Ye, K., Chern, F., Yu, F., Guo, R., and Kumar, S. The lazy neuron phenomenon: On emergence of activation sparsity in transformers. In The Eleventh International Conference on Learning Representations, 2023c. URL https://openrevie… find
- Meng, K., Bau, D., Andonian, A. J., and Belinkov, Y. Locating and editing factual associations in GPT. In Oh, A. H., Agarwal, A., Belgrave, D., and Cho, K. (eds.), Advances in Neural Information Processing Systems, 2022. URL https://openreview.net/forum? id=-h6WAS6eE4. ↗
- Merity, S., Xiong, C., Bradbury, J., and Socher, R. Pointer sentinel mixture models. In International Conference on Learning Representations, 2016.
- Hendel, R., Geva, M., and Globerson, A. In-context learning creates task vectors. In Bouamor, H., Pino, J., and Bali, K. (eds.), Findings of the Association for Computational Linguistics: EMNLP 2023, pp. 9318–9333, Singapore, December 2023. Association for Computational Linguistics. URL https://acla…
- Hendrycks, D. and Gimpel, K. Gaussian error linear units (gelus). arXiv preprint arXiv:1606.08415, 2016. find
- Hernandez, E., Sharma, A. S., Haklay, T., Meng, K., Wattenberg, M., Andreas, J., Belinkov, Y., and Bau, D. Linearity of relation decoding in transformer language models. arXiv preprint arXiv:2308.09124, 2023.
- Nanda, N., Lee, A., and Wattenberg, M. Emergent linear representations in world models of self-supervised sequence models. In Proceedings of the 6th BlackboxNLP Workshop: Analyzing and Interpreting Neural Networks for NLP, pp. 16–30, 2023.
- Nasr, M., Carlini, N., Hayase, J., Jagielski, M., Cooper, A. F., Ippolito, D., Choquette-Choo, C. A., Wallace, E., Tramèr, F., and Lee, K. Scalable extraction of training data from (production) language models. arXiv preprint arXiv:2311.17035, 2023.
- Nostalgebraist. Interpreting gpt: The logit lens, 2020. URL https://www.lesswrong.com/posts/AcKRB8wDpdaN6v6ru/ interpreting-gpt-the-logit-lens. ↗
- Ilharco, G., Ribeiro, M. T., Wortsman, M., Schmidt, L., Hajishirzi, H., and Farhadi, A. Editing models with task arithmetic. In The Eleventh International Conference on Learning Representations, 2022.
- Olah, C., Cammarata, N., Schubert, L., Goh, G., Petrov, M., and Carter, S. Zoom in: An introduction to circuits. Distill, 2020. doi: 10.23915/distill.00024.001. https://distill.pub/2020/circuits/zoom-in. — [doi:10.23915/distill.00024.001](https://doi.org/10.23915/distill.00024.001) ↗
- Jain, S., Kirk, R., Lubana, E. S., Dick, R. P., Tanaka, H., Grefenstette, E., Rocktäschel, T., and Krueger, D. S. Mechanistically analyzing the effects of fine-tuning on procedurally defined tasks. arXiv preprint arXiv:2311.12786, 2023. — [arXiv:2311.12786](https://arxiv.org/abs/2311.12786) find
- Ouyang, L., Wu, J., Jiang, X., Almeida, D., Wainwright, C., Mishkin, P., Zhang, C., Agarwal, S., Slama, K., Gray, A., et al. Training language models to follow instructions with human feedback. In Advances in Neural Information Processing Systems, 2022.
- Li, K., Hopkins, A. K., Bau, D., Viégas, F., Pfister, H., and Wattenberg, M. Emergent world representations: Exploring a sequence model trained on a synthetic task. In The Eleventh International Conference on Learning Representations, 2023a. URL https://openreview.net/ forum?id=DeG07_TcZvT. ↗
- Park, K., Choe, Y. J., and Veitch, V. The linear representation hypothesis and the geometry of large language models. In Causal Representation Learning Workshop at NeurIPS 2023, 2023.
- Li, K., Patel, O., Viégas, F., Pfister, H., and Wattenberg, M. Inference-time intervention: Eliciting truthful answers from a language model. arXiv preprint arXiv:2306.03341, 2023b.
- Qi, X., Zeng, Y., Xie, T., Chen, P.-Y., Jia, R., Mittal, P., and Henderson, P. Fine-tuning aligned language models compromises safety, even when users do not intend to! arXiv preprint arXiv:2310.03693, 2023.
- Rafailov, R., Sharma, A., Mitchell, E., Ermon, S., Manning, C. D., and Finn, C. Direct preference optimization: Your language model is secretly a reward model, 2023.
- X. (eds.), Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP), pp. 2153–2162, Hong Kong, China, November 2019. Association for Computational Linguistics. doi: 10.18653/v1/D19-… find
- Schulman, J., Wolski, F., Dhariwal, P., Radford, A., and Klimov, O. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347, 2017. find
- Sheng, E., Chang, K.-W., Natarajan, P., and Peng, N. The woman worked as a babysitter: On biases in language generation. In Inui, K., Jiang, J., Ng, V., and Wan, X. (eds.), Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conferen…
- Wei, A., Haghtalab, N., and Steinhardt, J. Jailbroken: How does LLM safety training fail? In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL https://openreview.net/forum? id=jA235JGM09. ↗
- Welleck, S., Kulikov, I., Roller, S., Dinan, E., Cho, K., and Weston, J. Neural text generation with unlikelihood training. In International Conference on Learning Representations, 2020. URL https://openreview.net/ forum?id=SJeYe0NtvH. ↗
- Stiennon, N., Ouyang, L., Wu, J., Ziegler, D., Lowe, R., Voss, C., Radford, A., Amodei, D., and Christiano, P. F. Learning to summarize with human feedback. Advances in Neural Information Processing Systems, 33: 3008–3021, 2020. find
- Xu, J., Lee, A., Sukhbaatar, S., and Weston, J. Some things are more cringe than others: Preference optimization with the pairwise cringe loss. arXiv preprint arXiv:2312.16682, 2023. — [arXiv:2312.16682](https://arxiv.org/abs/2312.16682) find
- Yang, X., Wang, X., Zhang, Q., Petzold, L., Wang, W. Y., Zhao, X., and Lin, D. Shadow alignment: The ease of subverting safely-aligned language models. arXiv preprint arXiv:2310.02949, 2023.
- Tenney, I., Das, D., and Pavlick, E. BERT rediscovers the classical NLP pipeline. In Korhonen, A., Traum, D., and Màrquez, L. (eds.), Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pp. 4593–4601, Florence, Italy, July 2019. Association for Computational Lin…
- Zhang, Z., Lin, Y., Liu, Z., Li, P., Sun, M., and Zhou, J. MoEfication: Transformer feed-forward layers are mixtures of experts. In Muresan, S., Nakov, P., and Villavicencio, A. (eds.), Findings of the Association for Computational Linguistics: ACL 2022, pp. 877–890, Dublin, Ireland, May 2022. Assoc…
- Todd, E., Li, M. L., Sharma, A. S., Mueller, A., Wallace, B. C., and Bau, D. Function vectors in large language models. arXiv preprint arXiv:2310.15213, 2023.
- Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.-A., Lacroix, T., Rozière, B., Goyal, N., Hambro, E., Azhar, F., et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971, 2023.
- Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L. u., and Polosukhin, I. Attention is all you need. In Guyon, I., Luxburg, U. V., Bengio, S., Wallach, H., Fergus, R., Vishwanathan, S., and Garnett, R. (eds.), Advances in Neural Information Processing Systems, v…
- Wallace, E., Feng, S., Kandpal, N., Gardner, M., and Singh, S. Universal adversarial triggers for attacking and analyzing NLP. In Inui, K., Jiang, J., Ng, V., and Wan, Zou, A., Wang, Z., Kolter, J. Z., and Fredrikson, M. Universal and transferable adversarial attacks on aligned language models. arXi…
- 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. find
- We start from Equation 2: d mlp ℓ ℓ MLP (x ) = find
Updated
This page in Markdown → https://md.research.elamur.ai/free/en/papers/lee-2024-mechanistic-understanding-alignment-algorithms-case/