Improving Robustness In Sparse Autoencoders via Masked Regularization

Vivek Narayanaswamy; Kowshik Thopalli; Bhavya Kailkhura; Wesam Sakla

cs.LGApr 7, 2026

Improving Robustness In Sparse Autoencoders via Masked Regularization

Authors: Vivek Narayanaswamy, Kowshik Thopalli, Bhavya Kailkhura, Wesam Sakla

Organizations: Lawrence Livermore National Laboratory, CA, USA

Abstract

Sparse autoencoders (SAEs) are widely used in mechanistic interpretability to project LLM activations onto sparse latent spaces. However, sparsity alone is an imperfect proxy for interpretability, and current training objectives often result in brittle latent representations. SAEs are known to be prone to feature absorption, where general features are subsumed by more specific ones due to co-occurrence, degrading interpretability despite high reconstruction fidelity. Recent negative results on Out-of-Distribution (OOD) performance further underscore broader robustness related failures tied to under-specified training objectives. We address this by proposing a masking-based regularization that randomly replaces tokens during training to disrupt co-occurrence patterns. This improves robustness across SAE architectures and sparsity levels reducing absorption, enhancing probing performance, and narrowing the OOD gap. Our results point toward a practical path for more reliable interpretability tools.

Figures & tables

Metric	Method	$\boldsymbol{\ell_{0}=20}$	$\boldsymbol{\ell_{0}=40}$	$\boldsymbol{\ell_{0}=80}$	$\boldsymbol{\ell_{0}=160}$
Mean Full Absorption ( $\uparrow$ )	w/o Masking	86.119	91.646	94.650	97.434
	w/ Masking	88.475	93.450	96.437	98.003
Explained Variance ( $\uparrow$ )	w/o Masking	72.172	77.421	82.303	87.841
	w/ Masking	71.266	76.874	82.419	88.823
Sparse Probing ( $\uparrow$ )	w/o Masking	73.483	75.099	78.705	79.574
	w/ Masking	75.574	77.112	77.749	79.639

Table 1: Performance comparison of the proposed masking strategy on the MatryoshkaBatchTopK SAE trained on Layer 8 of Pythia-160m-deduped across multiple evaluation metrics and sparsity levels. Across metrics, we see consistent improvements with our approach with a small trade-off in explained variance at lower sparsities. Overall, the method demonstrates robustness and strong performance across sparsity regimes.

Metric	Method	$\boldsymbol{\ell_{0}=20}$	$\boldsymbol{\ell_{0}=40}$	$\boldsymbol{\ell_{0}=80}$	$\boldsymbol{\ell_{0}=160}$
Mean Full Absorption ( $\uparrow$ )	w/o Masking	90.805	97.365	98.969	99.174
	w/ Masking	94.559	98.753	97.322	98.979
Explained Variance ( $\uparrow$ )	w/o Masking	53.516	58.984	64.453	69.922
	w/ Masking	53.125	59.375	64.844	71.094
Sparse Probing ( $\uparrow$ )	w/o Masking	74.473	73.300	74.659	75.659
	w/ Masking	74.206	77.341	74.243	75.363

Table 2: Performance comparison of the MatryoshkaBatchTopK SAE trained w/ and w/o masking on layer 12 of Gemma2-2-b. The masking strategy tends to provide steady improvements in mean absorption score and explained variance, with benefits visible across most sparsity levels. For other metrics such as sparse probing, TPP, and SCR, the w/ Masking approach also performs competitively, showing strong results overall with few trade-offs appear in a few settings with increase in

\ell_{0}

Figure 1: OOD probing results: Mean test AUC across 8 OOD datasets from [ 11 ] at $\ell_{0}=20$ with MatryoshkaBatchTopK SAEs on Pythia-160m-deduped (a) and Gemma-2-2b (b). While oracle probes (red, trained on raw LLM activations) outperform SAEs, our objective substantially improves OOD performance over the baseline.

Metric	$m=0.0$	$m=0.2$	$m=0.3$	$m=0.5$
Mean Full Absorption ( $\uparrow$ )	86.119	87.379	88.475	89.602
Explained Variance ( $\uparrow$ )	72.171	71.860	71.270	69.540

Table 3: Ablation on masking probabilities for Pythia-160m-deduped at

L_{0}=20

with Matryoshka SAEs. As a general trend, explained variance decreases as masking increases. We select

m=0.3

as it achieves the optimal trade-off, while

m=0.0

corresponds to no masking.

Sparse Autoencoders

Sparse autoencoders decompose neural network activations into sparse features to study, interpret, and control learned representations.

See all 217 papers

When Pruning Meets Interpretability: Preserving Sparse Autoencoder Robustness in LLMs
Aug 26, 2026Suchit Gupte, Xueru Zhang, Mohammad Mahdi KhaliliLLM Pruning Sparse Autoencoders
Department of Computer Science The Ohio State University Columbus, OH 43210, USA
C $^{2}$ R: Cross-sample Consistency Regularization Mitigates Feature Splitting and Absorption in Sparse Autoencoders
Jun 29, 2026Haoran Jin, Xiting Wang, Shijie Ren +2LLM Interpretability Representation Disentanglement
University of Science and Technology of China · State Key Laboratory of Cognitive Intelligence · Gaoling School of Artificial Intelligence, Renmin University of China
Towards Understanding the Robustness of Sparse Autoencoders
Apr 20, 2026Ahson Saiyed, Sabrina Sadiekh, Chirag AgarwalAdversarial Attacks on LLMs Jailbreak Robustness
University of Virginia · Independent Researcher

Aug 26, 2026J/K move · Enter open · S save

When Pruning Meets Interpretability: Preserving Sparse Autoencoder Robustness in LLMs

Suchit Gupte, Xueru Zhang, Mohammad Mahdi Khalili

Department of Computer Science The Ohio State University Columbus, OH 43210, USA

LLM Pruning Sparse Autoencoders Neural Network Pruning

Open paper AskPDF arXiv

Sparse autoencoders (SAEs) are widely used to interpret the internal representations of large language models (LLMs), yet their reliability under post-hoc model compression remains poorly understood. We present a systematic study of how pruning affects SAE behavior and theoretically show that, for a fixed SAE, its impact is governed by perturbation energy, a covariance-weighted norm. This perspective exposes a key limitation of magnitude pruning: by ignoring activation geometry, it distorts the learned representation space and degrades SAE functionality. Activation-aware methods such as Wanda and SparseGPT, in contrast, implicitly control perturbation energy and are therefore substantially more robust at preserving SAE behavior. We further reveal a consistent structural vulnerability across all pruning methods: middle layers are significantly more sensitive to pruning than early or late layers. Guided by this insight, we propose a layer-wise sparsity allocation strategy, achieving lower perplexity under the same average pruning sparsity. Experiments across four model architectures validate our theoretical findings. Code is publicly available at https://github.com/osu-srml/sae-robustness-under-pruning/tree/main.

Abstract

Sparse Autoencoders

Explore similar work

When Pruning Meets Interpretability: Preserving Sparse Autoencoder Robustness in LLMs

C2^{2}2R: Cross-sample Consistency Regularization Mitigates Feature Splitting and Absorption in Sparse Autoencoders

Towards Understanding the Robustness of Sparse Autoencoders

C $^{2}$ R: Cross-sample Consistency Regularization Mitigates Feature Splitting and Absorption in Sparse Autoencoders