[Paper Review] BatchTopK Sparse Autoencoders

Bart Bussmann, Patrick Leask+1 more|arXiv (Cornell University)|Dec 9, 2024

Generative Adversarial Networks and Image Synthesis4 citations

TL;DR

This paper introduces BatchTopK Sparse Autoencoders, a method that relaxes the per-sample top-k activation constraint of standard TopK SAEs by selecting the top-k activations across an entire batch, enabling variable latent activation per sample. The approach improves reconstruction performance without sacrificing average sparsity and allows direct specification of sparsity levels, outperforming TopK and matching JumpReLU SAEs in reconstruction while avoiding costly hyperparameter tuning.

ABSTRACT

Sparse autoencoders (SAEs) have emerged as a powerful tool for interpreting language model activations by decomposing them into sparse, interpretable features. A popular approach is the TopK SAE, that uses a fixed number of the most active latents per sample to reconstruct the model activations. We introduce BatchTopK SAEs, a training method that improves upon TopK SAEs by relaxing the top-k constraint to the batch-level, allowing for a variable number of latents to be active per sample. As a result, BatchTopK adaptively allocates more or fewer latents depending on the sample, improving reconstruction without sacrificing average sparsity. We show that BatchTopK SAEs consistently outperform TopK SAEs in reconstructing activations from GPT-2 Small and Gemma 2 2B, and achieve comparable performance to state-of-the-art JumpReLU SAEs. However, an advantage of BatchTopK is that the average number of latents can be directly specified, rather than approximately tuned through a costly hyperparameter sweep. We provide code for training and evaluating BatchTopK SAEs at https://github.com/bartbussmann/BatchTopK

Motivation & Objective

• To address the limitation of fixed per-sample sparsity in TopK SAEs, which assumes all inputs require the same number of active latents regardless of complexity.
• To improve reconstruction performance without increasing average sparsity by allowing variable activation counts per sample based on input complexity.
• To enable direct specification of average sparsity levels, eliminating the need for costly hyperparameter sweeps common in other SAE variants.
• To maintain interpretability while enhancing reconstruction by leveraging batch-level activation selection.
• To evaluate the method on GPT-2 Small and Gemma 2 2B, comparing performance to TopK and JumpReLU SAEs using standard metrics.

Proposed method

• BatchTopK replaces the per-sample top-k selection with a batch-level top-k selection, where the top-k activations across all samples in a batch are selected for reconstruction.
• The method uses a differentiable relaxation of the top-k operation to allow gradient flow during training, enabling end-to-end optimization of the latent dictionary.
• During inference, the model estimates a global threshold θ as the average of minimum positive activations across batches, replacing the batch-dependent selection with a threshold-based activation function.
• The threshold θ is computed as the expected value of the smallest positive activation across multiple batches, ensuring consistent inference behavior.
• The architecture uses a standard SAE encoder-decoder setup with ReLU activation and L2 reconstruction loss, augmented with a sparsity penalty and auxiliary loss to prevent dead latents.
• The auxiliary loss in BatchTopK is identical to TopK SAEs, using a small-scale reconstruction of dead latents to maintain their utility.

Figure 1: On GPT-2 Small activations, BatchTopK largely achieves better NMSE and CE than standard TopK across different dictionary sizes, for a fixed number of active latents of 32 (Left). JumpReLU SAEs are omitted from this comparison as their L0 cannot be fixed to a value. For fixed dictionary siz

Experimental results

Research questions

• RQ1 Can relaxing the per-sample top-k constraint to a batch-level constraint improve reconstruction performance in sparse autoencoders?
• RQ2 Does allowing variable numbers of active latents per sample lead to better reconstruction without increasing average sparsity?
• RQ3 Can BatchTopK achieve performance comparable to state-of-the-art JumpReLU SAEs while enabling direct specification of sparsity levels?
• RQ4 How does the distribution of active latents per sample differ between BatchTopK and fixed-top-k methods?
• RQ5 Does the improved flexibility of BatchTopK lead to better downstream performance in terms of cross-entropy degradation?

Key findings

• BatchTopK SAEs achieve significantly lower normalized mean squared error (NMSE) than TopK SAEs across all tested dictionary sizes and sparsity levels on GPT-2 Small and Gemma 2 2B.
• For a fixed number of active latents (L0=32), BatchTopK outperforms TopK in both NMSE and cross-entropy (CE) degradation on both models.
• On GPT-2 Small, BatchTopK outperforms JumpReLU SAEs in both NMSE and CE degradation across all sparsity levels, despite JumpReLU’s superior downstream robustness in high-activation regimes.
• On Gemma 2 2B, BatchTopK matches or exceeds JumpReLU in reconstruction (NMSE) but only outperforms it in CE degradation at the lowest sparsity level (k=16).
• The distribution of active latents per sample in BatchTopK shows wide variation, with some samples using only one latent and others using over 80, confirming the method’s adaptive activation strategy.
• BatchTopK enables direct specification of average sparsity levels without hyperparameter sweeps, unlike JumpReLU SAEs, which require tuning of multiple parameters to achieve target sparsity.

Figure 2: On Gemma 2 2B activations, BatchTopK outperforms TopK SAEs across different values of k. Although BatchTopK has a better reconstruction performance (left), it only outperforms JumpReLU in terms of downstream CE degradation in the setting where k=16 (right).