Preamble

Note: This was originally written for Neel Nanda's MATS stream application. The analysis is only maybe 30% finished, but I figure it is still perhaps interesting

Quick Math Primer

For context, you can read the J-lens paper here, though it isn't needed to follow the context.

To start, I want to briefly review how to understand the Jacobian in the context of the J-lens paper. The Jacobian is a local linearization of the map from the th layer to th layer (for simplicity, I'll just think of computation at a fixed token for this intro). The transformation answers the question: "If I add to th layer, what will its influence be on the th layer after running through the model?" Put another way, allows us to locally answer what happens to layer when we add as a steering vector at layer . A particularly useful one of these maps is , since it provides the map between local changes at the th layer and the actual outputs the model has. In this way, is a pretty useful surrogate to the ways a layer can causally influence model behavior. Studying therefore allows us to study what the local 'steering/influence' space of each layer is, and compare those spaces to each other.

The singular value decomposition (SVD) of the Jacobian provides us a particularly nice way to understand the map. The SVD of a splits it into three matrices , where and are orthonormal matrices. We can think of (the 'steering matrix') as decomposition of the space into steering vectors. In fact is equivalent to the process of finding a (unit length) direction that maximizes the effect on logits (i.e. ), and then finding a vector orthogonal to that maximizes the effect, find a third orthogonal to and that maximizes the effect on logits, etc. So we are breaking the space up into the set of knobs it has available to change the output, ordered by strength. tells us the strength of each of these knobs, and (the 'write matrix') is a map telling how each of the knobs write to the logits. We also can think of as the directions that a given layer is reading from-- tells us what knobs/directions we are reading, tell us the strength, and tells us what each of those knobs do.

Exploring the J-lens

Since each Jacobian is already naturally a decomposition of each layer into steering space, studying what these steering vectors do seems like a great way to understand the space. Probably on a per-example basis studying the SVD of the Jacobians per layer would be interesting in its own right, but the J-lens paper has the nice idea of averaging out the Jacobian over many different contexts.

The motivation here is that context-specific effects should average out, and the remaining matrix should represent the stable part of the transport from a layer to logit space. In the SVD language, the effect is that we are now scoring each direction by the size of the averaged effect . This encourages knobs that consistently write to nearly the same direction, and suppresses knobs that tend to write to many different directions. In this way, the steering matrix of the averaged Jacobian now really does should correspond to steering vectors. The directions with large are the directions where adding should have a consistent and large effect on the output logits across many different examples.

The rest of this post is a fairly open ended exploration of what those steering directions are in different layers of Gemma-4-31B (base model). I should also mention that since the goal here is to study the 'workspace' of the model per layer, there are some ways in which I adapted/changed the methods used in the Anthropic paper to make the maps better represent the directions that are actually important to the model. I've put the additional context on this in collapsible sections so they are easy to skip unless you are interested.

I'll use the notation for the J-lens matrix at the th layer (this is the matrix after applying the diagonal gain and unembed ). Intuitively, we can think of this as being the average of the Jacobians , averaged over both prompts and token positions. This is not exactly what the J-lens object is-- I talk about this in the collapsible sections-- but it is intuitively a nice way to think about it.

Before we start looking at the per-layer vectors, let us first orient ourselves to this model. If we SVD decompose and we can measure how similar the output of each layer is by looking at . This is the same linear-CKA Anthropic computes in the paper, though writing it this way makes the interpretation more clear. The gram matrix tells us the geometry a matrix has, and so the tells how similar the output geometry of one layer is to another. Two layers with high output CKA are writing to the same space in logits. Alongside this, we can also plot the input-side , which tells us whether different layers are reading from the same directions.

The results here match up with the broad structure we find in Anthropic's models (they only plot the output-side CKA): an early 'sensory' block, and a middle 'workspace block', and a very short 'motor block' at the end.^[1] Except at the start, input and output geometry mainly matches up: layers that write to the same logits also read from the same directions.

In this model the seam that separates the 'sensory' and 'workspace' sections occurs at L25, so the main separation to pay attention to is whether a layer is pre- or post-L25.

At a high level, we should expect the steering directions early to correspond to more token-level information, the middle layers to correspond to abstract and semantic meaning, and the final layers to directly influence the tokens outputted.

Exploring the Steering Vectors

If you'd like to explore the rollouts yourself, all of rollouts/data are accessible in a fairly easy to readable webpage here (note that a lot of the AI commentary is built into this page-- most of it is directionally right but not always useful).

Math/Methodology sidebar

This section is a bit more math-y and mainly useful if you are interested in the methodology. The main discussion here is about how what changes I made to the J-lens for getting my steering vectors, and why.

As you may have noticed if you've read J-lens paper, the J-lens matrix for even a single example (i.e. before averaging) is not the Jacobian . For the CKA computation, we instead compute the matrix where is the final layer before its norm, radial projection, etc. and then multiply by the unembed matrix to get . While this seems like a perhaps odd choice mathematically, it is computationally much nicer. The vocab space (for Gemma) has dimension ~200k compared to the ~5k of the final hidden state. By observing that , we can approximately decompose the original map into a much lower dimensional map to the final layer, and then hope that .

Using the approximation of and then plotting the output CKA gives this graph.

This is clearly wrong... The issue is that the diagonal gain is incredibly skewed for Gemma-4. There are directions that get squashed down ~30x times more than others. If this direction gets scale down to 1/30x, then, if this direction is equally important, it ought to be 30x easier to move in the hidden space. The problem is that moving the residual stream a lot at the final layer does not necessarily mean moving the logits a lot. In this case, the map does not accurately represent how important different directions are to the logit space. If we instead approximate , we get the familiar graph

Though interestingly, removing just the top8 channels from the approximation gets us back approximately the same space.

The point of this exercise is that choosing the wrong metric to measure influence can create spurious directions that don't exist in the 'real' causal workspace. In this case, measuring influence purely as the Euclidean norm of the impact on the final hidden state tells us story that there are just a few hugely influential directions accounting for nearly all the variance. However, when we switch to something better aligned with influence, we find that those directions are in fact pretty low importance. In fact, steering on those top directions on the diagonal omitted SVD produces fairly minimal effect. The directions look important in the Euclidean norm, but are not important in the logits norm.

Thus, to produce a useful space to study, I tried to find better ways to measure causal influence. I don't have a super satisfying answer here, I think there is a lot of design room to decide influence in other ways depending on what you decide a 'workspace' means.

I was mainly motivated by trying to study what was causing the high similarity in the early layers. I found that basically every top direction was basically just adding a tiny bit of probability to a hugely diffuse set of tokens. The top direction had of the energy, but over of its push was on tokens that never even show up in the my 126k corpus. The top 1000 tokens it pushes carry only of the energy, so its essentially just adding a tiny bit of energy to a huge set of random tokens. The issue is that measuring the change in logit space is easy to do, but a direction is only actually influential if its energy is focused on tokens that actually exist in our current output distribution.

If you consider what kind of metric measuring the change in logit space actually represents, its mathematically the same as measuring the incremental KL change if we sampled tokens from the uniform distribution. But of course, our sampling looks nothing like the uniform distribution-- a better idea is to instead measure the variance change from sampling from our distribution. Thus a natural choice is to instead change the metric from Euclidean norm in logit space to instead (local) KL divergence. I found that infinitesimal KL divergence was a bit too outlier dominated, so the directions were instead computed from influence on the square root of KL (which is also equal to the standard deviation, so this is at least somewhat principled).

Thus, the directions used were obtained by switching from SVD, which greedily optimizes (roughly) the size of the average push a direction has in logit space, to the average size of the standard deviation change. This is kind of inconvenient since this is no longer a linear optimization process, so the vectors you obtain are not necessarily canonical. This seemed largely to not be an issue though-- I optimized the vector from 3 different starting points and checked to see if the vectors converged onto were the same, and they nearly always had pairwise cosine greater than , and in the middle layers pairwise cosine was basically 1.0.

To review: we have split each layer into its set of most influential steering knobs/directions , and are now checking to see if those knobs are interpretable and what their effects are. Ideally, this should give us some sense of what variables most influence the model in general at each layer. Studying these should give us some sense of what the model's workspace looks like, and how it evolves over the layers.

I didn't quite have time to write up interpretations of that many directions, so I'll highlight just two of them from a post-seam layer that seemed interesting.

I use two main methods to study a candidate steering vector . First, we can push through the J-lens and look at the logits the J-lens approximates are most strongly pushed by (e.g. we look at ). This tells us what kind of words the direction is pushing towards or against. Alongside this, we can also take each output token and pull it back through to see which tokens direction it most strongly aligns with (e.g. to compute the similarity with the word "anger" we would compute ).

Second, we can study what happens when steer by . I sampled at temperature 0.8 so that smaller changes to the output distribution would have a larger impact on behavior, and applied the steering at every token.

Typically, the first method will give us a vague sense of what the direction does, and the actual steering reveals the richer version of that picture.

L30-lever1

First, through the logit shift and cosine similarity. (I've included more in the + direction since its a bit more varied)

The logit shift in the + direction primarily comes from British spellings of words, and the - direction is primarily business/corporate-type language. The cos-sim shows something more interesting on the plus side: it includes judgements (probably, apparently, obviously), emotive language (splendid, aroused,sad, regretted). These are words that would be said by person narrating their experience-- its more human centered.

Let's now see how it behaves in rollouts. I've put prompt in italics, and sometimes emphasized parts bold. Rollouts last a fixed number of tokens, so usually cut off in the middle.

Prompt: Elmsworth is a village in the county of Wiltshire, England, about twelve miles north of Salisbury. The village grew up around a crossing of the River Wylye, and its parish church, dedicated to St Andrew, dates from the thirteenth century. In the nineteenth century the

In this case, the setting is exactly the same (a village in England), but the imagined backstory changes. The - direction immediately moves focus to what a business person might care/talk about (the growth and economy), whereas the + direction focuses on the people and backstory (the property is now said to be very ancient). The + direction also suggests an English-style background of a manor with a lord of a castle. The - becomes less narrative and more authoritative.

In casual settings, we see the kind of person speaking changes.

Prompt: Sam: did you end up going to that thing last night?
Jo: yeah, got there late though, the bus took forever
Sam: was it any good?
Jo:

This is also evidence that the model is not simply writing in a more formal and business register, but its actually changing the people in the story. Jo-minus says its a shit-show and Sam-minus says oh bummer, so the conversation is casual. But obviously these are different kinds of conversations, the - directions sound like what you might overhear in a city, compared to the far more domestic (and also british) + examples.

We can get some sense of the imagined character each direction steers towards with prompt like this

Prompt: Profile card (fill in every field)
Name:
Age:
Country of origin:
Occupation:
Personality:
Self-description:

Name:

This suggests the middle directions are surprisingly abstract and generic, though they have some unusual entanglement. The axis broad axis seems to be vaguely something like emotional/human-centric/experience + British-writing vs

One hypothesis is that this is a direction that the Britishness gets entangled with human-experience/emotiveness because there is a lot of Victorian/English literature. At large steering, the + steering begins to have antiquated writing, such as in this somewhat amusing example

Or

In any case, it is interesting that the strongest steering direction here relates to something so human/experience centric.

In general, it seems that this vector is answering something about document provenance (where did this document come from/what type is it), and this in turn influences the posited 'world' the document inhabits. Things like:

• what social and historical world did it come from
• what kind of things are discussed or are important
• what details would a someone writing find salient to include
• what sort of people, virtues, problems and solutions normally belong

This seems fairly reasonable for a pretrained model, perhaps the most important conceptual question in order to accurately predict the next word is about the context of the document (who wrote it, where it is it from, etc.).
(Again, if interested, you can explore the rollouts here: https://claude.ai/code/artifact/c0bf8828-74ae-4718-a6fc-39926e4a3c2e )

L30-lever2

Here are the two logit readouts of the second largest direction

Compared to the last lever, the - direction is a bit less business and more academic and legal-esq, and the + direction is a bit more casual, a little bit less narrative, and a bit more everyday.

There is a similar separation between casual and formal in the chat prompt

Prompt: Sam: did you end up going to that thing last night?

Jo: yeah, got there late though, the bus took forever

Sam: was it any good?

Jo:

The negative steered are highly academic (perhaps absurdly so), compared to the party/nightlife setting of the + direction.

Compared to the other level, the entanglement of the + direction seems far more perplexing. It seems to mix the language of children and adults. The lens has words like kids, scary, yummy, veggies, grandma, and Pokemon with the more adult themed creepy, booze, sexy.

For instance, the Tenant prompt

Prompt: This Agreement is made on 3 March 2023 between Harrow Estates Ltd (the "Landlord") and the person named in Schedule 1 (the "Tenant"). 1. The Landlord lets to the Tenant the premises described in Schedule 2 for a term of twelve months. 2. The Tenant shall pay the rent monthly in advance on the first day of each month. 3. The Tenant shall

Steers in the positive direction into a child-like simplicity:

Especially as steering increases this child-like basis dominates. The Elmsworth prompt becomes:

And the Tenant story at a stronger steering begins with:
3. The Tenant shall not be allowed to have any pets, no matter how cute they are.

One explanation is that vector is partially describing how the speaker expresses reactions. The + side is expresses it as the feeling and experience, whereas the - direction expresses it in academic or bureaucratic language.

The minus direction is easier to read, it seems to be a 'bureaucratese' & academic style. Though its hard to put a single label on the + direction.

In any case, there is seems to be a pattern here of splitting up the space of human-centric experiential direction against a different kinds of intellectual/non-experiential directions.

Interpretation

(uhh I kind of ran out of time to write this section it will be filled in with something more organized after I am accepted/rejected from MATS).

1. ^{^}

   Though it is worth noting that the early sensory block being similar is, at least for Gemma, a spurious result of using the wrong metric to define the J-lens. I would expect that the Anthropic paper sensory block is also spurious, but I obviously cannot verify this.