Distributional Surgery for Language Model Activations (2025)

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 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.

## introduction language models (lms) have demonstrated a remarkable ability to understand and generate humanlike documents @xcite @xcite @xcite . however, inspection of their output often reveals undesirable generation, including inaccurate or toxic texts @xcite @xcite . meanwhile, developing effective strategies to control the generation process of lms remains a significant challenge @xcite . researchers have proposed numerous methods for controllable text generation in language models @xcite . these approaches primarily include model editing and supervised fine-tuning. both methods, however, require altering the model weights using a subset of text samples, which can result in unstable representations for other text instances @xcite . additionally, these approaches typically require substantial computational resources. to address these limitations, activation intervention has emerged as a promising alternative for controllable text generation @xcite @xcite . this approach involves altering the model activations responsible for undesirable output during inference. previous research has identified interpretable directions within the activation space of language models that play a causal role during inference. studies by @xcite and @xcite suggest that these directions can be manipulated to adjust model behavior in a controlled manner. this body of work indicates that the internal representations of language models are structured in ways that enable fine-grained control over generated text. drawing inspiration from these findings, activation intervention frameworks operate on the premise that the information needed to guide the model toward generating a target sentence is already encoded within the model. the hidden information is extracted as latent vectors and then used to steer the generation toward producing desirable outputs. the preliminary success of these activation intervention methods motivates our approach to improving the quality and controllability of language model generation. problem statement. we consider a language model consisting of l layers, each layer has h heads, each head has dimension d. for example, for llama-2, we have @xmath0, @xmath1, and @xmath2. the training dataset is denoted by @xmath3, y * i ) @xmath4,...,n , the i-th text is denoted by x i , and its ground truth label is y * @xmath5, 1}, where the label 1 (positive) represents the undesirable text, and the label 0 (negative) represents the desirable text. our goal is twofold: (i) detect an undesirable text, and (ii) modify an undesirable text into a desirable text. the activations for a text x i at layer ℓ ∈ {1, . . . , l} is denoted by a ℓ,i . the activation at layer ℓ + 1 is the output of the operation: a ℓ+1,@xmath6,i + ffn(a mid ℓ,i ), a mid ℓ,@xmath7,i + h @xmath8,i ). here, p ℓh ∈ r d×dh is the projection matrix that maps the output of each layer to the d dimensional head space, att is the attention operator @xcite , q ℓh ∈ r dh×d is the pull-back matrix and ffn is the feed-forward layer. each a ℓ,i is a concatenation of headwise activations a ℓh,i for @xmath9, . . . , h. inspired by @xcite , we aim to perform intervention at some selected a ℓh,i , the activations for head h of layer ℓ, if we detect that the activation is from undesirable content. contributions. we contribute an activation intervention method to detect and rectify the undesirable generation of lm. we call our method radiant (risk-aware distributional intervention policies for language models' activations). radiant comprises two components: ## layerwise risk-aware probes in the first step, we aim to find a classifier c ℓh : r @xmath10, 1} for each head @xmath11, . . . , h at each layer ℓ = 1, . . . , l to classify the activation value a ℓh of desirable and undesirable texts. we propose using a linear logistic classifier, parameterized by a slope parameter θ ℓh ∈ r d and a bias parameter ϑ ℓh ∈ r. the headwise classification rule is the training process of c ℓh must take into account two types of risk: (i) false-negative risk when an undesirable text is not detected, (ii) falsepositive risk when a desirable text is classified as undesirable, and is subsequently edited and loses its original semantics. therefore, a natural candidate for the loss function is a combination of the false positive rate (fpr) and the false negative rate (fnr). however, neither fpr nor fnr have smooth functions in optimizing variables. we, hence, resort to smooth surrogates of these two metrics that use the predicted probability of the classifier, similarly to @xcite . in detail, we use to conclude this step, we compute the classifier c ℓ for layer ℓ = 1, . . . , l by tuning the parameters α. the layer whose classifier c ℓ delivers the highest quality (accuracy or any risk-aware metric) will be the optimal layer to construct the probe. this optimal layer, along with the collection of headwise classifiers, is the final output of this step. ## headwise interventions with probabilistic guarantees we propose a distributional intervention to the activations of the samples predicted as undesirable by the layerwise classifier. in this section, we will focus on constructing a single headwise intervention, and in the next section, we will combine multiple headwise interventions into a layerwise intervention. a headwise intervention is a map ∆ ℓh : a ℓh → âℓh that needs to balance multiple criteria: (i) it should be easy to compute and deploy, (ii) it should be effective in converting the undesirable activations to the desirable regions, (iii) it should minimize the magnitude of the intervention to sustain the context of the input. intuitively, we propose solving an optimization problem with the loss and constraints that fit all the criteria listed. the details are as follows. to promote (i), we employ a simple linear map ∆ ℓh (a ℓh ) = g ℓh a ℓh + g ℓh parametrized by a matrix g ℓh ∈ r d×d and a vector g ℓh ∈ r d . this linear map can also be regarded as a pushforward map that transforms the undesirable-predicted activations to become desirable-predicted activations. let us now represent the undesirable-predicted activations as a d-dimensional random vector ãℓh . its distribution can be estimated using the training data after identifying the subset d+ ℓh of training samples that are predicted undesirable by c ℓh , that is, d+ ℓh ≜ {i : c ℓh (a ℓh,i ) = 1}. the activations of samples in d+ ℓh lead to an empirical distribution p ℓh . the linear map ∆ ℓh will pushforward the distribution p ℓh to the new distribution q ℓh = ∆ ℓh # p. using the pushforward distribution q ℓh , we can impose criteria (ii) and (iii) above in an intuitive method. to promote (ii), we require that the activations distributed under q ℓh should be classified as desirable by c ℓh with high probability. finally, to promote (iii), we require that the distributions q ℓh and p ℓh be not too far from each other. let γ ∈ (0, 0.5) be a small tolerance parameter, and let φ be a measure of dissimilarity between probability distributions, we propose to find ∆ ℓh by solving the following stochastic program the proof of theorem 1 can be found in appendix d. figure 1 illustrates the effect of the headwise intervention ∆ ℓh : the headwise classifier c ℓh is represented by the red linear hyperplane ϑ ℓh + θ ⊤ ℓh @xmath12; the undesirable-predicted (label 1) region is towards the top left corner, while the desirable-predicted (label 0) region is towards the bottom right corner. the activations of the undesirable-predicted samples are represented as a gaussian distribution with mean ( µ, σ), drawn as the red ellipsoid. the edit map ∆ ℓh pushes this distribution to another gaussian distribution q ℓh drawn as the green ellipsoid. the distribution q ℓh has a coverage guarantee on the desirable-predicted region with probability at least 1γ. one can also verify that q ℓh has mean µ ⋆ and covariance matrix (s ⋆ ) 2 . problem (4) can be solved using semidefinite programming solvers such as copt or mosek. the moment information µ and σ can be estimated from the subset d+ ℓh . one can intuitively expect a trade-off between the tolerance level γ and the magnitude of the headwise mapping. if γ is lower, the activations will be edited at a higher magnitude so that the edited activations will likely end up in the desirable-predicted region of the classifier c ℓh . in contrast, if γ is higher, the activations will be edited with a smaller magnitude due to the less stringent constraint to swap the predicted label. one can view the distribution q ℓh ∼ (µ ⋆ , (s ⋆ ) 2 ) as the counterfactual distribution of the undesirablepredicted activations with minimal perturbation. this distribution q ℓh is found by optimization, which is in stark contrast with the design of the counterfactual distribution in mimic @xcite , in which the intervention is computed based on the activations of the desirable-predicted activations. as a comparison to iti @xcite , we note that the headwise intervention of iti does not depend on the value of the activations: iti shifts the activations along the truthful directions for a stepsize multiplied by the standard deviation of activations along the intervention (truthful) direction. in contrast, our headwise intervention depends on the value a ℓh , and one can verify that the magnitude of the proposed shift amounts to ∥(g ⋆ ℓh -i)a ℓh + g ⋆ ℓh ∥ 2 . moreover, iti does not provide any (probabilistic) guarantee for the intervention, while the probabilistic guarantee is internalized in our method through the design of the mapping in equation ( 3 ). remark 1. we observe that the two following tricks increase the empirical performance of our intervention framework. first, to avoid the collapse of q ℓh into a dirac distribution and to ensure the similarity between the real and the constructed covariance matrix of desirable content, we can add the constraint s ⪰ σ 1 2 0 to the optimization problem (4), where σ 0 is the empirical covariance matrix of the desirable activations {i : y * @xmath13. second, to avoid taking the inverse cdf of the standard normal distribution, we use γ ← φ -1 (1γ) and finetune γ instead of γ. finally, given the input with activation a ℓ at layer ℓ, suppose that a ℓ is predicted undesirable by c ℓ , we propose to edit the activations of only the heads that are predicted undesirable by the headwise classifier c ℓh . more specifically, we edit the headwise activations a ℓh to a new headwise activations âℓh through the relationship ## experiments in this section, we present empirical evidence for the effectiveness of our method radiant. we evaluate radiant on the truthfulqa benchmark @xcite , consisting of two tasks: the main task is the generation, and the secondary task is multiple choice. the generation task requires the model to generate an entire answer for each question using greedy autoregressive decoding. the accuracy and helpfulness of the answer are best assessed by humans. however, in almost all recent works in the field, including @xcite and @xcite , this criterion is measured by an alternative large language model finetuned on the target dataset. the multiple-choice task contains candidate answers to each question, requiring the model to give probabilities. higher probabilities for truthful answers yield higher scores. ## conclusion we introduced radiant, an activation intervention method consisting of two components: (i) a layerwise probe to detect undesirable content and (ii) headwise interventions to rectify the head activations upon undesirably predicted outcomes. contrary to existing intervention methods, where the interventions can be scattered across different layers, our intervention is focused on a single layer of the network. this focus helps alleviate the distributional shifts of the activations in subsequent layers. moreover, our headwise intervention aims to minimize the perturbations to the activations while keeping a reasonable guarantee of the effectiveness of the intervention. this is further demonstrated in empirical results, where our method outperforms the baseline intervention methods for various lms. limitations and social impact. our paper focuses on improving the truthfulness of lms, and the results aim to improve trustworthy artificial intelligence. apart from language generation, our paper can also be implemented in other domains for activation editing. however, it is important to acknowledge the potential misuse of our method. there exists a risk that adversarial actors could exploit our approach to transform truthful outputs into misleading or false information. this dualuse nature underscores the importance of ethical guidelines and safeguards in developing artificial intelligence. by promoting transparency and accountability in using our framework, we want to raise awareness of the risks while maximizing the benefits of improved truthfulness in language generation.