Taming Polysemanticity in LLMs: Provable Feature Recovery via Sparse Autoencoders

Abstract: We study the challenge of achieving theoretically grounded feature recovery using Sparse Autoencoders (SAEs) for the interpretation of LLMs. Existing SAE training algorithms often lack rigorous mathematical guarantees and suffer from practical limitations such as hyperparameter sensitivity and instability. To address these issues, we first propose a novel statistical framework for the feature recovery problem, which includes a new notion of feature identifiability by modeling polysemantic features as sparse mixtures of underlying monosemantic concepts. Building on this framework, we introduce a new SAE training algorithm based on ``bias adaptation'', a technique that adaptively adjusts neural network bias parameters to ensure appropriate activation sparsity. We theoretically \highlight{prove that this algorithm correctly recovers all monosemantic features} when input data is sampled from our proposed statistical model. Furthermore, we develop an improved empirical variant, Group Bias Adaptation (GBA), and \highlight{demonstrate its superior performance against benchmark methods when applied to LLMs with up to 1.5 billion parameters}. This work represents a foundational step in demystifying SAE training by providing the first SAE algorithm with theoretical recovery guarantees, thereby advancing the development of more transparent and trustworthy AI systems through enhanced mechanistic interpretability.

• The paper presents a novel Group Bias Adaptation (GBA) algorithm that provably recovers monosemantic features from polysemantic LLM activations.
• It introduces a rigorous statistical framework with ε-identifiability to ensure reliable feature recovery under targeted sparsity and balance conditions.
• Empirical results on Qwen2.5-1.5B show that GBA achieves superior sparsity-loss trade-offs, robustness, and consistent feature discovery.

This paper introduces a novel approach to training Sparse Autoencoders (SAEs) for LLMs with the goal of recovering interpretable, monosemantic features from polysemantic activations. The authors address key limitations of existing SAE training methods, such as the lack of theoretical guarantees for feature recovery, sensitivity to hyperparameters, and training instability.

The core contributions include:

1. A new statistical framework for feature recovery, modeling polysemantic activations ($x$) as sparse linear combinations of underlying monosemantic features ($V$) with non-negative coefficients ($H$), i.e., $x \approx HV$. This framework introduces a rigorous notion of feature $\epsilon$-identifiability, accounting for ambiguities like permutation, scaling, and feature splitting.
2. A novel SAE training algorithm called Group Bias Adaptation (GBA). GBA directly controls neuron activation sparsity by adaptively adjusting bias parameters ($b_m$) for groups of neurons, aiming to meet pre-defined Target Activation Frequencies (TAFs) for each group.
3. Theoretical proof that a simplified version of GBA (Modified Bias Adaptation) can provably recover all true monosemantic features when input data is sampled from their proposed statistical model under specific conditions on network width, bias range, and feature balance.
4. Empirical demonstration of GBA's superior performance on LLMs up to 1.5 billion parameters (Qwen2.5-1.5B), achieving a better sparsity-loss trade-off and learning more consistent features across runs compared to L1 regularization and TopK activation methods.

Statistical Framework and Feature Identifiability

The paper models an LLM's internal activation vector $x \in \mathbb{R}^d$ as a sparse, non-negative linear combination of $n$ monosemantic feature vectors $v_i \in \mathbb{R}^d$ (rows of $V \in \mathbb{R}^{n \times d}$): $V$0, where $V$1 is the coefficient matrix for $V$2 data points, with each row being $V$3-sparse. The goal is to recover $V$4.

To address the inherent non-uniqueness of this factorization, $V$5-identifiability is defined. A feature matrix $V$6 is $V$7-identifiable if any alternative factorization $V$8 implies that $V$9 is equivalent to $H$0 up to permutation, feature splitting (where a feature in $H$1 is a positive linear combination of features in $H$2), and small cosine similarity deviations bounded by $H$3. Theorem 5.3 states that under certain conditions on $H$4 (row-wise sparsity, non-degeneracy, low co-occurrence $H$5) and $H$6 (incoherence), $H$7 is $H$8-identifiable with $H$9.

