Llama Scope: Extracting Millions of Features from Llama-3.1-8B with Sparse Autoencoders

Motivation

• Interpretability research on models exceeding 8 billion parameters remains early, while comprehensive analyses need SAEs across multiple sites and feature sizes.These resources support SAE-based circuit analysis, feature-splitting studies, and identification of ultra-rare features.Source: §1 · pp. 1–2

Contribution

• Llama Scope provides 256 open-source SAEs trained on Llama-3.1-8B to reduce the need for extensive retraining.The suite is intended to facilitate mechanistic interpretability research and provide a shared resource for researchers.Source: §1 · p. 2

Decomposing Every Activation Space of Llama-3.1-8B-Base

• The suite decomposes four activation spaces across all 32 layers of Llama-3.1-8B-Base, using 32K and 128K feature widths.The spaces are the post-MLP residual stream, attention output, MLP output, and Transcoders, yielding 256 SAEs.Source: §1 · p. 2

Improved TopK SAEs

• The authors modify Top-K SAEs with decoder-column 2-norms in TopK computation, JumpReLU post-processing, and a K-annealing training schedule.These modifications are described in Section 2.4.Source: §1 · p. 2

Comprehensive Evaluation

• The evaluation covers sparsity-fidelity trade-offs, latent firing frequency, feature interpretability, transfer to longer contexts and fine-tuned models, and latent geometry.The generalization tests concern SAEs trained on base models.Source: §1 · p. 2

Disk-IO Friendly Training Infrastructure

• A mixed-parallelism training approach reduces the memory bottleneck when training SAEs with many features.The infrastructure is designed to be disk-IO friendly.Source: §1 · p. 2

2.1 ATTACKING SUPERPOSITION WITH SPARSE AUTOENCODERS

• The superposition hypothesis posits independent features that are decomposable, linearly represented, sparse in activation, and overcomplete relative to representation dimensionality.Source: §2.1 · p. 2

2.2 VANILLA SPARSE AUTOENCODERS

• A vanilla SAE encodes inputs through a hidden layer, applies sparsity constraints such as L1 regularization while minimizing reconstruction error, and decodes features back into input space.Source: §2.2 · p. 3
• SAEs offer encoder and decoder perspectives: activations reveal which inputs elicit features, while decoder columns expose the geometry of the latent basis.Source: §2.2 · p. 3

2.3 TOP-K SPARSE AUTOENCODERS

• Top-K SAEs retain only the K most active features for reconstruction; prior work reported improved sparsity-fidelity trade-offs and performance on novel metrics (Gao et al., 2024; Templeton et al., 2024a).Source: §2.3 · p. 3

2.4 MODIFICATIONS TO TOP-K SAES

• Top-K scoring incorporates decoder-column 2-norms instead of requiring unit-norm columns and gradient pruning, though the norms remain unfixed during training and are addressed at inference.Source: §2.4 · p. 4
• Post-training thresholding retains activations above a threshold calibrated to average K active features, yielding JumpReLU-like, independently evaluated feature activation.Source: §2.4 · p. 4
• Gradually reducing active features from D to K over the first 10% of training steps improves convergence, while the authors omit auxiliary dead-feature revival loss because dead features are rarely an issue.Source: §2.4 · p. 4

3.1 COLLECTING ACTIVATIONS

• Training activations come from SlimPajama with subset proportions preserved, use bfloat16, and exclude document-boundary tokens after documents are truncated to 1024 tokens.A regularly refilled and shuffled buffer adds randomness to address correlations among positions within a document.Source: §3.1 · p. 5

3.2 TRAINING POSITION

• SAEs cover four positions per Transformer block: post-MLP residual stream, attention output, MLP output, and a transcoder that predicts MLP output from the layer-normalized residual stream.Across 32 layers, these yield 128 possible training positions.Source: §3.2 · p. 5, §3.3.1 · p. 6

3.3.1 INITIALIZATION

• Decoder weights use Kaiming-uniform initialization, encoder weights begin as their transpose and are then untied, and zero biases produce near-zero initial reconstruction loss.For transcoders, encoder weights are not initialized from the decoder transpose because their input and output distributions differ.Source: §3.3.1 · pp. 5–6

3.3.2 OPTIMIZATION

