Source-linked AI summary
Llama Scope: Extracting Millions of Features from Llama-3.1-8B with Sparse Autoencoders
Zhengfu He, Wentao Shu, Xuyang Ge, Lingjie Chen, Junxuan Wang, Yunhua Zhou, Frances Liu, Qipeng Guo, Xuanjing Huang, Zuxuan Wu, Yu-Gang Jiang, Xipeng Qiu
TL;DR
SAE research on industrial-scale language models is limited by the need to train many models across activation sites and feature sizes. Llama Scope addresses this with a 256-SAE Llama-3.1-8B suite, improved TopK training, and broad evaluation. The paper reports a reusable open-source foundation for mechanistic-interpretability research, while noting limits in evaluating low-activation features and longer-context adaptation.
Problem
Industrial-scale SAE research remains limited, and comprehensive analyses require models trained across multiple activation sites and feature sizes.
Method
The authors train 256 SAEs across Llama-3.1-8B layers and sublayers, improve TopK SAEs, and evaluate reconstruction, sparsity, interpretability, generalization, and latent geometry.
Results
The work introduces Llama Scope SAEs and reports a foundation for future SAE-training improvements and mechanistic-interpretability investigations.
Takeaways & Limitations
The open-source Llama Scope suite is intended to help researchers save time and effort by providing reusable SAEs for investigations.
Takeaways & Limitations
The evaluation focuses on top activations, so monosemanticity scores may not reflect feature behavior on low-activating samples.
Abstract
from arXiv · showhide
Sparse Autoencoders (SAEs) have emerged as a powerful unsupervised method for extracting sparse representations from language models, yet scalable training remains a significant challenge. We introduce a suite of 256 SAEs, trained on each layer and sublayer of the Llama-3.1-8B-Base model, with 32K and 128K features. Modifications to a state-of-the-art SAE variant, Top-K SAEs, are evaluated across multiple dimensions. In particular, we assess the generalizability of SAEs trained on base models to longer contexts and fine-tuned models. Additionally, we analyze the geometry of learned SAE latents, confirming that \emph{feature splitting} enables the discovery of new features. The Llama Scope SAE checkpoints are publicly available at~\url{https://huggingface.co/fnlp/Llama-Scope}, alongside our scalable training, interpretation, and visualization tools at \url{https://github.com/OpenMOSS/Language-Model-SAEs}. These contributions aim to advance the open-source Sparse Autoencoder ecosystem and support mechanistic interpretability research by reducing the need for redundant SAE training.
LLAMA SCOPE: EXTRACTING MILLIONS OF FEATURES FROM LLAMA-3.1-8B WITH SPARSE AUTOENCODERS
The paper is authored by Zhengfu He, Wentao Shu, Xuyang Ge, Lingjie Chen, Junxuan Wang, Yunhua Zhou, Frances Liu, Qipeng Guo, Xuanjing Huang, Zuxuan Wu, Yu-Gang Jiang, and Xipeng Qiu.
- Zhengfu He, Wentao Shu, and Xuyang Ge are listed among the authors.
- Lingjie Chen, Junxuan Wang, Yunhua Zhou, Frances Liu, Qipeng Guo, and Xuanjing Huang are listed among the authors.
- Zuxuan Wu, Yu-Gang Jiang, and Xipeng Qiu are listed among the authors.
1 INTRODUCTION
The paper addresses the difficulty of conducting comprehensive SAE research at industrial scale by releasing a broad Llama-3.1-8B suite and improving TopK SAEs. It evaluates these resources across reconstruction, sparsity, interpretability, generalization, and latent geometry.
- Motivation: The work targets the limited availability of rigorous interpretability research and reusable SAEs for models exceeding 8 billion parameters.Comprehensive analyses otherwise require training SAEs at multiple sites and feature sizes.
- Contributions: Llama Scope provides 256 SAEs covering every Llama-3.1-8B sublayer across 32 layers and 32K or 128K feature widths.The suite covers post-MLP residual streams, attention outputs, MLP outputs, and Transcoders.
- Contributions: The authors modify TopK SAEs with decoder-norm-aware selection, JumpReLU post-processing, and K-annealing training.
- Evaluation: Llama Scope SAEs are evaluated on sparsity-fidelity metrics, feature firing, interpretability, longer-context and fine-tuned-model generalization, and latent geometry.
- Infrastructure: A mixed-parallelism, disk-IO-friendly infrastructure reduces the memory bottleneck of training SAEs with many features.
2 CONCEPTUAL AND TECHNICAL BACKGROUND
The paper presents SAEs as sparse, linear, decomposable, overcomplete representations and describes vanilla and TopK variants. Its TopK modifications regulate activation selection, training convergence, decoder norms, and inference thresholding.
- SAE motivation: SAEs seek sparse, linear, decomposable features in an overcomplete latent space to address superposition in neural-network representations.
- Vanilla Sparse Autoencoders: A vanilla SAE encodes an input into a nonlinear hidden representation and decodes it back while minimizing reconstruction error with a sparsity penalty.
- TopK SAEs: TopK SAEs retain only the K most active features for reconstruction, while decoder columns are normalized after each minibatch.
- Modifications to Top-K SAEs: The authors incorporate decoder-column norms into TopK selection to enforce sparsity without pruning gradients parallel to decoder columns.
- Modifications to Top-K SAEs: Post-training thresholding creates a JumpReLU-like variant in which features activate independently when their values exceed a threshold.
- Modifications to Top-K SAEs: K-annealing gradually reduces active features from D to K during early training, improving convergence by easing adaptation to sparse activations.
3 TRAINING LLAMA SCOPE SAES
Llama Scope trains a broad SAE suite on Llama-3.1-8B activations using multiple Transformer positions, feature widths, and scalable preprocessing. The training procedure includes initialization, normalization, scheduling, and post-processing choices designed for stable reconstruction and analysis.
- Data: Training uses SlimPajama activations with preserved subset proportions, bfloat16 computation, 1024-token documents, excluded boundary-token activations, and buffer shuffling.
- Training positions and widths: The suite trains SAEs at four positions across all 32 layers, using 8x or 32x hidden-size expansions corresponding to 32K or 128K features.
- Initialization: The encoder and decoder are initialized with normalized Kaiming-based weights or transposed decoder weights, while biases start at zero to provide near-zero reconstruction loss.
- Training positions: The four positions are the post-MLP residual stream, attention output, MLP output, and Transcoder input-output mapping.Transcoders take layer-normalized residual streams as input and predict MLP outputs.
- Optimization: Optimization uses Adam with a learning rate warm-up and final linear decay, while K is annealed early to distribute feature activation more evenly.
- Normalization and post-processing: Input and output norms are normalized during training and weights are rescaled afterward so the final SAE operates on original values.Decoder columns are additionally normalized to unit 2-norm so they represent feature directions rather than strength.
4 EVALUATION
Llama Scope evaluates 256 SAEs across sparsity-fidelity, interpretability, activation frequency, out-of-distribution generalization, and feature geometry. TopK and wider SAEs generally improve efficiency while preserving reconstruction quality, and wider models can discover new features, but interpretability evaluation and long-context testing have stated limitations.
- 4.2 SPARSITY-FIDELITY PARETO EFFICIENCY: TopK SAEs reduce L0 sparsity from around 150 to 50 while maintaining or improving explained variance and Delta LM loss versus vanilla SAEs.The authors attribute this improvement likely to mitigating feature shrinkage and removing weakly firing features.
- 4.2 SPARSITY-FIDELITY PARETO EFFICIENCY: Wider SAEs outperform narrower ones in reconstruction at the same L0 sparsity, and their firing distributions shift toward lower frequencies because fixed top-k activates a smaller fraction of features.This pattern is reported for L7R-8x, L15R-8x, and L23R-8x TopK SAEs.
- 4.2 SPARSITY-FIDELITY PARETO EFFICIENCY: Across all 256 SAEs, TopK matches or slightly exceeds vanilla reconstruction quality with significantly better L0 sparsity, while wider SAEs improve Delta LM loss and explained variance at the same sparsity.Residual-stream SAEs generally outperform those trained at other positions across the three metrics.
- 4.3 INTERPRETABILITY OF FEATURES: TopK and vanilla SAEs show no significant difference in automated monosemanticity, while manual analysis finds about 10% of features are not interpretable.Manual evaluation does not distinguish reliably among monosemanticity scores 2–5.
- 4.3 INTERPRETABILITY OF FEATURES: Interpretability analysis focuses on top activations and may miss feature behavior on low-activating samples, limiting how fully the reported monosemanticity scores characterize features.The authors note that TopK and JumpReLU specifically remove weakly firing instances.
- 4.5.1 ACROSS SEQUENCE LENGTH: At 8192-token contexts, average reconstruction loss in the last 1024 tokens is 0.0086, a 12% increase over training data, while L0 sparsity rises from 50 to 55.Performance degradation converges around 8192 tokens, although the SAEs were trained on 1024-token sequences.
- 4.5.2 TO INSTRUCTION-FINETUNED MODELS: SAEs generally generalize to instruction-finetuned models without significant Delta LM loss or L0 degradation, except for L31R-8x and L31R-32x.For L15R-32x-TopK, Delta LM loss increases by 0.162 in the base model versus 0.090 in the instruction-finetuned model at the same L0 sparsity; the authors caution this may reflect task robustness.
- 4.6 FEATURE GEOMETRY: Feature-geometry analysis finds TopK and vanilla SAEs share a universal feature geometry, while wider SAEs can learn new features rather than only recombining existing ones.A distinct Brexit feature appears in L15R-32x-TopK, whereas the closest smaller-SAE feature represents broader historical movements; related clusters include wars, climate change, pandemics, and financial crises.
5 RELATED AND FUTURE WORK
The paper outlines future directions for scaling SAEs, extending neuron-level analysis, and making SAE latent spaces more interpretable.
- Scaling up Sparse Autoencoders: Scaling SAEs to larger models may require customized GPU acceleration and structured sparsity for mixture-of-experts SAEs.
- Extending Neuron-Level Analysis to SAE Features: Methods developed for studying MLP neurons could be applied to SAE features to investigate knowledge representation, universality, and multilingual capabilities.
- Revealing a More Interpretable Latent Space: Linear probes, decision trees, and similarity metrics such as CCA are proposed as tools for understanding SAE latent spaces.
6 CONCLUSION
The paper concludes that Llama Scope provides a foundation for future SAE training and mechanistic interpretability research. Its open-source resources are intended to reduce researchers’ time and effort.
- Llama Scope introduces TopK SAEs trained on the Llama-3.1-8B-Base model as a foundation for future SAE improvements.
- 40,96x higher latent-activation storage and disk-throughput requirements than text make data loading a major SAE training bottleneck.
- A producer-consumer buffer generates, shuffles, and serves activations to SAEs while limiting the need to store all activations in advance.
A.2 MIXED PARALLELISM
Mixed parallelism balances the memory demands of SAE training with the slower process of activation generation. It combines data parallelism for language-model activation generation with tensor parallelism for distributed SAE training.
- Each GPU loads a copy of the language model, distributes the SAE across GPUs, and maintains an independent activation buffer.
- Mixed parallelism combines data-parallel activation generation with tensor-parallel SAE training to manage memory efficiently.
- The strategy accelerates activation-buffer refilling because activation generation is slower and less memory-intensive than SAE training steps.
A.3 COMPARISON TO PRE-SAVING APPROACH
Online activation generation reduces storage requirements compared with pre-saving activations, but repeated SAE training at one position can increase computation.
- Online activation generation eliminates vast storage requirements and supports training an 8B-model SAE on a single NVIDIA A100 GPU.
- Training multiple SAEs with different widths at one position can require 1–2 times the redundant activation generation of pre-saving.
- Figure 10 evaluates all 256 SAEs using L0 sparsity, explained variance, and Delta LM loss.
C ABLATION STUDIES
The ablations examine post-processing, K-annealing, and decoder-norm choices in Top-K SAEs. JumpReLU preserves performance, K-annealing accelerates convergence, while the decoder-norm baseline comparison remains uncertain because replication bugs are suspected.
- Post-Processing: Nearly identical L0 sparsity and MSE efficiency show that JumpReLU inference causes negligible performance loss across 32 LXR-8x-TopK SAEs.JumpReLU replaces the TopK sparsity constraint with a threshold calibrated so K features fire per position in expectation.
- K-Annealing: K-annealing from 4096 to 50 during the first 10% of training activates features earlier and reaches the same MSE loss faster.Without annealing, features begin activating only around one-third of training; with annealing, almost all features fire at least once over 1e6 tokens.
- Decoder L2 Norm: The decoder-L2-norm ablation reports at least a 0.05 drop in variance explained for the replicated baseline, but suspected bugs make this comparison inconclusive.The authors leave replication of the baseline for future work rather than treating the result as definitive.