Sparse Autoencoders, Again?

Thanks for the comprehensive review of our work, particularly with respect to checking proof details and pointing out multiple valuable corrections.

Comment: The benchmarks seems sensible ... but the only comparison made is against other sparse autoencoder models.

Response: Actually, we also compare against diffusion in Section 5.3 which is definitively not an SAE. Additionally, we compare extensively against vanilla VAEs, which have not historically been considered as SAEs per se. Still, the predominant focus of our work is indeed on improving upon SAE models.

Comment: I did not check the experiments in detail ... Went through the proofs in appendix F...

Response: Thanks for checking the proofs, this is a great help. The diversity and complementary nature of our experiments also serve as a key part of our contribution.

Comment: (+)The idea for this paper is well motivated and the paper is well structured (+) The novel ideas in this paper warrant exploration and I believe this to be a good direction for research.

Response: We appreciate acknowledgement of the novelty of our work, and indeed no prior paper has repurposed and analyzed VAE models in the way that we have.

Comment: (-) The overall contribution in this paper feels a bit shallow ... comparisons against other sparse autoencoder models only ...

Response: We emphasize that VAEs represent one of the highest cited, influential ML frameworks of the last 10 years, and yet they have never been previously applied or analyzed in the context of SAE tasks as we have done. Nor have the specific attributes of adaptive sparsity, local minima smoothing, and guaranteed maximally sparse global optima been cohesively addressed in the context of widely-used SAE models. Still, as a contribution to an ML conference we do not match the mathematical depth or generality of a pure theoretical contribution on manifold geometry, etc. Instead, we produce a brand new VAE use case, that requires no additional model complexity, and elucidate unique attributes shared by no prior approaches. We also offer extensive supporting experimentation, including comparisons with diffusion models beyond SAEs alone.

Comment: (-) The theoretical contributions need clarification ...

Response: Please see our specific responses elsewhere. However, broadly speaking, our results establish that the proposed VAEase energy function uniquely and simultaneously achieves two desirable properties: (i) The global minima exactly reflect union-of-manifold structure in data, and (ii) the loss exhibits a selective smoothing mechanism capable of reducing local minimizers. No prior SAE or VAE approach achieves both; provably so in the sense described in Section 4 for the wide classes we consider as baselines.

Comment: Question about Definition 4.1 and and whether stated conditions can be relaxed, etc.

Response: Great suggestion. The issue here relates to an inadvertent typo in our submission, where we used $\mapsto$ instead of $\rightarrow$ in Definition 4.1. Crucially, we do not require the image set in $\mathbb{R}^r$ to be an open ball. This implies: (i) the manifold's image can be an arbitrary topology as long as a global diffeomorphism exists; (ii) manifolds with a boundary remain valid objects of study as the definition focuses on the dimensionality of image set rather than imposing openness constraints.

Comment: Regarding Definition 4.3, do we really need $\sigma_z(x;\phi_\gamma) = \Theta (1)$?

Response: In the proof from Section F.1 (Line 1003), the correct asymptotic requirement should be $\sigma_z^2(x;\phi_\gamma) = O(\gamma)$ for active dims or $1 - \sigma_z^2(x;\phi_\gamma) = O(\gamma)$ for inactive dims to eliminate constants in the inequality. We will clarify the revision accordingly; thanks for pointing out this oversight in our original draft.

Comment: The notation is not quite clear ... specifically brackets vs parentheses in equations (2) and (4).

Response: We adopt the convention (shared by some ML works) that parentheses and brackets are interchangeable, and are chosen purely for making nested cases more readable, i.e., to avoid multiple parentheses or brackets in a row that can be visually difficult to parse. But this is obviously a minor stylistic choice that can be changed if the reviewer prefers.

Comment: Why is 4.6 called a Corollary ...

Response: Corollary 4.6 involves analysis closely related to Theorem 4.5, and the final conclusions are complementary. Still, we can easily change this designation if preferable.

Comment: The first paragraph in F.1 is not clear, etc...

Response: Great catch. Indeed, the original text was backwards; we will correct it to state that we first construct a feasible solution and then analyze to form a necessary condition for optimality.

Thanks for the quite comprehensive and accurate summary of our work. Indeed this description captures the essence of what we hoped to convey.

Comment: ...additional empirical validation of language model data (e.g., on a more recent, larger model) would make the paper stronger.