Group Bias Adaptation (GBA) Algorithm

The GBA algorithm aims to overcome the limitations of traditional sparsity-inducing methods like L1 regularization (which causes activation shrinkage) and TopK activation (which can be sensitive to initialization).

Key Ideas:

1. Bias Adaptation: Instead of an explicit sparsity penalty in the loss function, GBA directly controls the activation frequency of each neuron. The bias $x \approx HV$0 of a neuron $x \approx HV$1 (in pre-activation $x \approx HV$2) is adjusted periodically.
   • If a neuron activates too frequently (actual frequency $x \approx HV$3 TAF $x \approx HV$4), its bias $x \approx HV$5 is decreased.
   • If a neuron activates too rarely ($x \approx HV$6), its bias $x \approx HV$7 is increased.
2. Neuron Grouping: Neurons are divided into $x \approx HV$8 groups, each assigned a different TAF ($x \approx HV$9). TAFs are typically set in an exponentially decaying sequence (e.g., $\epsilon$0). This allows the SAE to capture features with varying natural occurrence frequencies.

Algorithm 1: Group Bias Adaptation (GBA)

1. Input: Data $\epsilon$1, initial SAE parameters $\epsilon$2, neuron groups $\epsilon$3, optimizer Opt.
2. Hyperparameters: Iterations $\epsilon$4, batch size $\epsilon$5, buffer size $\epsilon$6, bias adaptation rates $\epsilon$7, rarity threshold $\epsilon$8.
3. Initialize buffers $\epsilon$9 for each neuron $b_m$0.
4. For $b_m$1:

   a. Sample mini-batch $b_m$2, normalize rows. b. Compute pre-activations $b_m$3. c. Compute reconstruction loss $b_m$4. d. Update SAE parameters $b_m$5 (except biases $b_m$6) using Opt. e. Add pre-activations $b_m$7 to buffers $b_m$8. f. If $b_m$9 (buffer full):

   i. Update biases $x \in \mathbb{R}^d$0 using Subroutine $x \in \mathbb{R}^d$1 (Algorithm 2). ii. Empty all buffers $x \in \mathbb{R}^d$2.

5. Return final SAE parameters $x \in \mathbb{R}^d$3.

Algorithm 2: GBA Subroutine $x \in \mathbb{R}^d$4 (Bias Adaptation)

1. Input: Current biases $x \in \mathbb{R}^d$5, buffers $x \in \mathbb{R}^d$6, groups $x \in \mathbb{R}^d$7, hyperparameters $x \in \mathbb{R}^d$8.
2. For each neuron $x \in \mathbb{R}^d$9: a. Compute activation frequency $n$0. b. Compute max pre-activation $n$1.
3. For each group $n$2: a. Compute average max pre-activation for active neurons in group $n$3.
4. For each group $n$4 and each neuron $n$5: a. If $n$6: $n$7. b. If $n$8: $n$9.
5. Return updated biases $v_i \in \mathbb{R}^d$0.

The biases are clamped to $v_i \in \mathbb{R}^d$1 to maintain sparsity and prevent over-sparsification. Efficient implementation involves iteratively updating $v_i \in \mathbb{R}^d$2 and $v_i \in \mathbb{R}^d$3 rather than storing all pre-activations.

Theoretical Recovery Guarantees

For a simplified "Modified BA" algorithm (single neuron group, fixed bias $v_i \in \mathbb{R}^d$4 implying a fixed TAF $v_i \in \mathbb{R}^d$5, smooth ReLU-like activation, and vanishingly small output scales $v_i \in \mathbb{R}^d$6), Theorem 6.1 provides provable feature recovery. It states that if:

