Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Distributional Surgery for Language Model Activations

Bao Nguyen, Binh Nguyen, Duy Nguyen, Viet Anh Nguyen

TL;DR

This work introduces RADIANT, a two-stage inference-time framework to reduce undesirable outputs from language models by intervening on activations rather than updating weights. It combines layerwise risk-aware probes (per-head linear logistic classifiers with a smooth risk surrogate over FPR and FNR) to detect problematic content and select an optimal layer, with headwise distributional interventions (linear maps) that transform undesirable activations into desirable ones using semidefinite programming to guarantee probabilistic effectiveness while minimizing perturbations. The method achieves strong improvements on TruthfulQA across a range of models and datasets, while maintaining small distribution shifts as evidenced by KL and CE metrics, and it generalizes to various architectures, including sparse Mixtures of Experts. Its inference-time, parameter-free edits offer a lightweight, scalable alternative to fine-tuning or weight editing, with broad applicability to truthfulness, toxicity mitigation, and other controllable-generation tasks. The combination of risk-aware probing and SDP-guided interventions provides principled guarantees and practical performance gains, making RADIANT a versatile tool for safer and more reliable text generation.

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.

Distributional Surgery for Language Model Activations

TL;DR

This work introduces RADIANT, a two-stage inference-time framework to reduce undesirable outputs from language models by intervening on activations rather than updating weights. It combines layerwise risk-aware probes (per-head linear logistic classifiers with a smooth risk surrogate over FPR and FNR) to detect problematic content and select an optimal layer, with headwise distributional interventions (linear maps) that transform undesirable activations into desirable ones using semidefinite programming to guarantee probabilistic effectiveness while minimizing perturbations. The method achieves strong improvements on TruthfulQA across a range of models and datasets, while maintaining small distribution shifts as evidenced by KL and CE metrics, and it generalizes to various architectures, including sparse Mixtures of Experts. Its inference-time, parameter-free edits offer a lightweight, scalable alternative to fine-tuning or weight editing, with broad applicability to truthfulness, toxicity mitigation, and other controllable-generation tasks. The combination of risk-aware probing and SDP-guided interventions provides principled guarantees and practical performance gains, making RADIANT a versatile tool for safer and more reliable text generation.

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.
Paper Structure (33 sections, 1 theorem, 15 equations, 2 figures, 21 tables)

This paper contains 33 sections, 1 theorem, 15 equations, 2 figures, 21 tables.

Table of Contents

  1. Introduction
  2. Related Works
  3. Layerwise Risk-aware Probes
  4. Headwise Interventions with Probabilistic Guarantees
  5. Experiments
  6. Experimental Settings
  7. Numerical Results
  8. Comparison between Finetuning-free Techniques
  9. Comparison against ITI with Instruction Finetuning Methods.
  10. Comparison against ITI with Representation Intervention Finetuning Methods.
  11. Ablation Studies
  12. Computational Cost
  13. Additional Experiments
  14. Conclusion
  15. Additional Results on TruthfulQA Benchmark
  16. ...and 18 more sections

Key Result

Theorem 1

Suppose that $\widehat{\mathbb P}_{\ell h} \sim \mathcal{N}(\widehat{\mu}, \widehat{\Sigma})$ and $\mathbb Q_{\ell h} \sim \mathcal{N} (\mu, \Sigma)$ and $\varphi$ admits the form Let $(\mu^\star, S^\star, t^\star)$ be the solution of the following semidefinite program Then, by defining $G^\star_{\ell h} = \widehat{\Sigma}^{-\frac{1}{2}} ( \widehat{\Sigma}^{\frac{1}{2}} (S^\star)^2 \widehat{\Sig

Figures (2)

  • Figure 1: At head $h$ of layer $\ell$, we learn a headwise intervention, linear mapping $\Delta_{\ell h}$ to transform the undesirable-predicted activations to desirable-predicted activations.
  • Figure 2: FNR and FPR metrics with different hyperparameters $\alpha$ across layers of Llama-7B.

Theorems & Definitions (3)

  • Theorem 1: Optimal headwise intervention
  • Remark 1
  • proof : Proof of Theorem \ref{['thm:intervene']}