Distributional Surgery for Language Model Activations
This addresses the issue of harmful or toxic outputs in language models, which is a concern for users and developers, though it appears incremental as it builds on existing detection and mitigation techniques.
The paper tackles the problem of undesirable content generation in language models by introducing a two-stage method that detects and mitigates such content through activation rectification, resulting in improved performance over baselines in reducing undesirable outputs.
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.