We thank the reviewer for their careful review and address the weaknesses they brought up as well as their questions here.

Weaknesses

Writing: We agree with the reviewer about the repetitiveness of the empirical observations and have moved the majority of the figures to the appendix. We have also removed unnecessary technical details in section 4. However we believe our informal style is appropriate for this type of empirically backed and exploratory theoretical work, as we believe the function of formalization is primarily to avoid errors or to enshrine well-established results.

Theory: We have addressed the reviewer's concerns regarding the confusing wording of this paragraph. Additionally, we emphasize Section 4.1 is meant to give a qualitative description of the algorithm before rigorous arguments are given in later sections.

Experiments: Our work demonstrates that models become easier to analytically analyze when their weights inherit symmetries in their training data. While we considered the maximal case of permutation symmetry, a similar simplification of the loss occurs if the weights inherit smaller symmetries. As a demonstration we have added\inc{ look at appendix. } Real world data is of course significantly more complicated, but we nevertheless expect model weights to reflect its structure. However, we believe a first principles understanding of how this occurs requires first the development of fundamental tools to analyze toy models of increasing sophistication. You may also see this discussion on the differences between fundamental and applied mechanistic interpretability.

Questions

Q1: An example of the trained weights can be seen in figure 3.

Q2: The size of the off-diagonal noise becomes zero element-wise in the limit, but does not go to zero averaged over a row. Therefore it is not appropriate to write that $W = aI$ in the limit. In fact, the off-diagonal elements encode enough correlations so that $W$ is indeed low-rank. (Note also that $r\to\infty$ was a typo, we meant $n_s\to\infty$ with $r$ fixed.)

Q3: $\mu$ referred to the mean value of the noise $\nu$. We have clarified the exposition by absorbing it into the bias $b$. The goal of the paragraph is to explain the context and give intuition for the calculation that follows, and we have edited it to improve the clarity.

Q4,5: We have more clearly defined a Hadamard matrix by its defining properties, which makes it more clear how the construction works. We have also included a new section in the appendix explaining why it satisfies the permutation symmetry for off-diagonal terms along with a step-by-step discussion of the other two conditions. Additionally, we fixed the $R$ versus $R_{ij}$ typo.

Minor

Mostly done, but will confirm these later. Update: Thanks for finding these issues. We made most of the proposed changes, except for a few minor stylistic ones. We also checked the maximum absolute difference and found it was converging to zero as well (see appendix).

Questions

