Adaptive Sparse Allocation with Mutual Choice \& Feature Choice Sparse Autoencoders

## Abstract

Abstract Sparse autoencoders (SAEs) are a promising approach to extracting features from neural networks,
enabling model interpretability as well as causal interventions on model internals.
SAEs generate sparse feature representations using a sparsifying activation function that implicitly defines
a set of token-feature matches.
We frame the token-feature matching as a resource allocation problem constrained by a total sparsity upper bound.
For example, TopK SAEs solve this allocation problem with the additional constraint that each token matches with at most k 𝑘 k italic_k features.
In TopK SAEs, the k 𝑘 k italic_k active features per token constraint is the same across tokens,
despite some tokens being more difficult to reconstruct than others.
To address this limitation, we propose two novel SAE variants, Feature Choice SAEs and Mutual Choice SAEs ,
which each allow for a variable number of active features per token.
Feature Choice SAEs solve the sparsity allocation problem under the additional constraint that each feature matches with at most m 𝑚 m italic_m tokens.
Mutual Choice SAEs solve the unrestricted allocation problem where the total sparsity budget can be allocated freely between tokens and features.
Additionally, we introduce a new auxiliary loss function, 𝚊𝚞𝚡 ⁢ _ ⁢ 𝚣𝚒𝚙𝚏 ⁢ _ ⁢ 𝚕𝚘𝚜𝚜 𝚊𝚞𝚡 _ 𝚣𝚒𝚙𝚏 _ 𝚕𝚘𝚜𝚜 \mathtt{aux\_zipf\_loss} typewriter_aux _ typewriter_zipf _ typewriter_loss ,
which generalises the 𝚊𝚞𝚡 ⁢ _ ⁢ 𝚔 ⁢ _ ⁢ 𝚕𝚘𝚜𝚜 𝚊𝚞𝚡 _ 𝚔 _ 𝚕𝚘𝚜𝚜 \mathtt{aux\_k\_loss} typewriter_aux _ typewriter_k _ typewriter_loss to mitigate dead and underutilised features.
Our methods result in SAEs with fewer dead features and improved reconstruction loss at equivalent sparsity levels
as a result of the inherent adaptive computation.
More accurate and scalable feature extraction methods provide a path towards
better understanding and more precise control of foundation models.

## 1 Introduction

Understanding the internal mechanisms of neural networks is a core challenge
in Mechanistic Interpretability.
Increased mechanistic understanding of foundation models could provide
model developers with tools to identify and debug undesirable model behaviour.

Dictionary learning with sparse autoencoders (SAEs) has recently
emerged as a promising approach for extracting sparse, meaningful, and
interpretable features from neural networks, particularly language
models (Cunningham et al., ; Sharkey et al., ).

One problem with wide SAEs for foundation models is that there are often many
dead features (Rajamanoharan et al., ; Templeton et al., ; Gao et al., ).
Dead features are features which are remain inactive across inputs, effectively
wasting model capacity and hampering efficient training.
Another problem is that approaches like TopK SAEs (Gao et al., )
don’t have a natural way to take advantage of Adaptive Computation: spending
more computation, and crucially more features, to reconstruct more difficult tokens.

We frame the problem of generating sparse feature activations corresponding to
some given neural activations as a resource allocation problem, allocating the
scarce total sparsity budget between token-feature matches to maximise the
reconstruction accuracy.
Within this framing, we naturally
motivate two novel SAE variants which can provide Adaptive Computation:
Feature Choice SAEs and Mutual Choice SAEs (FC and MC SAEs respectively).
Feature Choice SAEs solve the sparsity
allocation problem under the additional constraint that each feature matches with at most m tokens.
Mutual Choice SAEs solve the unrestricted allocation problem where the total sparsity budget can
be allocated freely between tokens and features.
These approaches combine the Adaptive Computation of Standard SAEs with
the simple optimisation and improved performance of TopK SAEs.

Figure: Figure 1: A comparison of the pre-activation affinities and the resulting feature activations following different sparsifying activation functions. Red and blue represent positive and negative affinities respectively, with deeper colours representing larger magnitudes. In the first three approaches we have a total sparsity budget of 6. Affinities (Far-left): The token-feature affinities Z’, before any sparsifying activation function. Token Choice/TopK (Center-left): We activate the top $k$ features corresponding to each token. Note that there are features that don’t fire in this batch, which could lead to dead features. Feature Choice (Center): For each feature, it activates corresponding to the top $m$ tokens with the highest affinity. Note that all features fire in this batch. Mutual Choice (Center-right): The elements with the largest magnitude affinities activate, regardless of their token or feature affiliations. ReLU/Standard (Far-right): All strictly positive elements activate. Here we allow low-magnitude feature activations which may be false positives and which cause the $L_{0}$ to be higher. Jump-ReLU SAEs can be seen as a special case of ReLU SAEs where the activation threshold is non-zero.
Refer to caption: extracted/5985782/new_sparse_affinity_matrices.png

Our contributions are as follows:

- •
We provide a framing for sparsifying activation functions in SAEs
as a solution to a resource allocation problem.
- •
We introduce two new SAE architectures: the Mutual Choice
SAE and Feature Choice SAE, which are Pareto improvements
on both standard SAEs and TopK SAEs. Additionally,
the Feature Choice approach is, to our knowledge, the first SAE training method
which reliably results in  zero dead features even at large scale.
- •
We show that our methods naturally enable Adaptive Computation: using
more features to reconstruct more difficult tokens. Instead of setting
the number of features per token as a fixed $k$, we fix $\mathbb{E}[k]$ as a hyperparameter
and allow the model to learn how to allocate the sparsity budget, without
increasing computational overhead.
- •
We propose a novel auxiliary loss function, $\mathtt{aux_zipf_loss}$, which
mitigates under-utilisation of features and better utilises the SAE’s full capacity.

We believe that with increasingly accurate approaches to feature extraction,
it will become possible to connect sparse features over many layers and understand
how models are computing outputs in a mechanistic, circuits-driven fashion.
In particular, given that Feature Choice SAEs generally have no dead features,
they can scale reliably to very large autoencoders, which are
likely to be necessary for effective reconstruction on large
foundation models such as GPT-4 (OpenAI et al., ) or Llama 3 (Dubey et al., ).