Response: We agree that further investigation into LLM applications of our work is desirable. However, because we are not domain specialists w.r.t. LLM interpretability in particular (which involves a number of qualitative investigations), we have deferred this application-specific dimension to future work. We also remark that prior SAE papers devoted to LLMs have proposed the same metrics that we adopt herein for quantification purposes.

That being said, as our emphasis is on diversity of application, especially very recent ones, we have added an additional LLM-driven experiment whereby the downstream goal is to generate hypotheses regarding relationships between text data (e.g., headlines) and a target variables (e.g., clicks); please see https://arxiv.org/pdf/2502.04382. While we defer details to this prior paper (no affiliation to us), the key optimization task involves learning sparse representations of text embeddings from Yelp reviews. Results are shown in the new table below, where we observe that VAEase again requires the fewest average active dimensions (AAD) even while maintaining the best reconstruction error (RE). This further expands the diversity of scenarios whereby VAEase outperforms existing alternatives. Thanks for the suggestion.

Comment: What does the reconstruction error/sparsity trade-off for language model data look like? A visualization similar to Figure 3 (possibly in the appendix) would be appreciated.

Response: Good question. While Openreview does not allow us to upload new figures, we have run this experiment and the trend closely matches that from Figure 3. This further solidifies the stable advantage of VAEase across quite distinct domains. And such a figure can easily be added to a revised appendix as the reviewer suggests.

We are appreciative of the detailed, constructive comments, and for pointing out the high potential of our work being used by those interested in sparse autoencoders (which is our intended audience). We address main reviewer points in turn below.

Comment: As far as I understand, the masking of active dimensions in the experiment illustrated in Fig. 3 follows the partitioning criteria described in lines 367--384 ("throughout this paper" is stated therein). Would it not be fairer, when viewing this against reconstruction error, to mask dimensions in order of least to most disturbing w.r.t. the reconstruction error? In other words, sequentially remove the next active dimension that would perturb the average reconstruction loss the least (maybe different dimensions are picked per model).

Response: Actually, the shape of Figure 3 does not appreciably change when enacting the reviewer's suggestion. This is because the magnitude of $\mu_z$ values for SAE models, or $\sigma_z$ values for VAE models, already closely reflect their importance w.r.t. reconstruction error by design. Moreover, using these magnitudes better reflects practical use cases, whereby we would ideally like to assess active dimensions without having to compute $O(\kappa^2)$ separate reconstruction errors.

Comment: In lines 420--423: This sentence in the analysis is faulty in my view, since their own method estimates $r \approx 15$ on Pseudo-MNIST. Since the setup upper bounds the ground-truth $r$, a simple (but very laborious) test could follow this setup from 16 dimension down to one dimension, stopping whenever the whole setup degenerates.

Response: Our intended point on Lines 420-423 and the surrounding text is the following: By design of the experiment, we know that 16 is a strict upper bound on the ground-truth manifold dimension. Hence any method that produces an estimate above 16 must be incorrect. In particular, both diffusion approaches in Table 4 produce estimates vastly exceeding 16. Meanwhile our VAEase approach produces an estimate below 16 despite no knowledge of the generative process. Hence the VAEase estimate is at least plausible, even while conceding that the true manifold dimension might still be a bit smaller (which would also be even further from the diffusion estimates). Of course intuition suggests it can't be dramatically smaller, as visual inspection of human-written MNIST digits indicates multiple non-trivial degrees of freedom.

Overall though, our main objective here is to show the superiority of VAEase relative to diffusion, not exactly pinpoint any exact manifold dimensions associated with MNIST. The reviewer's suggestion would indeed help with the latter, but this is computationally intensive and involves subjectively deciding when GAN samples deviate too far from MNIST. Still, the reviewer's comment is well-noted, and we will adjust the writing to ensure our desired points are more clearly conveyed.

Comment: I reviewed most of the supplementary material. It was surprising to me how the quality of the writing is considerably worse in the appendix sections compared to the main paper.

Response: Admittedly in the rush to submit, more attention was devoted to polishing the main text, which the reviewer agrees is well written and organized. For the revision we can definitely remedy this imbalance.

Comment: A worthy mention is perhaps the applicability of sparse autoencoders in general, coupled with the size of the change proposed to VAEs. While properly mathematically substantiated and experimentally verified, the main contribution is still arguably quite straightforward and that could be viewed as its weakness.

