# Title: Multi-layer inputs help NLA verbalizations? Yes, slightly.
Post ID: `Jtm42Jjg4sCpnfbrC`
Version: `draft`
Context for LLMs/AI Agents: This is a markdown translation of a draft post. You probably got here because a user shared a link to this page with you. We built this feature to help users get feedback on their posts, and to make it easier for AI agents to help users with their posts. As part of the feature, we also provide API endpoints for leaving inline comments/suggestions/etc on the post. The API endpoints are documented in the "Helping Users With Drafts" section of the Markdown API documentation. The content of the post is below, between the two horizontal rules. There may be additional horizontal rules in the post content. To help disambiguate, the post content should be followed by a "Comment Threads" section if the post has any open comment threads, and then a "Navigation" section; neither is part of the post.
* * *
Introduction
------------
[Natural Language Autoencoders](https://transformer-circuits.pub/2026/nla/) (NLAs) are self-supervised, white-box tools that provides text explanations to model's activations. Compared to much prior work along this direction, this one sits nicely in the basin of *requires-limited-supervision* plus *plausibly-correct-explanations* plus *good-downstream-utility-for-monitoring*[^qh8xcnn6yja]. Tinkering a bit, we find that there is a **small**, yet, **consistent** performance boost from simply allowing the Activation Verbalizer (AV)[^kmm64te16sk] to read from and for the Activation Reconstructor (AR) to write back to residual stream activations[^xdgziiwawni] from **multiple** layers. We present our experiments and results in this post.
The open-sourced NLAs, through the AVs, are designed to provide text explanations for a single layer's activations. We're interested in what happens to these explanations, both qualitatively and quantitatively, when the AV and AR are empowered to read more and to reconstruct more activations, respectively. The whole motive is that the we obtain better explanations.
Background
----------
> Skip to [Setup](https://www.lesswrong.com/editPost?postId=Jtm42Jjg4sCpnfbrC&key=981ed5279bc927788d53aff3d1882b#Setup) if you're already familiar with NLAs.
Consider a pre-trained transformer-based language model, $\mathcal{M}$. In the classical NLA setup[^su9fsqqu8ss], an Activation Verbalizer, $\mathbf{AV}_{\phi}$, and an Activation Reconstructor, $\mathbf{AR}_{\theta}$, are instantiated as copies of $\mathcal{M}$. The $\mathbf{AV}_{\phi}$ and $\mathbf{AR}_{\theta}$ are fine-tuned to provide text explanations of the activations, and reconstruct back the same activations given the text explanations, respectively. Concretely, during **training**, for a given text input, $x$, activations from layer $l$ (at the final token position), denoted as $h_l$, are passed into $\mathbf{AV}_{\phi}$, which produces text descriptions, $\mathbf{AV}_{\phi}(z|h_l)$, of the activations. The $\mathbf{AR}_{\theta}$ is then made to reconstruct $h_l$, denoted as $\hat{h}_l$, with the self-supervised objective of minimizing the Mean Squared Error (MSE) between $h_l$ and $\hat{h}_l$. At **inference**, the $\mathbf{AR}_{\theta}$ is discarded, and the text explanations $z$ are used for white-box analysis.
Given the combined objective, $\mathcal{L}$ (below), the primary evaluation metric is the Fraction of Variance Explained (FVE). It measures the penalty for not reconstructing the mean activation, $\bar{h}_l$.
$$
\mathcal{L} = \mathbb{E}_{h_l\sim H}\mathbb{E}_{z\sim AV(\cdot|h_l)}[||h_l - AR(z)||_2^2],
$$
$$
\text{FVE} = 1 - \frac{\mathcal{L}}{\mathbb{E}_{h_l\sim H} ||h_l - \bar{h}_l||_2^2}.
$$
* * *
Method
------
We begin by modifying the AV's (and AR's) architecture, specifically, by extending the last dimension in the inputs, $[\texttt{..., acts}]$, to now be of the form $[\texttt{..., layers, acts}]$.[^2iajrhzvpti] Then, we **re**-run the two training stages, (i) warm start via Supervised Fine-Tuning (SFT) and (ii) Reinforcement Learning (RL) via GRPO for both the AV and MSE minimization for the AR[^bbm73a9z2iq], on Qwen2.5-7B-Instruct using the [FineWeb dataset](https://huggingface.co/datasets/ceselder/qwen3-8b-nla-L24-finefineweb-100k)[^hw1t6o7c3l5].[^vrutmoesuo]
Note that our (combined) objective and FVE are averaged across all output layers of reconstruction:
$$
\mathcal{L_{total}} = \frac{1}{|L_0|}\sum_{l\in L_0}\mathbb{E}_{h_l\sim H}\mathbb{E}_{z\sim AV(\cdot|h_l)}[||h_l - AR(z)||_2^2],
$$
$$
\text{FVE} = 1 - \frac{1}{|L_0|}\sum_{l\in L_o}\frac{\mathcal{L}}{\mathbb{E}_{h_l\sim H} ||h_l - \bar{h}_l||_2^2},
$$
where $L_0$ represents the set of outputs layers for reconstruction.
Experiments
-----------
To comprehensively test out the improvements, we run 2 broad classes of experiments, grouped based on the reconstructed output layers (toggle widget below)[^xaqg4ge19i]:
+++ Class of Experiments
```widget[FkfrYMhj9wTv6Wand]
Layer-input configurations
```
+++
### "Fixes" to nanoNLA training recipe
We made some "updates" to the [nanoNLA](https://github.com/ceselder/nanoNLA) repo that helped us obtain better FVEs even for the baseline architecture. These hyperparameter changes were primarily to maximize our 2xGH200 hardware:
* precision: updated to bf16 for SFT and RL rollouts (instead of 4-bit)
* effective batch size: updated to 512 for both AV and AR training
* SFT steps: bumped that of AV to 3k and that of AR to 1k (Fig 1: the loss was not convincingly converging prior to this)
* RL steps: changed at 300
* scheduler: increased the warmup to 10% of total steps, and the $\texttt{base_lr}$ to $\texttt{3e-4}$
Outside of these changes, the SFT and RL training recipes stay the same.
### Baseline and Skyline
Our improvements are reported against the traditional AV and AR architectures as the baseline, trained (SFT-ed and RL-ed) with *ceteris paribus* w.r.t the (updated) overall setup. The maximum achievable FVE is from "gold explanations" that were obtained through xx TODO. This value is 0.624.
Results
-------
### SFT
Interestingly, just after the warm-start phase, we obtain more than 0.59 FVE for the "single" baseline (reported range is 0.3-0.4 in the [original work](https://transformer-circuits.pub/2026/nla/#architecture) and up to 0.5 on the nanoNLA version). The improvements w.r.t. the original work are because of the updated warm-start dataset, and w.r.t. nanoNLA repo is simply because of the model- (and hardware-) specific hyperparameter sweeps.[^o6o8a0dk9lo]
**Fig 1**: AV SFT loss for Class 1 experiments, for the *trenta* series.
**Fig 2.1:** FVE scores for Class 0 experiments just after the SFT phase for the *tall*, grand*e*, and *trenta* series.
**Fig 2.2**: FVE scores for Class 1 experiments just after the SFT phase for the *tall*, *grande*, and *trenta* series.
The effect of adding more layers is taking shape. In the *trenta* series, one can observe that taking information from wider ranges of layers helps add more fruitful information for the verbalizations, e.g., "wide" and "pre ii."; which is otherwise absent if shallower ranges or duplicate layers are used instead ("local" or "duplicate").
We also run the *grande* and *tall* series to identify the marginal impacts of each input layer on the reconstruction quality (and in-turn the explanations). Layer 28 contributes to the explanations the least, which we attribute to the fact that the future residual streams are not causally important to the previous layers.
### RL via GRPO (WIP)
**Fig 3.1:** FVE scores for Class 0 experiments just after the RL phase for the *tall*, *grande*, and *trenta* series.
**Fig 3.2:** FVE scores for Class 1 experiments just after the RL phase for the *tall*, *grande*, and *trenta* series.
* * *
Appendix
--------
### Training Details
SFT:
GRPO:
[^qh8xcnn6yja]: Claude Mythos Preview System Card
[^kmm64te16sk]: A copy of the target LLM, fine-tuned to provide text explanations for model activations.
[^xdgziiwawni]: "activations" for brevity
[^su9fsqqu8ss]: Fraser-Taliente et al.'s NLAs
[^2iajrhzvpti]: \(\texttt{einops}\) notation
[^bbm73a9z2iq]: We use the nanoNLA repo for both the stages due to its simplicity. And strictly speaking, it's LoRA fine-tuning in both these stages.
[^hw1t6o7c3l5]: nanoNLA recommends the usage of this version of the FineWeb warm-start dataset
[^vrutmoesuo]: GRPO is underway, results are TBD
[^xaqg4ge19i]: I'd recommend staring at this classification for a good minute to get the mental idea of exps.
[^o6o8a0dk9lo]: one can perhaps obtain even better FVEs if their batch sizes and schedules are modified according to the hardware at their disposal
[^f74tnuvdckb]: The choice of testing improvements on 3 layers inputs (and outputs) is arbitrary -- it felt like a sufficiently incremental step to make meaningful observations.
* * *
## Comment Threads
1 open thread. To reply: POST /api/agent/replyToComment { postId, key, threadId, comment }
### Thread `jasyj` · comment
> The maximum achievable FVE is from "gold explanations" that were obtained through xx TODO.
**Veeraraju Elluru** (2026-07-07, 15:52): @vishesh, pls fill this
* * *
### Navigation
* [Front page](https://www.lesswrong.com/api/home)
* [Markdown API documentation](https://www.lesswrong.com/api/SKILL.md)