Scale Sparse Autoencoder (Scale SAE)

• Scale SAE is a sparse autoencoder architecture that partitions latent space into multiple experts to improve computational efficiency and interpretability.
• It employs a multiple expert activation scheme with a global Top-K selection to ensure sparsity by activating only the highest magnitude features across selected experts.
• Adaptive feature scaling reduces redundancy by emphasizing high-frequency deviations, fostering specialist expert representations and enhanced feature diversity.

Scale Sparse Autoencoder (Scale SAE) is a sparse autoencoder architecture designed to resolve key efficiency and interpretability challenges that arise when scaling autoencoders for LLM analysis. Scale SAE leverages a multi-expert mixture-of-experts (MoE) framework, dual innovations in expert activation and feature scaling, and rigorous empirical validation to significantly improve both computational tractability and feature diversity relative to previous MoE-SAEs (Xu et al., 7 Nov 2025).

1. Standard Sparse Autoencoder Foundations

Sparse autoencoders (SAEs) target the decomposition of high-dimensional neural activations (e.g., from transformer residual streams) into a sparse set of interpretable directions, balancing reconstruction fidelity and feature sparsity according to the objective

$$\mathcal{L}(x) = \|x - (W^{dec} z + b_{dec})\|^2_2 + \lambda \|z\|_1,$$

where $x \in \mathbb{R}^{d_{model}}$ is a LM layer activation, $z \in \mathbb{R}^H$ is the sparse latent code (often via $\text{TopK}_k(W^{enc}(x-b_{pre}))$), and $\lambda$ tunes the sparsity-penalty. The TopK nonlinearity yields exactly $k$ nonzero latents per input, which greatly improves interpretability by ensuring every activation is represented by a concise subset of learned features (Gao et al., 2024).

2. Scale SAE Architecture: Multi-Expert Partitioning

Scale SAE organizes the $H = N \cdot h$ total latent dimensionality into $N$ independent experts, each of width $h$ (Xu et al., 7 Nov 2025). A router (parameterized by $W_{router} \in \mathbb{R}^{N \times d_{model}}$ and $b_{router}$) produces a gating distribution $$p(x) = \mathrm{softmax}(W_{router} (x - b_{router}))$$. For each token, only the top $e$ experts are selected (by $p_i(x)$), promoting load-balancing and computational savings.

Each expert $i$ comprises its own encoder $W^{enc}_i \in \mathbb{R}^{h \times d_{model}}$ and decoder $W^{dec}_i \in \mathbb{R}^{d_{model} \times h}$. This partitioning, if naively implemented, can suffer from redundancy: experts often converge to overlapping feature sets. Previous MoE-SAEs address efficiency but fail to deliver distinct specialist experts, limiting practical interpretability.

3. Multiple Expert Activation and Global Top-K Selection

Scale SAE introduces a “Multiple Expert Activation” scheme:

• For each input, activate the top $e$ experts (by gating probability).
• Compute pre-activation vectors $f_i = W^{enc}_i x$ for $i$ in the active set.
• Globally apply TopK across all $e \cdot h$ entries (Equation 3), ensuring that exactly $K$ nonzero features are distributed among the selected experts:

$$z_{ij} = \begin{cases} f_{ij}, & \text{if } (i,j) \in \mathcal{K} \ 0, & \text{otherwise} \end{cases}$$

where $\mathcal{K}$ indexes the $K$ largest entries over all active experts’ features.

The final reconstruction is a weighted sum: $$\hat{x} = \sum_{i \in \mathcal{T}} p_i(x) E_i(x), \quad E_i(x) = W^{dec}_i z_i + b_{dec,i}$$ where each expert’s decoded output is weighted by its gating score.

An auxiliary load-balancing term is added: $L_{aux} = N \sum_{i=1}^{N} f_i P_i$ with $f_i$ as the fraction of tokens routed to expert $i$ and $P_i$ its average probability. The full objective is

$\mathcal{L} = \|x - \hat{x}\|_2^2 + \alpha L_{aux}$

ensuring both reconstruction and balanced specialist usage.

4. Adaptive Feature Scaling for Expert Diversity

To combat redundancy and promote feature specialization, each expert encoder is decomposed into mean and deviation components: $$\bar{W}^{enc}_i = \frac{1}{h} \sum_{j=1}^h W^{enc}_{i,j}, \qquad \Delta W^{enc}_i = W^{enc}_i - \bar{W}^{enc}_i$$ A learnable scalar $\omega_i$ rescales deviation: $$\hat{W}^{enc}_i = \bar{W}^{enc}_i + (1 + \omega_i) \cdot \Delta W^{enc}_i$$ This process emphasizes “high-frequency” deviations relative to the mean, functioning as a differentiable high-pass filter that increases diversity in expert feature sets. Empirically, $\omega_i$ converges to positive values and scales with expert activation count.

5. Training Algorithm and Implementation

Scale SAE models are initialized with random router parameters and expert encoders/decoders. The training loop proceeds as follows:

for each minibatch {x^(b)}: for x in batch: p = softmax(W_router(x - b_router)) T = top-e indices in p for i in T: W_bar = row_mean(W_enc_i) Delta_W = W_enc_i - W_bar W_hat = W_bar + (1 + omega_i) * Delta_W f_i = W_hat @ x z = global_top_K(stack(f_i)) for i in T: E_i = W_dec_i @ z_i + b_dec_i x_hat = sum_i(p_i * E_i) L_recon = ||x - x_hat||_2^2 L_sparse = sum_i ||z_i||_1 accumulate expert usage statistics L_aux = N * sum_i f_i P_i L_total = mean(L_recon) + lambda mean(L_sparse) + alpha L_aux optimizer.step()

All parameter updates are driven by backpropagation through the gating, encoding, and rescaling operations.

6. Empirical Findings: Scalability, Diversity, and Interpretability

Extensive FLOPs-matched evaluations on OpenWebText and large transformer models demonstrate:

• Up to 24% lower reconstruction mean-squared error than TopK-SAE, Gated-SAE, and Switch-SAE architectures.
• 37–42% improvement at high sparsity ($L_0$ large).
• A 99% reduction in feature redundancy, measured as the fraction of dictionary vectors with cosine similarity $>0.9$.
• Automated interpretability and faithfulness metrics (e.g., Loss Recovered, LLM judge, max-activate) equal or surpass prior architectures.
• Aggressive reduction (25–31%) in neuronal activation similarity in high-sparsity MoE regimes, corroborating improved feature specialization.

Benchmarks confirm that Multiple Expert Activation enforces specialization and Feature Scaling suppresses redundancy, so Scale SAE delivers specialist, diverse, monosemantic features without degrading computational efficiency (Xu et al., 7 Nov 2025).

7. Practical Implications and Comparative Significance

Scale SAE substantially extends the scalability of SAE-based interpretability frameworks. Unlike conventional SAEs and standard MoE-SAEs—which often suffer from overlap, redundancy, and poor scaling—Scale SAE obtains computational savings proportional to the number of experts and active features. Simultaneously, it supports larger dictionaries and higher sparsity budgets without sacrificing disentanglement or monosemanticity. These advances enable the reliable analysis of LLM representations at extreme scale and offer rigorous statistical guarantees for feature diversity, making them suitable for transparent inspection and probing of LLM internals.