Enhancing Neural Network Interpretability with Feature-Aligned Sparse Autoencoders

## Abstract

Abstract Sparse Autoencoders (SAEs) have shown promise in improving the interpretability of neural network activations, but can learn features that are not features of the input, limiting their effectiveness. We propose Mutual Feature Regularization (MFR) , a regularization technique for improving feature learning by encouraging SAEs trained in parallel to learn similar features.
We motivate MFR by showing that features learned by multiple SAEs are more likely to correlate with features of the input. By training on synthetic data with known features of the input, we show that MFR can help SAEs learn those features, as we can directly compare the features learned by the SAE with the input features for the synthetic data.
We then scale MFR to SAEs that are trained to denoise electroencephalography (EEG) data and SAEs that are trained to reconstruct GPT-2 Small activations. We show that MFR can improve the reconstruction loss of SAEs by up to 21.21% on GPT-2 Small, and 6.67% on EEG data. Our results suggest that the similarity between features learned by different SAEs can be leveraged to improve SAE training, thereby enhancing performance and the usefulness of SAEs for model interpretability.

## 1 Introduction

Interpretability aims to explain the relationship between neural network internals and neural network outputs. Many interpretability techniques examine raw activations, equating proposed fundamental units of neural networks such as neurons or polytopes to human understandable concepts (Erhan et al., ; Nguyen et al., ; Bau et al., ; Olah et al., ; Black et al., ). These techniques often benefit from a clean correspondence between those fundamental units and concepts, and may fail if concepts are distributed over many units, or many concepts focused in a single unit, such as in the case of feature superposition (Elhage et al., ). We describe features of the input as the atomic, human-understandable concepts represented by input data.

To derive a representation of activations with a stronger one-to-one correspondence of features and neurons, sparse autoencoders (SAEs) have been trained on neural network activations. The decoders of SAEs trained on neural network activations have been shown to form dictionaries of features more easily explained than the neurons themselves, making SAEs potentially useful for understanding the internals of neural networks (Bricken et al., ; Cunningham et al., ; Gao et al., ).

Despite the recent popularity of SAEs, early results suggest they may learn features that are not features of the input, reducing their usefulness for interpretability (Till, ; Huben, ; Anders et al., ). One failure mode considers transformations on the space of inputs: features of the input that are ‘split’ over multiple decoder weights, or multiple features ‘composed’ in one decoder weight. Conceivably, the representation learned by the SAE could be so varied from the input as to contain features entirely incompatible with the input space. These failures are alarming, as studying an SAE would not be guaranteed to reveal information about the neural network that SAE was trained on activations from. In the worst case, if features were commonly split and composed, it is not obvious why studying SAEs would be more useful than studying the raw activations directly, although prior work has given evidence against this.

We hypothesize that if a feature is learned by multiple SAEs, that feature is more likely to be a feature of the input, and show that this is true for SAEs trained on synthetic data comprised of known features. Based on this result, we encourage multiple SAEs trained on the same data to learn common features through conditionally reinitializing SAE weights, and an auxiliary penalty calculated using the similarity of the SAE weights. We name this reinitialization technique and auxiliary penalty Mutual Feature Regularization (MFR).

Using SAEs trained with MFR, we learn more features of the input than baseline SAEs when training on synthetic data (Section 3). We then train SAEs with MFR on activations from GPT-2 Small (Radford et al., ) and on electronencephalography (EEG) data, showing that MFR improves SAEs at scale, and on real-world data (Section 4). Our findings indicate that MFR helps avoid features not in the input space and improves performance on key SAE evaluations, potentially increasing their usefulness for interpretability.

Figure: Figure 1: Our experimental pipeline for training SAEs with MFR. In step one, we extract activations from a neural network, represented by the interconnected nodes on the left. These activations are the inputs for our SAEs. In step two, we train multiple SAEs on the extracted activations. Each SAE learns to reconstruct the input activations through a sparsity constraint on the hidden layer. MFR involves several steps: We first check for inactive features in the SAE hidden state after applying the TopK activation function. If too many inactive features are detected, we reinitialize the weights of the affected SAE. We also include an auxiliary penalty to encourage the SAEs to learn similar features, shown by the final text box.
Refer to caption: x1.png

## 2 Background

### 2.1 Sparse autoencoders

Olshausen and Field () introduced unsupervised learning of sparse representations, capturing structure in data more efficiently than dense representations. Sparse autoencoders (SAEs) have since found wide application in domains such as representation learning Coates et al. (); Henaff et al. (), denoising Vincent et al. (); Duan et al. (), and anomaly detection Sakurada and Yairi (); Xu et al. ().

SAEs reconstruct an input $\mathbf{x}\in\mathbb{R}^{d}$ through a hidden representation $\mathbf{h}\in\mathbb{R}^{h}$, minimizing the reconstruction loss $\left\|\mathbf{x}-\hat{\mathbf{x}}\right\|_{2}^{2}$ while maintaining sparsity in $\mathbf{h}$, written as $\mathbf{\hat{x}}=\mathbf{W}^{\prime}\sigma(\mathbf{W}\mathbf{x}+\mathbf{b})$, where $\mathbf{W}\in\mathbb{R}^{h\times d}$ is the encoder weight matrix, $\mathbf{W}^{\prime}\in\mathbb{R}^{d\times h}$ is the decoder weight matrix, $\mathbf{b}\in\mathbb{R}^{h}$ is the encoder bias, and $\sigma$ is an activation function on $\mathbf{h}$. Ideally columns of $\mathbf{W}^{\prime}$ correspond to features comprising $\mathbf{x}$.

