Jumping Ahead: Improving Reconstruction Fidelity with JumpReLU Sparse Autoencoders

Thank you for your thoughtful and detailed feedback. We're particularly encouraged by your recognition of the non-trivial theoretical insights we introduce regarding training SAEs with parameter-discontinuous loss functions. Nevertheless, we appreciate your concerns and would like to address them directly.

First, regarding overall performance and significance: Our evaluations show that JumpReLU SAEs consistently match or exceed the reconstruction fidelity of both Gated and TopK SAEs across all tested combinations of layers, sites and models (including Pythia 2.8B results in the Appendix G). In particular, whereas TopK improves upon Gated at some sites and layers but is beaten by Gated at other sites and layers (looking at Figures 2, 14 and 15 together), JumpReLU consistently provides the best reconstruction fidelity across all sites and layers we have tested. This consistency is valuable: it means practitioners can confidently choose JumpReLU SAEs knowing they will get state-of-the-art or better performance without needing to experiment with different architectures for different use cases. Moreover, we show that this reliable performance advantage does not come with downsides – JumpReLU SAEs are comparably interpretable to TopK and Gated SAEs, faster to train, have similar or fewer high frequency features than TopK, and perform better on our disentanglement evaluation.

Regarding the intuition behind JumpReLU versus TopK: JumpReLU addresses a fundamental limitation of TopK that is explicitly noted in their own paper. TopK forces exactly k features to be active at every token, which is unnatural – we would expect the number of relevant semantic/causal features to vary across tokens. JumpReLU, by using thresholds rather than counts, allows this natural variation while maintaining average sparsity through the L0 penalty. This more flexible approach to sparsity likely contributes to JumpReLU's consistent performance advantages.

Regarding ablation studies: Thank you for this valuable suggestion. We have now added ablations in Appendix H.2 (referenced from Section 5.1) that isolate our two main contributions: (a) JumpReLU activation with traditional L1 sparsity penalty, and (b) ReLU activation with our L0 sparsity penalty. Neither configuration matches the performance of the complete JumpReLU SAE with L0 penalty, demonstrating that both components are essential for the observed improvements. Regarding the suggested TopK + L0 ablation – we respectfully point out that this combination doesn't make sense conceptually since TopK already enforces sparsity directly through its k parameter. Adding a L0 penalty (or any other sparsity penalty) to TopK would be redundant.

On computational efficiency: While the complexity difference between O(k log n) and O(1) might seem modest, it becomes significant in practice, particularly when training large SAEs or conducting extensive hyperparameter searches. In our setup, TopK (with AuxK) is approximately 50% slower to train at K=70 than JumpReLU, even using an approximate top-K algorithm. Moreover, TopK's computational cost scales linearly with K, while JumpReLU remains constant; this is important to keep in mind when there is still no consensus on the optimal level of sparsity required for SAEs to be useful for downstream applications.

Regarding the manual interpretability study: We understand your concerns about potential bias in our interpretability ratings. However, we implemented rigorous controls to ensure objectivity: the study was conducted double-blind, with features randomly assigned across architectures and raters unaware of which features came from which SAE architecture during assessment. This methodology was specifically designed to eliminate the kind of bias you have identified as a concern.

In conclusion, we believe our work makes substantial contributions to the field of mechanistic interpretability through SAE training. Unlike TopK or Gated SAEs, the JumpReLU SAEs introduced in this paper achieve consistently superior reconstruction fidelity across models, sites and layers, doing so through a more flexible and computationally efficient sparsity mechanism. Furthermore, our theoretical insights connecting straight-through estimators to kernel density estimation open new possibilities for training models with discontinuous loss functions, with potential applications beyond SAEs.

Thank you for your thorough review of our paper, we are heartened by your positive comments on the paper’s clarity and technical correctness. We have carefully considered the weaknesses you have listed and believe we can effectively address both your primary concerns: (a) that you feel our testing has been too selective, and (b) that the contribution of this paper is too incremental.

On (a), the selectiveness of our testing, we believe this is straightforward to address. Firstly, we do indeed compare JumpReLU against TopK and Gated SAEs on another model, Pythia 2.8B, in Appendix G (and signposted in the Conclusion section) precisely to address the concern that our results may not generalize beyond Gemma 2. Secondly, for each tested layer, we conduct evaluations across three different activation sites (residual stream, MLP output, and attention output), showing the JumpReLU SAEs consistently provide the most faithful reconstructions across multiple sites, i.e. not just for residual stream activations. Finally, our choice of layers 9, 20 and 31 of Gemma 2 9B was made because these are evenly spaced layers roughly 1/4, 1/2 and 3/4 through this 42-layer model, i.e. to check that our results generalise to different layers in the model. We have added text at the start of Section 5 to clarify why we chose these layers in the paper. Most importantly, we emphasize that we chose these to evaluate the SAEs on these layers prior to running the experiments, i.e. we did not cherry-pick these layers to improve our results.

