Matryoshka NLAs: training activation verbalizers to frontload reconstruction-relevant information

TLDR: We train a “matryoshka” NLA that, unlike standard NLAs, is trained to put the most important details at the start; it is trained by randomly truncating the verbalizer’s explanations before showing them to the reconstructor. We find that our matryoshka NLA frontloads claims that are important to reconstruction (more than standard NLAs), and that this can be used as a heuristic for saliency of the represented feature. We did not find strong evidence that matryoshka NLAs are significantly more useful than normal NLAs, though we didn’t look very hard. Natural language autoencoders (you should probably read that paper before this!) are pretty cool! Earlier we tried using length penalties while training them and found you could shorten NLA explanations a lot without losing much reconstruction ability. This led to thinking about training a kind of "matryoshka NLA"[1], where you randomly truncate the output of the AV (activation verbalizer) during training, which would force it to put the most important information at the start of the explanation[2]. This is conceptually similar to matryoshka SAEs. We have now trained such an NLA on Qwen3.6-27B (layer 42). Take a look at some NLA explanations (or try with your own text)!  Motivation The motivation for training matryoshka NLAs was to make the AV order its thoughts about the activation, and force it to put the most reconstruction-relevant information at the start of its response. With matryoshka NLAs we should be able to force the model to put the text that's actually relevant for increasing FVE (fraction of variance explained) at the top, so you can easily only look at the important items. You might also be able to do better experiments with matryoshka NLAs where you modify or remove lines, because the AR (activation reconstructor) is trained to handle explanations that are truncated at any point.  With matryoshka NLAs, you can look at the marginal per-line (or per-token) FVE to see how useful each line is (with normal NLAs truncating the explanation is out-of-distribution for the AR, so this is less meaningful). It might make sense to highlight marginal FVE prominently when showing matryoshka NLA outputs, as in the widget above. Method We use EasyNLA to train a Qwen 3.6 NLA. We create 450k warm-start explanations that are formatted as ten shortish bullet points, using Claude Sonnet 4.6. Here’s what that looks like: Final token "BI" begins an author name/affiliation entry Author list continuation after first author, second name expected Academic paper header/metadata section with author names Bibliographic format: surname in caps followed by given name Pattern of author names after keywords section in papers Domain: power systems/control engineering academic publication [...] This might be a better format for explanations because it describes the activation as a series of short independent bullet points, each of which can be mostly considered in isolation. Normal NLAs write 3-5 paragraphs, which makes it harder to do things like ablate individual claims. In the AR warm-start, the Claude-written explanations are randomly truncated (in the same way as in RL); otherwise the AR is really bad at the start of RL because it had never seen a truncated explanation before.[3] We do 400 RL steps where the AV explanation is truncated[4] to a uniform random number of tokens from 1 to 120 before giving it to the AR. (The truncation is the same for all explanations in each GRPO group). We triple the KL penalty from the EasyNLA default, and taper the KL penalty down significantly for later tokens, because otherwise later tokens (which can't contribute to FVE much, because earlier tokens gobbled up most of the FVE) end up having their reward being dominated by the KL penalty, and have a much weaker incentive to contribute to reconstruction. (The exact KL value was chosen in a kinda ad hoc manner, which partially confounds direct comparisons.) We previously used a lower KL value and got results that were less aesthetically pleasing to read, largely because they overused quotes. The hyperparameters chosen for the KL tapering probably weren't optimal (they were chosen mostly based on vibes), and there’s probably a better way to choose them. Matryoshka NLAs capture most of their FVE in a few lines We can plot how good the reconstruction (the FVE) of the matryoshka NLA is when truncating the explanation at each token index (versus a normally-trained Qwen3.6-27B NLA as a control) on random held-out pretraining documents (averaged over 100 explanations from separate documents). The standard NLA does better on shorter truncation lengths when you reverse the line order[5], so we show it both normally (grey dashed) and with reversed lines. The matryoshka NLA does better than the standard NLA for shorter truncation lengths (less than ~120 tokens); the standard NLA does better when there's no truncation. (This is somewhat confounded since it shows both the effect of the different warm