• The data $v_i \in \mathbb{R}^d$7 is decomposable with i.i.d. Gaussian features $v_i \in \mathbb{R}^d$8.
• Network width $v_i \in \mathbb{R}^d$9 is sufficiently large: $V \in \mathbb{R}^{n \times d}$0, where $V \in \mathbb{R}^{n \times d}$1 relates to coefficient concentration.
• Bias $V \in \mathbb{R}^{n \times d}$2 is in a specific range, implying a TAF $V \in \mathbb{R}^{n \times d}$3 such that $V \in \mathbb{R}^{n \times d}$4. This range depends on the superposition regime ($V \in \mathbb{R}^{n \times d}$5 vs $V \in \mathbb{R}^{n \times d}$6).
• A "Feature Balance" condition holds (all features appear sufficiently often with sufficiently large coefficients).

Then, Modified BA recovers all monosemantic features $V \in \mathbb{R}^{n \times d}$7 with high probability. The proof involves showing good initialization, approximately Gaussian pre-activations via Gaussian conditioning, and analyzing the dynamics of weight alignment using Efron-Stein inequalities.

Empirical Results

Experiments were conducted on the Qwen2.5-1.5B model, extracting MLP activations from layers 2, 13, and 26 on Pile Github and Wikipedia datasets. JumpReLU was used as the activation function for all methods.

• Sparsity-Loss Frontier: GBA performs comparably to the best TopK (post-activation sparsity) and significantly outperforms L1 regularization and a non-grouped Bias Adaptation (BA) baseline.
• Hyperparameter Robustness: GBA is nearly tuning-free. Its performance is robust to the number of groups $V \in \mathbb{R}^{n \times d}$8 and the specific TAFs, provided the Highest TAF (HTF) is adequately high (e.g., 0.1-0.5) and $V \in \mathbb{R}^{n \times d}$9 is sufficiently large (e.g., 10-20). This is a significant practical advantage.
• Feature Consistency: Measured by Maximum Cosine Similarity (MCS) across runs with different random seeds, GBA learns significantly more consistent features than TopK. L1 is generally more consistent, but GBA surpasses L1 for the most active (top 0.05%) features.
• Feature Analysis: Scatter plots of Z-scores vs. other metrics (max activation, activation fraction, MCS) and a feature dashboard example show that GBA learns sparse, selective, and consistent features. For example, high Z-score GBA neurons often correspond to specific, infrequent concepts and show high MCS.

Implementation Considerations

• Computational Cost: Training SAEs is generally expensive. GBA adds minimal overhead compared to standard SAE training; the bias adaptation step is efficient.
• Buffer Management: The bias adaptation step requires a buffer of pre-activations. The paper suggests updating biases every 50 gradient steps with the largest batch size hardware permits. Efficiently, only $V$00 and $V$01 need to be tracked, not the full buffer.
• Activation Function: While theory uses smooth ReLU-like functions, experiments use JumpReLU, which empirically works well with GBA, especially for decoupling neuron output magnitude from its bias.
• Deployment: Trained SAEs can be used to replace MLP layers in LLMs for interpretability or potentially to steer model behavior. The features learned by GBA are more consistent, making them more reliable for downstream interpretability tasks.

Practical Implications

This research offers a more robust and theoretically grounded method for training SAEs.

• Improved Interpretability: By recovering more consistent and potentially more monosemantic features, GBA can enhance our understanding of LLM internal workings.
• Reduced Tuning Effort: The near tuning-free nature of GBA makes it more practical for researchers and practitioners to apply SAEs without extensive hyperparameter searches.
• Reliable Feature Discovery: Higher consistency implies that the features discovered are less likely to be artifacts of random initialization, leading to more trustworthy interpretations.

The paper lays a strong foundation by bridging theoretical understanding and practical application of SAEs, paving the way for more transparent and trustworthy AI systems. Future work includes extending theoretical guarantees to more general settings and using the learned features for model interventions and circuit discovery.