week 11 / 12
Steering & function vectors
Once you've found a direction, you can push the model along it.
works through arena 1.3.2 · 1.3.2 Function Vectors & Model Steering
Reading to writing
Weeks 6 through 10 mostly asked what is inside the model. This week asks what happens when you edit it. Steering is a causal test: if adding a direction changes behavior in the predicted way, the direction is doing real work.
A minimal mental model is activation addition:
import torch
resid = torch.randn(1, 12, 768)
steering_vec = torch.randn(768)
resid[:, -1, :] = resid[:, -1, :] + 2.0 * steering_vec
print(resid.shape)
# torch.Size([1, 12, 768]) — the shape never changes; only the last position moved
The hard part is not the addition. The hard part is finding a direction worth adding and choosing where to add it.
Function vectors
A function vector is a task direction extracted from examples. For antonyms, you show prompts full of pairs like hot -> cold, average each of a chosen set of attention heads' outputs over many such prompts, sum those averages into one vector, and inject it into a new prompt. The model can then perform the task without the examples being present. The heads are chosen by measuring which ones have the largest causal effect on in-context task performance — the same "measure, then intervene" loop as week 9.
This is in-context learning made portable: a task inferred from context becomes a vector you can move between runs.
nnsight as an intervention tool
nnsight is another library for tracing and editing model internals, especially when the model is too large to run locally. Code inside a trace describes what to save or edit during the forward pass; it feels a bit like building a query plan.
# no-run
with model.trace(prompt):
hidden = model.transformer.h[9].output[0].save()
model.transformer.h[9].output[0][:] += steering_vec
out = model.lm_head.output.argmax(-1).save()
TransformerLens and nnsight share the same conceptual verbs: read activations, write activations, measure outputs.
Connection to SAEs
SAE-latent steering is a cleaner modern variant. Instead of adding a dense direction found by contrastive prompts, you turn up a sparse latent such as a "Golden Gate Bridge" feature and observe output changes. The evidence standard is the same: specific intervention, specific behavioral prediction.
Pair-session guide
Core work is ARENA 1.3.2 sections 1 through 3: introduction to nnsight, task-encoding hidden states, and function vectors. These sections run on GPT-J-6B (about 12 GB in bfloat16), which is the point of nnsight: you can either run it on a big local GPU or execute remotely on NDIF's servers. If you plan to use remote execution, request an NDIF API key before the session and check the model is live on the status page. Stretch work is section 4, steering vectors in GPT2-XL — a smaller model, but it is not hosted remotely, so it must fit on your own GPU with REMOTE = False. Pair rule: always write down the target layer, token position, and scale before running a steering experiment.
What you should see
You should see a model complete antonym-style tasks because you injected a function vector, not because the prompt contained examples. In the steering stretch, you should see generation shift in sentiment or topic when you add the steering vector at the right layer and scale — the classic demo is GPT2-XL veering into weddings mid-sentence.
Where to go next
Week 12 is the capstone: reverse-engineering a learned algorithm end to end. The work this notebook replicates, plus the follow-up ARENA points to, is all readable now:
- Function Vectors in Large Language Models by Todd et al. (project page at functions.baulab.info) is the paper behind sections 2-3.
- Steering GPT-2-XL by adding an activation vector by Turner et al. is the post behind section 4, including the "Love − Hate" and wedding demos.
- Steering Llama 2 via Contrastive Activation Addition by Panickssery et al. is the follow-up ARENA's bonus section suggests: the same idea, made systematic on a chat model.
this week's pair session
core
- 1-3: nnsight, task-encoding states, function vectors
stretch
- GPT2-XL steering