arena-in-12-weeks
glossary about

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:

this week's pair session

core

  • 1-3: nnsight, task-encoding states, function vectors

stretch

  • GPT2-XL steering