Recent work has shown that the TopK activation function (Makhzani and Frey, ) better approximates the L0 norm in training than alternative techniques such as L1 regularization (Gao et al., ), allowing more precise control over the sparsity of $\mathbf{h}$. The TopK activation function on $\mathbf{h}$ is defined as:

$$ $\sigma_{k}(\mathbf{h})_{i}=\begin{cases}h_{i}&\text{if }h_{i}\geq\tau_{k}( \mathbf{h})\\ 0&\text{otherwise}\end{cases}$ $$

where $\tau_{k}(\mathbf{h})$ is the $k$th largest activation in $\mathbf{h}$. We focus exclusively on SAEs with a TopK activation function, and use the transpose of the encoder weights as the decoder, halving the number of trainable parameters with minimal performance impact as shown by Cunningham et al. ().

### 2.2 Evaluating Sparse autoencoders

Evaluating SAEs is challenging due to the lack of a ground truth for the input features represented by large neural networks. Thus, SAE evaluations act as proxies the extent to which interpretable representation of these input features have been learned in a way that does not require access to them.

The reconstruction loss, measured as the Euclidean distance between the SAE input and output, is a widely used metric for the faithfulness of an SAE’s learned representations to the input features represented by a neural network. Although reconstruction loss does not account for the interpretability of the learned features, improved reconstruction loss has previously been accompanied by improved performance on interpretability evaluations Rajamanoharan et al. (); Gao et al. (); Rajamanoharan et al. (). However, the reconstruction loss is not itself sufficient to evaluate SAEs, as there may be solutions to optimizing reconstruction loss that do not preserve the structure of the input features or convey information about them, such as learning the identity function.

The interpretability of SAE features has been evaluated by the ability of humans and language models to accurately describe those features Bricken et al. (); Cunningham et al. (); Rajamanoharan et al. (). In this evaluation, humans and language models generate feature descriptions based on token sequences and their corresponding activations. The accuracy of these descriptions is then evaluated by predicting feature activations on unseen tokens. The correlation between predicted and true activations, typically quantified using the Pearson correlation coefficient, is used as a measure of description accuracy. However, recent work has critiqued this method, suggesting that even highly accurate feature descriptions may not faithfully represent the model being explained Huang et al. ().

Alternative SAE evaluations analyze the sparsity of SAE outputs through the L0 norm, the presence of consistently inactive features, and ’loss recovered’. Loss recovered refers to the discrepancy in model loss between zero ablation of a layer and the insertion of an SAE output as if it were the activations of that model. The motivation for loss recovered is that it more directly approximates the information preserved by the SAE output, as this may not be accurately measured by the reconstruction loss.

### 2.3 Semi-supervised learning with multiple models

Semi-supervised learning with multiple models involves training several models on both human-labeled and model-labeled data. Co-training, a semi-supervised technique, uses two distinct and conditionally independent views of the same data to iteratively improve the performance of two classifiers Blum and Mitchell (); Zhou and Li (). This approach aims to maximize agreement between the classifiers and has been shown to improve their accuracy Nigam et al. (). Similar techniques have been applied in deep learning for tasks such as machine translation Xia et al. () and image recognition Qiao et al. ().

‘Mutual learning’ and ‘co-teaching’ have been used to describe techniques where student models trained in parallel teach each other by minimizing the divergence between their predictions. These methods have shown superior performance for the size of the student model compared to distillations of larger models Zhang et al. (); Tarvainen and Valpola (); Han et al. (); Ke et al. (); Wu and Xia (). Related techniques include temporal ensembling, which improves model robustness by averaging predictions over multiple training epochs Laine and Aila (), and fraternal dropout, which encourages models trained in parallel to make similar predictions for the same data points, serving as a method of regularization to prevent overfitting Żołna et al. ().

Our work builds on the semi-supervised learning literature. However, our motivation differs from the motivation for most of these techniques, which often relates to a lack of training data rather than learning features not in the input space.

## 3 Experiments with synthetic data

Following Sharkey et al. (), we generate a synthetic dataset of vectors that represent more features than their dimensionality, allowing a direct measurement of how well an SAE learns the features of inputs with superposition. This simplifies our analysis by mitigating the problem of imperfect SAE evaluations (discussed in Section 2.2) as we can then directly compare the SAE and input features, but is only applicable when the input features are known, which is not typical for real-world data.

### 3.1 Generating synthetic data

We aim to create a synthetic dataset of vectors similar to activation vectors sampled from a neural network, but where the features of the input represented by the network are known and can be compared with the features learned by the SAE. These vectors should have similar properties to neural network activations, such as representing superposed features, and having correlations in the activation of features, but simultaneously be learnable by SAEs.

To do so, we generate a dataset $\mathcal{D}=\{\mathbf{x}^{(1)},\mathbf{x}^{(2)},\dots,\mathbf{x}^{(N)}\}$ of vectors $\mathbf{x}^{(i)}\in\mathbb{R}^{d}$. Each vector represents the activation of $G$ features in $d$ dimensions, where $G>d$, is intended to mimic feature superposition in neural networks. We define the feature matrix $\mathbf{F}\in\mathbb{R}^{d\times G}$, where each element is sampled from a standard normal distribution:

$$ $F_{ij}\sim\mathcal{N}(0,1)\quad\forall i\in{1,\ldots,d},j\in{1,\ldots,G}$ $$

We assign probabilities to each feature in $\mathbf{F}$ based on its index through exponential decay:

$$ $p_{j}=\frac{\lambda^{j}}{\sum_{k=1}^{G}\lambda^{k}}\quad\forall j\in{1,\ldots,G}$ $$

where $\lambda\in(0,1)$ is the decay rate hyperparameter, and is a specified constant. By raising $\lambda$ to the power of the index of the feature, we increase the decay for that feature such that the probability of a feature’s activation decreases exponentially with its index.

To introduce correlations in the activations of features, we partition the features into $E$ groups of equal size $S=\frac{G}{E}$. Let $\mathcal{G}_{e}$ be the set of indices for features in group $e$:

$$ $\mathcal{G}_{e}=\{(e-1)S+1,\ldots,eS\}\quad\forall e\in{1,\ldots,E}$ $$

To construct each data point $\mathbf{x}^{(i)}$, we randomly select an active group $e_{i}\in{1,\ldots,E}$ and choose $K$ active features within that group according to the probability $p_{j}$ of a feature. We denote this set $\mathcal{A}_{i}=\{\mathbf{f}_{1,i},\dots,\mathbf{f}_{j,i},\dots,\mathbf{f}_{K,
i}\}$, where $\mathcal{A}_{i}\subset\mathcal{G}{e_{i}}$. Finally, we sample a sparse feature coefficients $a_{ij}$ for each feature in each sample according to:

$$ $a_{ij}=\begin{cases}u_{ij}&\text{if }j\in\mathcal{A}_{i}\\ 0&\text{otherwise}\end{cases}$ $$

where $u_{ij}\sim\mathcal{U}(0,1)$ are the non-zero coefficients. $\mathcal{D}$ is created by linearly combining the ground truth features using the sparse feature coefficients:

$$ $\mathbf{x}^{(i)}=\sum_{j=1}^{G}a_{ij}\mathbf{f}_{j}$ $$

where $\mathbf{f}_{j}$ is the $j$-th column of the ground truth feature matrix $\mathbf{F}$.

### 3.2 Training sparse autoencoders with mutual regularization on synthetic data

We train SAEs with and without MFR on the synthetic dataset $\mathcal{D}$, generated with the parameters $G=512$, $d=256$, $E=12$, $K=3$ and $\lambda=0.99$ in accordance with Section 3.1 (Figure 7). Samples in $\mathcal{D}$ are then comprised of $512$ features represented in $256$ dimensions with $36$ active features, imitating feature superposition in neural networks. We train on $100$ million unique examples. For the MFR SAEs, we train two SAEs in parallel. A complete description of MFR is given in Section 3.3. All SAEs trained in this section have a hidden size of $512$, equalling the input feature count.

We train with a learning rate of $0.01$ with AdamW, and a batch size of $10,000$. On a $40\text{GB}$ A100 GPU, an SAE with these hyperparameters trains in approximately six minutes. ‘Baseline’ SAEs are trained only to minimize reconstruction loss, with sparsity enforced through the TopK activation function on the hidden state. When comparing baseline SAEs with MFR, we maintain an identical architecture and hyperparameter selection, excluding details specific to MFR. We use the exact value of total active features, 36, for the value of k in the TopK activation function for both

### 3.3 Analysis

We hypothesize that features learned by multiple SAEs trained on the same data will tend to correlate more strongly with a feature of the input than a feature learned by only one SAE. To test this hypothesis, we analyze the decoder weight matrices of two baseline SAEs, and compare them with the feature matrix $\mathbf{F}$ for their training dataset $\mathcal{D}$.

Let $\mathbf{W^{\prime}}^{(1)}$ and $\mathbf{W^{\prime}}^{(2)}$ represent the decoder weight matrices from two SAEs. For each feature $\mathbf{w}^{(1)}_{i}$ in $\mathbf{W^{\prime}}^{(1)}$, we find the corresponding feature $\mathbf{w}^{(2)}_{j^{*}}$ in $\mathbf{W^{\prime}}^{(2)}$ that maximizes cosine similarity:

$$ $j^{*}=\arg\max_{j}\cos\left(\mathbf{w}^{(1)}_{i},\mathbf{w}^{(2)}_{j}\right).$ $$

Likewise, for each $\mathbf{w}^{(2)}_{j}$ in $\mathbf{W^{\prime}}^{(2)}$, we find the most similar feature $\mathbf{w}^{(1)}_{i^{*}}$ in $\mathbf{W^{\prime}}^{(1)}$ using the same cosine similarity maximization, resulting in pairs of features between the two SAEs that are most similar. To ensure a one-to-one correspondence of features, we use the Hungarian algorithm to assign the pairs.

We again use the same cosine similarity maximization, this time between $\mathbf{W^{\prime}}^{(1)}$ and $\mathbf{F}$, as well as $\mathbf{W^{\prime}}^{(2)}$ and $\mathbf{F}$, finding pairs of features between a decoder and feature matrix that are most similar. We plot these similarities for all features in $\mathbf{W^{\prime}}^{(1)}$ and $\mathbf{W^{\prime}}^{(2)}$ in Figure 3, illustrating the positive relationship (correlation coefficient = 0.625) between feature similarity across SAEs, and feature similarity with the input features.