Response: With regard to SAEs, we believe their relevance is increasing along multiple fronts. For example, as we discuss in Section 5 and elsewhere, there is rapidly growing usage w.r.t. learning interpretable representations of LLM activation layers. Moreover, other representative SAE examples spanning vision and language are frequently appearing in the literature, e.g., as in https://arxiv.org/pdf/2502.04382, https://arxiv.org/pdf/2502.06755v1, and https://arxiv.org/pdf/2410.06981. And to better reflect this trend, we have added additional related experiments; please see our rebuttal response to Reviewer WZdZ for details.

Secondly, although our VAEase proposal is admittedly simple, we view this as a fundamental strength as it is easy for anyone to apply, portending larger downstream impact. Moreover, the underlying novelty is evidenced by 10 years of extensive VAE usage without prior work ever uncovering, let alone rigorously analyzing, the highly-effective modification we have proposed.

Comment: Some typos/mistakes ... Minor comments

Response: Thanks for pointing these out, it is extremely helpful for improving the revision. We will definitely correct each of them.

Thanks for acknowledging the many positive aspects of our work, including the novel design, comprehensive proofs, solid theoretical foundation, well-designed experiments, and the robustness and generalizability of our proposed model. We also appreciate the reviewer's statement that there is clear and convincing evidence of the claims made in our submission.

Comment: The proposed method make sense for the problem, but I do not find the advantages compared with traditional VAE.

Response: A traditional VAE model cannot readily learn adaptive sparse solutions (meaning solutions whereby the locations of informative latent dimensions vary from sample to sample). In fact, we rigorously prove this VAE limitation via Corollary 4.6, and verify it empirically; see for example Table 2, where the VAE is incapable of learning the correct manifold structure which requires adaptive sparsity to do so.

Comment: The benchmark dataset and evaluation metrics in this paper are not popular in the field of autoencoder research.

Response: SAE models are commonly applied to a wide variety of different tasks, such that we are not aware that there is any particular most popular benchmark. That being said, arguably the most high-profile recent application is to the analysis of LLM activation layers, and we have included such experimentation in our paper. We have also added another new application example to the rebuttal; please see our response to Reviewer WZdZ. Moreover, the metrics we adopt closely mirror prior work as well.

Comment: 1. ... Testing on more diverse image datasets like ImageNet or COCO.

Response: While we agree further testing can always be valuable, our focus is not specifically on computer vision tasks, which is why our submission also includes multiple non-image examples. Moreover, we do not commonly see SAE models applied to ImageNet or COCO. Still, application to broader image classes is a useful suggestion for future work.

Comment: 2. Although the paper claims improvements over traditional VAEs, reported gains might not be substantial enough to justify the additional complexity introduced by the VAEase model ...

Response: We respectfully disagree with this point for multiple key reasons. First and most importantly, our proposed VAEase approach does not introduce additional complexity over a regular VAE. Instead, it merely reuses the exact same encoder variances as a novel gating mechanism; hence no further justifications are needed with regard to complexity. And secondly, for a fixed reconstruction error, VAEase does in fact produce much greater sparsity across a wide range of tasks as desired. Of course the reported reconstruction errors are often similar, but this is specifically part of the experimental design to facilitate meaningful sparsity comparisons. Along these lines, we also note that under listed strengths, the reviewer specifically mentioned that our experiments are "well-designed and demonstrate clear improvements of VAEase over existing methods."

Comment: 3. The paper emphasizes the theoretical and experimental aspects of VAEase but lacks a clear discussion on its practical applications ... benefits of adaptive sparsity and reduced reconstruction error are not directly linked to specific use cases.

Response: Given the vast literature involving SAE models, we largely defer treatment of particular practical applications to prior work. That being said, we do link the benefits of adaptive sparsity and reduced reconstruction errors to the specific use case of LLM activation interpretability. In particular, to complement more subjective/qualitative analysis of SAE-learned representations, we follow common practice in the LLM literature that advocates specifically for comparing sparsity levels at a shared reconstruction error, with greater sparsity (meaning adaptive sparsity from sample to sample as our VAEase achieves) being associated with better interpretability. Even so, we are happy to include further references to this effect (although space in our submission is tight, so there is quite limited room for much additional discussion).

Comment: 4. ...The authors could benefit from a more concise summary of the key innovations and their implications, rather than focusing solely on detailed proofs.

Response: As an important clarification, our work does not focus solely on detailed proofs. In fact, all proofs are exclusively deferred to the very end of our paper deep in the appendices, and certainly not treated as a central focus. And for reference, we concisely list our central innovations and their consequences on the bulleted paragraphs beginning on Line 60 as well as pointers to other sections therein.