Thanks for pointing out the Matryoshka SAE connection.
The Matryoshka loss function could probably be used directly in the NLA case.
Hmm, yeah, if this works it would be much more efficient than what I described, since you can train all the scales simultaneously.
However, something about Matryoshka loss feels not quite right to me in the NLA case, specifically the fact that the prefix lengths it uses for grading an example are determined “in advance” and can’t depend on information about that specific example.
This is appropriate in the SAE case, where you want each feature to have some context-independent meaning (and—unlike with AV-sampled natural-language tokens—the features don’t have pre-existing interpretations attached to them before training). But it seems awkward in token space, where some examples will inherently require longer descriptions than others to reach any given reconstruction-error threshold with equal fluency/legibility. “Describe what’s going on here in exactly 16 tokens” is more feasible in some cases than others[1], and if our loss encourages the model to do well at tasks like that across all examples, I’d worry that this would encourage steganography / unnatural-sounding word-pretzels / that sort of thing.
In the setup described in this post, the model has the freedom to use different lengths on different examples, “spending” more tokens in cases where the improvement in reconstruction error justifies the expenditure. The Matryoshka loss doesn’t adapt like this.
I wonder whether the following would work: introduce one or more new tokens into the vocab which the AV will use to demarcate the end of each “graded prefix,” so that the AV samples would look like:
Then, have a term in the loss for each end-of-prefix separator, where each term is a sum of “reconstruction loss using the text up to that separator” and “length penalty on the tokens between it and the previous separator.”
This is a generalization of what happens with the EOS token in the setup described in the post, where the AV can sample that token earlier or later in the sequence and makes this decision by trading off reconstruction and length-penalty. With the end-of-prefix separators, it’s doing the exact same thing, but once for each nested prefix we’re going to grade.
As a stylized example, consider two activations, each of which is representing a bunch of things which have well-defined names in natural language, but where:
in one case, the model’s tokenizer encodes those names very inefficiently
in the other, the model’s tokenizer encodes the names much more efficiently
If we had a combined loss with a length penalty as in the post, the AV might end up reciting all of the names in both cases even though that results in a longer text in the first case; we just need the marginal reconstruction-error benefit of including each name to exceed the marginal length-penalty cost.
However, if we grade the model on reconstruction based on the first N tokens—where N does not adaptively vary with the content of the examples—then N might be too small to list all the names in the first case (but not in the second), and this could result in unnatural-sounding (or just undesirably imprecise) verbalizations of examples like that first case.
Yeah, I think the Matryoshka loss is more “mostly applicable” rather than “directly applicable”.
With Matryoshka SAEs there are two types of prefixes. The first is the dictionary subsets, which in our implementation were fixed and chosen in advance, which I agree doesn’t make sense for an NLA. The second is that the SAE could have K active features, and it could allocate the K active features across subsets as desired. This seems close to your “graded NLA prefixes” proposal.
There are also alternative forms of hierarchy that could be explored. For example, in Appendix C.1 of the Matryoshka paper there was an alternative implementation where the dictionary subset sizes where randomly sampled from a distribution when performing inference, which encouraged a more continuous hierarchy.
Another random idea (which may not be good) is to just pass the first K tokens of the AV verbalization to the AR (where K is chosen randomly), which may encourage the AV to put the more important stuff earlier in the generation, with the fine-grained details towards the end. But it could also encourage weird compression.
Thanks for pointing out the Matryoshka SAE connection.
Hmm, yeah, if this works it would be much more efficient than what I described, since you can train all the scales simultaneously.
However, something about Matryoshka loss feels not quite right to me in the NLA case, specifically the fact that the prefix lengths it uses for grading an example are determined “in advance” and can’t depend on information about that specific example.
This is appropriate in the SAE case, where you want each feature to have some context-independent meaning (and—unlike with AV-sampled natural-language tokens—the features don’t have pre-existing interpretations attached to them before training). But it seems awkward in token space, where some examples will inherently require longer descriptions than others to reach any given reconstruction-error threshold with equal fluency/legibility. “Describe what’s going on here in exactly 16 tokens” is more feasible in some cases than others[1], and if our loss encourages the model to do well at tasks like that across all examples, I’d worry that this would encourage steganography / unnatural-sounding word-pretzels / that sort of thing.
In the setup described in this post, the model has the freedom to use different lengths on different examples, “spending” more tokens in cases where the improvement in reconstruction error justifies the expenditure. The Matryoshka loss doesn’t adapt like this.
I wonder whether the following would work: introduce one or more new tokens into the vocab which the AV will use to demarcate the end of each “graded prefix,” so that the AV samples would look like:
Then, have a term in the loss for each end-of-prefix separator, where each term is a sum of “reconstruction loss using the text up to that separator” and “length penalty on the tokens between it and the previous separator.”
This is a generalization of what happens with the EOS token in the setup described in the post, where the AV can sample that token earlier or later in the sequence and makes this decision by trading off reconstruction and length-penalty. With the end-of-prefix separators, it’s doing the exact same thing, but once for each nested prefix we’re going to grade.
As a stylized example, consider two activations, each of which is representing a bunch of things which have well-defined names in natural language, but where:
in one case, the model’s tokenizer encodes those names very inefficiently
in the other, the model’s tokenizer encodes the names much more efficiently
If we had a combined loss with a length penalty as in the post, the AV might end up reciting all of the names in both cases even though that results in a longer text in the first case; we just need the marginal reconstruction-error benefit of including each name to exceed the marginal length-penalty cost.
However, if we grade the model on reconstruction based on the first N tokens—where N does not adaptively vary with the content of the examples—then N might be too small to list all the names in the first case (but not in the second), and this could result in unnatural-sounding (or just undesirably imprecise) verbalizations of examples like that first case.
Yeah, I think the Matryoshka loss is more “mostly applicable” rather than “directly applicable”.
With Matryoshka SAEs there are two types of prefixes. The first is the dictionary subsets, which in our implementation were fixed and chosen in advance, which I agree doesn’t make sense for an NLA. The second is that the SAE could have K active features, and it could allocate the K active features across subsets as desired. This seems close to your “graded NLA prefixes” proposal.
There are also alternative forms of hierarchy that could be explored. For example, in Appendix C.1 of the Matryoshka paper there was an alternative implementation where the dictionary subset sizes where randomly sampled from a distribution when performing inference, which encouraged a more continuous hierarchy.
Another random idea (which may not be good) is to just pass the first K tokens of the AV verbalization to the AR (where K is chosen randomly), which may encourage the AV to put the more important stuff earlier in the generation, with the fine-grained details towards the end. But it could also encourage weird compression.