Figure: Figure 2: The relationship between feature similarity across SAEs, and feature similarity with the input features for two baseline SAEs.
Refer to caption: x2.png

This correlation is weakened by a cluster of features with high similarity across SAEs, but low similarity with $\mathbf{F}$, potentially harming SAE performance due to the lack of similar input features for features in that cluster. We found that features in this cluster were significantly less likely to be active after the TopK activation function (Figure 5). By avoiding learning this cluster of features, we could improve SAE performance, as it comprises many of the features uncorrelated with those in $\mathbf{F}$. Additionally, it could increase the correlation between feature similarity across SAEs and feature similarity with features in $\mathbf{F}$, potentially allowing for further improvements in SAE performance by encouraging SAEs to learn features present in both their decoder, and the decoders of other SAEs trained on the same data.

Figure: Figure 4: The relationship between feature similarity across SAEs, feature similarity with the input features and the likelihood a feature is active after the TopK activation function on the hidden representation for two baseline SAEs.
Refer to caption: x4.png

Across multiple training runs, a subset would have a reduction in the size of this cluster, resulting in SAEs learning features that correlate more strongly with features in $\mathbf{F}$ (Figure 8). This variation was binary: either the cluster would be larger, at approximately 15 features per SAE, or smaller, at approximately 5. We did not observe other variations. As no hyperparameters were modified, we hypothesize that this is caused by differences in the random weight initializations, and found that we could reliably detect these superior weight initializations by the presence of features consistently not active after the TopK activation function, often in the first 100 training steps (Figure 7).

By reinitializing the SAE weights if a measure of these inactive features exceeded a threshold, we consistently find initializations that do not result in that cluster of features being learned. Doing so strengthens the correlation between the similarity of features learned across SAEs, and the similarity of features learned by the SAEs with $\mathbf{F}$ (correlation coefficient = 0.625 increased to 0.668) (Figure 3).

We found that the particular metric used to decide whether to reinitialization the SAE weights did not effect performance, as the behavior of initializations with smaller clusters of these features uncorrelated with features in the input feature matrix were easily identified by all metrics tested that measure feature inactivity after the TopK activation function. We give an example in Figure 7, plotting the deviation of features from the mean activation probability of a feature, calculated as the value of $k$ used for the TopK activation function divided by the decoder size.

Finally, to incentivize features present in the decoders of other SAEs trained on the same data, we add an auxiliary penalty to the SAE loss function. We define this auxiliary penalty as

$$ $\frac{\alpha}{\binom{N}{2}}\sum_{i=1}^{N-1}\sum_{j=i+1}^{N}(1-\text{MMCS}( \mathbf{W}^{(i)},\mathbf{W}^{(j)}))$ $$

where $\alpha$ is a constant that weights the penalty, $N$ is the number of SAEs, and MMCS is a function that returns the mean of the max cosine similarity pairs across the weight matrices $\mathbf{W}^{(i)}$ and $\mathbf{W}^{(j)}$. We calculate MMCS as

$$ $\text{MMCS}(\mathbf{W}^{(i)},\mathbf{W}^{(j)})=\frac{1}{|\mathbf{W}^{(i)}|} \sum_{w_{i}\in\mathbf{W}^{(i)}}\max_{w_{j}\in\mathbf{W}^{(j)}}\text{CosineSim} (w_{i},w_{j}).$ $$

We name the combined use of our reinitialization method and auxiliary penalty MFR. We find that MFR results in SAEs recovering more of $\mathbf{F}$ (Figure 8), and that SAEs trained with MFR did not have the cluster of features with high similarity across SAEs, but low similarity with $\mathbf{F}$ ‘(Figure 5).

Figure: Figure 6: We plot $\frac{1}{N}\sum_{i=1}^{N}\left(\frac{|\mathbf{W}x_{i}-(k/N)|}{k/N}\right)$, where $N$ is the neuron count of $\mathbf{W}$ and $k$ is the number of active neurons in the hidden layer after $\sigma_{k}$. $k/N$ is then the frequency each feature would be active if all features were equally likely to activate. Hyperparameters were identical across runs.
Refer to caption: x6.png

Figure: Figure 8: MMCS of the decoder weights with the input feature matrix of a baseline SAE, and two SAEs trained with MFR. The MFR and ’superior initialization’ SAEs are reinitialized if $\frac{1}{N}\sum_{i=1}^{N}\left(\frac{|\mathbf{W}x_{i}-(k/N)|}{k/N}\right)=1$, which serves as a threshold of feature inactivity. For the MFR SAE, the constant $\alpha$ that weights the auxiliary penalty is set to 3.
Refer to caption: x8.png

## 4 Scaling mutual regularization

In this section we scale MFR to larger models and real-world data. We train SAEs with MFR to reconstruction activations sampled from GPT-2 Small, or to reconstruct EEG data, showing improved performance compared to baselines. We choose these tasks because they demonstrate the results in Section 3 generalize to natural data from a neural network, and to a non-interpretability task: denoising.

### 4.1 GPT-2 Small

