We anchored a single concept in the residual stream of a transformer. It ended up where we wanted, with nearby colors graded sensibly, and without degrading task accuracy. Steering is next.
Earlier posts in this sequence introduced Sparse Concept Anchoring (SCA), tested in autoencoders. This post applies the technique to transformers. You don't need to have read the earlier posts to understand this one. Light revisions by Claude Fable 5, and experiments run with help from all the Claude 5s.
Information flows through a transformer from input to output via the residual stream — an internal state, kind of like a mental state. It's composed of pieces, addressable by slice (layer) and position (token). At the input, each state is a token embedding. At deeper layers, the hidden states take on abstract meaning, being progressively transformed until they reach the final layer, where they are used to predict the next token.
Those hidden states are the subject of a lot of mech interp research, because we (researchers) expect to find structure that reveals something about how the model "thinks" about the world[1].
In our experiments, we used a mini color-mixing language of expressions like "red + blue = purple". The models were trained to predict the answer, and they could do so accurately, even for pairs of colors they had never seen before.
What would you expect to see, if you looked at the embeddings for the color tokens? We fit a linear probe on all the colors except the one being plotted, and decoded the held-out color with the probe. That was repeated for all colors in the vocabulary. Here's what we found:
All 216 colors of our 6-subdivision RGB cube, decoded from embeddings with linear probes. Open rings ○ show where the colors should be (ideally); filled circles ● show where they were measured to be.
This model has never actually seen a color; all we showed it was a bunch of opaque tokens in color-mixing equations, and it inferred the geometry of the RGB cube.[2]
However, we had to go looking for that structure. It would be much more convenient if we could just tell the model where to put the colors, e.g. we might say "red must be in dimension 1." We can't do that for all concepts — that would leave no room for superposition, and the capacity and accuracy of the model would quickly degrade. But perhaps we can do it for a handful of concepts. That's what Sparse Concept Anchoring (SCA) aims to do.
Suppose red was actually represented in dimension 1, at (1, 0, …). If we measured the cosine similarity from that direction to every color, we should see something like this:
That is, colors that are not red should be approximately orthogonal to (1, 0, …), while red should be fully aligned , and red-ish colors like orange should be partially aligned. And what do we actually find in this un-anchored model?
Per-color alignment of the first operand in unanchored models, mean over all slices and seeds. Each pixel is the true color of the input (op1), dithered to show the cube internals.
... Fair enough! It turns out that, even though the states in this model are normalized (vectors of length 1), they are all approximately orthogonal to each of the dimensions of the residual stream. Without constraint, the model learns arbitrary concept vectors, and random vectors in a high-dimensional space () are almost guaranteed to be near-orthogonal to all others (not merely unaligned). That's great actually, because it means we should be able to anchor a couple of concepts wherever we like, and the model should still have plenty of freedom to arrange the other concepts and the relationships between them.
Anchoring
We will use two of the geometric activity regularizer terms from our autoencoder experiments: anchor and anti-subspace.[3]
Anchor pulls the hidden states toward a direction — in this case, (1, 0, ...) — but only for color equations that have a red operand (labeling and targeting is described below). Anti-subspace pushes all activations away from that same direction and its antipode.
Let's see what that does to the states at the first operand.
Per-color alignment of the first operand, mean over all slices and seeds. The gray bands are sampled from individual slices and seeds.
All anchored conditions have exact-match accuracy within 0.02 of the un-anchored baseline. But the anchor term seems to have an excessive global effect: it pulls red to the anchor direction, but it drags the other colors along with it. Adding the anti-subspace term restores the separation.
Labeling
To isolate certain concepts, we're going to need some labeled samples. For safety-relevant concepts in LLMs, realistic labels might be along the lines of "this tweet contains a lie" or "this paragraph is about protein synthesis". To make our toy scenario as transferable as possible, we'll use a comparable labeling scheme:
Entire sequences are labeled, not tokens. So in "red + blue = purple", the equation is labeled "red", not just the red token.
Imprecise and incomplete: we don't give the model a graded curve of redness. Red is labeled as such 8% of the time; the labels are binary (true/false), and nearby colors may sometimes be labeled "red" too.
Pure red accounts for 31% of all labeled lines, and the ten reddest colors account for 85%. Overall, about 0.12% of training lines are labeled. The model needs to learn to position, separate, and grade colors in spite of the noisy, coarse signal.
Targeting
Our equation-level labels don't say which states to pull toward the anchor. In the experiments above the pull was uniform: the anchor term was averaged over all hidden states in a batch. That means that in equations labeled "red", a non-red operand and the + and = positions would all be pulled just as strongly.
The figure below shows how aligned each state is given the color of the first operand (whatever it may be).[4]
Per-(slice, position) alignment with the anchor. The height of the red lines shows measured alignment when the first operand is pure red; the grading clouds show the same thing for all colors.
Notice that all states have become aligned with the anchor: when operand 1 is red, all other positions point toward the nominated red direction. The model still has high task accuracy, but the anchor term seems to have had an excessive global effect.
Let's try to make it more targeted. We can't tell the regularizer which states to pull, but we can let it choose. Instead of using a uniform average over positions, we'll use mellowmax pooling within color equations (lines) before averaging over the batch.
It's a soft minimum function — a bit like softmax, but rather than focusing next-token probability, think of it as concentrating the regularizer loss on one or two positions. Since the anchor term is a distance, the soft minimum focuses on the states that are already closest to the anchor. Here's what pooling does to the measurements:
Alignment when using mellowmax pooling over positions. Left: per-(slice, position), keyed on the first operand. Right: mean over slices and seeds at op1.
That seems better: now when the first operand is red, most of the alignment toward the anchor is in fact at the first operand. It's unclear whether that's optimal — perhaps a different position should have been favored at deeper slices. But at least the effect isn't dominating over all positions.
Grading is also improved, with non-reds roughly as unaligned as in the un-anchored models. Again it's unclear if the grading is optimal for reddish colors. Should the relationship between redness and alignment to the anchor be linear? Quadratic? Sigmoid? We scored grading against a target, similar to the steering effect we saw in autoencoders, but the plots above look more like a sigmoid. We didn't spend too long on this question, because what actually matters is how steerable it is.
Future and related work
Soon we will try intervening on the anchored models by projecting out the anchored direction or ablating the related weights, as we did in autoencoders. We expect to see reduced (near-random) task accuracy for equations involving red, while non-reds should be unaffected.
We will also expand the mini language to have more mixing operations, and we'll anchor those as well as the color operands.
At that point, we would like to compare SCA to other methods like Gradient Routing and steering with SAEs.
Conclusion
We have shown that Sparse Concept Anchoring "works" in transformers: we picked a single concept and nudged it onto an axis of our choice. Nearby concepts became partially aligned, non-red colors stayed as unaligned as in the un-anchored models, and linear probes still decode the RGB cube from anchored models roughly as well as they did before (although that last comparison is noisy). Task accuracy was unchanged.
We are looking forward to progressing to interventions (steering) and more complex models.
The hidden states and embeddings in this model are 64-dimensional, and normalized to length 1. So the embedding space is a hypersphere, and yet the RGB cube is right there, linearly decodable.
An activity regularizer is a loss term, calculated from the hidden states (activations) rather than the model output. Ours are formulated as before: anchor is simply the cosine distance to the concept vector, and anti-subspace is the squared distance of intrusion into the subspace:
where is the state vector, is the concept vector, is the th component of the state vector, and is the subspace (dimensions) being reserved for the concept.
We anchored a single concept in the residual stream of a transformer. It ended up where we wanted, with nearby colors graded sensibly, and without degrading task accuracy. Steering is next.
Information flows through a transformer from input to output via the residual stream — an internal state, kind of like a mental state. It's composed of pieces, addressable by slice (layer) and position (token). At the input, each state is a token embedding. At deeper layers, the hidden states take on abstract meaning, being progressively transformed until they reach the final layer, where they are used to predict the next token.
Those hidden states are the subject of a lot of mech interp research, because we (researchers) expect to find structure that reveals something about how the model "thinks" about the world[1].
In our experiments, we used a mini color-mixing language of expressions like "red + blue = purple". The models were trained to predict the answer, and they could do so accurately, even for pairs of colors they had never seen before.
What would you expect to see, if you looked at the embeddings for the color tokens? We fit a linear probe on all the colors except the one being plotted, and decoded the held-out color with the probe. That was repeated for all colors in the vocabulary. Here's what we found:
All 216 colors of our 6-subdivision RGB cube, decoded from embeddings with linear probes. Open rings○ ●
This model has never actually seen a color; all we showed it was a bunch of opaque tokens in color-mixing equations, and it inferred the geometry of the RGB cube.[2]
However, we had to go looking for that structure. It would be much more convenient if we could just tell the model where to put the colors, e.g. we might say "red must be in dimension 1." We can't do that for all concepts — that would leave no room for superposition, and the capacity and accuracy of the model would quickly degrade. But perhaps we can do it for a handful of concepts. That's what Sparse Concept Anchoring (SCA) aims to do.
Suppose red was actually represented in dimension 1, at (1, 0, …). If we measured the cosine similarity from that direction to every color, we should see something like this:
That is, colors that are not red should be approximately orthogonal
Per-color alignment of the first operand in unanchored models, mean over all slices and seeds. Each pixel is the true color of the input (op1), dithered to show the cube internals.
... Fair enough! It turns out that, even though the states in this model are normalized (vectors of length 1), they are all approximately orthogonal to each of the dimensions of the residual stream. Without constraint, the model learns arbitrary concept vectors, and random vectors in a high-dimensional space (
Anchoring
We will use two of the geometric activity regularizer terms from our autoencoder experiments: anchor and anti-subspace.[3]
Anchor pulls the hidden states toward a direction — in this case, (1, 0, ...) — but only for color equations that have a red operand (labeling and targeting is described below). Anti-subspace pushes all activations away from that same direction and its antipode.
Let's see what that does to the states at the first operand.
Per-color alignment of the first operand, mean over all slices and seeds. The gray bands are sampled from individual slices and seeds.
All anchored conditions have exact-match accuracy within 0.02 of the un-anchored baseline. But the anchor term seems to have an excessive global effect: it pulls red to the anchor direction, but it drags the other colors along with it. Adding the anti-subspace term restores the separation.
Labeling
To isolate certain concepts, we're going to need some labeled samples. For safety-relevant concepts in LLMs, realistic labels might be along the lines of "this tweet contains a lie" or "this paragraph is about protein synthesis". To make our toy scenario as transferable as possible, we'll use a comparable labeling scheme:
Pure red accounts for 31% of all labeled lines, and the ten reddest colors account for 85%. Overall, about 0.12% of training lines are labeled. The model needs to learn to position, separate, and grade colors in spite of the noisy, coarse signal.
Targeting
Our equation-level labels don't say which states to pull toward the anchor. In the experiments above the pull was uniform: the anchor term was averaged over all hidden states in a batch. That means that in equations labeled "red", a non-red operand and the
+and=positions would all be pulled just as strongly.The figure below shows how aligned each state is given the color of the first operand (whatever it may be).[4]
Per-(slice, position) alignment with the anchor. The height of the red lines shows measured alignment when the first operand is pure red; the grading clouds show the same thing for all colors.
Notice that all states have become aligned with the anchor: when operand 1 is red, all other positions point toward the nominated red direction. The model still has high task accuracy, but the anchor term seems to have had an excessive global effect.
Let's try to make it more targeted. We can't tell the regularizer which states to pull, but we can let it choose. Instead of using a uniform average over positions, we'll use mellowmax pooling within color equations (lines) before averaging over the batch.
It's a soft minimum function — a bit like softmax, but rather than focusing next-token probability, think of it as concentrating the regularizer loss on one or two positions. Since the anchor term is a distance, the soft minimum focuses on the states that are already closest to the anchor. Here's what pooling does to the measurements:
Alignment when using mellowmax pooling over positions. Left: per-(slice, position), keyed on the first operand. Right: mean over slices and seeds at op1.
That seems better: now when the first operand is red, most of the alignment toward the anchor is in fact at the first operand. It's unclear whether that's optimal — perhaps a different position should have been favored at deeper slices. But at least the effect isn't dominating over all positions.
Grading is also improved, with non-reds roughly as unaligned as in the un-anchored models. Again it's unclear if the grading is optimal for reddish colors. Should the relationship between redness and alignment to the anchor be linear? Quadratic? Sigmoid? We scored grading against a
Future and related work
Soon we will try intervening on the anchored models by projecting out the anchored direction or ablating the related weights, as we did in autoencoders. We expect to see reduced (near-random) task accuracy for equations involving red, while non-reds should be unaffected.
We will also expand the mini language to have more mixing operations, and we'll anchor those as well as the color operands.
At that point, we would like to compare SCA to other methods like Gradient Routing and steering with SAEs.
Conclusion
We have shown that Sparse Concept Anchoring "works" in transformers: we picked a single concept and nudged it onto an axis of our choice. Nearby concepts became partially aligned, non-red colors stayed as unaligned as in the un-anchored models, and linear probes still decode the RGB cube from anchored models roughly as well as they did before (although that last comparison is noisy). Task accuracy was unchanged.
We are looking forward to progressing to interventions (steering) and more complex models.
Our experiments are published at z0u.github.io/sca2 (source).
Or "models the data distribution", if you prefer.
The hidden states and embeddings in this model are 64-dimensional, and normalized to length 1. So the embedding space is a hypersphere, and yet the RGB cube is right there, linearly decodable.
An activity regularizer is a loss term, calculated from the hidden states (activations) rather than the model output. Ours are formulated as before: anchor is simply the cosine distance to the concept vector, and anti-subspace is the squared distance of intrusion into the subspace:
where
Figures showing the same thing but keyed by the second operand are available in this notebook.