The remainder of this paper is organized as follows: In Section 2, we review the sparse autoencoder literature and the challenges in training SAEs.
Section 3 details our resource allocation framing for the sparsifying activation function and introduces the Mutual Choice and Feature Choice SAE architectures (MC and FC respectively).
In Sections 4 & 5, we describe our methods and experimental setup and in Section 6 we present results from training our sparse autoencoders.
We discuss our findings and the limitations of our work in Section 7 and conclude in Section 8 with directions for future research.

## 2 Related Work

### 2.1 Adaptive Computation

In Adaptive Computation, neural networks decide how much compute (and/or which parameters) to apply to a given
input example
(Graves, ; Xue et al., ).
Ideally, the model should learn to apply less compute to easier examples and
more compute to more difficult examples in order to maximise
performance within a compute budget.
In our setting, we consider the token-feature matches to be the scarce quantity to allocate,
where we say that a token matches with a feature if the
feature is activated on that token.

### 2.2 Sparse Autoencoders

Sparse Autoencoders (SAEs) (Lee et al., ; Le et al., ; Mairal et al., ) learn an over-complete basis, or dictionary, of sparsely activating features.
The feature activations, $z$, correspond to their associated neural activations, $x$, via the feature dictionary. In particular, we can write an SAE as:

$$ $\displaystyle z$ $\displaystyle=\sigma_{s}(\text{Enc}(x))\in\mathbb{R}^{F}$ (1) $\displaystyle\hat{x}$ $\displaystyle=\text{Dec}(z)\in\mathbb{R}^{N}$ (2) $$

where $\sigma_{s}$ is a sparsifying activation function (e.g. ReLU), Dec is an affine map and $x\in\mathbb{R}^{N}$
(^1^11In practice, there may be additional pre-processing and post-processing of the neural activations, $x$.).

SAEs are trained to minimize the Reconstruction Error (Mean Squared Error) between $x$ and $\hat{x}$.
This reconstruction error term is combined with an optional
Sparsity Loss term (for example, an $L_{1}$ penalty to induce sparsity)
and an optional Auxiliary Loss term to reduce dead features:

$$ $\mathcal{L}(x)=|x-\hat{x}|_{2}^{2}+\lambda_{1}\mathcal{L}_{\text{sparsity}}(z) +\lambda_{2}\mathcal{L}_{\text{aux}}(x,z,\hat{x})$ (3) $$

Rajamanoharan et al. (); Templeton et al. (); Gao et al. ()
have shown that decomposing neural activations using the SAE feature dictionary
allows for increased human interpretability of models even at
model sizes comparable to frontier foundation models.

### 2.3 TopK SAEs

TopK SAEs (Gao et al., ) use a TopK activation function instead of the $L_{1}$
penalty to induce sparsity, as in Makhzani & Frey ().
Though in the standard $L_{1}$ SAE formulation, the number of features-per-token is variable,
in the TopK formulation the features-per-token is fixed at the same $k$ for all tokens.
We hypothesize that having a fixed $k$ is a key drawback of the TopK method.
Variable $k$ values introduce Adaptive Computation which can focus
more of the token-feature matching budget on more difficult tokens.

In concurrent work, Bussmann et al. () introduce BatchTopK
which is closely analogous to our Mutual Choice SAEs and also provides adaptive computation.
However, they do not deal with the problem of underutilised features.

### 2.4 Dead Features

SAE features which remain inactive across many inputs are known as dead features.
Bricken et al. () declare a feature to be dead
when it hasn’t fired for at least 1e7 tokens.
Dead features present a challenge especially when scaling to
larger models and wider autoencoders.
For example, Templeton et al. () find 64.7% of
features are dead for their autoencoders with 34M features.
Our Feature Choice approach naturally results in  zero dead features by ensuring that each feature activates for every batch.

### 2.5 Auxiliary k loss function

Gao et al. () propose the auxiliary loss function, $\mathtt{aux_k_loss}$
to reduce the proportion of dead features.
Given the SAE residual $e=x-\hat{x}$, they define
the auxiliary loss $\mathcal{L}_{\text{aux}}$ = $|e-\hat{e}|^{2}$,
where $\hat{e}$ = Dec($z_{\text{dead}}$) is the reconstruction using the top $k_{\text{aux}}$ dead features.
Gao et al. () find fewer dead features when using the $\mathtt{aux_k_loss}$ for SAE training.

However, the $\mathtt{aux_k_loss}$ is only applied to features which qualify as dead.
We apply a similar auxiliary loss, the $\mathtt{aux_zipf_loss}$, to underutilised, but not yet dead, features.

## 3 Background

### 3.1 Sparsifying Activation Functions as Resource Allocators

We consider the following many-to-many matching problem:

- •
We have $F$ features and a batch of $B$ tokens.
We would like to have at most $M$ token-feature matches,
where we say that a token matches with a feature if the feature is activated on that token
(^2^22Here we follow the Mixture of Expert literature in abusing notation slightly to refer to neural activations corresponding to a given token position as a ”token”.)
.
- •
We would like to allocate our budget of $M$ token-feature matches such that the reconstruction error is minimised.

Formally, we seek a reconstruction-error optimal weighted subgraph
$\mathcal{H}\subseteq\mathcal{G}=\{\{1,...,B\}\times\{1,...,F\},\textbf{E}\}$
where $\mathcal{H}$ has at most $M$ edges, for $M\ll BF$.
The edge weights can be viewed as the token-feature affinities:
the pre-sparsifying activation feature magnitudes z’.
(^3^33As an illustrative example of the same problem, we can imagine a university which is able to confer at most $M$ degrees where a student may take many degree subjects and a degree subject may admit many students.)

This problem doesn’t immediately admit an efficient solution because it is currently
unspecified how the edge weights contribute to the token reconstruction error.
We make a simplifying assumption that we denote the Monotonic Importance Heuristic -
the edges with the largest edge weights are likely to represent the most important
contributions to the reconstruction error.
With this heuristic, we can solve the problem of allocating token-feature matches
by choosing the M edges with the
largest magnitude edge weights as the edges for our subgraph $\mathcal{H}$.

