# Learning Distribution-Wise Control in Representation Space for Language Models

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

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Builds on 49 of its 70 references that are in the map.
- First posted on arXiv 7 Jun 2025.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Training Verifiers to Solve Math Word Problems](https://mlsanity.com/paper/2110.14168/) (2021)
- [Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challenge](https://mlsanity.com/paper/1803.05457/) (2018)
- [Parameter-Efficient Transfer Learning for NLP](https://mlsanity.com/paper/1902.00751/) (2019)
- [HellaSwag: Can a Machine Really Finish Your Sentence?](https://mlsanity.com/paper/1905.07830/) (2019)
- [Prefix-Tuning: Optimizing Continuous Prompts for Generation](https://mlsanity.com/paper/2101.00190/) (2021)
- [Locating and Editing Factual Associations in GPT](https://mlsanity.com/paper/2202.05262/) (2022)
- [BoolQ: Exploring the Surprising Difficulty of Natural Yes/No Questions](https://mlsanity.com/paper/1905.10044/) (2019)
- [Can a Suit of Armor Conduct Electricity? A New Dataset for Open Book Question Answering](https://mlsanity.com/paper/1809.02789/) (2018)
- [GLU Variants Improve Transformer](https://mlsanity.com/paper/2002.05202/) (2020)
- [Transformer Feed-Forward Layers Are Key-Value Memories](https://mlsanity.com/paper/2012.14913/) (2020)
- [Inference-Time Intervention: Eliciting Truthful Answers from a Language Model](https://mlsanity.com/paper/2306.03341/) (2023)
- [Interpretability in the Wild: a Circuit for Indirect Object Identification in GPT-2 small](https://mlsanity.com/paper/2211.00593/) (2022)
- [Deep Variational Information Bottleneck](https://mlsanity.com/paper/1612.00410/) (2016)
- [Steering Language Models With Activation Engineering](https://mlsanity.com/paper/2308.10248/) (2023)
- [Are NLP Models really able to Solve Simple Math Word Problems?](https://mlsanity.com/paper/2103.07191/) (2021)
- [Deep Learning and the Information Bottleneck Principle](https://mlsanity.com/paper/1503.02406/) (2015)
- [LLaMA-Adapter: Efficient Fine-tuning of Language Models with Zero-init Attention](https://mlsanity.com/paper/2303.16199/) (2023)
- [Towards Automated Circuit Discovery for Mechanistic Interpretability](https://mlsanity.com/paper/2304.14997/) (2023)
- [The Linear Representation Hypothesis and the Geometry of Large Language Models](https://mlsanity.com/paper/2311.03658/) (2023)
- [Program Induction by Rationale Generation : Learning to Solve and Explain Algebraic Word Problems](https://mlsanity.com/paper/1705.04146/) (2017)
- [MAD-X: An Adapter-Based Framework for Multi-Task Cross-Lingual Transfer](https://mlsanity.com/paper/2005.00052/) (2020)
- [Causal Abstractions of Neural Networks](https://mlsanity.com/paper/2106.02997/) (2021)
- [UltraFeedback: Boosting Language Models with Scaled AI Feedback](https://mlsanity.com/paper/2310.01377/) (2023)
- [Solving General Arithmetic Word Problems](https://mlsanity.com/paper/1608.01413/) (2016)
- [Towards Best Practices of Activation Patching in Language Models: Metrics and Methods](https://mlsanity.com/paper/2309.16042/) (2023)
- [LLM-Adapters: An Adapter Family for Parameter-Efficient Fine-Tuning of Large Language Models](https://mlsanity.com/paper/2304.01933/) (2023)
- [Emergent Linear Representations in World Models of Self-Supervised Sequence Models](https://mlsanity.com/paper/2309.00941/) (2023)
- [LEACE: Perfect linear concept erasure in closed form](https://mlsanity.com/paper/2306.03819/) (2023)
- [Patchscopes: A Unifying Framework for Inspecting Hidden Representations of Language Models](https://mlsanity.com/paper/2401.06102/) (2024)

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

- [Spherical Steering: Geometry-Aware Activation Rotation for Language Models](https://mlsanity.com/paper/2602.08169/) (2026)
- [Steering Information Utility in Key-Value Memory for Language Model Post-Training](https://mlsanity.com/paper/2507.05158/) (2025)
- [The Generalization Ridge: Information Flow in Natural Language Generation](https://mlsanity.com/paper/2507.05387/) (2025)
- [BOReFT: Manifold Steering of Language Models for Black-box Optimization](https://mlsanity.com/paper/2609.33722/) (2026)

## Abstract (arXiv)

Interventions in language models (LMs) are applied strategically to steer model behavior during the forward pass. Learnable interventions, also known as representation fine-tuning, aim to apply pointwise control within the concept subspace and have proven effective in altering high-level behaviors. In this work, we extend this approach to the distribution level, enabling the model to learn not only pointwise transformations but also the surrounding regions of the concept subspace. We demonstrate that these methods perform effectively in early layers, with larger standard deviations correlating strongly with improved performance. Across eight commonsense reasoning and seven arithmetic reasoning benchmarks, our distribution-wise interventions consistently outperform pointwise interventions in controllability and robustness. These results illustrate that distribution-wise interventions provide a more comprehensive method for steering model behavior and enabling finer-grained control over language models. The code is at: hrefhttps://github.com/chili-lab/D-Interventionhttps://github.com/chili-lab/D-Intervention.