Power Steering: Behavior Steering via Layer-to-Layer Jacobian Singular Vectors

Summary

cross-posted from my blog TLDR The map of how the activations of one ‘source’ layer in an LLM impact the activations in some later ‘target’ layer can provide vectors for steering LLM behavior. Computing this map, or the Jacobian, is costly but the top high rank components can be determined in just ~15 forward passes in a process called power iteration. This method is cheap enough that every source/target pair in the model can be examined producing a sensitivity map. The use of power iteration to find steering vectors gave the natural name ‘Power Steering’. The resulting power steering vectors produce comparable performance to similar but more costly non-linear optimization techniques that find vectors for maximizing source layer to target layer impacts. The cheap computation of power steering allows one to map all source/target pairs in a model to find interesting steering behavior. Steering behavior is mostly easily found using prompts that have decision forks for the model but it can also induce latent behaviors. Introduction and Background on Steering Vectors Steering LLM Behavior Controlling LLM behavior without low level mechanistic interpretability is an attractive facet of the technical AI safety field. Certain methods such as Steering Vectors work by adding a direction in representation space to shift model behavior toward a target concept. This is often performed at a specific layer and token position during inference to elicit specific behaviors. Steering techniques are underpinned by the Linear Representation Hypothesis, that model internals are composed of many salient linear directions in representation space. Steering vectors are primarily formed via Contrastive Activation Addition (CAA). The LLM is given pairs of prompts that elicit different opposing concepts (e.g. love vs hate). Run both through the model, capture the model representation at some points (often the residual stream), and take the difference. That’s your steering vector. At inference, add or subtract this vector (scaled by some coefficient) to shift model behavior. This can work with as little as a single prompt pair, though more recent approaches average over many pairs. For a deeper introduction, see Steering Llama 2 via Contrastive Activation Addition. CAA has known reliability issues: prior work demonstrated that steerability is highly variable across inputs, and spurious biases can inflate apparent effectiveness. Unsupervised Steering One interesting approach for finding steering vectors in an ‘unsupervised’ manner is MELBO, which finds vectors that elicit latent behaviors in language models by “learning unsupervised perturbations of an early layer of an LLM.” The goal of MELBO is to maximize the change in a specified target layer activations , through perturbation at a source layer. This perturbation, a bias vector of specific size , is added to the source layer output. The bias vector is found via non-linear optimization using the following loss function which has been simplified for this discussion: Additional vectors can be found via orthogonalization during optimization, and the process can exploit nonlinearities to find off-manifold directions. MELBO could find interesting steering vectors to undo safety refusals and reveal latent behaviors such as chain-of-thought. These results would generalize to other contexts and could be elicited using single prompts. A natural question following MELBO may be, how does a linear approximation of the same effect perform? What does the gradient between layers already tell us about these interactions for a given prompt? Comparison of CAA, MELBO, and Power Steering methods of producing steering vectors for a given model and prompt(s). Power Steering: The Local Linear Approximation An alternative to MELBO is the local linear approximation through the layer-to-layer Jacobian, which maps how perturbations at some source layer activations will impact the activations at some target layer, : . directly describes how the target representation changes through small perturbations at the source layer. The right singular vectors of are the directions at the source layer that produce the largest response at the target layer, i.e. the directions the network is already most sensitive to for a given prompt. Stated another way, these vectors are the unit norm directions at the source which produce the largest linear response at the target layer. The vectors are the solution to the linearized version of the MELBO objective. The right singular vectors can be found without explicitly forming the Jacobian, which for an LLM would be a matrix of dimension . Power iteration only needs the matrix-vector product , and repeated application converges to the top right singular vector1. Computing requires both vector-Jacobian products (VJPs, ) and Jacobian-vector products (JVPs, ). Standard backprop via Pytorch autodiff gives VJPs directly. JVPs can be obtained via reverse-over-reverse: differentiating the VJ