JumpReLU SAE training on the roadmap?

saprmarks / dictionary_learning

MIT License

151 stars 40 forks source link

JumpReLU SAE training on the roadmap? #19

Open ejmichaud opened 3 months ago

ejmichaud commented 3 months ago

Just curious if anyone is thinking about implementing a training pipeline for JumpReLU SAEs! They have a couple of properties which are really desirable for something I'm working on.

ejmichaud commented 3 months ago

Ah whoops I see the jump_trainer branch now and that @canrager is working on it :)

adamkarvonen commented 3 months ago

This was just merged: https://github.com/saprmarks/dictionary_learning/pull/22

I'm not sure how well tested it is though. Let us know if you have any problems / how it works for you.