Describing the computation performed in a stack of transformer layers
Anwen Hao, mentored by Adrians Skapars
Anthropic’s natural language autoencoders (NLAs) is a promising method to automatically generate explanations of activations. But what if we want to explain the computation that occurs over a stack of layers?
To address this question, I propose natural language transcoders (NLTs), a tool that, if successful, will automatically generate explanations of the computation performed in a stack of layers in a transformer. This post discusses my progress on this idea.
Architecture
Below is the architecture of the natural language transcoder. It follows closely from the NLA architecture with some key changes:
where
: output from layer N.
: output from layer M. M > N.
AV: activation verbalizer.
AR: activation reconstructor.
So, the natural language transcoder reads the delta between the input and output of a stack of layers in a transformer, verbalizes it into an explanation of the computation over that stack, and reconstructs the delta from that explanation alone. Instead of reconstructing an activation back to itself, the NLT reconstructs the delta back to itself.
Experimental Design
For my first experiment, I chose to train an NLT to describe what an entire transformer is doing because I believe that it is easier to make high-level claims about the behavior of a transformer than to describe exactly what is happening over a stack of layers. So, in this case,
where
: residual stream before entering the first transformer layer (after the embedding).
: residual stream leaving the final layer.
I use Qwen2.5-0.5B-Instruct as the frozen target model and the base model for the AV and AR, so this covers all 24 transformer blocks and produces an 896-dimensional vector.
For my training process to train the NLT, I closely follow the original NLA training procedure, while adapting key steps of the procedure accordingly. The training process consists of two phases.
SFT Warm Start
A teacher model is prompted to generate explanations that will be paired with delta vectors and serve as labels for the delta vectors. Both the AV and AR are SFT’ed using these pairs.
Originally, in the NLA training procedure, the teacher model is prompted to explain the information in an activation given the prefix (all the text before the token position of the activation). However, since our goal is to explain computation, I adapted the prompt to the teacher model.
In addition to the prefix, I also apply the norm and unembedding matrix (logit lens) to and , and we give the teacher model:
Original text prefix
The highest-probability tokens from unembedding
The highest-probability next tokens from unembedding
Largest logit increases and decreases
Entropy change: whether the distribution after unembedding became more concentrated or more diffuse than the distribution after unembedding .
Margin change: whether the final top candidate became more or less separated from the runner-up.
Note that the AV does not receive the prefix and this much information; it only receives the delta. The teacher model was asked for exactly two or three short bullets describing:
What the transformer strengthened
What it suppressed, when supported
What ambiguity it resolved
What immediate next-token behavior it prepared for
Joint RL
After SFT, both the AV and AR are trained simultaneously with RL. Given a delta, the AV generates 8 candidate explanations. We score each explanation through the current AR. We update the AV from within group relative rewards and update the AR with backpropagation.
Joint RL also motivates the necessity of a warm start. Without a warm start, the AR does not know how to transform text to a vector and cannot provide a useful reward signal to the AV from the beginning, which prevents the training from converging.
Results
Below is the final evaluation on all 5,000 held-out examples with greedy AV generation.
Metric
After SFT
After RL
Raw-delta FVE
0.119
0.503
Raw-delta MSE
3.810
2.150
Delta cosine
0.667
0.818
Prediction KL
5.571
2.221
Target-model top-1 agreement
11.9%
29.5%
Target-model top-5 overlap
18.2%
39.6%
I sample 100 explanations from the NLT and audit the explanations for correctness based on the following criteria:
The NLT’s prediction for the model’s next-token behavior is plausible based on context
The NLT’s claim about which tokens are promoted vs. suppressed is correct
Below are the results of the audit:
Contextually plausible
Token information correct
Count
Yes
Yes
20
Yes
No
28
No
Yes
28
No
No
24
The NLT is strongest in these cases:
The continuation has few reasonable possibilities:
“ranks” shifting toward “of” in “ranks of 1A athletics”
numerical tokens shifting toward “inch” in a pizza-size list
Clear section and discourse transitions: The model is leaving one segment and beginning the next.
Narrow, repeatedly established topics: When the prefix strongly and repeatedly establishes one topic, the resulting delta may contain a clearer, more redundant signature of that topic. This makes it easier for the AV to recover the right interpretation from the delta alone.
Related tokens promoted together: The tokens that are promoted are related and suggest the same kind of continuation, making it easier to generate a plausible explanation.
The NLT struggles in cases of:
Partial tokens and when exact continuation is required: It struggles with BPE fragments, dates, timestamps, numbers, and partially completed words. It often recognizes the broad category (“a date comes next”) but misses the exact continuation or its direction.
Naming specific entities: It frequently detects a general domain but names the wrong entity or subdomain.
Confusing final predictions with delta changes: The NLT sometimes calls a token promoted because it appears in the final top candidates, even when its actual logit change is negative.
Noisy or composite text: Forum metadata, malformed web text, concatenated pages, multiple unrelated headlines, and similar inputs produce unstable explanations. There may be no single coherent computation for the verbalizer to summarize.
Weakly identifiable global deltas: Because the AV receives the delta without the original prefix, several different contexts can produce similar directions in activation space. It can recover coarse information such as “prepare for a noun” or “increase a numerical continuation,” but not necessarily which noun, number, entity, or event is intended.
Next Steps
The above audit suggests that the AV struggles on more difficult continuation tasks, which suggests that a larger model may be helpful. For my next steps, I plan to repeat this experiment using a 7B model as the base model for the AV and AR. For example, I can use Qwen2.5-7B-Instruct, which is the size of the smallest released NLA checkpoint.
Training these models will require significant amounts of compute, so if you have compute resources, please reach out!
See the complete results and code of my experiment here:
Describing the computation performed in a stack of transformer layers
Anwen Hao, mentored by Adrians Skapars
Anthropic’s natural language autoencoders (NLAs) is a promising method to automatically generate explanations of activations. But what if we want to explain the computation that occurs over a stack of layers?
To address this question, I propose natural language transcoders (NLTs), a tool that, if successful, will automatically generate explanations of the computation performed in a stack of layers in a transformer. This post discusses my progress on this idea.
Architecture
Below is the architecture of the natural language transcoder. It follows closely from the NLA architecture with some key changes:
where
So, the natural language transcoder reads the delta between the input and output of a stack of layers in a transformer, verbalizes it into an explanation of the computation over that stack, and reconstructs the delta from that explanation alone. Instead of reconstructing an activation back to itself, the NLT reconstructs the delta back to itself.
Experimental Design
For my first experiment, I chose to train an NLT to describe what an entire transformer is doing because I believe that it is easier to make high-level claims about the behavior of a transformer than to describe exactly what is happening over a stack of layers. So, in this case,
where
I use Qwen2.5-0.5B-Instruct as the frozen target model and the base model for the AV and AR, so this covers all 24 transformer blocks and produces an 896-dimensional vector.
For my training process to train the NLT, I closely follow the original NLA training procedure, while adapting key steps of the procedure accordingly. The training process consists of two phases.
SFT Warm Start
A teacher model is prompted to generate explanations that will be paired with delta vectors and serve as labels for the delta vectors. Both the AV and AR are SFT’ed using these pairs.
Originally, in the NLA training procedure, the teacher model is prompted to explain the information in an activation given the prefix (all the text before the token position of the activation). However, since our goal is to explain computation, I adapted the prompt to the teacher model.
In addition to the prefix, I also apply the norm and unembedding matrix (logit lens) to
Note that the AV does not receive the prefix and this much information; it only receives the delta. The teacher model was asked for exactly two or three short bullets describing:
Joint RL
After SFT, both the AV and AR are trained simultaneously with RL. Given a delta, the AV generates 8 candidate explanations. We score each explanation through the current AR. We update the AV from within group relative rewards and update the AR with backpropagation.
Joint RL also motivates the necessity of a warm start. Without a warm start, the AR does not know how to transform text to a vector and cannot provide a useful reward signal to the AV from the beginning, which prevents the training from converging.
Results
Below is the final evaluation on all 5,000 held-out examples with greedy AV generation.
Metric
After SFT
After RL
Raw-delta FVE
0.119
0.503
Raw-delta MSE
3.810
2.150
Delta cosine
0.667
0.818
Prediction KL
5.571
2.221
Target-model top-1 agreement
11.9%
29.5%
Target-model top-5 overlap
18.2%
39.6%
I sample 100 explanations from the NLT and audit the explanations for correctness based on the following criteria:
Below are the results of the audit:
Contextually plausible
Token information correct
Count
Yes
Yes
20
Yes
No
28
No
Yes
28
No
No
24
The NLT is strongest in these cases:
The NLT struggles in cases of:
Next Steps
The above audit suggests that the AV struggles on more difficult continuation tasks, which suggests that a larger model may be helpful. For my next steps, I plan to repeat this experiment using a 7B model as the base model for the AV and AR. For example, I can use Qwen2.5-7B-Instruct, which is the size of the smallest released NLA checkpoint.
Training these models will require significant amounts of compute, so if you have compute resources, please reach out!
See the complete results and code of my experiment here: