by Zachary Baker, Yuxiao Li, Maxim Panteleev, Maxim Finenko, Eslam Zaher

August 2025 | SPAR Spring '25

A post in our series "Feature Geometry & Structured Priors in Sparse Autoencoders"

Introduction

Here, we compare the difference of variational sparse autoencoders (vSAE) against a standard sparse autoencoder (SAE). We hypothesize that adding the KL divergence term into the loss function and sampling the activations from a Gaussian distribution should cause a pressure for features to 'move apart' leading to better representations of the language model's latent space, leading to increased performance and interpretability. We ultimately find that this addition does not benefit the interpretability of the language model and degrades performance.

Standard Sparse AutoEncoder (SAE)

The SAE learns a sparse representation of a language model's (LM) activations through deterministic encoding with sparsity regularization. Here, we use the residual stream of a 1-layer GeLU transformer decoder block as the LM activations (the blue input). The SAE is constructed in a very similar manner to an MLP with one hidden layer, as follows:

Encoder:

f = ReLU (W_{e n c} (x) + b_{e n c})

Decoder:

^x = W_{d e c} f + b_{d e c}

Loss Function:

L_{S A E} = ∥ x -^x ∥_{2}^{2} + λ ∥ f ∥_{1}

where λ controls the sparsity-reconstruction tradeoff and the L1 penalty ( $∥ f ∥_{1}$ ) enforces sparse activations.

Variational Sparse AutoEncoder (vSAE)

In the vSAE, the activations of the SAE dictionary features are sampled from a gaussian (mean=activation, variance=1), represented by the horizontal bars in the 'dictionary features' box. This changes the feature as a pointer in some direction representing the state of the LM into an area about some direction. We hypothesize that this will cause a 'pressure' for other features to move away from each other, leading to a more organized space of features for the LM state representations.

The vSAE adds in a probabilistic framework using a Gaussian prior:

Encoder (Mean, as variance is always set to 1):

μ = W_{e n c} (x - b_{d e c}) + b_{e n c}

Reparameterization:

z = μ + ϵ ⊙ σ, ϵ \sim N (0, I)

Decoder:

^x = W_{d e c} z + b_{d e c}

Loss Function:

L_{V S A E} = ∥ x -^x ∥_{2}^{2} + β \cdot KL [q (z | x) ∥ p (z)] + λ ∥ f ∥_{1}

where

q (z | x) = N (μ, σ^{2}),

p (z) = N (0, I)

KL [q (z | x) ∥ p (z)] = \frac{1}{2} d \sum i = 1 (μ_{i}^{2} + σ_{i}^{2} - 1 - log σ_{i}^{2})

The fundamental distinction lies in sparsity mechanisms: standard SAEs achieve sparsity through deterministic ReLU gating and explicit L1 regularization, while vSAEs achieve sparsity through stochastic s...

Posts

Wikitag Contributions

Comments