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.
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.