Activation Scaling for Steering and Interpreting Language Models (2024)

Abstract: Given the prompt “Rome is in”, can we steer a language model to flip its prediction of an incorrect token “France” to a correct token “Italy” by only multiplying a few relevant activation vectors with scalars? We argue that successfully intervening on a model is a prerequisite for interpreting its internal workings. Concretely, we establish a three-term objective: a successful intervention should flip the correct with the wrong token and vice versa (effectiveness), and leave other tokens unaffected (faithfulness), all while being sparse (minimality). Using gradient-based optimization, this objective lets us learn (and later evaluate) a specific kind of efficient and interpretable intervention: activation scaling only modifies the signed magnitude of activation vectors to strengthen, weaken, or reverse the steering directions already encoded in the model. On synthetic tasks, this intervention performs comparably with steering vectors in terms of effectiveness and faithfulness, but is much more minimal allowing us to pinpoint interpretable model components. We evaluate activation scaling from different angles, compare performance on different datasets, and make activation scalars a learnable function of the activation vectors themselves to generalize to varying-length prompts.

## introduction understanding which components of a language model play which roles in which tasks is a core aim of mechanistic interpretability. given the prompt rome is in, for instance, one might ask which components of the model most influence it to favor italy over some incorrect answer token, such as france. in addressing this question, a natural axiom is that a given component can only be understood as influential for a given task if intervening on it meaningfully alters the model's task-specific behavior. 1 code to experiment with our method is available at https://github.com/niklasstoehr/activationscaling . l,i for a model to favor one answer token over another token. this could be the mlp output at layer l for token position i. we learn multiplicative scalars, α (s) l,i , using gradient-based optimization. these correspond to interpretable interventions that generalize to test set prompts while relying on fewer parameters than additive steering vectors. building on this basic axiom, a growing literature seeks to both generate and test hypotheses about where certain behaviors are localized in a model by employing targeted interventions with methods such as activation patching @xcite @xcite , among others. studies in this literature often produce a set of attribution scores associated with various locations in the model which represent how much the model's output changed after editing the activation vectors at each location. although an effective intervention may be necessary to believe a given localization hypothesis, it is not sufficient, as interventions to other model locations may be similarly effective. indeed, recent work has questioned the relationship between interpretability and intervention on this basis, and has advocated for more rigorous and deliberate methodology for connecting the two @xcite @xcite . a parallel literature on model steerability also seeks to develop effective interventions, not for the primary purpose of interpretability, but to steer models toward desirable behaviors, like factuality @xcite , or away from undesirable behaviors, like toxicity @xcite . methods in this literature are typically designed to be maximally effective for steering and thus tend to involve model-wide (i.e., not necessarily localized) interventions to the weights @xcite @xcite or activation vectors @xcite . this paper seeks to synthesize the goals of model-wide interventions while still offering mechanistic insights into the model. specifically, we seek interventions that are effective, and not localized a priori, but are nevertheless minimal and faithful. we define an intervention to be effective if it flips the prediction of the correct answer token (e.g., italy) to an incorrect token (e.g., france). we define an intervention to be faithful if it does not substantially alter the probabilities of tokens unrelated to the given task. and, we define an intervention to be minimal if it alters activation vectors sparsely, in only a small number of locations. our notions of minimality and faithfulness follow from recent work in the transformer circuits literature @xcite . we operationalize these three desiderata via a threeterm objective allowing us to learn model-wide interventions using gradient-based optimization. more specifically, we propose a kind of modelwide yet parsimonious intervention that associates a single scalar parameter with each of many locations in a model. the scalar parameters are learned so that the model can be effectively steerede.g., away from italy and toward france-by simply scaling the activation vector at each location. we call this approach activation scaling (activscalar), inspired by the idea that some model components are highly specialized for certain task-relevant computations @xcite . existing work @xcite @xcite has identified individual components, e.g., name mover heads @xcite , in a first step and successfully steered a model by scaling only the contributions of those individual components in a second step. to evaluate our approach, we construct a baseline method that is the same in all aspects except it learns additive vectors at all locations, rather than single multiplicative scalars. we refer to this approach as steering vectors (steervec). in a suite of experiments, we find overall that activation scaling learns interventions which are as effective and faithful as steering vectors, while requiring far fewer learnable parameters. our results suggest that merely scaling the signed magnitude of activation vectors, without further affecting their direction, is sufficient for effective steering on simple tasks. moreover, we find activation scalars are highly interpretable. they are easy to understand as simply strengthening or weakening the steering directions already encoded in the model @xcite , and when visualized, they provide sparse and localized descriptions of important model components. finally, to extend our approach beyond fixed prompt templates, we develop a dynamic version of activation scaling. dynscalar makes the activation scalars learned functions of the activation vectors themselves, thus allowing learned interventions to transfer to test set prompts of varying length. ## transformer language models let σ be an alphabet of tokens, a finite, non-empty set, and let σ * be the set of all strings with tokens drawn from σ. a language model p is a probability distribution over σ * . as is current practice, most language models are defined autoregressively. each local conditional distribution p(• | @xmath0, where eos ̸ ∈ σ is the end-of-string token. in the context of an autoregressive language model, we call @xmath1. let x @xmath2. a common way to define the local conditional is via the softmax function σ which maps from r |σ| to the probability simplex ∆ |σ|-1 : dataset prompt answer toks country-capital conflict (ccc) the capital of germany is paris. q: what is the capital of germany? a: berlin / paris indirect object identification (ioi) when anne met with tom, tom gave the book to anne / tom ## experiments we fit the two intervention methods, activs-calar and steervec, on the three-term objective in eq. ( 9 ) and conduct evaluations for effectiveness, faithfulness and minimality. ## extension to variable-length prompts activscalar and steervec, as well as most existing interpretability methods such as activation patching, learn intervention parameters that are tied to specific token positions i in the prompt x n . to reuse learned parameters for intervention with a new prompt, the test prompt must match the length i n of the training prompt, and ideally match it syntactically. existing work circumvents the requirement of matched prompts by intervening only on the last token position @xcite @xcite , or on token positions that can be easily aligned across prompts, like the main verb, or the last token of the subject @xcite @xcite . ## conclusion we show that scaling the signed magnitude of a few relevant activation vectors is often sufficient to flip a model's prediction between a correct and a wrong answer token. besides being effective at steering, activation scaling requires many fewer parameters than additive steering vectors which intervene both on the magnitude and direction of activation vectors. our gradient-based multi-objective learning scheme can be understood as reversing the interpretability pipeline, putting steering performance based on clearly defined objectives first and interpretability as a natural by-product second. limitations steervec, activscalar and dynscalar are controllable via different hyperparameters: the margin m in the effectiveness objective, λ f weighing the faithfulness term and λ m weighing the strength of the ℓ 1 -regularization. there are additional training-related hyperparameters such as the learning rate, the number of epochs, the batch size, the number of data instances, and the standard deviation of the gaussian noise initialization of the intervention parameters, that have a strong influence on the results. for instance, increasing the noise or the margin, training for more epochs, or weakening the ℓ 1 -regularization, results in activation scalars that deviate more from zero. access to many hyperparameters can make a method more difficult to deploy. on the other hand, hyperparameters with well-founded semantic can offer desirable, fine-grained controls. for instance, a larger λ m hyperparameter leads to sparser activation scalars that are easier to interpret, a level of control not offered by many existing methods. yet, methods like activpatch or attrpatch also require finicky hyperparameter choices such as how to corrupt the prompt, e.g., deciding on the standard deviation of the gaussian noise @xcite to obtain a second corrupted prompt for patching. a limitation of this work is the size of models studied and the small size and synthetic character of the tasks. the two largest models considered in this work are gpt2-xl (1.5 billion parameters) and pythia-1.4b (1.4 billion parameters). beyond the required compute, we do not anticipate problems applying activation scaling to larger models. testing activscalar on more real-world datasets with longer prompts is another future avenue. to boost the performance of dynscalar, one could expand the computational expressivity of the activation vector-to-scalar function g (s) l . finally, more work is needed in extending our evaluation based on effectiveness to existing methods such as activation patching. for instance, a promising direction is to fix the attribution scores obtained from activation patching and then post-hoc learning a suitable β parameter that facilities the answer token flipping behavior.