Time-Aware Feature Selection: Adaptive Temporal Masking for Stable Sparse Autoencoder Training

AI summary

Overview

Research area: Mechanistic interpretability of large language models, specifically sparse autoencoder (SAE) training methods for decomposing neural activations into interpretable features.

Technical level: Intermediate. The paper assumes familiarity with sparse autoencoders, L1/L0 sparsity penalties, reconstruction loss, and probing techniques, but its core idea (tracking which features matter over time instead of applying fixed thresholds) is explainable without deep mathematics.

Scope: The paper proposes Adaptive Temporal Masking (ATM), a training method that reduces "feature absorption" in SAEs, and evaluates it against three baselines on layer 12 of Gemma-2-2B.

What This Paper Is About

Sparse autoencoders are used to break down a language model's internal activations into individual, human-interpretable features. A known failure mode called feature absorption occurs when one feature swallows another that implies it (for example, the feature for "short" absorbing the feature for "starts with S"), which improves the sparsity score that SAE training optimizes but destroys interpretability. The paper's goal is a training procedure that selects which features stay active based on how their importance evolves over time, rather than using rigid hard thresholds, in order to keep absorption low while preserving reconstruction quality.

Key Contributions

1. Adaptive Temporal Masking (ATM) — a training mechanism that tracks feature activation magnitudes, frequencies, and reconstruction contributions through exponential moving averages, combining them into an importance score that changes as training proceeds.
2. A probabilistic masking scheme that replaces hard thresholding: masking probability is computed from how far a feature's importance score falls below a statistically derived threshold (mean plus a multiple of the standard deviation), producing a soft rather than abrupt cutoff, with a guaranteed minimum number of features always retained.
3. A principled importance-based selection rule that combines two signals — how large a feature's activations are (magnitude EMA) and how much it contributes to reducing reconstruction error (the gradient magnitude of reconstruction loss with respect to the feature).
4. A comprehensive evaluation on Gemma-2-2B covering feature absorption, reconstruction metrics, unsupervised SAE metrics, and sparse probing across 35 binary tasks and five datasets.

Main Findings

• Absorption score: ATM achieves a mean absorption score of 0.0068, compared with 0.1402 for TopK SAEs and 0.0114 for JumpReLU. The paper's conclusion describes this as a 40% reduction in feature absorption compared to the best existing method on this metric.
• Hard cases: On 'e'-words, ATM scores 0.0069 versus 0.3364 for TopK and 0.0138 for JumpReLU.
• Reconstruction quality: ATM records MSE 0.5508 and cosine similarity 0.9727, better than TopK (MSE 2.53125, cosine 0.875) and JumpReLU (MSE 1.6719, cosine 0.9297), but worse than the Standard SAE baseline (MSE 0.0898, cosine 0.9961).
• Information retention: ATM scores 0.9967 on the cross-entropy loss score and 0.9965 on KL divergence, close to the Standard SAE (1.0 and 0.9996) and above TopK (0.9556 and 0.9565) and JumpReLU (0.9951 and 0.9945).
• Explained variance: 0.9102 for ATM, versus 0.6016 for TopK, 0.7344 for JumpReLU, and 0.9844 for the Standard SAE.
• Sparsity: ATM reports L0 sparsity of 3280 and L1 sparsity of 1704. TopK reports L0 40 and L1 366; JumpReLU reports L0 2666 and L1 4832; the Standard SAE reports L0 8724 and L1 12544. The paper states ATM's sparsity performance "roughly matches JumpReLU."
• Sparse probing: ATM reaches top-1 test accuracy of 0.7161, below TopK (0.7698) and above JumpReLU (0.7154) and the Standard SAE (0.6379). The paper reports ATM "particularly excelling on bias detection (0.796) and sentiment analysis (0.853)."
• Baseline comparison set: Table 1 also includes a Standard SAE baseline with absorption score 0.0161, in addition to the TopK and JumpReLU comparisons emphasized in the abstract.

Methodology in Plain English

The researchers start from a transformer layer's activation vectors (hidden dimension 2304) and train an encoder that expands them into 16384 latent features, plus a decoder that maps back to the original space. Rather than allowing every feature to stay on and penalizing the total activation size, ATM decides which features to keep using a running record of how each feature has behaved.

Three statistics per feature are updated as exponential moving averages: how large its activations typically are, how often they occur, and how strongly the reconstruction loss depends on it (measured by gradient magnitude). The first two are multiplied into a single importance score. The method then computes a threshold from the current mean and standard deviation of importance scores across features, and any feature below that line becomes progressively more likely to be masked off — but probabilistically, not deterministically, so features can recover if they become important again. A minimum number of the most important features are always kept, so information can still flow. Training starts with a warmup period where nothing is masked, letting initial importance patterns form, and later includes periodic pruning phases where the threshold multiplier is raised to push for more aggressive sparsity. The final loss combines reconstruction error on the masked activations with an L1 penalty on them.