Regarding (b) your concern about incremental contribution, we believe our work represents a meaningful advance in both theoretical and practical aspects of the field. Theoretically, we introduce a principled approach to training through discontinuous loss functions using straight-through estimators, which has broader implications beyond just SAEs. Empirically, our extensive evaluations demonstrate that JumpReLU SAEs deliver superior or equivalent reconstruction fidelity compared to existing approaches across all tested scenarios – spanning different layers, activation sites, and model architectures. Although there may be some sites and layers where the difference between JumpReLU and TopK SAEs isn’t massive, the key strength of our approach lies in its consistency and reliability: as Figures 2, 8, 14 and 15 show, JumpReLU SAEs reliably provide the best reconstruction fidelity across models, sites and layers. This cannot be said of TopK and Gated SAEs, which trade places between different sites and layers. This reliability of JumpReLU SAE performance is particularly valuable from a practical standpoint: researchers can adopt JumpReLU SAEs with confidence, knowing they will achieve at least state-of-the-art performance without the need to experiment with different architectures for different applications. Furthermore, these improvements come with significant practical benefits – JumpReLU SAEs are more efficient to train than alternatives (avoiding both the auxiliary terms needed for Gated SAEs and expensive TopK operations), demonstrate comparable or better interpretability metrics, and show stronger performance on disentanglement tasks. We believe this combination of theoretical insight, reliable performance improvements, and practical advantages represent a meaningful contribution to the problem of training faithful SAEs to decompose LM activations, a problem that has received much attention this year in the field of mechanistic interpretability.

Addressing your specific questions and suggestions:

• Improvements to figures Thank you for your feedback, We have updated all figures accordingly and added an example of the dashboards used in the interpretability study (Figure 12).

• Why is JumpReLU restricted to a positive threshold? While negative thresholds are possible, we follow the convention from Bricken et al. (2023) of looking for non-negative linear combinations of features, stemming from the superposition hypothesis and empirical results suggesting LMs could use negative feature directions to represent unrelated features. Exploring negative thresholds could nonetheless be interesting future work.

• The preliminary definition of SAE is incorrect: We respectfully disagree. While traditional autoencoders often reduce dimensionality, all work on LM SAEs since their introduction has used an overcomplete basis (m>>n) to discover interpretable features. For example, our experiments use m=131,072 features for n=3,584 dimensional activations.

• Presentation of algorithms as pseudocode: While we appreciate algorithms are typically better in CS-style pseudocode, in this case we believe the implementation details are more valuable. The mathematical formulation is already specified in Equations 8-12. The code specifically demonstrates how to implement custom gradients in an autograd framework – a crucial detail that would be obscured by traditional pseudocode notation.

Thank you for your thorough and insightful review of our paper. We particularly appreciate your recognition of our paper's technical contributions, including both the state-of-the-art performance and the principled approach to training with discontinuous loss functions.

We acknowledge and appreciate the constructive criticisms you have raised. Below we address each of your main points and questions:

• Downstream evaluations. While we would certainly have liked to include more evaluations, like the two that you suggested, our preliminary experiments suggested that these two specific evaluations lack the resolution needed to meaningfully distinguish between SAE architectures (although they can be useful, for example, for distinguishing between SAEs of differing widths within a single architecture, or e.g. for showing that SAEs of any type provide better decompositions than PCA or the neuron basis). We consider it to still be an open problem how to evaluate SAEs on downstream tasks in a way that provides sufficient resolution to distinguish between architectures, a problem which we wish to tackle in future work.
• Hyperparameter introduction. While we agree that having a principled way to select ε would be valuable, we note that this parameter has a clear interpretation as a kernel bandwidth parameter (or, in the case of a rectangular kernel, also as a finite difference step-size parameter), making it more theoretically grounded than, for example, the auxiliary loss terms required by Gated SAEs, and hence easier to reason about when performing hyperparameter selection. Moreover, we strongly suspect it should be possible to develop automatic bandwidth selection methods by building on the extensive literature on bandwidth selection for kernel density estimation. (With reference to your question about alternative formulations of the pseudo-derivative, we cannot think of any formulation that doesn’t involve some sort of bandwidth parameter that controls the bias-variance trade-off that is inevitably required for estimating gradients in expectation, but are open to suggestions.) However, we did not prioritize further investigation along these lines after finding that, after applying dataset normalization, our specific choice of hyperparameters seems to transfer well across different models, sites and layers (and moreover, training is relatively insensitive to these hyperparameters). Nevertheless, we do think this would be an interesting line of further work to improve the JumpReLU training methodology.
• High frequency features. We respectfully point out that we in fact find that JumpReLU SAEs typically have slightly fewer high frequency features than TopK SAEs (Figure 4), although we agree with the general point that the presence of such high frequency features (in both JumpReLU and TopK SAEs) is undesirable. Regarding your suggestion about resampling high-frequency features – while this is an interesting direction, we suspect that unlike dead features, high-frequency features often play an important role in reconstruction fidelity, making their resampling during training potentially disruptive. However, we believe there may be more promising approaches to addressing this issue through modifications to the sparsity penalty. Indeed, this highlights another advantage of JumpReLU SAEs over TopK SAEs – our approach allows us to explore different potentially discontinuous sparsity penalties beyond L0 (cf. Appendix F) to incorporate additional desirable properties. We are actively investigating this direction in ongoing work.