We train five baseline and five MFR SAEs for 2,000,000 tokens on the first layer MLP outputs of GPT-2 Small, constraining the active neurons in a hidden layer of size 3072 to 6, 12, 18, 24 and 30 respectively. We use a batch size of 500, and a learning rate of 0.001 with AdamW. On a single V100, this takes approximately 2 hours to train both baseline and MFR SAEs.

For the MFR SAEs, we set the coefficient that weights the auxiliary penalty $\alpha$ such that the initial reconstruction loss and auxiliary penalty are equivalent, and use a 100 training step cosine warmup for the penalty. The penalty is applied to all five SAEs trained with MFR, such that they are all encouraged to learn similar features in training. We found that the warmup could prevent features becoming too similar early in training, and would allow setting $\alpha$ large enough to cause convergence later in training without increasing the reconstruction loss.

The three SAEs with the smallest values of $k$ (6, 12, 18) achieved superior reconstruction loss using MFR (Figure 9). For the two remaining SAEs ($k=24$ and $k=32$), we found equivalent or inferior reconstruction loss. Over the five SAEs, we found a mean reduction in the reconstruction loss of 5.66%. The most significant improvement was in the $k=6$ SAE, with a reduction of 21.21%, and the most significant degredation was in the $k=30$ SAE, with an increase of 7.89% (Table 1).

**Table 1: Comparison of the final reconstruction accuracy of SAEs trained with and without MFR.**
| k | MFR | Baseline | Relative Difference |
| --- | --- | --- | --- |
| 6 | 0.00132 | 0.00160 | -21.21% |
| 12 | 0.00121 | 0.00135 | -11.57% |
| 18 | 0.00116 | 0.00122 | -5.17% |
| 24 | 0.00114 | 0.00112 | 1.75% |
| 30 | 0.00114 | 0.00105 | 7.89% |

MFR consistently results in superior loss recovered compared to baselines (Figure 9). For this metric we extract the layer 0 MLP outputs of GPT-2 Small and reconstruct them using SAEs. We then insert the SAE outputs as though they were the MLP layer outputs, and measure the cross-entropy loss of the model on 10,000 randomly selected sequences from OpenWebText (Gao et al., ). We find a mean improvement of 5.45% in the MFR SAEs, a maximum improvement of 8.58%, and a minimum improvement of 3.51% (Table 2). With no modifications, GPT-2 Small’s cross-entropy loss on this dataset is 3.12, and 132.27 with the first MLP layer zero ablated.

**Table 2: The cross-entropy loss of GPT-2 Small on a subset of OpenWebText2 with SAE outputs inserted as MLP outputs.**
| k | MFR | Baseline | Relative Difference |
| --- | --- | --- | --- |
| 6 | 10.121 | 10.798 | -6.27% |
| 12 | 9.782 | 10.300 | -5.03% |
| 18 | 9.367 | 9.742 | -3.85% |
| 24 | 10.136 | 10.505 | -3.51% |
| 30 | 9.624 | 10.527 | -8.58% |

We believe these results suggest MFR causes SAEs to learn more information about the features that underly their training dataset. Specifically, the reduced loss recovered indicates the SAEs preserve more information about their inputs, and the improvements in reconstruction loss show more accurate reconstructions in terms of Euclidean distance to the input, but that this depends on the value of $k$ in the TopK activation function relative to the other SAEs being trained.

Figure: Figure 9: The reconstruction loss and loss recovered of various SAEs trained on activations from the first MLP layer of GPT-2 Small. We plot the reconstruction loss on a logarithmic scale.
Refer to caption: x9.png

### 4.2 Electroencephalography Data

We train five baseline and five MFR SAEs for 3,500,000 tokens at a learning rate of 0.001 with AdamW on vectorized EEG data from the TUH EEG Corpus (Obeid and Picone, ). We use a hidden size of 4096 and values of 12, 24, 36, 48 and 60 for the TopK activation function. We preprocess the EEG data with a low-cut frequency of 0.5Hz, a high-cut frequency of 45Hz and a filter order of 5. $\alpha$ is set to equal the initial reconstruction loss, and we use a cosine warmup of 100 steps. The auxiliary MFR penalty considers all five MFR SAEs. On a single V100 GPU, with the above hyperparameters, the five baseline and MFR SAEs train in one hour with a training batch size of 1024.

We train on EEG data to show that MFR can be applied to SAEs trained on non-neural network data. SAEs have been applied to EEG denoising in the past (Qiu et al., ; Li et al., ), and in both finding more interpretable representations of neural network activations and denoising accurate feature learning is beneficial, so plausibly MFR is useful for denoising EEG data.

For the reconstruction loss, we find a mean improvement of 1.8%, a maximum improvement of 6.67%, and a maximum degradation of 4.04% (Figure 10). The benefits of MFR on this dataset are reduced significantly from Section 4.1. We hypothesize that this is because MFR is designed to encourage SAEs to learn accurate representations of input features in which features are represented with superposition in the training data. Although there is evidence that features are superposed in neural network activations (Elhage et al., ; Jermyn et al., ), the same evidence is not present for EEG data.

Figure: Figure 10: The reconstruction accuracy and loss recovered of various SAEs trained on vectorized EEG data from the TUH EEG Corpus.
Refer to caption: x10.png

