Distributional Surgery for Language Model Activations

TL;DR

This paper introduces a two-stage method to detect and mitigate undesirable content in language models by rectifying activations through distributional steering, improving safety without significantly altering model behavior.

Contribution

It proposes a novel ensemble of classifiers for detection and a distributional steering approach via semidefinite programming for mitigation, advancing safety in language model outputs.

Findings

Outperforms baselines in reducing harmful outputs

Effective across multiple language models and datasets

Minimally perturbs attention distributions

Abstract

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.