Natural language autoencoders are pretty interesting; I'm a big fan of using meta-models to interpret what's going on inside LLMs and NLAs seem like a promising route to understanding them better.
One thing I've been interested in for a while is using NLAs to interpret the differences between activations: given an activation from an earlier and later layer, can we verbalize the changes between those two layers? It seems like it would be pretty cool if we could verbalize the activation diff between each layer in a transformer, since it would let us better understand how the model does computations.
Unfortunately training an activation-diff-NLA that actually works well is pretty hard[1], so I've done something easier: create an NLA that can interpret the outputs of the attention heads of a single layer (L24) of Qwen3-8B. It, uh, kinda works? I can get explanations that are (by the way the NLA is trained) text describing the attention head outputs! It's not clear to me how useful these explanations are though, and my NLA only works well for some attention heads.
Here's a widget where you can look at some example outputs from my NLA:
This is pretty early work, and I'll probably make better versions of this kind of thing in the future! I'm not really sure how good head NLAs are because I'm not sure how to evaluate them well. Please keep in mind that I've probably made some mistakes here, and you should generally take any results here with a grain of salt.
NLA recap
[You can skip this section if you're already familiar with NLAs.]
A traditional natural language autoencoder (NLA) is two models that are fine-tuned from the model we want to interpret: an activation verbalizer (AV) that takes in an activation at a specific layer as input and outputs text, and an activation reconstructor (AR) that takes that explanation as input and outputs an activation (specifically the AR is truncated at the layer the activation is from[2]).
First the AV and AR are SFTed on a bunch of Claude-written explanations on random tokens from random documents, that are written by giving Claude some context and asking it to make up something that seems like the model could plausibly be thinking about. Then the AV and AR are trained together to try to make it so when you take an activation h (from a random token from a random document), [3] (where ), with a KL penalty to keep the AV from straying too far from the SFTed AV. The AV is RLed with GRPO, and AR is directly optimized.
Matryoshka NLAs are my variant of NLAs where you randomly truncate the explanations during training, to try to make the most important information be at the beginning.
defining attention head NLAs
Qwen3 concatenates the per-head outputs (attn_weights @ v) of each attention head into one big vector, then multiplies that by the output matrix. Effectively the output matrix is projecting each attention head's output from a d_head-vector to a d_model-vector and summing over all the heads. But we can also use it to project just a single head into d_model-space and use that to interpret the output of a single attention head (Qwen3 doesn't have a bias for the output projection).
An attention head NLA is like a normal NLA, but instead of being given a normalized activation, it's given the normalized output of a single attention head (specifically, what you would get if you zeroed out all the other attention heads when doing output projection). Also, the AR is told what attention head it's looking at in its prompt, to try to make it so the AV doesn't start by encoding that information using some code to get cheap FVE (fraction of variance explained).[4]
Actually, instead of using just the attention head outputs, I took a bunch of attention head outputs over many tokens to get the average output, and subtracted the head's mean projected output from the attention head outputs. (I calculated the mean value once and reused it for everything.) In hindsight I think this was bad and I shouldn't have centered the outputs like that.
grouped query attention
Qwen3 models use grouped query attention (GQA) to make inference faster - specifically, the key and value matrices are tied for the 4 attention heads in each GQA group. The query matrices aren't tied, so the attention patterns among attention heads in a GQA group will differ: each head has its own query but shares the keys that the query looks for.
I didn't investigate the impact of GQA on my NLAs too much. It seems that often if my head NLA is good at explaining one head in a group it's roughly as good as the other heads in the group; I don't know why this is.
the explanations
Mostly when I look at explanations from my attention head NLA, I think "Yeah, that seems like something vaguely like what the attention head might be doing". I don't have much evidence on how useful this NLA is for understanding what attention heads do; I hope to have more information about that in the future.
evaluating them
It's pretty hard to evaluate attention head NLAs well, since we don't really have a good ground truth for what individual attention heads do. I'm pretty unsure what exactly is going on in the AV that makes it write the explanations the way it does (aside from the warmstart pushing it to write in a certain format initially). Often the explanations don't seem to be very relevant to the input text (despite having high FVE). I don't think this is that surprising - probably the model really is thinking about that thing for some reason, and it just ultimately doesn't end up being useful. I slightly prefer looking at the matryoshka NLA explanations since they can frontload the important stuff into the first several tokens and so I don't need to read the whole thing.
It would be interesting to try seeing the impact of ablating attention heads, and seeing if you can predict that impact from attention-head-NLA explanations.
quoting
The head NLA will often quote tokens that the attention head attended to. The head NLA doesn't get to directly see what tokens the attention head attended to (just the centered attention head output), so this is some weak evidence that the head NLA is actually doing something meaningful.
is it a hidden code?
I did some cursory analysis to see if the NLA was writing in some kind of hidden code, instead of just plainly describing the head's write. Having Claude Haiku 5.5 rewrite explanations (on held-out chat samples) to keep the substance but change the wording resulted in only losing 3.9 points of FVE (2.7 points for the matryoshka one). So it's probably not the case that the AV is writing in some shared code language for the AR.
how I trained these NLAs
I started with an SFT warm-start with explanations written by Claude Haiku 5.5, given a list of tokens that the attention head boosted, for each attention head, for 100k tokens (10 tokens from 10k FineFineWeb documents) when the centered attention head write was injected into the model as a soft token. This was a kinda arbitrary way of doing it but it didn't need to be great; I just needed a starting point for the NLA that pointed it in the right direction.
These explanations weren't intended to be very good, but they were enough to give the NLA a starting point that wasn't awful and made it so the explanations were written in a vaguely sensible way. I did ~0.35 H100 hours of SFTing on Qwen3-8B. I do everything on all attention heads of layer 24 (with one NLA shared across all heads), and train everything as unquantized (bf16) LoRAs.
Once I had the normally trained NLA, I tried doing some extra training on top of it to convert it into a matryoshka NLA. I mostly did this because I wanted to make it easier to understand the explanations I was getting from the NLA, since I wasn't sure how much of the explanation text was actually relevant (versus filler text that wasn't actually used by the AR).
I did 240 steps with the matryoshka objective (truncating the explanations to U[1,120] tokens) on top of the normal-head-objective NLA to make the matryoshka head NLA variant for 2.7 H100 hours.[6] For the first 40 steps I only trained the AR so it learned to handle truncated explanations first. The switch to the matryoshka objective caused the FVE to drop initially because the AR wasn't expecting truncated explanations, but it caught up quickly.
Here's the FVE on held-out data, for both the standard and matryoshka NLAs (x-axis is log(step)):
Here's the train-time FVE for the standard NLA; if you plot it against training steps on a log plot the FVE grows with log(steps):
I'm pretty good at this "scaling laws" thing
The KL penalty turned out to have been too low; the explanations drifted pretty far from the warm-start.
(Note that the way FVE is computed for NLAs is a little weird, and not what I expected at first. Mostly this is because the AR can't guess the mean vector, because the mean vector isn't unit length.)
reconstruction quality
The head NLA is better at some heads than others, and roughly tends to be about as good among heads in a GQA group:
It might be that the NLA kinda learns to specialize more in heads that are easier to verbalize.
misc graphs
These ones aren't that interesting but I had Claude make some other graphs in case you wanted to see them.
On the matryoshka NLA here's the FVE when truncating to various token counts on held-out pretraining documents:
Here's a better way to see how similar the FVE is among heads in the 8 GQA KV groups:
Here's how much more quoted tokens are attended to:
attention sinks
Qwen3 doesn't have an explicit beginning-of-sequence token, so instead attention heads rest unused attention on the first token. For the chat-tuned model this isn't really a problem though because the first token is always <|im_start|>, although the NLA was trained on web text which wasn't chat-formatted. I did some quick checks and it doesn't seem like attention sinks contribute to the FVE of my NLA much.
what should I think of activation diffs as
I kinda think of the activations after a layer in the residual stream as the "mental state" of an LLM at that point, and so I guess attention head outputs should be thought of as a change to that "mental state".
I'm not sure if this is the best way to think of the attention head output vectors though. For example, it might be that a model splits something I would think of as one logical change across multiple layers, and so there might not be a good explanation for what a single attention head (or layer) does - you might need to think about things more holistically, which is harder.
future directions
Ways you could improve head NLAs
It's kinda weird for the model to go from having normal activations in the residual stream to having diffs there. There's almost certainly a better way to give the attention head output vectors to the model.
Improve the way attention head outputs are centered (or maybe just don't center them?)
Comparing them to other ways of interpreting what attention heads do
Probably you could twiddle the hyperparameters (e.g. learning rate) to make it better
Try training one for a model that doesn't use GQA, so each attention head is more individually meaningful
Increase the KL penalty so that it's large enough to meaningfully affect training
Train it for more steps, use a bigger model, do full-parameter fine-tuning instead of a LoRA, etc.
Looking at a bunch of explanations and trying to understand them deeply
Have some explicit handling for attention sinks
Reward it for being good at all attention heads, instead of just effectively averaging the reconstruction error across all attention heads
Train on the base model instead of the chat one
I think coming up with a better way to interpret changes in activations is more important than just making attention head NLAs better. IMO you should think of attention head NLAs as more of a stepping stone to an actually good mechanism for understanding individual changes to LLM activations.
Random other ideas:
Try giving attention patterns to a meta-model
Consider using a better way to score the NLA than just looking at the reconstruction error
conclusion
I'm pretty unsure if methods like NLAs will actually end up being useful for solving LLM interpretability, but it seems like they might be pretty useful! Understanding changes between activations seems pretty useful to me for figuring out what the heck LLMs are doing.
I'm trying to write about my research more frequently by giving more updates, and this post is a part of that. There's still a lot of stuff related to attention NLAs (e.g. evaluating them in an Actually Good way) to explore, and I'll be writing about what I'm doing more!
Thanks to Alibaba for helping advance the study of LLM interpretability by releasing Qwen3-8B! I'm also grateful for EasyNLA, which is a fork of nanoNLA, which in turn is a reimplementation of Anthropic's NLA code. I've published the source code I used and training run logs.
Please let me know if you have any advice/ideas for me; I'd love to hear what you think. :) Please experiment with new kinds of NLA yourself - there's a lot to discover here and it would be really great if we could solve interpretability!
Technically there's also a value head tacked on the end that maps the activation at the truncated layer to the AR's output. But in practice the value head is very close to the identity matrix and doesn't do anything; it should probably be omitted if you're creating a new NLA architecture (unless you have a really strange model architecture where d_model varies between layers or something like that).
The NLA paper doesn't directly say if you normalize the AR output or not, but the open source code for training/inferencing NLAs normalizes the AR output. I define as a function that normalizes an activation.
It's probably pretty easy for the AV to figure out what attention head it's looking at from the input vector, and knowing what attention head you're looking at is really useful for the AR. If the AR wasn't given the attention head index, the AV will just learn to encode which attention head wrote the vector in some way, and I want to avoid the AV/AR developing a coded language.
The time-per-step for matryoshka NLAs is lower, because there are less tokens in each rollout, so with matryoshka NLAs you get more steps in the same time period. The downside is that the model doesn't work very well for tokens near the end of the truncation range.
Natural language autoencoders are pretty interesting; I'm a big fan of using meta-models to interpret what's going on inside LLMs and NLAs seem like a promising route to understanding them better.
One thing I've been interested in for a while is using NLAs to interpret the differences between activations: given an activation from an earlier and later layer, can we verbalize the changes between those two layers? It seems like it would be pretty cool if we could verbalize the activation diff between each layer in a transformer, since it would let us better understand how the model does computations.
Unfortunately training an activation-diff-NLA that actually works well is pretty hard[1], so I've done something easier: create an NLA that can interpret the outputs of the attention heads of a single layer (L24) of Qwen3-8B. It, uh, kinda works? I can get explanations that are (by the way the NLA is trained) text describing the attention head outputs! It's not clear to me how useful these explanations are though, and my NLA only works well for some attention heads.
Here's a widget where you can look at some example outputs from my NLA:
This is pretty early work, and I'll probably make better versions of this kind of thing in the future! I'm not really sure how good head NLAs are because I'm not sure how to evaluate them well. Please keep in mind that I've probably made some mistakes here, and you should generally take any results here with a grain of salt.
NLA recap
[You can skip this section if you're already familiar with NLAs.]
A traditional natural language autoencoder (NLA) is two models that are fine-tuned from the model we want to interpret: an activation verbalizer (AV) that takes in an activation at a specific layer as input and outputs text, and an activation reconstructor (AR) that takes that explanation as input and outputs an activation (specifically the AR is truncated at the layer the activation is from[2]).
First the AV and AR are SFTed on a bunch of Claude-written explanations on random tokens from random documents, that are written by giving Claude some context and asking it to make up something that seems like the model could plausibly be thinking about. Then the AV and AR are trained together to try to make it so when you take an activation h (from a random token from a random document),
Matryoshka NLAs are my variant of NLAs where you randomly truncate the explanations during training, to try to make the most important information be at the beginning.
defining attention head NLAs
Qwen3 concatenates the per-head outputs (
attn_weights @ v) of each attention head into one big vector, then multiplies that by the output matrix. Effectively the output matrix is projecting each attention head's output from a d_head-vector to a d_model-vector and summing over all the heads. But we can also use it to project just a single head into d_model-space and use that to interpret the output of a single attention head (Qwen3 doesn't have a bias for the output projection).An attention head NLA is like a normal NLA, but instead of being given a normalized activation, it's given the normalized output of a single attention head (specifically, what you would get if you zeroed out all the other attention heads when doing output projection). Also, the AR is told what attention head it's looking at in its prompt, to try to make it so the AV doesn't start by encoding that information using some code to get cheap FVE (fraction of variance explained).[4]
Actually, instead of using just the attention head outputs, I took a bunch of attention head outputs over many tokens to get the average output, and subtracted the head's mean projected output from the attention head outputs. (I calculated the mean value once and reused it for everything.) In hindsight I think this was bad and I shouldn't have centered the outputs like that.
grouped query attention
Qwen3 models use grouped query attention (GQA) to make inference faster - specifically, the key and value matrices are tied for the 4 attention heads in each GQA group. The query matrices aren't tied, so the attention patterns among attention heads in a GQA group will differ: each head has its own query but shares the keys that the query looks for.
I didn't investigate the impact of GQA on my NLAs too much. It seems that often if my head NLA is good at explaining one head in a group it's roughly as good as the other heads in the group; I don't know why this is.
the explanations
Mostly when I look at explanations from my attention head NLA, I think "Yeah, that seems like something vaguely like what the attention head might be doing". I don't have much evidence on how useful this NLA is for understanding what attention heads do; I hope to have more information about that in the future.
evaluating them
It's pretty hard to evaluate attention head NLAs well, since we don't really have a good ground truth for what individual attention heads do. I'm pretty unsure what exactly is going on in the AV that makes it write the explanations the way it does (aside from the warmstart pushing it to write in a certain format initially). Often the explanations don't seem to be very relevant to the input text (despite having high FVE). I don't think this is that surprising - probably the model really is thinking about that thing for some reason, and it just ultimately doesn't end up being useful. I slightly prefer looking at the matryoshka NLA explanations since they can frontload the important stuff into the first several tokens and so I don't need to read the whole thing.
It would be interesting to try seeing the impact of ablating attention heads, and seeing if you can predict that impact from attention-head-NLA explanations.
quoting
The head NLA will often quote tokens that the attention head attended to. The head NLA doesn't get to directly see what tokens the attention head attended to (just the centered attention head output), so this is some weak evidence that the head NLA is actually doing something meaningful.
is it a hidden code?
I did some cursory analysis to see if the NLA was writing in some kind of hidden code, instead of just plainly describing the head's write. Having Claude Haiku 5.5 rewrite explanations (on held-out chat samples) to keep the substance but change the wording resulted in only losing 3.9 points of FVE (2.7 points for the matryoshka one). So it's probably not the case that the AV is writing in some shared code language for the AR.
how I trained these NLAs
I started with an SFT warm-start with explanations written by Claude Haiku 5.5, given a list of tokens that the attention head boosted, for each attention head, for 100k tokens (10 tokens from 10k FineFineWeb documents) when the centered attention head write was injected into the model as a soft token. This was a kinda arbitrary way of doing it but it didn't need to be great; I just needed a starting point for the NLA that pointed it in the right direction.
These explanations weren't intended to be very good, but they were enough to give the NLA a starting point that wasn't awful and made it so the explanations were written in a vaguely sensible way. I did ~0.35 H100 hours of SFTing on Qwen3-8B. I do everything on all attention heads of layer 24 (with one NLA shared across all heads), and train everything as unquantized (bf16) LoRAs.
After that I did 400 RL steps (with the same mechanism as for normal NLAs (based on the EasyNLA self-contained config[5])) for 8 H100 hours.
Once I had the normally trained NLA, I tried doing some extra training on top of it to convert it into a matryoshka NLA. I mostly did this because I wanted to make it easier to understand the explanations I was getting from the NLA, since I wasn't sure how much of the explanation text was actually relevant (versus filler text that wasn't actually used by the AR).
I did 240 steps with the matryoshka objective (truncating the explanations to U[1,120] tokens) on top of the normal-head-objective NLA to make the matryoshka head NLA variant for 2.7 H100 hours.[6] For the first 40 steps I only trained the AR so it learned to handle truncated explanations first. The switch to the matryoshka objective caused the FVE to drop initially because the AR wasn't expecting truncated explanations, but it caught up quickly.
Here's the FVE on held-out data, for both the standard and matryoshka NLAs (x-axis is log(step)):
Here's the train-time FVE for the standard NLA; if you plot it against training steps on a log plot the FVE grows with log(steps):
I'm pretty good at this "scaling laws" thing
The KL penalty turned out to have been too low; the explanations drifted pretty far from the warm-start.
(Note that the way FVE is computed for NLAs is a little weird, and not what I expected at first. Mostly this is because the AR can't guess the mean vector, because the mean vector isn't unit length.)
reconstruction quality
The head NLA is better at some heads than others, and roughly tends to be about as good among heads in a GQA group:
It might be that the NLA kinda learns to specialize more in heads that are easier to verbalize.
misc graphs
These ones aren't that interesting but I had Claude make some other graphs in case you wanted to see them.
On the matryoshka NLA here's the FVE when truncating to various token counts on held-out pretraining documents:
Here's a better way to see how similar the FVE is among heads in the 8 GQA KV groups:
Here's how much more quoted tokens are attended to:
attention sinks
Qwen3 doesn't have an explicit beginning-of-sequence token, so instead attention heads rest unused attention on the first token. For the chat-tuned model this isn't really a problem though because the first token is always
<|im_start|>, although the NLA was trained on web text which wasn't chat-formatted. I did some quick checks and it doesn't seem like attention sinks contribute to the FVE of my NLA much.what should I think of activation diffs as
I kinda think of the activations after a layer in the residual stream as the "mental state" of an LLM at that point, and so I guess attention head outputs should be thought of as a change to that "mental state".
I'm not sure if this is the best way to think of the attention head output vectors though. For example, it might be that a model splits something I would think of as one logical change across multiple layers, and so there might not be a good explanation for what a single attention head (or layer) does - you might need to think about things more holistically, which is harder.
future directions
Ways you could improve head NLAs
I think coming up with a better way to interpret changes in activations is more important than just making attention head NLAs better. IMO you should think of attention head NLAs as more of a stepping stone to an actually good mechanism for understanding individual changes to LLM activations.
Random other ideas:
conclusion
I'm pretty unsure if methods like NLAs will actually end up being useful for solving LLM interpretability, but it seems like they might be pretty useful! Understanding changes between activations seems pretty useful to me for figuring out what the heck LLMs are doing.
I'm trying to write about my research more frequently by giving more updates, and this post is a part of that. There's still a lot of stuff related to attention NLAs (e.g. evaluating them in an Actually Good way) to explore, and I'll be writing about what I'm doing more!
Thanks to Alibaba for helping advance the study of LLM interpretability by releasing Qwen3-8B! I'm also grateful for EasyNLA, which is a fork of nanoNLA, which in turn is a reimplementation of Anthropic's NLA code. I've published the source code I used and training run logs.
Please let me know if you have any advice/ideas for me; I'd love to hear what you think. :) Please experiment with new kinds of NLA yourself - there's a lot to discover here and it would be really great if we could solve interpretability!
I've written about this a bit in the past.
Technically there's also a value head tacked on the end that maps the activation at the truncated layer to the AR's output. But in practice the value head is very close to the identity matrix and doesn't do anything; it should probably be omitted if you're creating a new NLA architecture (unless you have a really strange model architecture where d_model varies between layers or something like that).
The NLA paper doesn't directly say if you normalize the AR output or not, but the open source code for training/inferencing NLAs normalizes the AR output. I define
It's probably pretty easy for the AV to figure out what attention head it's looking at from the input vector, and knowing what attention head you're looking at is really useful for the AR. If the AR wasn't given the attention head index, the AV will just learn to encode which attention head wrote the vector in some way, and I want to avoid the AV/AR developing a coded language.
Except that I train in full bf16 instead of quantizing to 4 bits.
The time-per-step for matryoshka NLAs is lower, because there are less tokens in each rollout, so with matryoshka NLAs you get more steps in the same time period. The downside is that the model doesn't work very well for tokens near the end of the truncation range.