[Question] Best way to reimplement T-SAE and related methods within SAELens

[xXCoolinXx]

Questions

T-SAE paper: https://openreview.net/pdf?id=bojVI4l9Kn (confusingly, this is different from Temporal SAE)

I am working on developing an architecture similar to but slightly different from the above paper. In the end, I would like to be able to implement my architecture as well as T-SAE as a baseline. The T-SAE architecture itself doesn't seem too hard to implement, just using an InfoNCE loss between tokens.

My big issue is figuring out how to properly modify the activation store class, in a way that is both general enough to accommodate other new architectures and also functional for both my architecture and T-SAE.

My main idea to implement the T-SAE activation store would be to return pairs of activations instead of just single activations, so that, during the training loop, one could use the pair as a positive example and the rest of the batch as a negative example.

I have already done a preliminary implementation in my fork of SAELens here: https://github.com/xXCoolinXx/SAELens/blob/context-sae/sae_lens/training/activations_store.py

The main changes are that I added a context_processor function argument to the activation store, that can be called before the full context activations are returned. For example, the near_pairs context processor can find pairs between different activations in the context. The issues I've had with this approach is that it doesn't work with concatenated batches, and I've additionally had to pass in a mask to avoid padding. It also gets a bit messy trying to work with the paired activations in the downstream SAE.

I would like to find a method that is general enough to meet this use case and also be suitable to commit the changes to the library. My question would then be, what is the preferred approach for this goal?

[chanind]

This is a bit tricky, I'm hesitant to build-in support for something like this into the base library since it adds a lot of complication and it only works for training a single type of SAE that may never become mainstream (e.g. if you were to pass the option to do this grouping for any other SAE it would break the SAE). It seems like a lot of the complication here comes from the fact that it's not enough to simply turn off activation shuffling, you still want to shuffle activations but keep pairs of activations together. You should already be able to disable concatenating sequences in the same batch by setting disable_concat_sequences=True.

I'd be open to the idea that you can pass in an arbitrary context_processor as a user-defined lambda function to the activation store, or maybe having a registry like in sae_lens.registry.py where users can register named context processors and associated configs to use in the LLM training runner (e.g. you could register a near_pairs context processor, and then set context_processor="near_pairs" in the runner config to enable it) - that could keep things flexible enough that it's not too hard for users to do whatever they want to LLM activations without needing to modify the activations store class itself and write their own runner. What do you think?