SAE features¶
A sparse autoencoder turns a residual into a few named features. Reading them needs the whole SAE; acting on one needs only its two rows — an encoder row to compute the feature's activation from the live residual, a decoder row to add or remove its contribution — so the intervention runs in the worker on every step with a few kilobytes shipped, and the decision (is the feature on, how strongly) is made from the model's actual state rather than fixed in advance.
The SAE here is Qwen-Scope for
layer 24: TopK-50, 65,536 features, trained on the residual stream leaving the block of
Qwen3-8B-Base (its card says applying it to the post-trained checkpoint is reasonable; the
reconstruction error below is the price).
Find the features: full SAE, once, on the client¶
import torch
from huggingface_hub import hf_hub_download
from nnsight.modeling.vllm import VLLM
LAYER, K = 24, 50
sae = torch.load(hf_hub_download("Qwen/SAE-Res-Qwen3-8B-Base-W64K-L0_50", f"layer{LAYER}.sae.pt"),
map_location="cuda")
W_enc, W_dec, b_enc, b_dec = sae["W_enc"], sae["W_dec"], sae["b_enc"], sae["b_dec"] # (65536, 4096), (4096, 65536)
model = VLLM("Qwen/Qwen3-8B", dispatch=True)
PROMPT = "The capital of France is"
with model.trace(PROMPT, temperature=0.0):
resid = sum(model.model.layers[LAYER].output).save() # the stream leaving block 24, [pos, 4096]
pre = resid.float() @ W_enc.T + b_enc # the encoder takes the raw residual
top = pre.topk(K, dim=-1)
acts = torch.zeros_like(pre).scatter_(-1, top.indices, top.values)
recon = acts @ W_dec.T + b_dec
print([round(e, 3) for e in ((recon - resid.float()).norm(dim=-1) / resid.float().norm(dim=-1)).tolist()])
# [0.291, 0.363, 0.421, 0.467, 0.488]
last = acts[-1]
feats = last.nonzero().flatten()
feats = feats[last[feats].argsort(descending=True)][:8].tolist()
print([(f, round(last[f].item(), 1)) for f in feats])
# [(16957, 42.8), (51672, 35.2), (41823, 31.4), (37656, 23.4), (25446, 21.3), (32674, 20.5), (61073, 20.0), (53363, 18.4)]
sum(layers[24].output) is the hook point the SAE was trained on (the block's output, as a
HuggingFace forward hook sees it). One trace, one tensor home; the 2 GB SAE never leaves the
client.
Ablate each feature in the worker¶
Two rows per feature ship with the block. The feature's activation is recomputed from the live residual at the last position, and its decoder direction is subtracted in proportion — nothing here is a fixed vector.
def rows(f):
return W_enc[f].to(torch.bfloat16), b_enc[f].item(), W_dec[:, f].to(torch.bfloat16)
paris = model.tokenizer.encode(" Paris")[0]
with model.trace(temperature=0.0, max_tokens=1) as tracer:
with tracer.invoke(PROMPT):
base = model.logits.save()
for f in feats:
w_in, b_in, w_out = rows(f)
with tracer.invoke(PROMPT):
hidden, residual = model.model.layers[LAYER].output
x = hidden[-1] + residual[-1]
a = torch.relu(x @ w_in.to(x.device) + b_in) # this feature, on this forward
hidden[-1] -= a * w_out.to(x.device) # remove what it contributes
logits = model.logits.save()
p = lambda l: round(l.float().softmax(-1)[0, paris].item(), 3)
print(p(base), [(f, p(l)) for f, l in zip(feats, logits)])
# 0.537 [(16957, 0.351), (51672, 0.431), (41823, 0.499), (37656, 0.402), (25446, 0.491), (32674, 0.649), (61073, 0.56), (53363, 0.547)]
Feature 16957 — the most active at is — carries a third of P(Paris) on its own; 37656
and 51672 carry some; 32674 suppresses it.
Clamp a feature through generation¶
Scale the feature's activation to a target multiple, at every step, from its live value.
w_in, b_in, w_out = rows(16957)
for target in (0.0, 3.0):
with model.trace(PROMPT, temperature=0.0, max_tokens=8) as tracer:
seen = list().save()
for _ in tracer.iter[:8]:
hidden, residual = model.model.layers[LAYER].output
x = hidden[-1] + residual[-1]
a = torch.relu(x @ w_in.to(x.device) + b_in)
seen.append(round(a.item(), 1))
hidden[-1] += (target - 1) * a * w_out.to(x.device) # a -> target * a
out = tracer.result.save()
print(target, seen, repr(out.outputs[0].text))
# 0.0 [42.8, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 40.0] ' Paris. The capital of France is Paris'
# 3.0 [42.8, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 34.5] ' Paris. The capital of Italy is Rome'
The feature fires on the prompt's is and again on the generated is seven steps later, and
nowhere between — seen is the feature's own timeline, read for free while acting on it.
Removing it changes what follows the second is; tripling it leaves the model's own
continuation alone.
Rules¶
- Encoder on the raw residual, as the SAE's card does — no
b_decsubtraction before encoding for this family; check yours. - Base-model SAE, post-trained model. The reconstruction error (0.29–0.49 by position) reflects that mismatch and the 50-feature budget; use the features as handles, not as ground truth.
- Whole-SAE reads stay on the client. Shipping a 2 GB
W_encwith a block would serialise it on every request; ship rows. - Under tensor parallelism the residual is whole on every rank and the write is re-sharded, so the same block runs unchanged (Tensor parallelism).