• The training schedule warms the learning rate over 10K steps and linearly decays it during the final 20% of training, while gradually reducing active features from D to K in the first 10%.The feature-count schedule is intended to promote more even feature activation.Source: §3.3.2 · p. 6

3.3.3 INPUT AND OUTPUT NORMALIZATION

• Inputs are normalized to 2-norm D before entering the SAE, and reconstruction loss uses the same normalization; transcoders normalize both input and output.The input and output normalization factors differ only for transcoders.Source: §3.3.3 · p. 6

3.3.4 POST-TRAINING PROCESSING

• After training, weights are rescaled to undo training normalization and decoder columns are set to unit 2-norm, separating each feature’s direction from its encoded magnitude.The processed SAE operates on the original input and reconstructs the original output.Source: §3.3.4 · pp. 6–7

4.1 DATA AND METRICS

• Evaluation uses L0, explained variance, and Delta LM loss; unlike reconstruction metrics, Delta LM loss measures the SAE’s impact on overall language-model performance.Evaluations use 50 million held-out tokens sampled from the training distribution.Source: §4.1 · pp. 7–8

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 reconstruction at the same L0.On quartile-layer post-MLP residual streams, TopK reduced L0 from around 150 to 50 while maintaining or improving explained variance and Delta LM loss; residual-stream SAEs generally performed best on the three metrics.Source: §4.2 · p. 8, Figure 2 · p. 8

4.3 INTERPRETABILITY OF FEATURES

• Among 768 sampled features, automated monosemanticity scores showed no significant TopK–Vanilla difference, and manual analysis found about 10% were uninterpretable.The analysis used top-activating contexts, so it may not capture behavior on low-activating samples; this limits the monosemanticity assessment.Source: §4.3 · p. 9, Figure 3 · p. 9

4.4 ACTIVATION FREQUENCY

• Wider SAEs’ firing-frequency distributions shift toward lower frequencies, as expected with more features and a fixed top-k sparsity of 50.Training aimed for around 10% inactive features and fewer than 2% ultra-active features, defined as firing more than 0.1 of the time.Source: §4.4 · pp. 9–10, Figure 4 · p. 10

4.5.1 ACROSS SEQUENCE LENGTH

• On 8192-token contexts, degradation converged: last-1024-token reconstruction loss was 0.0086, 12% above training data, and L0 rose from 50 to 55.SAEs were trained on 1024-token sequences and were not fine-tuned on longer sequences.Source: §4.5.1 · p. 10, Figure 5 · p. 10

4.5.2 TO INSTRUCTION-FINETUNED MODELS

• On instruction-finetuned models, all but L31R-8x and L31R-32x showed no significant Delta LM loss degradation or L0 increase.For L15R-32x-TopK, language-model loss increased by 0.162 on the base model versus 0.090 on the instruction-finetuned model at the same L0; the smaller increase may reflect greater downstream-task robustness, not better SAE fit.Source: §4.5.2 · p. 11, Figure 6 · p. 11

4.6 FEATURE GEOMETRY

• A wider SAE learned a distinct Brexit feature where the smaller SAE had a broader Historical Movements feature, supporting discovery of new rather than merely composed features.The analysis also found a coherent cluster of features about wars, pandemics, natural disasters, and other harmful events; local UMAP clusters included climate change and financial crises.Source: §4.6 · pp. 11–13, Figure 7 · p. 12
• TopK and Vanilla features had higher decoder-column cosine similarity than a random baseline, indicating shared feature geometry; cosine similarity is not itself a direct measure of semantic similarity.The 36-feature example’s farthest cosine similarity was 0.28, above the random-projection maxima of 0.09 for 8x and 0.10 for 32x SAEs.Source: §4.6 · pp. 11–13, Figure 8 · p. 13

Scaling up Sparse Autoencoders

• Future SAE scaling may use customized GPU acceleration (Gao et al., 2024) and structured sparsity in MoE SAEs to reduce training costs while maintaining performance (Mudide et al., 2024).Structured sparsity in MoE SAEs is identified as a promising avenue for exploration.Source: §5 · p. 13

Extending Neuron-Level Analysis to SAE Features

• Neuron-level findings on knowledge, universality, and multilinguality may inform SAE-feature analysis (Dai et al., 2022; Gurnee et al., 2023; 2024; Tang et al., 2024); SAE features may better reflect cross-architecture universality than MLP neurons (Wang et al., 2024).Source: §5 · p. 13