1. In this context, assuming sparse data of this form is standard (see Elhage et al. 2022, Ganguli and Sompolinksy 2010, and also the references in reviewer ru6M's comments on our work). Changing to $[-1, 0]$ won't matter in a practical case as this corresponds to flipping the sign of the embedding vectors.

• Our work demonstrates that models become easier to analytically analyze when their weights inherit symmetries in their training data. While we considered the maximal case of permutation symmetry, a similar simplification of the loss occurs if the weights inherit smaller symmetries. As a demonstration we have added such an example to the appendix. Real world data is of course significantly more complicated, but we nevertheless expect model weights to reflect its structure. However, we believe a first principles understanding of how this occurs requires first the development of fundamental tools to analyze toy models of increasing sophistication. See also this discussion on the differences between fundamental and applied mechanistic interpretability.

2. The rank condition need not be considered when deriving equation (4) since it is just an identity which arises from the definition of $\nu_i$. The $\nu_i$ certainly correlate with each other, but because each output is purely a function of $x_i + \nu_i$ this fact doesn't matter. No assumptions are made about the off diagonals.

• The current version of the paper should be more clear on this point now: the empirical observations of Section 3 are used as assumptions in Section 4. See the discussion regarding the Barry-Esseen theorem above as well.

3. One way to recover $W_{\text{in}}$ and $W_{\text{out}}$ from $W$ (generally) is to perform an SVD so that $W=UDV$. $W_{\text{out}}$ can then be written as $U\sqrt{D}$ where $\sqrt{D}$ is an $n_s \times n_d$ matrix with diagonals equal to the (positive) square-roots of the non-zero singular values of $W$. Similarly $W_{\text{in}} = \sqrt{D}V$. This is not unique because $W = W_{\text{out}} W_{\text{in}} = (W_{\text{out}}M)(M^{-1} W_{\text{in}})$ for any invertible $M$. This shows that we could equally well have taken different $W_{\text{in}}$ and $W_{\text{out}}$. The pair on line 398 would indeed work.

4. For this section we restricted to $W$ which have 1 on the diagonal rather than $a$ for ease of notation. It's correct that this equation is particular to our setting, in particular the facts that the variance of $x_i$ is $(4p - 3p^2)/12$ and that the $x_i$ are i.i.d.

We thank the referee for the review. Many of these points helped us clarify the writing.

Weaknesses: The theoretical findings of Section 4 should be correct, especially given that our newly derived lower bounds are in agreement with the upper bounds. Nevertheless, we can also add plots that confirm the derived bounds. We will confirm here when this is complete. Update: this is complete!

We have attempted to better explain the jargon and remove unnecessary verbosity. If there are any specific instances of this the reviewer would still like us to address, please let us know.

The toy model of Elhage et al. 2022 is a well-studied model for understanding superposition. It is not a perfect model of LLMs because it only captures how features are linearly encoded in activations. While in practice language models also contain non-linearly encoded features, we believe the methods developed in our work on the linear case will generalize to future toy models that capture these features.

Major

Abstract

In our updated version of the paper we derive a lower bound on the loss valid for any elementwise or gating activation function in the appendix. This clarifies the point made in the abstract.

Section 1

By "permutation symmetric data" we meant that the distribution of $x\in\mathbb{R}^n$ is such that $(x_1,...,x_n)$ is equal in distribution to $(x_{\pi(1)},..., x_{\pi(n)})$ for any permutation $\pi$. By ``thermodynamic limit'' we mean the limit as $n_s\to\infty$. These terms are standard in some parts of the machine learning community, but we have now also defined them explicitly in our text.

Section 2

We trained using SGD online on randomly generated samples. More experimental details can be found in the appendix. (Once the authors are made public, we will also submit a version which has links to a public Github, for easy reproducibility.)

Section 3

1. Update: We have changed the wording. See also the appendix with empirical tests.
2. Multi-line answer:

Thank you for pointing this out, it should indeed be $n_s\to\infty$. We have fixed the typo.

• We believe 8192 shows convergences to Gaussianity for most of the parameter space (see Figure 7 in the new paper or Figure 6 in the original one).
• Yes, thanks for the pointer. We mention this fact in the definition of the model (the beginning of Section 2) now.

3. Thank you, we want the paper to be as precise and readable as possible. We have edited this portion of the paper (beginning of Section 3.2). (In this case, we say a vector is permutation invariant if $x_{i} = x_{\pi(i)}$ for any permutation $\pi$. Similarly, as above, we say a square matrix is permutation invariant $A_{ij}=A_{\pi(i)\pi(j)}$ for any permutation $\pi$.)
4. This sentence was indeed unclear and we removed it in the rewrite of this section.
5. As above.
6. We have made the references to plots more precise.
7. Multi-line answer:

Indeed, we should have defined $\nu$ (or $\nu_i$ for the $i$'th row) earlier. We have done this now.

Indeed the $\Delta$ notation may have been confusing since it's effectively a sample variance (e.g. you were correct to think $\Delta\text{var}(\nu)$ was the variance of the variances of $\nu_i$ over $i$). In the paper, we reserved the terms mean and variance for random variables, to help keep clear what is random and what is considered fixed in our theory. We rewrote this section and so it should be clearer now.

8. Multi-line answer:

The quantity $\Lambda$ measures the distance of the error term $\eta_i$ to a Gaussian distribution. This being small justifies the use of the Gaussian approximation in Section 4.

As the referee is likely aware, such Gaussian approximations are typically justified by central limit theorems (taking $n_s$ to infinity) or else some Barry-Esseen-type (finite $n_s$) bound. The quantity we give is actually such a finite $n_s$ bound and we give the relevant pointer to the theorem in the text. The quantity $\Lambda$ also arises in Lyapunov's condition for the central limit theorem; see Billingsley, "Probability and Measure" or Durrett, "Probability: Theory and Examples" (the latter is open access and easily accessible online, but the theorem is left as an exercise in that book).

9. Done!

Section 4

1. The new version makes clear that $\sigma^2$ is the variance of $\nu$.
2. We show how the Persian Rug construction achieves a newly derived lower bound on the loss in the appendix now, proving optimality.
3. Equation 8 does not imply that we can fix any value of $b$, but rather that for a given value of $b$ we can find the optimal value of $a$. In general fitting the optimal value of $b$ must be done numerically, either by optimizing the expected or empirical loss.
4. In our updated draft we show a general lower-bound which, in combination with the upper bound, shows that the error is controlled by the level of noise, $\sigma$.

We thank the reviewer for their review and their insight into connections with the literature. Here we respond to the weaknesses brought up.

Interpretability and Autoencoders: It is a fair point that the main substance of the paper is the analysis of the autoencoder model. But this autoencoder model was put forth as a toy model of superposition, which is an important concept in the science of interpretabability. This is why we discuss the implications of our results to interpretability in the abstract, introduction, and conclusion. If there are any specific sections or sentences you think should be removed or further qualified, please let us know.

Main Concern: We agree with the reviewer that the trained neural networks do not converge to the Persian Rug and indeed we tried to emphasize this in the paper. So we are in the "latter'' case mentioned by the reviewer. However, our work does in fact interpret the learned algorithm as much as possible because only certain statistics of the weights are important. The Persian Rug has additional structure, which we warn against interpreting. While this is just for the toy model of compressing independent features, we think this meta point on mechansitic interpretability is worth emphasizing for future researchers wanting to do circuit analysis. We give the exact conditions as conditions on the spectrum and the diagonal for optimal $W$, and the Persian Rug as an example.

Centering of Data: In this context, assuming the sparse data is positive is standard (see Elhage et al. 2022, Ganguli and Sompolinksy 2010, and also the references in reviewer ru6M's comments on our work).

Choice of ReLU Nonlinearity: We have added to the paper a lower bound for the loss with the optimal elementwise activation function. The fact that the upper bound we derive for the ReLU case is equal (up to logarithmic factors) to this lower bound shows that the ReLU belongs to a class of models which all share the same problem with respect to reconstruction -- they don't accurately predict the sparse reconstructed variables. It's therefore a choice to show a general principle.

Connection to Literature: We have addressed the referees' concerns regarding the lack of attention towards two parts of the literature mentioned in their review. We have included a discussion of the connection between these subjects and our paper, which was not immediately apparent to us and we thank the referee for bringing them up. The first group of papers primarily deal with autoencoders of type [1] $\hat{x} = A \sigma(Bx)$. On the other hand, we work with autoencoders of the form [2] $\hat{x} = \sigma(ABx)$ where the distribution over $x$ is highly non-Gaussian. In our setting [2], we find that training with $x$ following a Gaussian distribution matching the first two moments does not recover similar $A$ or $B$, for example the overall scale $|AB|$ is significantly different between Gaussian and non-Gaussian $x$. This is unlike setting [1] where Gaussian and non-Gaussian $x$ behave similarly.

Our work also has an oblique relationship with message passing (and convex optimization) algorithms for signal recovery. While these are powerful algorithms for solving the general problem of signal recovery, we are interested in studying the specific recovery algorithm a network learns when it is simultaneously allowed to choose the encoding scheme. The algorithm is constrained by the network's architecture: a 1-layer network which decodes the sparse signal by applying a linear layer followed by a gating. This is similar in structure to a single step of a message passing algorithm (and also similar to [1] from reviewer ru6M). However what our analysis shows is that even in the best case dictionary the network can find, this model cannot efficiently decode the sparse signal without error, which is not the main area of interest in the current literature. Finally, our aim is to speak to the interpretability community, who are widely using the model (or very similar models) to the one we consider, and not to push forward the general theory of autoencoders. The ideas in our manuscript are not well known in this area, nor can they be derived from the literature cited by the reviewer.

We thank the reviewer for their careful review and answer their questions and address the weaknesses brought up.

Weaknesses: While the original work by Elhage et al. (2022) empirically reveals a number of important phenomena appearing in the toy model of superposition, the work does not explore the network's exploitation of randomness, which is reflected in their focus on smaller models. More explicitly, our experiments show permutation-symmetry emerges only in sufficiently large values of $n_s$. The utility of randomness in large systems is familiar to physicists, and it is our hope to convey this to the mechanistic interpretability community, who we feel often overlook this effect by focusing purely on small-scale, detailed structures. Thus while it is true that any matrix of the form in (original submission 7) is similarly optimal (so long as it leads to equal variances for all $\nu_i$), the purpose of the explicit construction is to demonstrate that such matrices do exist, and to highlight that the detailed structure of a matrix may not be important when doing interpretability.

Question 1: Yes, we now have empirical evidence that our study can be extended to non-isotropic cases in two ways: smaller symmetry groups and perturbations around a fixed symmetry. We empirically observe that model weights inherit statistical symmetries even for smaller symmetry groups, such as permutations within individual blocks. This implies a similar simplification to the loss. The second way concerns small perturbations around a given symmetry, for which we find the solution is stable. We have added a discussion of this to the the appendix in the updated manuscript.

Question 2: A linear model will not learn the same set of weights. Most importantly, a linear model will learn any (unscaled) rank $n_d$ projection, while the non-linear model will learn a scaled projection with equal diagonals. More generally, if some features are more important than others, the linear model's projection will be along the largest eigen-directions of the sample covariance matrix, with the bias $\mathbf{B}$, fitting the mean of the data. The way to interpret this is that a linear model cannot store more directions than $n_d$ while adding a ReLU allows the model to leverage the sparsity of the input and store more information about the input.

Question 3: Randomly drawing an orthogonal matrix will not allow us to satisfy the condition $W_{ii}$ are constant. A random projection works reasonably well, but the diagonal elements still fluctuate, and this leads to a greater error than necessary. Intuitively this is because the diagonal elements control the strength of the signal term $x_i$ which is particularly important. The purpose of the particular construction in 4.2.2 is to demonstrate that such matrices do exist, and to highlight the algorithm's insensitivity to small-scale structure in the $W$ matrix and in particular its exploitation of randomness.

We thank the reviewer for their time and their review. Here are the answers to the questions brought up.

Q1: When the input data is not sparse then it cannot be compressed into a small subspace. Mathematically when $p n_s > n_d$ then the number of sparse features which are on is typically greater than the number of dimensions available to store them. This is explained via our theory because the variance of $\nu$ is large and therefore $x + \nu$ has little information about $x$. More precisely, this means that the conditional distribution $P[x | x + \nu]$ becomes broad and is therefore not informative about the true value of $x$. Therefore the most interesting regime is the one where $p n_s < n_d$, where compression is possible (though still non-trivial). The symmetry property is preserved even in the non-sparse case, however we do not focus on this regime.

Q2: Our interest in this model is specifically related to understanding the behavior of sparse autoencoders, which take a distribution over vectors $x \in \mathbb{R}^{n_d}$ (which is not sparse) and find an equivalent representation $y(x) \in \mathbb{R}^{n_s}$ where $y(x)$ is sparse and there is some linear map $A$ such that $x = Ay(x)$. We consider the problem where the ground truth of the $y$ distribution is known, and as such that distribution is naturally sparse.

We thank the reviewer for an in-depth and constructive analysis of our work, and for pointing us to related literature on sparse-coding and autoencoders. We believe it is indeed worth engaging with this work, and have added their relationship to our related works section. However, in contrast to the reviewer's suggestion, we believe this work can not be trivially adapted to arrive at our conclusions, and furthermore that the techniques by which we solve the architecture are useful for a wider class of problems in mechanistic interpretability. We justify this claim in the following answers to the reviewer's questions.

The references [1-5] have to do with the sparse coding problem, which (as the reviewer is aware) is very related to but not the same as the toy model of superposition we study in this paper. For easy reference, we re-state these problems now.

Problem 1 (Sparse Coding; from [3]): Given a set of vectors $y^{(1)}, y^{(2)}, ..., y^{(p)} \in\mathbb R^{n}$, find a sparse coding matrix $A\in\mathbb R^{n\times m}$ and sparse coefficient vectors $x^{(1)},...,x^{(p)}\in\mathbb R^{m}$ that minimize the the following reconstruction error with sparsity penalty

$$\mathcal E (A,X) = \sum_{i=1}^p \left(||y^{(i)} - Ax^{(i)}||_2^2 + S(x^{(i)})\right).$$

Problem 2 (Toy Model of Superposition): Given a set of sparse vectors $x^{(1)},.., x^{(p)}\in\mathbb R^m$, find a sparse coding matrix $A\in\mathbb R^{n\times m}$ and decoder matrix $B\in\mathbb R^{m\times n}$ that minimizes the reconstruction error

$$\mathcal{F}(A,B,X) = \sum_{i=1}^p ||\text{ReLU}(BAx) - x||_2^2.$$

(Note that we write a slightly different variant of the toy model problem compared to the one in our paper for easy comparison with the Problem 1.)

The papers [1-5] assume the $y^{(i)}$ are generated by some true dictionary $A^*$ and coefficients $x^{(i)}$ that satisfy certain properties (e.g. that $A^*$ is incoherent) and then show that these assumptions imply various desirable properties for various algorithms for learning $A$ and $x^{(i)}$ (for example, convergence of the optimizer to the true generating dictionary $A^*$).

In contrast, the neural networks in our work search over the space of dictionaries to find ones that encode sparse information in a way particularly suitable for reconstruction by a single linear + ReLU layer. As a result, the dictionary our network finds contains additional structure optimized for a recovery process consisting of a single linear layer followed by local activation function. In particular we find that the relevant error parameter is not the incoherence $(\max_{i\neq j} |W_{ij}|)$ but rather the variance of off diagonal elements in each row $( \sigma=\sum_j W_{ij}^2)$, and that it is this parameter that needs to be minimized for a given compression ratio, and this leads to the additional non-trivial structure such as $W$ having being a projector with equal diagonals. This structure is not present in standard constructions of incoherent dictionaries. Choosing unit vectors randomly will fail to suppress $\sigma$ when $n_d\sim O(n_s)$, while a random projector will fail to suppress diagonal fluctuations sufficiently quickly as $n_d/n_s\rightarrow 0$.

The referee has proposed that the methods that are developed in the papers [1-5] for problem 1 should easily transfer to problem 2, even if the problems are somewhat different. We investigated this question and believe this is not the case for the following reasons.

The key difficulty in adapting these methods to our case is that, in our case, the corresponding matrix to $A^*$ (namely $W_{in}$) is changing over the learning process, and it is not clear the assumptions [1-5] put on $A^*$ will be satisfied. (For example, [3] assumes that $A^*$ is incoherent and $||A^* ||$ is suitably bounded, [4] requires a $A^*$ to satisfy a bound on its RIP constant and has bounded norm, and [5] requires $A^*$ to be square (which is definitely not true in our setting)). Establishing that these conditions hold over the course of an optimization process seems like interesting future work, but is not the point of this paper. It is also not immediately apparent that the theoretical machinery developed to leverage these assumptions will apply when the dictionary $A^*$ is changing, even if the assumptions hold over the course of training (or eventually).

We here address the most crucial concern, [2]. The point we wish to make is that there is additional structure beyond the fact that $A^*$ is an incoherent dictionary. While it is true that choosing $A^*$ to have unit norm columns will give diagonal elements of $1$, it will not result in a $W$ which is a projector. The latter condition gives a pattern of off diagonals which optimally interfere when acting on the input vector.

To make this point clearer, let us consider two incoherent dictionaries $A^1$ and $A^2$ such that the corresponding $W^1$ and $W^2$ both satisfy ($|W^1_{ij}| = |W^2_{ij}| = \epsilon$) for all $i\neq j$. However, for $W^2$, in each row half of the off diagonal elements are positive and the other negative, while for $W^1$ all elements are positive. Consider now the action of the first row of each $W$, $W^1_{1 \bullet}$, and $W^2_{1 \bullet}$, onto the all ones vector $\vec{1}$. In the first case we get $W^1_{1 \bullet}\cdot\vec{1} = 1+\epsilon (n_s-1)$, while in the latter case we obtain $W^2_{1 \bullet}\cdot\vec{1} = 1$. We see that in the latter case the oscillation of signs allows for the total cancellation of off diagonal elements, and that knowledge of the coherence parameter was not sufficient to determine how distorted the signal would be.

In summary, our model does not simply make use of the off diagonals having small absolute value as would be true for a typical incoherent dictionary, but also optimizes the distribution of signs in order to maximally suppress the noise. Since we are interested in recovering continuous signals as accurately as possible, this makes a large difference in certain parameter regimes.

The example was meant to illustrate the conceptual point that one must specify additional information beyond the fact that $A^*$ is an incoherent dictionary in order to obtain the optimal solution. Addressing the critical concern:

2. That our solution is the preferred symmetric solution for this architecture is proved in E.2 where we show a lower bound on the loss for architectures of this form together with 4.2.1 and appendix D where we derive the form of the optimal matrix. Because the loss is lower bounded above zero the reconstruction is necessarily imperfect, and there is no exact sense in which any particular construction does or does not satisfy a restricted isometry property, but only the degree to which it minimize the loss. The form we derive is provably optimal among symmetric solutions, which are the preferred solutions of these neural networks. We can illustrate the difference between our construction and the reviewers proposed construction (random matrices) by considering the extreme case of $n_d = n_s$. Our construction gives $W = I$, while the reviewers construction will give a $W$ with small off diagonal terms. The performance of these two solutions will converge in the limit $p\to 0$ and differ drastically near $p = 1$. Our solution correctly characterizes the network's learned behavior across the entire parameter space of $p\in$ 0,1 $$. Note that the difference in performance between the two constructions remains as we move slightly away from the extremal limit, such as $n_d = n_s - k$ for $k \ll n_s$, where the explicit form of the optimal $W$ is no longer obvious.