Jumping Ahead: Improving Reconstruction Fidelity with JumpReLU Sparse Autoencoders
← Back to all posts
Thread (53 posts)
JumpReLU SAEs only require a single forward and backward pass during a training step and have an elementwise activation function (unlike TopK, which requires a partial sort), making them more efficient to train than either Gated or TopK SAEs.
Can you elaborate on the training efficiency point here in more detail? Only focus comparison on TopK
Claude 4 Opus (T) via @hdevalence.bsky.social
Thinking Summary
The user is asking me to elaborate on the training efficiency comparison between JumpReLU and TopK SAEs, specifically focusing on TopK. The paper mentions that JumpReLU is more efficient to train than TopK, and I should explain why. From the paper, the key points are: 1. JumpReLU uses a single forward and backward pass 2. JumpReLU has an elementwise activation function 3. TopK requires a partial sort operation Let me think about the computational complexity implications: - TopK requires finding the top K values among M features for each input - This is typically O(M log K) or O(M) with partial sorting algorithms - JumpReLU just applies an elementwise threshold comparison, which is O(M) The paper mentions they use `jax.lax.approx_max_k` for their TopK implementation, which is an approximate algorithm that's faster than exact TopK but still more expensive than elementwise operations.
The training efficiency advantage of JumpReLU over TopK SAEs comes down to the fundamental difference in their activation mechanisms:
TopK's computational overhead:
- Partial sorting requirement: For each input, TopK must identify and retain only the K largest pre-activations out of M total features (where M is the dictionary size, e.g., 131k in the paper's experiments)
- Complexity: Even with efficient algorithms, this requires O(M log K) operations for exact TopK, or O(M) for approximate methods
- Non-local operation: The activation of any feature depends on the values of all other features, requiring global information
JumpReLU's efficiency:
- Simple elementwise comparison: Each feature activation is determined by a single comparison:
if pre_activation > threshold - Complexity: O(M) with a very low constant factor - just M comparisons
- Local operation: Each feature's activation is independent of others
Practical implications:
The paper mentions using jax.lax.approx_max_k (based on Chern et al., 2022) for TopK implementation, which is faster than exact sorting but still involves:
- Multiple passes through the data to approximate the top-K selection
- Additional memory access patterns that are less cache-friendly
- More complex gradient computation in the backward pass
For a 131k-width SAE trained on billions of tokens, these differences compound significantly. The elementwise nature of JumpReLU also makes it more amenable to parallelization and hardware acceleration compared to the inherently sequential nature of sorting operations.
This efficiency gain is particularly important when training multiple SAEs with different sparsity levels or when scaling to larger models and dictionary sizes.