**Table 3: The reconstruction loss of SAEs trained with and without MFR on vectorized EEG data from the TUH EEG Corpus, and the L2 distance of the decoder weights of those SAEs from the decoder weights of the $k$=6 SAE. We plot of the reconstruction loss on a logarithmic scale.**
| k | MFR | Baseline | Relative Difference |
| --- | --- | --- | --- |
| 12 | 0.42 | 0.44 | -4.76% |
| 24 | 0.30 | 0.32 | -6.67% |
| 36 | 0.26 | 0.27 | -3.85% |
| 48 | 0.23 | 0.23 | 1.80% |
| 60 | 0.21 | 0.20 | 4.04% |

### 4.3 Weight analysis

One concern with MFR may be that in encouraging the SAE decoder features to be similar, the decoder weight matrices end up more similar than without MFR. This could be problematic, as SAEs with lower values of $k$ for the TopK activation that were trained with MFR alongside SAEs with higher values of $k$ could end up less sparse by becoming more similar to the SAEs with higher values of $k$. To test this, we measure the L2 distance of the decoder weight matrices for the baseline and MFR SAEs trained in Section 10.

We find that SAEs trained with MFR tend to be more different in terms of the L2 distance, but that as $k$ increases they trend toward lower L2 distances (Figure 10). This is in contrast to the baseline SAEs, which are more similar at smaller values of $k$, but trend towards larger L2 distances as $k$ increases. At the values of $k$ we train at, we do not consider this problematic, as the L2 distances suggest the decoder weight matrices are more different on average rather than less.

## 5 Conclusion

We proposed a method for training SAEs designed to recover more features of the input. We first establish a motivating hypothesis for MFR, that feature similarity across SAEs is correlated with feature similarity to the input features, showing that this hypothesis is true for SAEs trained on synthetic data, and that MFR improves the fraction of features of the input recovered (Section 3). We then scale MFR to both language model activations and a realistic denoising task, and show that it improves SAE performance on key metrics at scale (Section 4). We believe our method encourages SAEs to learn more features of the input, increasing their usefulness for interpretability.

### Limitations

Although we show improved performance of SAEs with MFR, this comes at a relative increase in computational cost, as the auxiliary penalty used in MFR requires training additional SAEs. As all of our experiments are easily completed on a single GPU, this is not problematic in our work. However, larger models can require SAEs with very large hidden dimensions, making this cost unmanageable if SAEs need to be trained for many layers. Training a multiple of the SAEs that would need to be trained without the auxiliary penalty may not be justifiable depending on the scale of the experimentation.

Despite the increased computational requirements, we believe that the auxiliary penalty is still valuable due to the small computational budget for SAEs relative to the models they are trained on. For smaller models (where the cost of training SAEs is less significant), it may be worth the increase in training compute for more accurate SAEs. For example, in the case of GPT-2 Small, the additional compute may not be of concern, as training is manageable on a single GPU or a small cluster making the additional information about the input features worth the additional compute.

We hope that future work will investigate efficient mutual learning-based approaches for SAEs that can benefit from the positive relationship between feature similarity across SAEs, and feature similarity with the input features without requiring more compute.

### Reproducibility Statement

