# Distributional Surgery for Language Model Activations

arXiv: https://arxiv.org/abs/2501.15758 · HTML: https://mlsanity.com/paper/2501.15758/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 31 of its 49 references that are in the map.
- First posted on arXiv 27 Jan 2025.

## Builds on (most cited first, 30 of 31)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [Training a Helpful and Harmless Assistant with Reinforcement Learning from Human Feedback](https://mlsanity.com/paper/2204.05862/) (2022)
- [TruthfulQA: Measuring How Models Mimic Human Falsehoods](https://mlsanity.com/paper/2109.07958/) (2021)
- [TriviaQA: A Large Scale Distantly Supervised Challenge Dataset for Reading Comprehension](https://mlsanity.com/paper/1705.03551/) (2017)
- [Self-Instruct: Aligning Language Models with Self-Generated Instructions](https://mlsanity.com/paper/2212.10560/) (2022)
- [Locating and Editing Factual Associations in GPT](https://mlsanity.com/paper/2202.05262/) (2022)
- [Don't Stop Pretraining: Adapt Language Models to Domains and Tasks](https://mlsanity.com/paper/2004.10964/) (2020)
- [Understanding intermediate layers using linear classifier probes](https://mlsanity.com/paper/1610.01644/) (2016)
- [Inference-Time Intervention: Eliciting Truthful Answers from a Language Model](https://mlsanity.com/paper/2306.03341/) (2023)
- [RealToxicityPrompts: Evaluating Neural Toxic Degeneration in Language Models](https://mlsanity.com/paper/2009.11462/) (2020)
- [Probing Classifiers: Promises, Shortcomings, and Advances](https://mlsanity.com/paper/2102.12452/) (2021)
- [Plug and Play Language Models: A Simple Approach to Controlled Text Generation](https://mlsanity.com/paper/1912.02164/) (2019)
- [Discovering Latent Knowledge in Language Models Without Supervision](https://mlsanity.com/paper/2212.03827/) (2022)
- [The Geometry of Truth: Emergent Linear Structure in Large Language Model Representations of True/False Datasets](https://mlsanity.com/paper/2310.06824/) (2023)
- [DExperts: Decoding-Time Controlled Text Generation with Experts and Anti-Experts](https://mlsanity.com/paper/2105.03023/) (2021)
- [OLMoE: Open Mixture-of-Experts Language Models](https://mlsanity.com/paper/2409.02060/) (2024)
- [Supervised Contrastive Learning for Pre-trained Language Model Fine-tuning](https://mlsanity.com/paper/2011.01403/) (2020)
- [Does Localization Inform Editing? Surprising Differences in Causality-Based Localization vs. Knowledge Editing in Language Models](https://mlsanity.com/paper/2301.04213/) (2023)
- [Relative representations enable zero-shot latent space communication](https://mlsanity.com/paper/2209.15430/) (2022)
- [The Troubling Emergence of Hallucination in Large Language Models -- An Extensive Definition, Quantification, and Prescriptive Remediations](https://mlsanity.com/paper/2310.04988/) (2023)
- [Pretrained Language Models for Text Generation: A Survey](https://mlsanity.com/paper/2201.05273/) (2022)
- [LoFiT: Localized Fine-tuning on LLM Representations](https://mlsanity.com/paper/2406.01563/) (2024)
- [Representation Surgery: Theory and Practice of Affine Steering](https://mlsanity.com/paper/2402.09631/) (2024)
- [Unified Detoxifying and Debiasing in Language Generation via Inference-time Adaptive Optimization](https://mlsanity.com/paper/2210.04492/) (2022)
- [sigmoidF1: A Smooth F1 Score Surrogate Loss for Multilabel Classification](https://mlsanity.com/paper/2108.10566/) (2021)
- [Goodtriever: Adaptive Toxicity Mitigation with Retrieval-augmented Models](https://mlsanity.com/paper/2310.07589/) (2023)

## Built on by (most cited first, 2 of 2)

- [Structured Pruning for Diverse Best-of-N Reasoning Optimization](https://mlsanity.com/paper/2506.03978/) (2025)
- [Task-driven Layerwise Additive Activation Intervention](https://mlsanity.com/paper/2502.06115/) (2025)

## Abstract (arXiv)

Language models, while capable of generating remarkably coherent and seemingly accurate text, can occasionally produce undesirable content, including harmful or toxic outputs. In this paper, we present a new two-stage approach to detect and mitigate undesirable content generations by rectifying activations. First, we train an ensemble of layerwise classifiers to detect undesirable content using activations by minimizing a smooth surrogate of the risk-aware score. Then, for detected undesirable contents, we propose layerwise distributional steering policies that transform the attention heads. These policies are computed through principled semidefinite programming, which aims to minimally perturb the attention distribution while probabilistically guaranteeing the effectiveness of the editions. Empirical evaluations across multiple language models and datasets show that our method outperforms baselines in reducing the generation of undesirable output.