Speeding Up JumpReLU SAE Inference with Custom Triton Kernels (2–14× on Real SAEs)
Motivation Sparse Autoencoders (SAEs) have become a central tool in mechanistic interpretability research, providing a way to decompose a model's internal activations into sparse, interpretable features. However, extracting these features often requires running the SAE over large volumes of activations across many layers and tokens. This makes SAE inference efficiency...
Jun 149