To ensure the reproducibility of our results, code for all our experiments and instructions for running them are available at [https://github.com/luke-marks0/mutual-feature-regularization/](https://github.com/luke-marks0/mutual-feature-regularization/).

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 3
presentation: 2
questions: - Have you ensured that splits do not have overlapping subjects for EEG data?
- Looking at Figure 5, the model performs as Multi-View models would perform (e.g., SimSiam, SimCLR). The difference is you have a separate Encoder and Decoder, like in DCCAE, instead of multi-view input. Looking at Figure 4, you have a much more diverse representation, or it is because the SAE learned different subspaces. Do you think your regularization will reduce the diversity of the features learned from the data?
soundness: 2
strengths: - Authors performed experiments with synthetic, EEG, and GPT2 weights datasets.
summary: The manuscript introduces Mutual Feature Regularization (MFR), a technique to improve Sparse Autoencoder training by encouraging multiple SAEs trained in parallel to learn similar features. Using synthetic data with known features, the authors validated that MFR helps SAEs learn actual input features more effectively than SAE. Furthermore, the technique improves reconstruction loss by up to 21.21% on GPT-2 Small activations and 6.67% on EEG data.
weaknesses: - Tables 1, 2, and 3 do not report variability across folds or multiple runs. 
- The regularization has to be compared with baselines. For example, DCCAE (Wang et al., 2015) is an autoencoder with CCA that maximizes the correlation between different views. DCCAE also uses different encoders and decoders for each view.
- It needs to be made clear how the splits have been performed. The validation set should be used only for hyperparameter search or checkpoint selection. The final performance must be reported on the hold-out test set.
- It needs to be shown how improvements in reconstruction with regularization improved interpretability. Consider performing analysis with qualitative examples of latents.

Wang, Weiran, et al. "On deep multi-view representation learning." International conference on machine learning. PMLR, 2015.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: Why did the authors decide to use EEG data, and specifically the TUH dataset? \
Line 427: Why is the author’s model only better in terms of reconstruction error for small k values? See table 1 and 3 \
Line 320/321: Do the authors return the mean of the max cosine similarity pairs based on the Hungarian algorithm output? \
Line 244: Why do the authors use cosine similarity maximization and not another distance metric? Since cosine similarity does not take the magnitude of the two matrices into account.
soundness: 4
strengths: This work is of high quality and largely original. Moreover, reading the paper is easy, and the authors do a great job building up their intuitions, related work, and simulation results to verify their model. Results are explained incredibly well, examples of this are Figures 2-6, and I think their method can be significant. I think everything up to the real world results are done incredibly well, and I think the authors have a good base for a strong paper.
summary: The authors propose a regularization framework that both determines whether sparse auto-encoders need to be re-trained and enforces sparse autoencoders trained with different random initializations are learning the same features. The authors validate their results on a simulated dataset, and then apply their model on GPT-2 and EEG data.
weaknesses: Major weakness(es): \
As good as the paper is up to its real world results, I believe the real world results are not as convincing. One of the main issues with the real world results is that the authors do not provide any deeper insights about what their improvements tell us about the data they are analyzing. One main feature of their method in the title is to improve interpretability, but when they apply the model on real world data, they do not actually verify new interpretations, i.e. what does the author’s model tell us about EEG data that the baseline did not? How do the improvements in reconstruction accuracy translate into actionable or interesting results for both the GPT-2 and EEG results? In fact, in their introduction (Section 2.2) the authors mention in Lines 128-130 that “… the reconstruction loss itself is not sufficient to evaluate SAEs, …” which I tend to agree with, but Tables 1-3 and Figure 9&10 only evaluate exactly that (and cross-entropy for GPT-2, but in some sense this is a different measure of reconstruction accuracy: prediction preservation accuracy). I think the authors can make this paper very impactful and significant if they explain to the reader what the added improvements in reconstruction accuracy/cross entropy loss translate to in terms of interpretability, but at this moment the paper seems unfinished. I think the idea the authors propose is important, and I think that with substantial improvements in the regard that I have just described, I am open to significantly increasing my score.

Minor weakness(es):
1. Figure 8 does not show standard deviations, and the authors do not report significance tests for the numbers in their tables (1-3).
2. The authors show many training curves, I believe these should not be added in the main text unless they add something, moving them to the Appendix would provide the authors enough space to dig into interpretability results from their model.
3. The performance of the authors’ model is not better for 24 and 30, 48 and 60 for the GPT-2 and EEG results, respectively. The authors should explain better why this is the case, and potentially research why this phenomenon happens. Since just from a reconstruction error perspective, out of the models the authors have tested their baseline performs the best for both the GPT-2 and EEG data for k=30 and k=60, respectively. If the authors can show that smaller ks help with interpretability, then this is not perse an issue, but since most of the results currently do not actually dive into interpreting the results and mostly focus on reconstruction accuracy, it hurts the impact of the paper.
4. The authors cite two papers on Lines 466-467 that perform EEG denoising, but do not add these papers as baselines.
5. Some of the hyperparameters are not explained, e.g. a hidden size of 4096 and learning rate of 0.001 when training on the EEG data. To improve the confidence in their paper, I would suggest running atleast a small hyperparameter study using the validation set, and selecting the best model for the test set (this applies to all of the authors’ results, including the simulation results).

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 2
questions: Question/comment Set 1:
-"features that are not features of the input"  Rather directly correlate with known features of the input.
- The expression  "features that correlate with features of the input" should be clarified with the fact that the latter corresponds to human interpretability.
-I assume reconstruction loss in the abstract is at the same sparsity level.
-Are polytopes "fundamental units of neural networks" like unit activations?
-The premise of "studying an SAE would not be guaranteed to reveal information about the neural network that SAE was trained on activations from" is strange. Just because the SAE activations aren't as interpretable they are still a function of the activations of the neuron network. Of course by the information processing inequality they can contain less information (or possibly zero if they don't pass anything). However, I do agree with the premise that "studying SAEs would [not always] be more fruitful than studying the raw activations directly".
- It should be made clear that the 'model loss' in the definition of 'loss recovered' refers to a task loss, perhaps the training loss, but could also be a loss function used in validation. References for this definition should be given.



Comment 2:
While not studied formally the significance of the problem reminds me of work studying whether LASSO is a consistent estimator of the true support. The authors considered whether random initialization of the weights is necessary if bootstrap sampling could be used to create different training samples for the same data. This is explored in Francis Bach's work on bootstrap LASSO, which uses the intersection of the support of sparse linear models to find a reliable support. It would be interesting if the proposed regularization works from simply different sampling of the same data without random initialization of the weights. In this case the weights may correspond without the quadratic cost. 
 

Comment 3:
The Hungarian algorithm is used to assign pairs to maximize correspondence with results in Figures 2, 3, 4 and 5; however, the linear assignment problem assumes one-to-one so if there is duplication of a neuron/dictionary element it will lead other pairs to be improperly matched.  So it is not clear to me that when the result shows low similarity, which is interpreted there wasn't a correspondence, it is in fact the case that the assignment didn't correspond.  Partial optimal transport methods that don't require the marginals to be matched exactly can be used to give a more robust matching.

Comment 4:
The synthetic data has some arbitrary choices of simulation (groups all being the same size). Furthermore, after $p_j$ is chosen as active, wouldn't it make sense to bound that sparse features coefficients away from 0? Choosing uniform over [a,1] with a>0 or some other distribution.

Comment 5:
It is not clear to me why a cluster of features that are similar between SAEs but are to the input feature matter if they are not active in the TopK activation? While I agree that reinitializing those features (which to my understanding is common for SAEs) can make the network more efficient.

