Steering Vectors

The residual stream is a sum of contributions, and the unembedding reads it linearly. If a behaviour corresponds to a direction in that space, adding the direction should produce the behaviour — no fine-tuning, no prompt.

Finding the direction

The crude version subtracts one prompt’s activations from another’s at a chosen layer. The reliable version is difference-in-means over a dataset of contrastive pairs — same question, opposite answer — which averages away everything the two sets share.

Read the activations at the answer token, not at the end of the prompt: the prompt is identical across the pair by construction, so anything measured there is noise.

Negative $\alpha$ steers the other way, which is a useful check: a direction that only works in one sign has probably found something other than the axis you meant.

Choosing the layer

Sweep $\ell$ and $\alpha$ together and plot behaviour change against perplexity on unrelated text. The usable region is where behaviour moves and perplexity does not — beyond it the vector is overwhelming the computation rather than biasing it.

Steering is a causal claim about a direction and is properly evaluated as one. Sparse autoencoder latents give candidate directions with an interpretation attached, which is a better starting point than difference-in-means when a dictionary is available.

import torch
from contextlib import contextmanager
from torch import Tensor


@torch.no_grad()
def contrastive_vector(model, pairs, layer: int) -> Tensor:
    """pairs: [(positive_prompt, negative_prompt), ...] identical but for
    the answer token. Reads the residual stream at the final position."""
    deltas = []
    for positive, negative in pairs:
        _, pos = model.run_with_cache(positive)
        _, neg = model.run_with_cache(negative)
        deltas.append(
            pos["resid_post", layer][0, -1] - neg["resid_post", layer][0, -1]
        )
    return torch.stack(deltas).mean(0)


@contextmanager
def steer(block, vector: Tensor, alpha: float):
    def hook(_mod, _inp, out):
        hidden = out[0] if isinstance(out, tuple) else out
        hidden = hidden + alpha * vector.to(hidden.dtype)
        return (hidden, *out[1:]) if isinstance(out, tuple) else hidden

    handle = block.register_forward_hook(hook)
    try:
        yield
    finally:
        handle.remove()

Normalising $v$ to unit length makes $\alpha$ comparable across layers, which matters because residual-stream norm grows with depth — an $\alpha$ tuned at layer 10 is a much weaker intervention at layer 30. Decoder blocks commonly return a tuple whose first element is the hidden state; returning a bare tensor from the hook silently drops the rest of it.