Figure: Figure 2: $\mathcal{G}$ (left) is a weighted bipartite graph $\mathcal{G}=\{\{T_{1},T_{2},T_{3}\}\times\{F_{1},F_{2},...,F_{6}\},\textbf{E}\}$. Edge weights represent pre-activation affinities with red and blue representing positive and negative values respectively. We are seeking a subgraph $\mathcal{H}\subseteq\mathcal{G}$ with $M=6$ edges. Here we have defined $\mathcal{H}$ (right) by the TopK method for $k=2$; we select the 2 edges from each token with the largest edge weights.
Refer to caption: extracted/5985782/subgraph.png

We can equivalently view this allocation problem as choosing a binary mask
$\textbf{S}\in\mathbb{R}^{B\times F}$ with at most M non-zero elements
which maximises reconstruction accuracy.
This mask is to be element-wise multiplied with a token-feature affinity matrix $\textbf{Z'}\in\mathbb{R}^{B\times F}$.
Applying the Monotonic Importance Heuristic, we are looking for the mask S such that
$\sum_{i,j}\textbf{Z}_{i,j}=\sum_{i,j}\textbf{Z'}_{i,j}\odot\textbf{S}_{i,j}$ is maximised.

TopK SAEs:
We can now see the TopK SAE approach as a special case of the above allocation problem,
with the additional constraint that the number of features per token
is at most $k$ for each token.
In other words, $\sum_{i}(\textbf{S}_{i,j})=k$ $\forall j$, where $M=kB$.
This leads to the solution of $\textbf{S}=\text{TopKIndices}(\textbf{Z'},\text{dim}=-1)$,
i.e. S picks out the $k$ features with the highest affinity for each token.
Here $\sigma_{s}(\textbf{z'})=\textbf{S}\odot\textbf{z'}$;
element-wise multiplication with S defines our
sparsifying activation function $\sigma_{s}$.

We now consider two other variants of this problem displayed in fig. 1:

Feature Choice SAEs:
Whilst TopK SAEs require each token to match with at most $k$ features,
we instead add the constraint that each feature matches with at most $m$ tokens.
This is equivalently a constraint on the columns of S rather than its rows:
$\sum_{j}(\textbf{S}_{i,j})=m$ $\forall i$, where $M=mF$.
This leads to the solution of $\textbf{S}=\text{TopKIndices}(\textbf{Z'},\text{dim}=-1)$,
i.e. S picks out the $m$ tokens with the highest affinity for each feature.

Mutual Choice SAEs:
Here we don’t add any additional constraints and allow any choice of token-feature matching.
This leads to the solution of $\textbf{S}=\text{TopKIndices}(\textbf{Z'},\text{dim}=(0,1))$,
i.e. S picks out the largest elements of the Z’ affinity matrix, regardless of their position.

The Feature Choice (FC) and Mutual Choice (MC) sparsifying activation functions
can be seen as having the desirable properties of the TopK activation
(for example, preventing activation shrinkage,
reducing the impact of noisy, low magnitude activations,
allowing for a progressive recovery code,
enabling simple model comparison,
not requiring sparsity losses which are in conflict with the reconstruction loss etc.)
but whilst allowing for Adaptive Computation, see fig. 1.
Since we do not have constraints on the number of features per token in
FC or MC, it’s possible for one token to activate (i.e. match with)
more features than another token.

#### 3.1.1 Analogy To Mixture of Experts Routing

We choose our naming here to align with the study of token-expert matching in the Mixture of Experts paradigm, where there is a close analogy.
Token Choice MoE routing strategies (Shazeer et al., ; Fedus et al., ) have the constraint that each token can be
routed to at most $k$ experts allowing for expert imbalances.
On the other hand, Expert Choice routing strategies (Zhou et al., ) have the
constraint that each expert processes at most m tokens,
which eliminates the possibility of underutilised experts
and allows tokens to be routed to a variable number of experts.

Here the intuitions are “each token picks
the $k$ experts to be routed to” and “each expert picks
the $m$ tokens to be routed to that expert” for Token Choice and Expert Choice respectively.
The variety of MoE routing algorithms is explored in Liu et al. ().
The Feature Choice approach we propose is directly analogous to the Expert Choice approach in MoEs
and TopK SAEs are directly analogous to Token Choice MoEs. For this reason, we will call TopK SAEs, Token Choice SAEs to unify our notation.

### 3.2 Distribution of Feature Densities

When analysing the distribution of feature densities in open-source SAEs from
Gao et al. (),
we find that the distributions typically follow a power law
described by the Zipf distribution with $R^{2}$ = 0.982, see fig. 3.

Figure: Figure 3: The Feature Density Curve fits a Zipf curve with $R^{2}=0.982$. The middle part of the feature density distribution (feature 100-20,000) fit the Zipf curve with $R^{2}$=0.996
Refer to caption: extracted/5985782/zipf_r_squared.png

We note that although the Zipf distribution well fits most of the distribution,
there is a considerable residual at the lower end of the distribution (the 20k+ rank tokens).
We suggest that these features are underutilised.
Underutilised features see fewer gradient updates than other features
leading to a self-reinforcing cycle
of these features being less useful and hence further underutilised until they die.

We will refer to these underutilised features as dying features.
We define a dying feature as a feature which is one of the 25% least prevalent features
which is also <60% as prevalent as we might predict from the prevalence rank
of the feature using the fitted Zipf curve.

Previous approaches to dealing with dead features either resampled dead features
(Bricken et al., )
or applied gradients to dead features
(for example $\mathtt{aux_k_loss}$ (Gao et al., ))
but they didn’t address dying features.
We hypothesise that many of the revived dead features were still not appropriately utilised.

To address the problem of dying features,
we add an additional auxiliary loss for dying features,
which is a natural generalisation of the $\mathtt{aux_k_{l}oss}$.
Given the SAE residual $e=x-\hat{x}$, we define
the auxiliary loss $\mathcal{L}_{\text{aux_zipf}}=|e-\hat{e}|^{2}$,
where $\hat{e}$ = Dec($z_{\text{dying}}$) is the reconstruction
using the top $k_{\text{aux}}$ dying features.
We can think of this $\mathtt{aux_zipf_loss}$ acting preventatively on features which could be at risk of
becoming dead and acting rehabilitatively on features which have been recently revived.
In this way, we reduce the proportion of both dead and dying features.

### 3.3 Choosing the Feature Choice Constraint

In the Feature Choice approach, there remains the question of how to distribute the sparse feature
activations across the feature dimension.

The simplest approach to this is to take $m_{i}=\nicefrac{{M}}{{F}}$ for
all i, where F is the number of features.
In other words, each feature can pick exactly m tokens to
process. We call this approach Uniform Feature Choice.

Uniform prevalence is a natural way to organize the features
so that a feature firing provides maximal information about the token,
under the assumption that all features provide approximately equal information.
However, we have seen that in existing open-source SAEs,
all features are not equivalently prevalent.
Instead, they are approximately Zipf-law distributed.
To maintain this distribution of feature density
we choose $m_{i}\sim$ Zipf($\alpha$, $\beta$),
where Zipf represents a truncated Zipf distribution
and i is the rank of a given feature in terms of feature density.

$$ $m_{i}=\text{Zipf}(i)\propto\frac{1}{(i+\beta)^{\alpha}}$ (4) $$

We call the Feature Choice approach where the $m_{i}$ are Zipf-distributed, Zipf Feature Choice, henceforth simply Feature Choice.

## 4 Methods

Our approach is as follows:

- •
Given the Zipf exponent and bias hyperparameters, $\alpha$ and $\beta$,
(^4^44We may obtain these hyperparameters by performing a hyperparameter sweep. Alternatively, if we have trained SAE of the same dimensions, we may run inference with this SAE over a large dataset (100 million tokens) and track the number of times each feature activates as its feature density and each feature’s relative feature density rank. We can then fit the (rank, feature density) pairs to a Zipf distribution and estimate the exponent and bias parameters, $\alpha$ and $\beta$. For GPT-2 sized residual stream activations (n=768), we find $\alpha\approx 1.0,\beta\approx 6.8$. In our experiments, we fix the exponent to exactly 1 for simplicity and we find that we may reuse similar $\alpha$ and $\beta$ parameters across SAE widths.)
, we use algorithm 1 to determine the estimated feature density
for each ranked feature.
We use the estimated feature densities to define the threshold
for dying features for the $\mathtt{aux_zipf_loss}$.
- •
We then train Mutual Choice SAEs with both the $\mathtt{aux_zipf_loss}$ and $\mathtt{aux_k_loss}$.
- •
Finally, we optionally fine-tune these SAEs with the Feature Choice activation function,
adding the constraint on the number of tokens that each feature should process.
Here there are no auxiliary loss terms.

The sparsifying activation functions $\sigma_{s}$ for each approach are all TopK activations where the
TopK is taken over the feature dimension for Token Choice (i.e. TopK),
the batch dimension for Feature Choice and
all dimensions for Mutual Choice.

## 5 Experimental Setup

Inputs: We train our sparse autoencoders on the layer 6 residual stream activations of GPT-2 small (Radford et al., ).
We use a context length of 64 tokens for all experiments.
We preprocess the activations by subtracting the mean over the $d_{\text{model}}$
dimension and normalize all inputs to unit $L_{2}$ norm.
We shuffle the activations for training our SAEs (as in Nanda ()).
All experiments are performed without feature resampling, unless otherwise specified.

Hyperparameters:
We tune learning rates based on Gao et al. ()
suggestion that the learning rate scales like $\sqrt{n}$.
We use the AdamW optimizer (Kingma & Ba, ) and a batch size of 1,536
(^5^55For the Feature Choice approach, it’s important to have sufficiently large minibatch sizes so that each feature is expected to activate every few minibatches.).
We train each SAE for 10,000 steps or until convergence.
We use a weight decay of 1e-5 and apply gradient clipping.
We analyse SAE widths from 4x to 32x the width of the neural activations with 0.8% feature sparsity.
We use gradient accumulation for larger batch sizes.
We do not perform extensive hyperparameter sweeps.

Evaluation: After training, we evaluate autoencoders on sparsity $L_{0}$, reconstruction (MSE) and the difference on the model’s final (Cross-Entropy) loss.
We report a standard normalized version of the loss recovered (%).
We additionally evaluate our SAEs’ interpretability using Juang et al. ()’s automated interpretability (AutoInterp) process.
We report the percentage of dead features across models.

Baselines: We compare our SAEs against Standard (ReLU) SAEs and TopK SAEs
(^6^66We do not explicitly test against Gated SAEs, but Gao et al. () find that TopK SAEs perform similarly or better than Gated SAEs with $1.5\times$ less compute to convergence.).

## 6 Results

Figure: Figure 4: SAEs trained with the Mutual Choice activation function, and those finetuned with the Feature Choice activation function have better downstream loss recovered at equivalent sparsity levels.
Refer to caption: extracted/5985782/new_pareto.png

We find that Feature Choice SAEs are a Pareto improvement upon the Token Choice TopK SAEs,
as in fig. 4.

Figure: Figure 5: The Mutual Choice SAEs with the $\mathtt{aux_zipf_loss}$ applied and the Feature Choice SAEs have fewer dead features than SAEs trained without the $\mathtt{aux_zipf_loss}$.
Refer to caption: extracted/5985782/dead_features_placeholder.png

We find that both Mutual Choice and Feature Choice SAEs provide better utilisation of the SAE
capacity with fewer dead features
than comparable SAE methods fig. 5.

**Table 1: Percentage of Dead Features for Different SAE Types and Model Sizes**
| SAE Type | 6k latent dim | 34m dim | 16m dim |
| --- | --- | --- | --- |
| Standard SAE (Bricken et al., ) | 9.0% | - | >90.0% |
| SAE++ (w/ resampling) (Templeton et al., ) | 5.1 % | 64.7% | - |
| Token Choice (TopK) SAE (Gao et al., ) | 0.0% | - | 7.0% |
| Feature Choice SAE (Ours) | 0% | 0% | 0% |

Templeton et al. ()’s 34 million latent SAEs have a dead feature rate of 64.7% (with resampling);
Gao et al. ()’s 16 million latent SAEs have a dead feature rate of 90% without mitigations and 7% with mitigations.
We find that Feature Choice SAEs can be trained with a 0% dead feature rate for both 16m and 34m latent SAEs as in table 1.

## 7 Discussion

We summarize the comparison between our approach and related approaches
in table 2:

**Table 2: Comparison of SAE Types**
| SAE Type | Performance | Dead Features | Auxiliary Losses Required |
| --- | --- | --- | --- |
| Standard (Bricken et al., ) | Weak | Many | $L_{1}$ Sparsity loss |
| SAE ++ (Templeton et al., ) | Better | Many | Scaled $L_{1}$ sparsity loss |
| Jump-ReLU (Rajamanoharan et al., ) | Better | Fewer | $L_{1}$ Sparsity loss |
| Token Choice (TopK) (Gao et al., ) | Better | Fewer | $\mathtt{aux_k}$ |
| Mutual Choice (ours) | Best | Fewer | $\mathtt{aux_k}$ and $\mathtt{aux_zipf}$ |
| Feature Choice (ours) | Best | None | None |

To our knowledge, we present the first SAE training process which
explicitly contains two separate phases:
first *Mutual Choice training*, then *Feature Choice training*.
We speculate that there may be performance benefits to phased training.

We believe that conceiving of the role of SAE encoders as defining matching/routing
algorithms (similar to other Adaptive Computation work, for
example, within the Mixture of Experts literature) could be a valuable
intuition pump for further improvements to SAE architectures.

Our approaches can also be combined with the MDL-SAE (Ayonrinde et al., )
formulation which treats conciseness (Description Length)
as the relevant quantity for evaluation and model selection rather than sparsity ($L_{0}$).
The Description Length of a set of feature activations is a function of the sparsity,
the distribution of activation patterns for each feature and the SAE width.

### 7.1 Zipf Distributed Features

We find that constraining the number of tokens per feature with the Zipf distribution
outperforms using the uniform distribution by >10% model loss recovered. The large drop in performance
using the uniform distribution compared to the Zipf distribution gives additional evidence
that naturally occurring features are not uniformly distributed.

We hypothesise that the reason that features appear to be Zipf distributed may be strongly analogous to the
reason that the frequency of words in natural languages like English are also Zipf distributed.
Words are the semantic units of sentences.
Since features are the semantic units of computation within language models,
a similar mechanism could explain the empirical tendency for the densities of SAE features
to tend towards the Zipf distribution.

In the Computational Linguistics literature, it is well known that the distribution of words
in natural languages approximately follows a Zipf distribution (Zipf, ).
For example, in written English text, the empirical distribution over words
(treated as a categorical variable) can be modelled as Zipf($\alpha$, $\beta$) = Zipf(1, 2.7).

We speculate that the Preferential Attachment Theory explanation
(Zhu et al., ; Chen, )
for the tendency of words to be Zipf-distributed,
which states that frequently used words (features) tend to be used more often,
may be analogously applicable here.
At initialisation, there is some variance in feature prevalence.
The tokens which are initially most highly activated early in training
receive the most gradient signal and are most refined,
leading to a virtuous cycle where they are more effective and useful
for a larger part of the feature density spectrum.
fig. 6 illustrates this dynamic over time.
We would be excited about further work detailing an explicit mechanism for
why features in Neural Networks tend towards being Zipf distributed.

We also note that, for humans, some concepts appear more prevalent in the world than others.
Hence, if features do map to human-interpretable concepts,
we might also expect variance in the prevalence of SAE features.

### 7.2 Adaptive Computation for Varying Token Difficulty

One benefit of the Mutual Choice and Feature Choice approaches is that
they allow for Adaptive Computation - difficult to reconstruct tokens
(which may represent complex or rare concepts) can be reconstructed using more features,
whereas other tokens which are more straightforward can use fewer features.
Where Token Choice (TopK) SAEs suggest that all tokens are equal;
MC and FC SAEs suggest that, in fact, some tokens are more equal
than others Singer ().

As an extreme example, we might expect that the activations resulting from the $\mathtt{<BOS>}$
token are relatively easy to reconstruct, considering they are both very common and
have exactly the same value every time. We might expect that an effective SAE could learn to
productively reallocate the sparsity budget that would have been spent on the $\mathtt{<BOS>}$
token to more difficult tokens, thus increasing the SAE’s effective capacity.
We provide an example of the features per token distribution in appendix D.

### 7.3 Dead Features

Our methods, especially Feature Choice SAEs, have many fewer dead features than other approaches
(typically zero) without complex and compute-intensive resampling procedures (Bricken et al., ).
This is especially important for large SAEs where the problem of dead features is typically more
significant.
We hypothesize that this might be an additional reason for the improved performance: all features
receive some gradient signal at every step.

We note that this simplified approach eliminates the need for complex
resampling procedures. This approach suggests a
simplified procedure for SAE training and model selection than previous
SAEs with fixed sparsity budget.

### 7.4 Philosophical Notes

In 7.1, we discussed the fact that features which map to
human-understandable concepts and are learned from human-generated data
(like internet text) are likely approximately Zipf distributed.
We would be interested in how this varies for data which is less obviously
human-generated, such as images or video.
We speculate that an information theoretic analysis of the features
of the natural world which have predictive power could
provide evidence for ontological theories.
We suggest that exploring the features apparent in the natural world
information-theoretically could be a valuable lens for
theorists working in the Naming & Necessity paradigm (Kripke, )
or within Natural Abstractions Theory (Wentworth, ).

### 7.5 Limitations

Appeals to the Monotonic Importance Heuristic:
In section 3.1 we defined the Monotonic Importance Heuristic (MIH) -
the assumption that the importance of a feature is monotonically increasing in feature activation magnitude.
We use this assumption when we choose the feature activations with the largest magnitudes with our TopK-style
activation functions.
TopK SAEs also implicitly assume the MIH (Monotonic Importance Heuristic).
We can think of Jump-ReLU SAEs and Gated SAEs as relaxing this assumption slightly
to a weak MIH.
For Jump-ReLU and Gated SAEs, the importance of activations is still related
to their magnitude and they still filter out any low magnitude feature activations;
however, the filtering threshold varies for each feature.
So Jump-ReLU SAEs do not have to make magnitude comparisons
across features which may have different natural scales.
It may be, however, that there are low magnitude activations which,
within a certain context, are nonetheless critically important in capturing
information which is useful for reconstruction and/or downstream model performance.
These important but low magnitude activations are difficult to
capture with our current SAE approaches
(^7^77Standard ReLU SAEs do allow low magnitude feature activations but at the expense of failing to filter out noisy low magnitude activations, which can be seen as false positives (Rajamanoharan et al., ).).

Though the weight-sharing form of Gated SAEs (Rajamanoharan et al., ) implicitly encodes the weak MIH prior,
the non-sharing form does not.
Weight-sharing Gated SAEs, however, tend to perform better.
The improved performance of approaches which encode the MIH prior could
be considered as evidence for the truth of the claim.
Alternatively, we might note that the Monotonic Importance Heuristic
acts as an inductive bias for our models.
Good inductive biases often allow models to perform better at first;
however, with increased scale we may not need such inductive biases
and may prefer allowing the model to learn more of the solution independently
(Xiao, ), (Sutton, ).

Generalisation Across Modalities:
We tested our SAEs within the domain of language.
We currently don’t know to what extent our results generalise across modalities,
especially to inherently continuous modalities like audio and images.
We would be excited about future work applying similar techniques to
Interpretability problems in a wider range of modalities.

Evaluation:
The Mechanistic Interpretability field doesn’t currently have widely agreed upon metrics
for evaluating Sparse Autoencoders.
Disentanglement benchmarks like Huang et al. () have been proposed as well as evaluation on tasks
where the ground truth is known (Karvonen et al., ).
Developing a more comprehensive suite of benchmarks for SAEs would
help us to have higher confidence in our comparisons between SAE variants.

## 8 Conclusion

We introduce the Feature Choice and Mutual Choice SAEs as a simple drop-in change
to the sparsifying activation function in Sparse Autoencoders. We also provide a
new auxiliary loss, $\mathtt{aux_zipf_loss}$, which prevents dying features and hence allows the SAE to
more fully utilize all of its features without wasting capacity.

We would be excited about future work that defines scaling laws for
Features Choice SAE performance on even larger models than the ones we use in this
work.
We believe a promising direction for future work is dissecting the
phenomenology of features (Bricken et al., ) conditioning on
both feature prevalence and feature activation magnitude.
We might ask
the question “what kinds of features appear very often in naturally occurring
datasets and what does this imply about the predictive elements of the natural world?”.
We believe this direction could prove fruitful for both
Machine Learning researchers and philosophers of science.

### Acknowledgments

Thanks to Callum McDougall, Matthew Wearden, Jacob Drori, Alex Cloud, Joseph Miller and Evžen Wybitul for comments on early drafts.

Thanks to Lee Sharkey, Andrew Gritsevskiy, Derik Kauffman, Hans Gundlach, and Jan
Kulveit for useful conversations. Thanks to Sandy Tanwisuth for advice and encouragement.

Thanks to MATS for providing the space to do this work and to EleutherAI for providing compute.

## Appendix A Feature Density Moves Towards A Zipf Distribution Throughout Training

Figure: Figure 6: At initialisation, the feature distribution diverges from the Zipf distribution significantly at a smaller feature rank than after training the feature dictionary.
Refer to caption: extracted/5985782/feature_distribution_at_initialisation.png

We can compare this with fig. 3 where the drop-off comes much later, at 24k vs at 3k.

## Appendix B Inference with Feature Choice and Mutual Choice SAEs

To run inference on our SAE variants, you can do batch inference with the method exactly as in the training setup.
However to do single token (or single sequence) it may be beneficial to instead impute a threshold value and swap out the activation function to use this value instead with a JumpReLU style approach Erichson et al. ().

## Appendix C Neural Feature Matrix Loss

One problem with SAEs is undesirable feature splitting.
Feature splitting occurs when an SAE finds a sparse combination of existing directions that allows for a smaller $L_{0}$ (Ayonrinde et al., ). For example, Bricken et al. () note that a model may learn dozens of features which all represent the letter "P" in different contexts in order to maintain low sparsity.

In order to reduce feature splitting we propose two additional auxiliary losses: $\mathtt{nfm_loss}$ and $\mathtt{nfm_inf_loss}$.

The Neural Feature Matrix (NFM) is defined as $\mathtt{nfm}(W)=\hat{W}\hat{W}^{T}$ for a weight matrix $W$.
The NFM is a symmetric square matrix which describes the correlation of different rows in the matrix $W$.

We define $\mathtt{nfm_loss}=|\mathtt{nfm}(W_{Dec})|_{F}$ and $\mathtt{nfm_inf_loss}=\frac{1}{F}\sum_{i}{max({\mathtt{nfm}(W_{Dec})}_{i}})$.

For our largest runs we apply both of these auxiliary losses with small weight.
Empirically we find this seems to reduce undesirable feature splitting and avoid a failure mode we call "dictionary collapse" when many features of the decoder dictionary start to align with each other.

## Appendix D Adaptive Computation Allows For A Variable Number of Features Per Token

Figure: Figure 7: The number of features per token is variable as a result of the Adaptive Computation.
Refer to caption: extracted/5985782/features_per_token.png

We note that there is a bimodal distribution of features per token. The SAE allocates close to the main mode of features to each token, attempting to allocate more features to harder tokens and less features to easier tokens. The $\mathtt{<BOS>}$ token is responsible for the lower peak - this token is much easier to reconstruct and so it is prudent for the network to allocate fewer features that token. In the ideal SAE, the $\mathtt{<BOS>}$ token could be reconstructed with a single feature, rather than the 10 features required here. Future work might look into optimisations that would allow the $\mathtt{<BOS>}$ token to be single-feature reconstructed as a test case of allowing even greater variance in features per token.

## Appendix E Progressive Codes

Gao et al. () describe learning a progressive code using their Multi TopK
loss. In their setting, the Multi TopK loss (a weighted sum of TopK losses for different values
of $k$) is required because the SAE generally "overfits" to the value of $k$ which harms a progressive code.
In our case, the SAE is robust to having variable $k$ values even for the same token depending on the context of the batch.
Empirically, we obtain progressive codes for a greater range of values of $k$ than in the TopK case.

## Appendix F Determining The Feature Distribution

Given the Zipf exponent and bias hyperparameters, $\alpha$ and $\beta$, we use Algorithm 1 algorithm 1 to determine the estimated feature density
for each ranked feature.
We use the estimated feature densities to define the threshold
for dying features for the $\mathtt{aux_zipf_loss}$.

Figure: Algorithm 1 Calculate Zipf Feature Distribution

## Appendix G Monotonic Importance Heuristic

In section 3 we appeal to the Monotonic Importance Heuristic (MIH), as in TopK SAEs, in order to simplify our allocation problem. Empirically we find that this works well though we discuss the case for not using the MIH in section 7.5.

One theoretical (though informal) motivation for the MIH is as follows.
Since the decoder dictionary is fixed to unit norm,
the norm of any feature’s contribution to the output is exactly the
magnitude of the feature activation to which it corresponds.
Hence if there’s limited cancellation between features (which is likely in an N dimensional space where the per-token sparsity is much less than N) then we might expect the component of $\hat{x}$ in any feature direction to be very close to the feature activation for that feature.
In particular, consider a feature with a small magnitude of $\varepsilon$.
This feature can only possibly influence the reconstruction loss
by at most $\varepsilon|e|$ where $e=x-\hat{x}$.
Features corresponding to larger activations can plausibly influence the reconstruction loss by a greater amount.

## Peer Reviews

confidence: 2
contribution: 2
initial_rating: 5
presentation: 2
questions: See weakness
soundness: 3
strengths: The proposed methods effectively reduce dead features and allow for flexible feature allocations, which allows the model to utilize features better and potentially more interpretable.
summary: The paper proposed Feature Choice and Mutual Choice Sparse Autoencoder (SAE). The feature choice version enforces sparsity by restricting each feature to match with at most $m$ tokens; the mutual choice version allows any choice of token-feature matching. zipf distribution is used to set the constraints for feature choice, and an additional auxiliary loss is used to prevent dead features. Numerical experiments show the proposed methods outperform existing SAE methods in terms of reconstruction at different sparsity levels and percentage of dead features
weaknesses: 1. The experiments are rather limited, only one setting with GPT2 activations is performed. 
2. To demonstrate the significance of the results, I think it would be better to include the average and standard deviation of the evaluated metrics across different runs
3. SAE is described as a promising approach for extracting sparse, meaningful, and interpretable features from neural networks. But under the current experiment setup, it is hard to tell how the extracted features are meaningful and interpretable.
4. The problem is framed as a resource allocator in section 3, but I think it is not very clear how the edge weights are defined. If I understand correctly, it is somehow defined as related to the magnitude of feature activation, e.g. in line 483. 
5. There are some typos, e.g. in lines 210 and 221 $S$ is defined the same way, but shouldn't it be different?  in line 266, l is written in as a subscript in aux_k_loss

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 2
questions: 1. In Line 344, the authors mentioned that they evaluated the interpretability of the SAEs using Juang et al. (2024)'s automated interpretability (AutoInterp) process. Could the authors provide results of this evaluation?

2. Could the authors provide more details about the experimental setup, including the datasets and the exact procedures used, to ensure the reproducibility of the results? Ideally, releasing the codebase would greatly benefit the community.

3. Have the authors consistently observed the discrepancy between the feature density and the Zipf curve in different settings, or is this observation limited to the GPT-2 experiments presented here?
soundness: 2
strengths: Clear Visualization: Figure 1 effectively visualizes the design of the sparsifying activation functions and the differences between the proposed methods and existing approaches.

Adaptive Computation: Feature Choice SAE and Mutual Choice SAEs demonstrate improved performance through adaptive computation, effectively allocating more resources to more complex tokens.

Mitigating Dead Features: The introduction of the aux_zipf_loss and the proposed definition of dying features successfully mitigate the occurrence of dead and underutilized features, leading to better overall model performance.
summary: The paper proposes two new variants of sparse autoencoders (SAEs), namely Feature Choice SAEs and Mutual Choice SAEs, to address the limitations of existing TopK SAEs. A key improvement is that the proposed models adaptively allocate the total sparsity budget between tokens and features, resulting in improved reconstruction performance (Figure 4). Additionally, the paper introduces a new auxiliary loss function, aux_zipf_loss, to mitigate underutilized features based on the observation that the Zipf distribution generally fits most of the feature densities, except at the lower end. The proposed aux_zipf_loss enhances feature utilization and reduces dead features (Table 1).
weaknesses: Organization of Sections: Sections 3 and 4 appear somewhat disorganized, with some concepts in Section 3 that would be better suited for the Methods section. This restructuring could improve the overall flow and clarity of the paper.

Limited Evaluation Scope: The analysis of the Zipf distribution fit (Figure 2) is an interesting discovery, but it appears that the authors only conducted this analysis on GPT-2. Broader validation on different models or datasets would strengthen the claims.

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 2
questions: 1. The paper claims improved reconstruction accuracy, but how does this translate to downstream task performance beyond the provided loss metrics?
2. Could the authors elaborate on the computational overhead of Feature Choice and Mutual Choice SAEs compared to TopK SAEs?
3. How well do these methods generalize to other types of neural networks beyond language models?
4. What are the practical implications of these improvements for real-world applications?
soundness: 3
strengths: 1. The resource allocation perspective on sparsifying activation functions provides an intuitive framework for understanding and improving SAEs.
2. The Feature Choice and Mutual Choice variants offer clear improvements over existing methods. For example, the Feature Choice SAEs reliably achieve zero dead features even at large scale and the scaling results are good that is able to handle up to 34M dimension models.
3. The discussion on the Zipf distribution in feature usage is also thoughtful.
summary: The paper introduces two novel variants of Sparse Autoencoders (SAEs) - Feature Choice SAEs and Mutual Choice SAEs - for extracting interpretable features from neural networks. The key innovation is framing the sparsifying activation function as a resource allocation problem, where a fixed sparsity budget must be distributed across token-feature matches. The authors also introduce aux_zipf_loss, which addresses underutilized features by generalizing the existing aux_k_loss. The proposed methods show improvements in reconstruction accuracy and feature utilization compared to existing approaches like TopK SAEs, particularly in handling dead features.
weaknesses: 1. One of my main concerns is about the presentation. Some technical details are scattered throughout different sections and the Figures and Tables are not clear or elegant which are beyond the width of the paper. The figures could also be better explained. The paper would benefit from clearer organization and improved writing.

2. The experiment results are somewhat not sufficient. Experiments primarily focus on GPT-2 small. More diverse model architectures would strengthen claims.

3. The contribution of this work is somehow incremental with limited significance and impact. The core problem about dead features in SAEs handled by this paper has limited impact in the interpretability community.

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 2
questions: - **Lack Comparison with Concurrent Work:** In addition to the baseline SAE types listed in Table 2, the authors should include a comparison with Bussmann et al. (2024) [1], which introduced BatchTopK-a method that also offers adaptive computation and could serve as a useful benchmark agaist Mutual Choice SAEs.

References:

[1] Bart Bussmann, Patrick Leask, and Neel Nanda. BatchTopK: A Simple Improvement for TopK-SAEs. July 2024. https://www.alignmentforum.org/posts/Nkx6yWZNbAsfvic98/batchtopk-a-simple-improvement-for-topk-saes

- **Line 195**: Could the authors clarify the definition of the pre-sparsifying activation feature magnitudes $\mathbf{z}'$? I assume this refers to the token-feature affinities, but how is this value determined for real applications?

- **Line 258-260**: The authors should provide a more detailed definition of "dying features" to improve clarity.

- **Table 1**: Could the authors clarify the term "w/ resampling"? Is it a fair comparison to evaluate results with and without resampling? Additionally, in practical applications, which approach yields a lower percentage of Dead Features?

- **Line 855-857:** The authors mentioned the failure mode "dictionary Collapse", which I personally think it is analogous the the collapsing problem in sparse mixture of experts models. The authors should demonstrate this via visualization or quantitative empirical results.

- **Line 855-857**: The authors mention the failure mode "dictionary collapse," which I believe is analogous to the collapsing problem seen in sparse mixture of experts models [1,2]. To strengthen this point, the authors should consider demonstrating it through visualizations or quantitative empirical results as in [1,2].

References:

[1] Chi, Z., Dong, L., Huang, S., Dai, D., Ma, S., Patra, B., Singhal, S., Bajaj, P., Song, X., Mao, X.L. and Huang, H., 2022. On the representation collapse of sparse mixture of experts. Advances in Neural Information Processing Systems, 35, pp.34600-34613.

[2] Do, T.G., Pham, Q., Nguyen, T., Doan, T.N., Nguyen, B.T., Liu, C., Ramasamy, S., Li, X. and Steven, H.O.I., HyperRouter: Towards Efficient Training and Inference of Sparse Mixture of Experts. In The 2023 Conference on Empirical Methods in Natural Language Processing.
soundness: 3
strengths: This work has potential for applications needing efficient feature allocation and reduced model complexity, like sparse representation learning. The adaptive computation strategy to reduce dead features could improve SAE design, making it more efficient for high-dimensional settings. Though incremental, the contributions offer valuable insights into current SAE limitations. The paper provides a basic experimental setup with empirical validation on real-world datasets. The originality of this work lies in framing SAEs as resource allocation problems and introducing the Feature Choice SAE and Mutual Choice SAE variants, offering a new approach to adaptive computation.
summary: This paper addresses challenges in sparse autoencoders (SAEs) by introducing Feature Choice SAEs and Mutual Choice SAEs, designed to improve adaptive computation by tackling dead features and feature sparsity. These SAEs are framed as resource allocation problems via the Monotonic Importance Heuristic for adaptive feature selection, where features are adaptively assigned to tokens. Empirical validation of the proposed SAEs demonstrated that they achieve fewer dead features and improved reconstruction loss at similar sparsity levels, attributed to their intrinsic adaptive computation mechanism. Key contributions include new SAE architectures, a heuristic theoretical foundation for adaptive feature allocation, and valuable insights into efficient sparse neural models.
weaknesses: - **Lack of Reproducible Code**: The paper does not include reproducible code for the numerical experimental results, making it difficult to verify the empirical findings. Consequently, my comments are based solely on the results as presented in the paper.

- **Monotonic Importance Heuristic:** The theoretical framing of models such as TopK SAEs, Feature Choice SAEs, and Mutal Choice SAEs as resource allocation problems is coherent. However, it is important to verify the mathemtical rigor in their proofs, even though the authors offer an informal theoretical motivation in Appendix G.

- **Lack of Extensive and Comprehensive Experimental Setup**: Although the comparisons with standard SAEs aim to demonstrate improvements in dead features and reconstruction loss, the evaluation lacks thoroughness and would benefit from a more comprehensive approach (see additional details in Questions).

- **ICLR 2025 Citation Style**: The paper should adhere to the ICLR 2025 citation style: when the authors or publication are part of the sentence, the citation should appear without parentheses using \citet{} (e.g., “See Hinton et al. (2006) for more information.”). Otherwise, the citation should be placed in parentheses using \citep{} (e.g., “Deep learning shows promise to advance AI (Bengio \& LeCun, 2007).”) The authors should revise all citations to ensure they adhere to this standard format.

- **Inconsistent Mathematical Notations**: The notation for multivariate variables in equations (1), (2), (3), and beyond is inconsistent, alternating between bold and regular font. Following a consistent notation style, as recommended by the ICLR guidelines, would enhance clarity.

- **Missing Information on Citations**: There are several instances of missing conference or journal names in the references. The authors should address these omissions to ensure completeness. For example, on lines 648-650, the correct citation format should be: "William Fedus, Barret Zoph, and Noam Shazeer. Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity, 2022. *Journal of Machine Learning Research*, 23.120 (2022): 1-39."

- **Typos:**   **Line 206**: The mask $S$ should be written in bold as $\mathbf{S}$. **Line 266:** aux\_$k_l$oss should be aux\_k\_loss.  **Line 903:** "Algorithm 1 algorithm 1" should be "Algorithm 1". **Line 452:** "fig. 6" should be "Fig. 6".

Meta-review: This work considers the token-feature matching problem derived from sparse autoencoder. The authors cast it into resource allocation framework, and considered two variants distinguished by their own constraints. Empirical results show that the methods result in more active feature participation and lower reconstruction loss.

The reviewers agree that this is an interesting work, and is well motivated.

However, the reviewers raised concerns that the empirical study is insufficient to justify the main claim, due to testing on small scale models and lack of important evaluation metric such as mean and standard deviation. In addition, reviewers also found the paper is not well organized.

Decision: Reject