Specific settings reported: layer 12 of Gemma-2-2B, WikiText-103 data in 128-token sequences with a buffer size of 2048, 5M tokens of training, batch size 2048 for SAE training and 32 for model inference, Adam optimizer with learning rate 3e-4 and a 1000-step warmup, probabilistic masking decay rate r = 0.5, sparsity weight λ_sparse = 0.001, and unit-norm decoder weights maintained by gradient projection. Note that the paper lists the EMA decay rate as β = 0.9 in the method section and as β = 0.99 in the experimental setup; the value of the threshold multiplier c is not reported.

Absorption is measured following the first-letter classification task of Chanin et al., with a feature-split threshold τ_fs = 0.03, a probe-similarity threshold τ_ps = 0.025, and a probe-absorption threshold τ_pa = 0.4. Sparse probing follows Gurnee et al.: 35 binary tasks across five datasets (bias_in_bios, Amazon Reviews, Europarl, GitHub, AG News), with 4000 training and 1000 test samples per task, inputs truncated to 128 tokens, the first 150 characters removed for GitHub to avoid license headers, up to five classes per dataset, and at least a 0.2 positive instance ratio per binary split.

Why This Matters

Feature absorption is a concrete obstacle to trusting SAE-based interpretability: if features silently merge, an analyst reading a "short" neuron may be unknowingly reading a combined "short" and "starts with S" signal. Reducing absorption while keeping reconstruction quality high means the individual features extracted from a model are more likely to correspond to genuinely separate concepts.

Real-world applications implied by the paper:

• Safety and bias auditing: The paper frames interpretability as crucial for reliability and safety and cites bias work; ATM's reported bias-detection probing accuracy of 0.796 supports this use case.
• Model debugging and behavior analysis: More stable features make it easier to consistently identify and analyze model behaviors across layers.
• Knowledge editing: The authors list targeted interventions as a future direction, suggesting temporal feature statistics could support more precise editing while preserving model capabilities.
• Interpretability tooling for researchers: A training recipe that reduces absorption without sacrificing reconstruction gives practitioners a more trustworthy base for feature-level analysis.

Industry relevance: Any organization deploying or auditing large language models has an interest in methods that make internal representations more legible. The paper's emphasis on stability, safety, and bias detection places it in the interpretability tooling pipeline rather than in model capability improvement.

Future Directions

• Theoretical foundations: Develop a formal framework relating temporal dynamics to feature stability, and clarify the optimal balance between short-term responsiveness and long-term stability.
• Scaling behavior: Test ATM on larger models than Gemma-2-2B and adapt the temporal tracking mechanisms for better computational efficiency.
• Hyperparameter optimization: Explore a wider range of hyperparameters (the authors specifically mention batch size and dictionary width as unexplored).
• Targeted interventions and cross-layer integration: Use temporal feature statistics for more precise knowledge editing, and extend ATM to capture temporal patterns across multiple network layers to reveal hierarchical feature relationships.

The paper also notes practical constraints: all experiments ran on a single NVIDIA RTX 4090 GPU, and the authors state they could not evaluate models larger than Gemma-2-2B, train deeper networks or use more tokens, or systematically investigate hyperparameters. The conclusions section states that the experiments "reveal important limitations" but does not enumerate them in the provided content.

Target Audience

Interpretability researchers and engineers working with sparse autoencoders, especially those studying feature absorption and dictionary learning. It is also relevant to safety and alignment teams who use SAE-derived features for auditing model behavior, and to graduate students or practitioners with prior exposure to SAEs who want a concrete, reproducible training modification. Readers without background in sparse coding or probing methodology will need supplementary reading, since the paper assumes familiarity with L0/L1 sparsity, reconstruction metrics, and linear probing.

Authors’ abstract

Understanding the internal representations of large language models is crucial for ensuring their reliability and safety, with sparse autoencoders (SAEs) emerging as a promising interpretability approach. However, current SAE training methods face feature absorption, where features (or neurons) are absorbed into each other to minimize $L_1$ penalty, making it difficult to consistently identify and analyze model behaviors. We introduce Adaptive Temporal Masking (ATM), a novel training approach that dynamically adjusts feature selection by tracking activation magnitudes, frequencies, and reconstruction contributions to compute importance scores that evolve over time. ATM applies a probabilistic masking mechanism based on statistical thresholding of these importance scores, creating a more natural feature selection process. Through extensive experiments on the Gemma-2-2b model, we demonstrate that ATM achieves substantially lower absorption scores compared to existing methods like TopK and JumpReLU SAEs, while maintaining excellent reconstruction quality. These results establish ATM as a principled solution for learning stable, interpretable features in neural networks, providing a foundation for more reliable model analysis.