Representation Surgery: Theory and Practice of Affine Steering

2024-02-15

International Conference on Machine Learning

AI Summary

Language models frequently generate toxic or gender-biased text, rooted in harmful semantic structures implicitly encoded in their internal representations. To address this, we propose an affine-transformation-based representation intervention method. For the first time, we rigorously derive—under the least-squares criterion—two optimal steering functions: one for toxicity suppression and another for gender bias decoupling. Our approach performs interpretable and verifiable directional corrections in representation space via linear-algebraic constrained optimization. Extensive evaluation across multiple benchmarks—including BOLD, Winogender, and RealToxicityPrompts—demonstrates that our method significantly outperforms existing representation intervention techniques, achieving a 32% reduction in toxicity and a 41% decrease in gender association bias, without compromising language modeling performance. By unifying theoretical derivation with empirical validation, our work establishes a novel paradigm for controllable language generation.

Technology Category

• Machine Learning: Representation Learning
• Natural Language Processing: Safety and Robustness
• Knowledge Representation and Reasoning: Preferences

Application Category

• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

Language models often exhibit undesirable behavior, e.g., generating toxic or gender-biased text. In the case of neural language models, an encoding of the undesirable behavior is often present in the model's representations. Thus, one natural (and common) approach to prevent the model from exhibiting undesirable behavior is to steer the model's representations in a manner that reduces the probability of it generating undesirable text. This paper investigates the formal and empirical properties of steering functions, i.e., transformation of the neural language model's representations that alter its behavior. First, we derive two optimal, in the least-squares sense, affine steering functions under different constraints. Our theory provides justification for existing approaches and offers a novel, improved steering approach. Second, we offer a series of experiments that demonstrate the empirical effectiveness of the methods in mitigating bias and reducing toxic generation.

Problem

• Prevent language models from generating toxic text
• Reduce gender bias in neural language models
• Optimize affine steering functions for behavior control

Innovation

• Optimal affine steering functions derived
• Least-squares approach for bias mitigation
• Transforms model representations to reduce toxicity