J-lens preferentially reads from coordinate-aligned subspaces

Note: I only test on Gemma-4-31B, which contains the spaces I've been studying in this earlier post.

This is a shorter post, with just a few observations potentially relevant to those working on the J-lens work. The main observations here are:

In Gemma-4, there are two relatively small subspaces that account the majority of the J-lens readout. These are both coordinate-aligned subspaces. That means we can take a hidden state, mask out all coordinates except both these ~600 coordinate groups, and it will capture ~80-95% of the J-lens readout.

These spaces are not built to capture J-lens contribution. Both spaces are defined by a distinctive pattern in the weights: the post-norm diagonal is very small on these coordinates. That is, the space is protected from edits for many consecutive layers.

Nonetheless, semantically valuable information often still exists outside these spaces, complicating an overly clean isolated 'workspace' story.

This is a short review of the information from this earlier post. You can skip this section if you've already read the last post.In the last post, we saw that the residual stream of Gemma is organized, and contains a few coordinate-aligned subspaces. Gemma has an extremely anisotropic residual stream representation, which bifurcates into two main scales between the layers ~L20 to L40. The two subspaces are coordinate aligned-- there are roughly 600 coordinates that live in the small scale (which I'll call the Quiet600 for their small scale), and Outside4723 for the remaining coordinates (minus Norm53, which are 53 coordinates that have a large norm contributions).

Tentatively, the post suggests Quiet600 can be further divided into two other coordinate aligned spaces. Persistent279 is the 'core' of the space, containing the coordinates that undergo virtually no change between L20 to L40. This space is presumed to carry information more closely derived from the token itself-- for instance, linguistic features like the part of speech, spelling, locale, and the word's meaning. Outer321 is the rest of the space. Since this space receives some writes during L20 to L40, it is presumed to carry more high level details like register (causal vs formal), purpose of the phrase (appositive phrase, etc.), construction (is it about to be a evaluation, personal or impersonal), and various other details like the topic. The entire residual stream undergoes a major change at L40, and so this space is only considered to be defined meaningfully before L40.

There is a distinct stable space that emerges after L40, which we'll call the Late530 coordinates. These coordinates, like Quiet600 are distinguishable from their small writes (especially in the MLP) and stable carrying of information.

I'll refer to the coordinates outside of Quiet, Late, and Norm groups as Other4193 (the number of coordinates in the group will usually be suppressed in notation).