Comment 6: Asymmetry in penalty across models. See comment in Weaknesses. 

Comment 7:
In Figure 6, the definition of the average of the absolute value is not clear. Why would the activation be k/N in the top-K. The top-K doesn't ensure sums to 1, but rather that it is above the threshold. Otherwise the metric is not meaningful. Its not clear if the reset threshold in caption Figure 8 makes sense.

Question set 8:
Line 363, it is not clear what the initial reconstruction loss was? Is it the first batch, the first epoch? What does the "100 training step cosine warmup for the penalty" mean?

Comment 9:
The lower loss on the task loss in Table 2 is interesting. Especially since the lowest loss is achieved at at k=18 which isn't the lowest reconstruction loss. However, it is hard to tell if this difference is meaningful as the loss is much higher than the no modification loss, while much lower than the zeroing out ablation.

Question 10: 
For the EEG data it is not clear what is meant by vectorized. Is it multichannel windows that is vectorized? If so what is the size of the windows, sampling rate, and number of channels?

Minor points:
-Lines 97–100, need parenthetical references.  Also on lines 126-127, lines 132-133, 138–139 and elsewhere.
-Line 129 "to optimizing"
-Lines 140–141 is not a complete sentence.
-Figure 10 doesn't convey any information to the reader (especially without log y-axis).
soundness: 2
strengths: Directly ensuring robustness of SAE through an alignment loss appears original and seems to provide some gains. It addresses a problem in the use of SAE-based interpretation of networks.
summary: The paper investigates a mutual alignment regularization to train a set of sparse auto-encoders (SAE) in the context enhancing efficiency and interpretability. It does so by penalizing the elements in one dictionary (decoder basis) from not being matched with at least one  element in another dictionary. The method is validated on synthetic data created with a sparse, group-structured model. Only one group is active and $K$ elements in the group are active, based on underlying rates. The method is also tested on a small trained transformer GPT2 small and finally on EEG data. The results show that with fewer active components the regularized networks outperform the unregularized network in terms of reconstruction cost, and for the GPT2 corpus in terms of the cross-entropy loss.
weaknesses: Based on the premise I was expecting more resulting showing improved interpretation. While quantitative results are given, the significance of them is questionable.

The paper's organization could be improved by stating initially the formulation. Figure 1 does not convey much information.  It says on line 81 that "auxiliary penalty calculated using the similarity of the SAE weights" is used, but it is not introduced until Section 3.3 titled "Analysis" there is finally a formulation (in numbered equations) at the bottom of page 6.  Jumping directly to the description of the synthetic data seems out of place.

The formulation of the auxiliary penalty/regularization itself is cryptic. It appears to involve asymmetries (perhaps due to the choice of cardinalities) that are not motivated.  If this is used to regularize all networks, due to the asymmetry, it seems that the index of the network adjusts how it is applied.

The choice of number of SAEs in the regularization isn't precisely given. I take it that the SAEs trained at different levels cardinality (TopK) are co-learned, but this isn't stated explicitly. 

Some of the statements early on are not clear (noted in question/comment set 1 below).

The use of EEG data is not well motivated nor are the details of the data well explained. If there was a downstream task like in GPT2 case that would be more interesting.

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 1
questions: Page 5, lines 225-226: “We see the exact value of total active features, 36, ~~.” What do the authors mean by the exact value, and where does the value 36 come from?

How exactly does the proposed MFR is defined? It is not explicitly entailed.
soundness: 2
strengths: The paper addresses the important issue of interpreting an internal representation of a trained network.

Three datasets, one simulated and two real, were used to support the validity of the proposed method.
summary: This paper introduces a method of enhancing the feature interpretability of a learned neural network via sparse auto-encoders. Specifically, a mutual feature regularization technique has been proposed by correlating multiple SAEs. The experiments were conducted on a simulated dataset and two real-world datasets.
weaknesses: The paper needs to reorganize its structure to describe its ideas and contributions better. In its current form, it isn't easy to understand.
The technical contribution is marginal and is not well supported by the experiments.

This paper is about enhancing the interpretability of an internal representation of a neural network. However, there are quantitative values compared with baseline models, but there are no interpretation-related results or explanations in the experiments. 

When training SAEs, the hidden size plays an important role, which may cause the feature superposition or feature distribution to occur.
In this respect, there should be a comparison among varying hidden sizes in SAEs.  

The current experimental results do not support many hypothetical statements well in Section 3.3. The validity of their arguments needs rigorous efforts.

Page 6, line 307: A statistical significance test is needed to justify the increase in correlation with reinitialization better.

Meta-review: This paper presents a regularization method, referred to as "mutual feature regularization for improving feature learning by encouraging sparse autoencoders trained in parallel to learn similar features. The main claim here is to enhance neural network interpretability.  The paper is easy to read such that intuitions and related work are well described.  The most critical concern raised by all reviewers is that while enhancing interpretability is the main claim, it is not clear how the improved reconstruction improves the interpretability. In addition, experiments on real-world data is not convincing. When the model is applied to real-world data, the authors did not verify new interpretations, i.e.,  what does the proposed model tell us about EEG data that the baseline did not. Therefore, the paper is not recommended for acceptance in its current form. I hope authors found the review comments informative and can improve their paper by addressing these carefully in future submissions.

Decision: Reject