Thank you for your detailed feedback, which has helped us identify important areas for improvement. We have made concrete changes to address your concerns about motivation, ablation studies, and technical presentation, as detailed below.

Addressing your questions first:

• What are the motivations behind the JumpReLU activation? The motivation for replacing ReLU with JumpReLU is illustrated in Figure 1, which shows why this change is necessary: ReLU cannot effectively separate active from inactive features without introducing false positives or underestimating activation magnitude (shrinkage). When a feature's pre-activations are high where the feature is present but low (but not always negative) where the feature is absent, ReLU either allows false positives or requires bias adjustments that systematically underestimate feature magnitudes. JumpReLU solves this by providing a threshold that screens out false positives without affecting magnitude estimation. We have added a sentence to the introduction, signposting that the motivation behind the JumpReLU activation is explained in Figure 1.

• Can the authors conduct the ablation study as stated in the weakness? We have added an ablation study in Appendix H.2 (referenced from Section 5.2 in the main text) that separates our two main changes. Specifically, we evaluate:

  1. Using JumpReLU activation while keeping the L1 sparsity penalty from prior work
  2. Keeping the ReLU activation while introducing our L0 sparsity penalty

  Neither ablation matches the performance of JumpReLU SAEs trained with a L0 penalty, demonstrating both components are necessary for the improvements shown.

• Can the author give more context about the problem to solve, for example, why we need the sparse autoencoder to evaluate LM? Our Introduction and Related Work sections provide extensive context for SAEs in LM interpretability. Specifically, we explain that SAEs help identify interpretable directions in LM activations, which are valuable for tasks like circuit analysis and model steering. We also discuss how many concepts appear to be linearly represented in LM activations, making dictionary learning approaches like SAEs particularly promising for understanding and controlling these models. These sections cite the key foundational works that established the value of SAEs for LM interpretability. Since our paper's contribution is a methodological improvement to SAE training rather than introducing SAEs for LMs, we believe this level of context appropriately situates our work while allowing us to focus on our technical contributions.

Additionally, to address the weaknesses you have listed (not already covered above):

• Notation. Regarding notation, we acknowledge that $L^p$ is standard mathematical notation. However, our use of L0 and L1 follows established convention in the language model SAE literature (e.g. Bricken et al, 2023; Templeton et al, 2024) where this notation is widely used and understood by the community. For Equation 1, we use $\\mathbf{f}$ to denote both the function and its output for clarity of exposition and to maintain consistency with prior work.
• The idea of using JumpReLU activation for SAEs is already presented in Gated SAEs. While it is true that Gated SAEs can be shown to implement a JumpReLU activation function (as noted in the Gated SAEs paper), our key contribution is not the activation function itself but rather introducing a principled method for training SAEs with discontinuous activation functions. Instead of relying on auxiliary loss terms to provide indirect gradients for training the threshold (as Gated SAEs do), we derive and directly minimize the expected loss of the SAE, despite the discontinuous nature of the JumpReLU function. As our evaluation shows, this more principled training approach leads to consistently better reconstruction fidelity than Gated SAEs at any given level of sparsity.
• The performance improvement of the method is not evident. Our evaluations show that JumpReLU SAEs consistently match or exceed the reconstruction fidelity of both Gated and TopK SAEs across all tested combinations of layers, sites and models (including Pythia-2.8B results in the appendix). While TopK sometimes beats Gated and sometimes loses to Gated (considering all sites and layers shown in Figures 2, 14 and 15), JumpReLU reliably provides the best sparsity-fidelity trade-off regardless of where in the network it is applied. This consistency is valuable: it means practitioners can confidently choose JumpReLU SAEs knowing they will get state-of-the-art or better performance without needing to experiment with different architectures for different use cases. Moreover, we show that this reliable performance advantage does not come with downsides – JumpReLU SAEs are comparably interpretable to TopK and Gated SAEs, faster to train, have similar or fewer high frequency features, and perform better on our disentanglement evaluation.

Before this phase of the rebuttal period ends, we wanted to ask the reviewer whether we have addressed your concerns with our work?