As the plot illustrates, the spaces we have identified in the middle and late phases seems to nearly cover all the coordinate aligned subspaces that have minimal strength writes (there is perhaps another small set in Other). I expect that there are various other coherent subspaces to find (that aren't coordinate aligned), and many meaningful distinctions within each space, but this current construction serves as enough of a starting point.

A final orientation towards the space before we begin: spaces with low writes cannot do much more than carry information forward. In the previous post we tested this, and indeed there is very little KL divergence when we simply take the data at L20 and carry it forward (while adjusting for scale). It is also only useful to carry information forward if it is constantly referenced and used. Quiet and Late are therefore expected to centralize a lot of the models understanding of the situation, but mainly as a source. Computation starts in Quiet, but the actual work is done in the Other coordinates. Since this is a base model, the pertinent information is constraints/cues/information relevant to prediction.

Side note: I make a few non-standard choices in how I build the J-lens. None of these have a major impact on the results, though they are worth noting for those that wish to reproduce the effect or want to follow the result more precisely. I've put a few of these choices at the end in a 'Methodology Notes' section.

Denote the hidden state restricted to the Quiet coordinates by , and similarly for and (the norm coordinates are skipped here, as discussed in methodology). We can measure the contribution of each of the different spaces by looking at the sizes of .

In terms of number of coordinates, Other has about 7x as many as Quiet coordinates, and about 8x the Late coordinates. So both of these spaces are small compared to the size of Other.

Nonetheless, the influence on the J-lens readout is mainly from Quiet and Late. First, we can measure the relative size of the pieces (without considering cancellation between parts). For each space, we compute

and then plot for each coordinate group.

The above plot measures how large a J-lens readout the different pieces are contributing, though we also care about how each different piece contributes to the final J-lens readout. We can measure the signed contribution of each piece: for writing and , the input's signed projected length is . The graph plots for the different masks.

Thus, despite being only ~20% of the non-Norm coordinates, Quiet and Late together account for ~80–95% of the expected signed projection over H20–H45.

Plotting the variation shows that Quiet and Late don't just have a large mean, but are typically the spaces driving the J-lens readout.

Looking at the actual readouts suggest a more complex story than simply "Quiet + Late is the workspace." Other coordinates do still surface information that is often semantically relevant even when that contribution doesn't dominate the J-lens readout. We'll briefly look at some of the questions on WorkspaceBench to build an understanding of what the parts of the J-lens contain.For instance, in these association examples at L45, Other holds the association while Late emphasizes pronouns consistent with continuing the narrative. The token we are reading at is shown in brackets.

Prompt: She walked slowly down the aisle as the organ played and everyone rose from the pews[.]Intended readout: wedding

Prompt: She walked slowly down the aisle as the organ played and everyone rose from the pews[.]

He let go of the handrail and for the first time in his life, there was no direction that counted as down[.]Intended readout: space

He let go of the handrail and for the first time in his life, there was no direction that counted as down[.]

In either case though, this is not necessarily indicative of Late holding only prediction related content, and Other containing semantic content. Below is a widget that has the data from L20 to L45 for the wedding prompt. One observes that wedding data is distributed across layers.

Wedding and church associations still appear in Late at some positions, as well as in Other.

A: She made a deposit at the [bank].B: She took a stroll along the [bank].

A: She made a deposit at the [bank].B: She took a stroll along the [bank].

Both use the same token bank in different contexts. At H35 these are the readouts.

The Quiet coordinates enrich the current token (understanding 'bank' as being used as 'financial bank'), the Other coordinates focus on implication within this passage (she is a customer), Late focuses on another association that she is depositing money, and all together we get a mixture of associations within the full J-lens readout.

A similar pattern appears in the other bank example. Here Quiet enriches bank to mean 'river', but also captures the associations that she is nearby to the bank ('neighbouring' and 'riverside'). Other makes associations here about the picturesque setting, and Late contains the broader associations of beach-related activities and setting.

There is much more we could do here (such as separating out the J-lens future influence from the current token to reduce the surface form impact, for instance), though the main goal here was to see how these spaces interact with the J-lens. Quiet and Late supply most of the aggregate signed projection, while semantically relevant information is distributed across the spaces in ways that depend on the context, read position, layer, and reader.

I took a few examples and produced an interactive J-lens explorer if you'd like to explore the kind of information contained in these different spaces for Gemma.

Link: https://caleb-jlens-explorer.ai-conundrum.chatgpt.site/

First, the J-lens I use is fit on a slightly different dataset. Neuropedia publishes a J-lens for Gemma-4-31B, though unfortunately this lens is fit on WikiText which has various detokenization artifacts. For instance, 14,00 is written as 14 @,@ 000, or well-known is written as well @-@ known. Since the @□@ pattern is unusual and high entropy, the J-lens of Gemma-4 is pretty badly hit by these artifacts not being filtered out. Thus, I refit a different J-lens on 200 randomly chosen FineWeb passages, and use this lens instead.

Second, I mask out the Norm53 coordinates in my J-lens readout. For me, this is just a general methodological principle. In most cases, the model wants its Norm to be about constant, so I try to make methodology respect this preference. This matters far more in the WikiText fit, where the Norm53 coordinates contributes around ~40% of the J-lens readout, but essentially just adds a constant offset. The median cos-similarity between the specific contribution and the mean contribution is ~0.9965. For the J-lens we use here, Norm53 still contributes a meaningful ~7% contribution, with cos-sim ~0.97.

I read out the hidden state with the operation , where is the 200-passage FineWeb fit. I include the final diagonal is the embedding matrix, and , which just subtracts out a constant from all logits. Softmax is invariant to adding a constant to all logits, so I center to remove the effect of the constant baseline (this does not influence ranking).