Revealing a More Interpretable Latent Space

• Linear probes and decision trees predict harmful content from SAE features, while CCA offers insights into SAE latent spaces (Lan et al., 2024).Source: §5 · p. 14

• Llama Scope comprises TopK SAEs trained on Llama-3.1-8B-Base and offers a foundation for future improvements in SAE training.Source: §6 · p. 14
• The authors hope Llama Scope and other open-source SAEs will save researchers time and effort, and describe SAE research as a long-term direction.Source: §6 · p. 14

A.1 ACTIVATION BUFFER

• Latent activations require 8K bytes per token versus 2 bytes for text, a 4,096x increase in storage and disk throughput that makes data loading a training bottleneck.Source: §A.1 · p. 19
• The on-the-fly, buffered approach addresses the data-loading bottleneck created by the storage and throughput demands of latent activations.Source: §A.1 · p. 19
• The authors generate activations on the fly rather than pre-saving them as in Lieberum et al. (2024) and Templeton et al. (2024b).Source: §A.1 · p. 19
• A producer-consumer setup has language models generate activations into a buffer while SAEs consume them in random order.Source: §A.1 · p. 19
• Training begins when the activation buffer is full, and language models refill it after half its contents have been consumed.Source: §A.1 · p. 19
• The buffer is shuffled after each refill to randomize training data without saving and shuffling all activations at once.Source: §A.1 · p. 19

A.2 MIXED PARALLELISM

• Mixed parallelism uses data parallelism for activation generation and tensor parallelism for SAE training, with LM copies and independent activation buffers on each GPU.Each GPU samples a minibatch from its buffer; the minibatches are all-gathered as SAE input.Source: §A.2 · p. 19, Figure 9 · p. 19
• Activation generation is slower but uses less memory, while SAE training steps are fast but its models and gradients consume significant memory.Combining the two parallelism techniques accelerates buffer refilling while managing the training memory bottleneck.Source: §A.2 · p. 19

A.3 COMPARISON TO PRE-SAVING APPROACH

• Online generation trains SAEs on an 8B model with one NVIDIA A100 and no pre-saved activations, versus pre-saving’s approximately 2TB storage and at least 500MB/s disk throughput.This approach is described as more suitable for academic research than prior large-SAE training reliant on distributed disk reading and extensive storage resources (Lieberum et al., 2024; Templeton et al., 2024b).Source: §A.3 · p. 19
• Training different-width SAEs at the same position can require generating activations repeatedly, causing 1–2 times the redundant activation generation compared with pre-saving.This is a limitation of online activation generation when multiple SAEs share a position.Source: §A.3 · p. 19

C ABLATION STUDIES

• JumpReLU inference produces nearly identical L0 sparsity and MSE efficiency across 32 LXR-8x-TopK SAEs, with a negligible performance drop.The comparison evaluates SAEs with and without JumpReLU inference mode.Source: Figure 11 · p. 21, C ABLATION STUDIES · p. 21
• 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 do not begin activating until about one-third of training; with annealing, almost all features activate at least once over 1e6 tokens.Source: C ABLATION STUDIES · p. 21, Figure 12 · p. 21
• With K-annealing, almost all L15R-8x-TopK SAE features fire at least once over 1e6 tokens, while without it they begin activating only after about one-third of training.K is annealed from 4096 to 50 in the first 10% of training steps.Source: C ABLATION STUDIES · p. 21
• Fixing decoder L2 norm to 1 after each training step yields at least a 0.05 drop in variance explained; the authors suspect baseline-replication bugs and defer further work.Other hyperparameters were fixed in this ablation.Source: C ABLATION STUDIES · p. 22

D PROMPT DESIGN FOR FEATURE INTERPRETABILITY

• The feature-interpretability prompt asks for a monosemanticity score based on activation consistency, using a five-level rubric from no discernible pattern to a clear pattern with no deviating examples.Intermediate ratings distinguish broad themes, overall patterns with mismatches, and clear patterns with one or two deviating examples.Source: D PROMPT DESIGN FOR FEATURE… · p. 22
• The prompt presents non-negative token activation values, with higher values indicating stronger feature connections, and instructs the evaluator to return only a number.The prompt’s context begins with a sentence marker and <START>.Source: D PROMPT DESIGN FOR FEATURE… · p. 22