Author order randomized. Authors contributed roughly equally — see attribution section for details.

Update as of July 2024: we have collaborated with @LawrenceC to expand section 1 of this post into an arXiv paper, which culminates in a formal proof that computation in superposition can be leveraged to emulate sparse boolean circuits of arbitrary depth in small neural networks.

What kind of document is this?

What you have in front of you is so far a rough writeup rather than a clean text. As we realized that our work is currently highly relevant to recent questions posed by interpretability researchers, we put together a lightly edited version of private notes we've written over the last ~4 months. If you'd be interested in writing up a cleaner version, get in touch, or just do it. We're making these notes public before we're done with the project because of some combination of (1) seeing others think along similar lines and wanting to make it less likely that people (including us) spend time duplicating work, (2) providing a frame which we think provides plenty of concrete immediate problems for people to independently work on^[1] (3) seeking feedback to decrease the chance we spend a bunch of time on nonsense.

1 minute summary

Superposition is a mechanism that might allow neural networks to represent the values of many more features than they have neurons, provided that those features are present sparsely in the dataset. However, until now, an understanding of how computation can be done in a compressed way directly on these stored features has been limited to a few very specific tasks (for example here). The goal of this post is to lay the groundwork for a picture of how computation in superposition can be done in general. We hope this will enable future research to build interpretability techniques for reverse engineering circuits that are manifestly in superposition.

Our main contributions are:

1. Formalisation of some tasks performed by MLPs and attention layers in terms of computation on boolean features stored in superposition.
2. A family of novel constructions which allow a single layer MLP to compute a large number of boolean functions of features entirely in superposition.
3. Discussion of how these constructions could be leveraged:
   1. to emulate arbitrary large sparse boolean circuits entirely in superposition
   2. to allow the QK circuit of an attention head to dynamically choose a boolean expression and attend to past token positions where this expression is true.
   3. To explain the tentative observation that transformers may store arbitrary XORs of features
4. A construction which allows the QK circuit of an attention head to check for the presence of surprisingly many query-key feature pairs simultaneously in superposition, on the order of one pair per parameter^[2].

10 minute summary

Thanks to Nicholas Goldowsky-Dill for producing an early version of this summary/diagrams and generally for being instrumental in distilling this post.

Central to our analysis of MLPs is the Universal-AND (U-AND) problem:

• Given $m$ input boolean features $f_1,f_2,\dots ,f_m$. These features are sparse, meaning on most inputs only a few features are true, and encoded as directions in the input space $R^d$.
• We want to compute all $(\frac{m}{2})$ possible binary conjunctions of these inputs ($f_1\wedge f_2,f_1\wedge f_3,\dots$), and output them in different linear directions. Some small bounded error in these output values is tolerated.
• We want to compute this in a single MLP layer ($RReLU(W\overrightarrow{x}+b)$) with as few neurons as possible, for weight matrix $W$ with shape , bias $B$, and ‘readoff’ matrix $R$

This problem is central to understanding computation in superposition because:

• Many features that people think of are boolean in nature, and reverse engineering the circuits that are involved in constructing them consists of understanding how simpler boolean features are combined to make them.  For example, in a vision model, the feature which is 1 if there is a car in the image may be computed by combing the ‘wheels at the bottom of the image’ feature AND the ‘windows at the top’ feature ^[3].
• We will be focusing on the part of the network before the readoff with the matrix $R$. In an analogous way to the toy model of superposition, we consider the first two layers to represent  If we can do this task with an MLP with fewer than $(\frac{m}{2})$ neurons, then in a sense we have computed more boolean functions than we have neurons, and the values of these functions will be stored in superposition in the MLP activation space.
• Any boolean function can be written as a linear combination of ANDs with different numbers of inputs. For example$$\text{XOR}(A,B,C)=A+B+C-2A\wedge B-2A\wedge C-2B\wedge C+4A\wedge B\wedge C$$
• Therefore, if we can compute and linearly represent all the ANDs in superposition, then we can do so for any boolean function.

If $m=d_0$ (the dimension of the input space), then we can store the input features using an orthonormal basis such as the neuron basis. A naive solution in this case would be to have one neuron per pair which is active if both inputs are true and 0 otherwise. This requires $(\frac{m}{2})=\Theta \left(d_0^2\right)$ neurons, and involves no superposition:

On this input $x_1,x_2$ and $x_5$ are true, and all other inputs are false.

We can do much better than this, computing all the pairwise ANDs up to a small error with many fewer neurons. To achieve this, we have each neuron care about a random subset of inputs, and we choose the bias such that each neuron is activated when at least two of them are on. This requires $d=\Theta \left(\text{polylog}\right(d\left)\right)$ neurons:

Importantly:

• A modified version works even when the input features are in superposition. In this case we cannot compute all ANDs of potentially exponentially many features. Instead, we must pick up to $\tilde{\Theta }\left(d^2\right)$ logical gates to calculate at each stage.
• A solution to U-AND can be generalized to compute many ANDs of more than two inputs, and therefore to compute arbitrary boolean expressions involving a small number of input variables, with surprisingly efficient asymptotic performance (superpolynomially many functions computed at once). This can be done simply by increasing the density of connections between inputs and neurons, which comes at the cost of interference terms going to zero more slowly.
• It may be possible to stack multiple of these constructions in a row and therefore to emulate a large boolean circuit, in which each layer computes boolean functions on the outputs of the previous layer. However, if the interference is not carefully managed, the errors are likely to propagate and eventually become unmanageable. The details of how the errors propagate and how to mitigate this are beyond the scope of this work.
• We study the performance of our constructions asymptotically in $d$, and expect that insofar as real models implement something like them, they will likely be importantly different in order to have low error at finite $d$.
• If the ReLU is replaced by a quadratic activation function, we can provide a construction that is much more efficient in terms of computations per neuron. We suspect that this implies the existence of similarly efficient constructions with ReLU, and constructions that may perform better at finite $d$.

Our analysis of the QK part of an attention head centers on the task of skip feature-bigram checking:

• Given residual stream vectors ${\overrightarrow{a}}_1,\dots ,{\overrightarrow{a}}_T$ (for sequence length $T$) storing boolean features in superposition .
• Given a set $B$ of skip feature-bigrams (SFBs) which specify which keys to attend to from each query in terms of features present in the query and key. A skip feature-bigram is a pair of features such as $({\overrightarrow{f}}_6,{\overrightarrow{f}}_{13})$, and we say that an SFB is present in a query key pair if the first feature is present in the key and the second in the query.
• We want to compute an attention score which contains, in each entry, the number of SFBs in $B$ present in the query and key that correspond to that entry. To do so, we look for a suitable choice of the parameters in the weight matrix $W_{QK}$, a $d_{\text{resid}}\times d_{\text{resid}}$ matrix of rank $d_{\text{head}}$. Some small bounded error is tolerated.

This framing is valuable for understanding the role played by the attention mechanism in superposed computation because:

• It is a natural modification of the ‘attention head as information movement’ story that engages with the many independent features stored in residual stream vectors in parallel, rather than treating the vectors as atomic units. Each SFB can be thought of as implementing an operation corresponding to statements like ‘if feature ${\overrightarrow{f}}_{13}$ is present in the query, then attend to keys for which feature ${\overrightarrow{f}}_6$ is present’. 
• The stories normally given for the role played by a QK circuit can be reproduced as particular choices of $B$. For example, consider the set of ‘identity’ skip feature-bigrams: $B_{\text{Id}}=\{$‘if feature ${\overrightarrow{f}}_i$ is present in the query, then attend to keys for which feature ${\overrightarrow{f}}_i$ is also present’$|\forall i\}$. Checking for the presence of all SFBs in $B_{\text{Id}}$ corresponds to attending to keys which are the same as the query. 
• There are also many sets $B$ which are most naturally thought of in terms of performing each check in $B$ individually. 

A nice way to construct $W_{QK}$ is as a sum of terms for each skip feature-bigram, each of which is a rank one matrix equal to outer product of the two feature vectors in the SFB. In the case that all feature vectors are orthogonal (no superposition) you should be thinking of something like this:

where each of the rank one matrices, when multiplied by a residual stream vector on the right and left, performs a dot product on each side: 

$${\overrightarrow{a}}_s^T{W}_{QK}{\overrightarrow{a}}_t=\sum _i({\overrightarrow{a}}_s\cdot {\overrightarrow{f}}_i^k)({\overrightarrow{f}}_i^q\cdot {\overrightarrow{a}}_t)$$

where $(f_1^k,f_1^q),\dots ,(f_{|B|}^k,f_{|B|}^q)$ are the feature bigrams in $B$ with feature directions $({\overrightarrow{f}}_i^k,{\overrightarrow{f}}_i^q)$, and ${\overrightarrow{a}}_s$ is a residual stream vector at sequence position $s$. Each of these rank one matrices contributes a value of $1$ to the value of ${\overrightarrow{a}}_s^T{W}_{QK}{\overrightarrow{a}}_t$ if and only if the corresponding SFB is present. Since the matrix cannot be higher rank than $d_{\text{head}}$, typically we can only check for up to $\tilde{\Theta }\left(d_{\text{head}}\right)$ SFBs this way.

In fact we can check for many more SFBs than this, if we tolerate some small error. The construction is straightforward once we think of $W_{QK}$ as this sum of tensor products: we simply add more rank one matrices to the sum, and then approximate the sum as a rank $d_{\text{head}}$ matrix, using the SVD or even a random projection matrix $P$. This construction can be easily generalised to the case that the residual stream stores features in superposition (provided we take care to manage the size of the interference terms) in which case $W_{QK}$ can be thought of as being constructed like this: 

When multiplied by a residual stream vector on the right and left, this expression is $${\overrightarrow{a}}_s^T{W}_{QK}{\overrightarrow{a}}_t=\sum _i({\overrightarrow{a}}_s\cdot {\overrightarrow{f}}_i^k)(P{\overrightarrow{f}}_i^q\cdot {\overrightarrow{a}}_t)$$

Importantly:

• It turns out that the interference becomes larger than the signal when roughly one SFB has been checked for per parameter: $|B|=\tilde{\Theta }\left(d_{\text{resid}}{d}_{\text{head}}\right)$
• When there is structure to the set of SFBs that are being checked for, we can exploit this to check for even more SFBs with a single attention head.
• If there is a particular linear structure to the geometric arrangement of feature vectors in the residual stream, many more SFBs can be checked for at once, but this time the story of how this happens isn’t the simplest to describe in terms of a list of SFBs. This suggests that our current description of what the QK circuit does is lacking. In fact, this example exemplifies computation performed by neural nets that we don’t think is best described by our current sparse boolean picture. It may be a good starting point for building a broader theory than we have so far that takes into account other structures.

Indeed, there are many open directions for improving our understanding of computation in superposition, and we’d be excited for others to do future research (theoretical and empirical) in this area.

Some theoretical directions include:

• Fitting the OV circuit into the boolean computation picture
• Studying error propagation when U-AND is applied sequentially
• Finding constructions with better interference at finite $d$
• Making the story of boolean computation in transformers more complete by studying things that have not been captured by our current tasks
• Generalisations to continuous variables

Empirical directions include:

• Training toy models to understand if NNs can learn U-AND and related tasks, and how learned algorithms differ.
• Throwing existing interp techniques at NNs trained on these tasks and trying to study what we find. Which techniques can handle the superposition adequately?
• Trying to find instances of computation in superposition happening in small language models. 

Structure of the Post

In Section 1, we define the U-AND task precisely, and then walk through our construction and show that it solves the task. Then we generalise the construction in 2 important ways: in Section 1.1, we modify the construction to compute ANDs of input features which are stored in superposition, allowing us to stack multiple U-AND layers together to simulate a boolean circuit. In Section 1.2 we modify the construction to compute ANDs of more than 2 variables at the same time, allowing us to compute all sufficiently small^[4] boolean functions of the inputs with a single MLP. Then in Section 1.3 we explore efficiency gains from replacing the ReLU with a quadratic activation function, and explore the consequences.

In Section 2 we explore a series of questions around how to interpret the maths in Section 1, in the style of FAQs. Each part of Section 2 is standalone and can be skipped, but we think that many of the concepts discussed there are valuable and frequently misunderstood.

In section 3 we turn to the QK circuit, carefully introducing the skip feature-bigram checking task, and we explain our construction. We also discuss two scenarios that allow for more SFBs to be checked for than the simplest construction would allow.

We discuss the relevance of our constructions to real models in Section 4, and conclude in Section 5 with more discussion on Open Directions.

Notation and Conventions

• $d$ is the dimension of some activation space. $d_0$ may also be used for the dimension of the input space, and $d$ for the number of neurons in an MLP
• $m$ is the number of input features. If the input features are stored in superposition, $m>d$, otherwise $m=d$
• ${\overrightarrow{e}}_1,{\overrightarrow{e}}_2,\dots ,{\overrightarrow{e}}_d$ denotes an orthogonal basis of vectors. The standard basis refers to the neuron basis.
• All vectors are denoted with arrows on top like this: ${\overrightarrow{f}}_i$
• We use single lines to denote the size of a set like this: $|S_i|$ or the $L^2$ norm of a vector like this: $|{\overrightarrow{f}}_i|$
• We say that a boolean function $g$ has been computed $ϵ$-accurately for some small parameter $ϵ$ if the computed output never differs from $g$ by more than $ϵ$. That is, whenever the function has the output $1$, the computation outputs a number between $1\pm ϵ$ and whenever the function outputs $0$, the computation outputs a number between $\pm ϵ$.
• We say that a pair of unit vectors is $ϵ$-almost orthogonal (for a fixed parameter $ϵ$) if their dot product is $<ϵ$ (equivalently, if they are orthogonal to $ϵ$-accuracy). We say that a collection of unit vectors is $ϵ$-almost-orthogonal if they are pairwise almost orthogonal. We assume $ϵ$ to be a fixed small number throughout the paper (unless specified otherwise).
• It is known that for fixed $ϵ,$ one can fit exponentially (in $d$) many almost orthogonal vectors in a $d$-dimensional Euclidean space. Throughout this paper, we will assume present in each NN activation space a suitably “large” collection of almost-orthogonal vectors, which we call an overbasis.
• Vectors in this overbasis will be called f-vectors^[5], and denoted ${\overrightarrow{f}}_1,{\overrightarrow{f}}_2,\dots ,{\overrightarrow{f}}_m$. We assume they correspond to binary properties of inputs relevant to a neural net (such as “Does this picture contain a cat?”). When convenient, we will assume these f-vectors are generated in a suitably random way: it is known that a random collection of vectors is, with high probability, an almost orthogonal overbasis, so long as the number of vectors is not superexponentially large in $d$^[6].

In this post we make extensive use of Big-O notation and its variants, little o, $\Theta ,\Omega ,\omega$. See wikipedia for definitions. We also make use of tilde notation, which means we ignore log factors. For example, by saying a function $f\left(n\right)$ is $\Theta \left(g\right(n\left)\right)$, we mean that there are nonzero constants $c_1,c_2>0$ and a natural number $N$ such that for all $n>N$, we have $c_{1}g\left(n\right)\le f\left(n\right)\le c_{2}g\left(n\right)$. By saying a quantity is $\tilde{\Theta }\left(f\right(d\left)\right)$, we mean that this is true up to a factor that is a polynomial of $\log d$ — i.e., that it is asymptotically between $f\left(d\right)/\text{polylog}\left(d\right)$ and $f\left(d\right)\text{polylog}\left(d\right)$.

1 The Universal AND

We introduce a simple and central component in our framework, which we call the Universal AND component or U-AND for short. We start by introducing the most basic version of the problem this component solves. We then provide our solution to the simplest version of this problem. We later discuss a few generalizations: to inputs which store features in superposition, and to higher numbers of inputs to each AND gate. More elaboration on U-AND — in particular, addressing why we think it’s a good question to ask — is provided in Section 2.

1.1 The U-AND task

  The basic boolean Universal AND problem: Given an input vector which stores an orthogonal set of boolean features, compute a vector from which can be linearly read off the value of every pairwise AND of input features, up to a small error. You are allowed to use only a single-layer MLP and the challenge is to make this MLP as narrow as possible.

More precisely: Fix a small parameter $ϵ>0$ and let $d_0$ and $\ell$ be integers with $d_0\ge \ell$^[7]. Let ${\overrightarrow{e}}_1,\dots ,{\overrightarrow{e}}_{d_0}$ be the standard basis in $R^{d_0}$, i.e. ${\overrightarrow{e}}_i$ is the vector whose $i$th component is $1$ and whose other components are $0$. Inputs are all at most $\ell$-composite vectors, i.e., for each index set $I\subseteq \left[d\right]$ with $|I|\le \ell$, we have the input ${\overrightarrow{x}}_I={\sum }_{i\in I}{\overrightarrow{e}}_i\in R^{d_0}$. So, our inputs are in bijection with binary strings that contain at most $\ell$ ones^[8]. Our task is to compute all $(\frac{d_0}{2})$ pairwise ANDs of these input bits, where the notion of ‘computing’ a property is that of making it linearly represented in the output activation vector $\overrightarrow{a}\left(\overrightarrow{x}\right)\in R^d$. That is, for each pair of inputs $i,j$, there should be a linear function $r_{i,j}:R^d\to R$, or more concretely, a vector ${\overrightarrow{r}}_{i,j}\in R^d$, such that ${\overrightarrow{r}}_{i,j}^T\overrightarrow{a}\left(x\right){\approx }_{ϵ}{\text{AND}}_{i,j}\left(x\right)$. Here, the ${\approx }_{ϵ}$ indicates equality up to an additive error $ϵ$ and ${\text{AND}}_{i,j}$ is $1$ iff both bits $i$ and $j$ of $x$ are 1. We will drop the subscript $ϵ$ going forward.

We will provide a construction that computes these $\Theta \left(d_0^2\right)$ features with a single $d$-neuron ReLU layer, i.e., a $d_0\times d$ matrix $W$ and a vector $\overrightarrow{b}\in R^d$ such that $\overrightarrow{a}\left(x\right)=\text{ReLU}(W\overrightarrow{x}+\overrightarrow{b})$, with $d\ll d_0$. Stacking the readoff vectors ${\overrightarrow{r}}_{i,j}$ we provide as the rows of a readout matrix $R$, you can also see us as providing a parameter setting solving $\overrightarrow{\text{ANDs}}\left(\overrightarrow{x}\right){\approx }_{ϵ}R\left(\text{ReLU}\right(W\overrightarrow{x}+\overrightarrow{b}\left)\right)$, where $\overrightarrow{\text{ANDs}}\left(\overrightarrow{x}\right)$ denotes the vector of all $(\frac{d_0}{2})$ pairwise ANDs. But we’d like to stress that we don’t claim there is ever something like this large, size $(\frac{d_0}{2})$, layer present in any practical neural net we are trying to model. Instead, these features would be read in by another future model component, like how the components we present below (in particular, our U-AND construction with inputs in superposition and our QK circuit) do.

There is another kind of notion of a set of features having been computed, perhaps one that’s more native to the superposition picture: that of the activation vector (approximately) being a linear combination of f-vectors — we call these vectors f-vectors— corresponding to these properties, with coefficients that are functions of the values of the features. We can also consider a version of the U-AND problem that asks for output vectors which represent the set of all pairwise ANDs in this sense, maybe with the additional requirement that the f-vectors be almost orthogonal. Our U-AND construction solves this problem, too — it computes all pairwise ANDs in both senses. See the appendix for a discussion of some aspects of how the linear readoff notion of stuff having been computed, the linear combination notion of something having been computed, and almost orthogonality hang together.

1.2 The U-AND construction

We now present a solution to the U-AND task, computing $(\frac{d_0}{2})$ new features with an MLP width that can be much smaller than $(\frac{d_0}{2})$. We will go on to show how our solution can be tweaked to compute ANDs of more than 2 features at a time, and to compute ANDs of features which are stored in superposition in the inputs.

To solve the base problem, we present a random construction: $W$ (with shape $d_0\times d$) has entries that are iid random variables which are $1$ with probability $p\left(d\right)\ll 1$, and each entry in the bias vector is $-1$. We will pin down what $p$ should be later. 

We will denote by $S_i$ the set of neurons that are ‘connected’ to the $i$th input, in the sense that elements of the set are neurons for which the $i$th entry of the row of the weight vector that connects to that neuron is $1$. $\overrightarrow{S_i}$ is used to denote the indicator set of $S_i$: the vector which is $1$ for every neuron in $S_i$ and $0$ otherwise. So $\overrightarrow{S_i}$ is also the $i$th column of $W$.

Then we claim that for this choice of weight matrix, all the ANDs are approximately linearly represented in the MLP activation space with readoff vectors (and feature vectors, in the sense of Appendix B) given by

$$v_{(x_i\wedge x_j)}=v_{ij}=\frac{\overrightarrow{S_i\cap S_j}}{|S_i\cap S_j|}$$

for all $i,j$, where we continue our abuse of notation to write $S_i\cap S_j$ as shorthand for the vector which is an indicator for the intersection set, and $|S_i\cap S_j|$ is the size of the set.

We preface our explanation of why this works with a technical note. We are going to choose $d$ and $p$ (as functions of $d_0$) so that with high probability, all sets we talk about have size close to their expectation. To do this formally, one first shows that the probability of each individual set having size far from its expectation is smaller than any $1/\text{poly}\left(d_0\right)$ using the Chernoff bound (Theorem 4 here), and one follows this by a union bound over all only $\text{poly}\left(d_0\right)$ sets to say that with probability $1-o\left(1\right)$, none of these events happen. For instance, if a set $S_i\cap S_j$ has expected size ${\log }^4{d}_0$, then the probability its size is outside of the range ${\log }^4{d}_0\pm {\log }^3{d}_0$ is at most $2e^{-\mu {\delta }^2/3}=2e^{-{\log }^2{d}_0}=2e{d}_0^{-\log d_0}$ (following these notes, we let $\mu$ denote the expectation and $\delta$ denote the number of $\mu$-sized deviations from the expectation — this bound works for $\delta <1$ which is the case here). Technically, before each construction to follow, we should list our parameters $d,p$ and all the sets we care about (for this first construction, these are the double and triple intersections between the $S_i$) and then argue as described above that with high probability, they all have sizes that only deviate by a factor of $1+o\left(1\right)$ from their expected size and always carry these error terms around in everything we say, but we will omit all this in the rest of the U-AND section.

So, ignoring this technicality, let’s argue that the construction above indeed solves the U-AND problem (with high probability). First, note that $|S_i\cap S_j|\sim \text{Bin}(d,p^2)$. We require that $p$ is big enough to ensure that all intersection sets are non-empty with high probability, but subject to that constraint we probably want $p$ to be as small as possible to minimise interference^[9]. We'll choose $p={\log }^2{d}_0/\sqrt{d}$, such that the intersection sets have size $|S_i\cap S_j|\approx {\log }^4{d}_0$. We split the check that the readoff works out into a few cases:

• Firstly, if input features $i$, $j$, and at most $\ell -2$ other input features are present (recall that we are working with $\ell$-composite inputs), then letting $\overrightarrow{a}$ denote the post-ReLU activation vector, we have ${\overrightarrow{f}}_{{\text{AND}}_{ij}}\cdot \overrightarrow{a}=1$ plus an error that is at most $\ell$ times [the sum of sizes of triple intersections involving $i,j$ and each of the $k-2$ other features which are on, divided by the size of the $S_i\cap S_j$]. This is very likely less than $O\left(1/{\log }^2{d}_0\right)$ for all polynomially many pairs and sets of $\ell -2$ other inputs at once^[10], at least assuming $d=\omega \left({\log }^8{d}_0\right)$. The expected value of this error is ${\log }^2{d}_0/\sqrt{d}$.
• Secondly, if only one of $i,j$ is present together with some at most $\ell -1$ other features, then we get nonzero terms in the sum that expanding the dot product ${\overrightarrow{f}}_{{\text{AND}}_{ij}}\cdot \overrightarrow{a}$ precisely for neurons in a triple intersection of $i,j$, and one of the $\ell -1$ other features, so the readoff $\approx 0$ — more precisely, $O\left(1/{\log }^2{d}_0\right)$ (again, assuming $d=\omega \left({\log }^8{d}_0\right)$), and $\frac{{\log }^2{d}_0}{\sqrt{d}}$ in expectation).
• Finally, if neither of $i,j$ is present, then the error corresponds to quadruple intersections, so it is even more likely at most $O\left(1/{\log }^2{d}_0\right)$ (still assuming $d=\omega \left({\log }^8{d}_0\right)$), and $\frac{{\log }^4{d}_0}{d}$ in expectation.

So we see that this readoff is indeed the AND of $i$ and $j$ up to error $ϵ=O\left(1/{\log }^2{d}_0\right)$.

To finish, we note without much proof that everything is also computed in the sense that 'the activation vector is a linear combination of almost orthogonal features' (defined in Appendix B). The activation vector being an approximate linear combination of pairwise intersection indicator vectors with coefficients being given by the ANDs follows from triple intersections being small, as does the almost-orthogonality of these feature vectors.

U-AND allows for arbitrary XORs to be efficiently calculated

A consequence of the precise (up to $ϵ$) nature of our universal AND is the existence of a universal XOR, in the sense of every XOR of features being computed. In this post by Sam Marks, it is tentatively observed that real-life transformers linearly compute XOR of arbitrary features in the weak sense of being able to read off tokens where XOR of two tokens is true using a linear probe (not necessarily with $ϵ$ accuracy). This weak readoff behavior for AND would be unsurprising, as the residual stream already has this property (using the readoff vector ${\overrightarrow{f}}_i+{\overrightarrow{f}}_j$ which has maximal value if and only if $f_i$ and $f_j$ are both present). However, as Sam Marks observes, it is not possible to read off XOR in this weak way from the residual stream. We can however see that such a universal XOR (indeed, in the strong sense of $ϵ$-accuracy) can be constructed from our strong (i.e., $ϵ$-accurate) universal AND. To do so, assume that in addition to the residual stream containing feature vectors ${\overrightarrow{f}}_i$ and ${\overrightarrow{f}}_j$, we’ve also already almost orthogonally computed universal AND features ${\overrightarrow{f}}_{{\text{AND}}_{i,j}}$ into the residual stream. Then we can weakly (and in fact, $ϵ$-accurately) read off XOR from this space by taking the dot product with the vector ${\overrightarrow{f}}_{{\text{XOR}}_{i,j}}:={\overrightarrow{f}}_i+{\overrightarrow{f}}_j-2{\overrightarrow{f}}_{{\text{AND}}_{i,j}}$. Then we see that if we had started with the two-hot pair ${\overrightarrow{f}}_{i^{\prime }}+{\overrightarrow{f}}_{j^{\prime }},$ the result of this readoff will be, up to a small error $O\left(ϵ\right),$ 

$$\{\begin{array}{ccc}0=0-0,& |\{i,j\}\cap \{i^{\prime },j^{\prime }\}|=0& \text{(neither coefficient agrees)}\\ 1=1-0,& |\{i,j\}\cap \{i^{\prime },j^{\prime }\}|=1& \text{(one coefficient agrees)}\\ 0=2-2,& \{i,j\}=\{i^{\prime },j^{\prime }\}& \text{(both coefficients agree)}\end{array}$$

This gives a theoretical feasibility proof of an efficiently computable universal XOR circuit, something Sam Marks believed to be impossible.

1.3 Handling inputs in superposition: sparse boolean computers

Any boolean circuit can be written as a sequence of layers executing pairwise ANDs and XORs^[11] on the binary entries of a memory vector. Since our U-AND can be used to compute any pairwise ANDs or XORs of features, this suggests that we might be able to emulate any boolean circuit by applying something like U-AND repeatedly. However, since the outputs of U-AND store features in superposition, if we want to pass these outputs as inputs to a subsequent U-AND circuit, we need to work out the details of a U-AND construction that can take in features in superposition. In this section we explore the subtleties of modifying U-AND in this way. In so doing, we construct an example of a circuit which acts entirely in superposition from start to finish — nowhere in the construction are there as many dimensions as features! We consider this to be an interesting result in its own right.

U-ANDs ability to compute many boolean functions of inputs features stored in superposition provides an efficient way to use all the parameters of the neural net to compute (up to a small error) a boolean circuit with a memory vector that is wider than the layers of the NN^[12]. We call this emulating a ‘boolean computer’. However, three limitations prevent any boolean circuit from being computed:

1. An injudicious choice of a layer executing XORs applied to a sparse input can fail to give a sparse output vector. Since U-AND only works on inputs with sparse features, this means that we can only emulate circuits with the property than on sparse inputs, their memory vector is sparse throughout the computation. We call these circuits ‘sparse boolean circuits’.
2. Even if the outputs of the circuit remain sparse at every layer, the $ϵ$ errors involved in the boolean read-offs compound from layer to layer. We hope that it is possible to manage this interference (perhaps via subtle modifications to the constructions) enough to allow multiple steps of sequential computation, although we leave an exploration of error propagation to future work.
3. We can't compute an unbounded number of new features with a finite-dimensional hidden layer. As we will see in this section, when input features are stored in superposition (which is true for outputs of U-AND and therefore certainly true for all but possibly the first layer of an emulated boolean circuit), we cannot compute more than $\tilde{\Theta }\left(d_{0}d\right)$ (number of parameters in the layer) many new boolean functions at a time.

Therefore, the boolean circuits we expect can be emulated in superposition (1) are sparse circuits (2) have few layers (3) have memory vectors which are not larger than the square of the activation space dimension.

Construction details for inputs in superposition

Now we generalize U-AND to the case where input features can be in superposition. With f-vectors ${\overrightarrow{f}}_1,\dots ,{\overrightarrow{f}}_m\in R^{d_0}$, we give each feature a random set of neurons to map to, as before. After coming up with such an assignment, we set the $i$th row of $W$ to be the sum of the f-vectors for features which map to the $i$th neuron. In other words, let $F$ be the $m\times d_0$ matrix with $i$th row given by the components of ${\overrightarrow{f}}_i$ in the neuron basis:

$$F=\left(\begin{array}{c}{\overrightarrow{f}}_1\to \\ ⋮\\ {\overrightarrow{f}}_m\to \end{array}\right)$$

Now let \hat{W} be a sparse matrix (with shape $d\times m$) with entries that are iid Bernoulli random variables which are $1$ with probability $p\left(d\right)\ll 1$. Then:

$$W=\hat{W}F$$

Unfortunately, since the ${\overrightarrow{f}}_1,\dots ,{\overrightarrow{f}}_m$ are random vectors, their inner product will have a typical size of $1/\sqrt{d_0}$. So, on an input which has no features connected to neuron $i$, the preactivation for that neuron will not be zero: it will be a sum of these interference terms, one for each feature that is connected to the neuron. Since the interference terms are uncorrelated and mean zero, they start to cause neurons to fire incorrectly when $\Theta \left(d_0\right)$ neurons are connected to each neuron. Since each feature is connected to each neuron with probability $p=\frac{{\log }^2{d}_0}{\sqrt{d}})$ this means neurons start to misfire when $m=\tilde{\Theta }\left(d_0\sqrt{d}\right)$^[13]. At this point, the number of pairwise ANDs we have computed is $(\frac{m}{2})=\tilde{\Theta }\left(d_0^{2}d\right)$.

This is a problem, if we want to be able to do computation on input vectors storing potentially exponentially many features in superposition, or even if we want to be able to do any sequential boolean computation at all:

Consider an MLP with several layers, all of width $d_{\text{MLP}}$, and assume that each layer is doing a U-AND on the features of the previous layer. Then if the features start without superposition, there are initially $d_{\text{MLP}}$ features. After the first U-AND, we have $\Theta \left(d_{\text{MLP}}^2\right)$ new features, which is already too many to do a second U-AND on these features!

Therefore, we will have to modify our goal when features are in superposition. That said, we're not completely sure there isn't any modification of the construction that bypasses such small polynomial bounds. But e.g. one can't just naively make $\hat{W}$ sparser — $p$ can't be taken below $d^{-1/2}$ without the intersection sets like $|S_i\cap S_j|$ becoming empty. When features were not stored in superposition, solving U-AND corresponded to computing $d_0^2$ many new features. Instead of trying to compute all pairwise ANDs of all (potentially exponentially many) input features in superposition, perhaps we should try to compute a reasonably sized subset of these ANDs. In the next section we do just that. 

A construction which computes a subset of ANDs of inputs in superposition

Here, we give a way to compute ANDs of up to $d_{0}d$ particular feature pairs (rather than all $(\frac{m}{2})$ ANDs) that works even for $m$ that is superpolynomial in $d_0$^[14]. (We’ll be ignoring $\log$ factors in much of what follows.)

In U-AND, we take $\hat{W}$ to be a random matrix with iid 0/1 entries with probability $p=\frac{{\log }^2{d}_0}{\sqrt{d}}$. If we only need/want to compute a subset of all the pairwise ANDs — let $E$ be this set of all pairs of inputs $\{i,j\}$ for which we want to compute the AND of $i$ and $j$ — then whenever $\{i,j\}\in E$, we might want each pair of corresponding entries in the corresponding columns $i$ and $j$ of the adjacency matrix $\hat{W}$, i.e., each pair ${\left(\hat{W}\right)}_{ki}$, ${\left(\hat{W}\right)}_{kj}$ to be a bit more correlated than an analogous pair in column $i^{\prime }$ and $j^{\prime }$ with $\{i^{\prime },j^{\prime }\}\notin E$. Or more precisely, we want to make such pairs of columns $\{i,j\}$ have a surprisingly large intersection for the general density of the matrix — this is to make sure that we get some neurons which we can use to read off the AND of $\{i,j\}$, while choosing the general density in $\hat{W}$ to be low enough that we don’t cross the density threshold at which a neuron needs to care about too many input features.

One way to do this is to pick a uniformly random set of ${\log }^4{d}_0$ neurons for each $\{i,j\}\in E$, and to set the column of $\hat{W}$ corresponding to input $i$ to be the indicator vector of the union of these sets (i.e., just those assigned to gates involving $i$). This way, we can compute up to around $|E|=\tilde{\Theta }\left(d_{0}d\right)$ pairwise ANDs without having any neuron care about more than $d_0$ input features, which is the requirement from the previous section to prevent neurons misfiring when input f-vectors are random vectors in superposition with typical interference size $\Theta \left(1/\sqrt{d_0}\right)$.

1.4 ANDs with many inputs: computation of small boolean circuits in a single layer

It is known that any boolean circuit with $k$ inputs can be written as a linear combination (with possibly exponential in $k$ terms, which is a substantial caveat) ANDs with up to $k$ inputs (fan-in up to $k$)^[15]. This means that, if we can compute not just pairwise ANDs, but ANDs of all fan-ins up to $k$, then we can write down a ‘universal’ computation that computes (simultaneously, in a linearly-readable sense) all possible circuits that depend on some up to $k$ inputs.

The U-AND construction for higher fan-in

We will modify the standard, non-superpositional U-AND construction to allow us to compute all ANDs of a specific fan-in $k$.

We'll need two modifications:

1. We're now interested in $k$-wise intersections between the $S_i$. The size of these intersections is smaller than double intersections, so we need to increase $p$ to guarantee they are nonempty. A sensible choice for fan-in $k$ is $p=\frac{{\log }^2{d}_0}{d^{1/k}}$.
2. We only want neurons to fire when $k$ of the features that connect to them are present at the same time, so we require the bias to be $-k+1$.

Now we read off the AND of a set $I$ of input features along the vector ${\bigcap }_{i\in I}{S}_i$.

We can straightforwardly simultaneously compute all ANDs of fan-ins ranging from $2$ to $k$ by just evenly partitioning the $d$ neurons into $k-1$ groups — let’s label these $2,3,\dots ,k$ — and setting the weights into group $i$ and the biases of group $i$ as in the fan-in $i$ U-AND construction.

A clever choice of density can give us all the fan-ins at once

Actually, we can calculate all ANDs of up to some constant fan-in$k$ in a way that feels more symmetric than the option involving a partition above^[16] by reusing the fan-in $2$ U-AND with (let’s say) $d=d_0$ and a careful choice of $p=\frac{1}{{\log }^2{d}_0}$ . This choice of  $p$ is larger than $\frac{{\log }^2{d}_0}{d^{1/k}}$ for any $k$, ensuring that every intersection set is non-empty. Then, one can read off ${\text{AND}}_{i,j}$ from $S_i\cap S_j$ as usual, but one can also read off ${\text{AND}}_{i,j,k}$ with the composite vector

$$-\frac{S_i\cap S_j\cap S_k}{|S_i\cap S_j\cap S_k|}+\frac{S_i\cap S_j}{|S_i\cap S_j|}+\frac{S_i\cap S_k}{|S_i\cap S_k|}+\frac{S_j\cap S_k}{|S_j\cap S_k|}$$ In general, one can read off the AND of an index set $I$ with the vector $$\sum _{I^{\prime }\subseteq I\text{s.t.}|I^{\prime }|\ge 2}(-1{)}^{|I|-|I^{\prime }|+1}{v}_{I^{\prime }}$$ where  $$v_{I^{\prime }}=\frac{{\bigcap }_{i\in I^{\prime }}{S}_i}{\left|{\bigcap }_{i\in I^{\prime }}{S}_i\right|}$$One can show that this inclusion-exclusion style formula works by noting that if the subset of indices of $I$ which are on is $J$, then the readoff will be approximately $$\sum _{I^{\prime }\subseteq I\text{s.t.}|I^{\prime }|\ge 2}(-1{)}^{|I|-|I^{\prime }|+1}max(0,|I^{\prime }\cap J|-1)$$. We’ll leave it as an exercise to show that this is $0$ if $J\ne I$ and $1$ if $J=I$. 

Extending the targeted superpositional AND to other fan-ins

It is also fairly straightforward to extend the construction for a subset of ANDs when inputs are in superposition to other fan-ins, doing all fan-ins on a common set of neurons. Instead of picking a set for each pair that we need to AND as above, we now pick a set for each larger AND gate that we care about. As in the previous sparse U-AND, each input feature gets sent to the union of the sets for its gates, but this time, we make the weights depend on the fan-in. Letting $K$ denote the max fan-in over all gates, for a fan-in $k$ gate, we set the weight from each input to $K/k$, and set the bias to $-K+1$. This way, still with at most about $\tilde{\Theta }\left(d^2\right)$ gates, and at least assuming inputs have at most some constant number of features active, we can read the output of a gate off with the indicator vector of its set.

1.5 Improved Efficiency with a Quadratic Nonlinearity

It turns out that, if we use quadratic activation functions $x\mapsto x^2$ instead of ReLU's $x\mapsto \text{ReLU}\left(x\right),$ we can write down a much more efficient universal AND construction. Indeed, the ReLU universal AND we constructed can compute the universal AND of up to $\tilde{\Theta }\left(d^{3/2}\right)$ features in a $d$-dimensional residual stream. However, in this section we will show that with a quadratic activation, for $\ell$-composite vectors, we can compute all pairwise ANDs of up to $m=\Omega \left(\exp \right(\frac{1}{2\ell }{ϵ}^2\sqrt{d}\left)\right)$^[17] features stored in superposition (this is exponential in $\sqrt{d}$, so superpolynomial in $d$(!)) that admit a single-layer universal AND circuit.

The idea of the construction is that, on the large space of features $R^m,$ the AND of the boolean-valued feature variables $f_i,f_j$ can be written as a quadratic function $q_{i,j}:\{0,1{\}}^m\mapsto \{0,1\}$; explicitly, $q_{i,j}(f_1,\dots ,f_m)=f_i\cdot f_j.$ Now if we embed feature space $R^m$ onto a smaller $R^r$ in an $ϵ$-almost-orthogonal way, it is possible to show that the quadratic function $q_{i,j}$ on $R^m$ is well-approximated on sparse vectors by a quadratic function on $R^r$ (with error bounded above by $2ϵ$ on 2-sparse inputs in particular). Now the advantage of using quadratic functions is that any quadratic function on $R^r$ can be expressed as a linear read-off of a special quadratic function $Q:R^r\to R^{r^2}$ given by the composition of a linear function $R^r\to R^{r^2}$ and a quadratic element-wise activation function on $R^{r^2}$ which creates a set of neurons which collectively form a basis for all quadratic functions. Now we can set $d=r^2$ to be the dimension of the residual stream and work with an $r$-dimensional subspace $V$ of the residual stream, taking the almost-orthogonal embedding $R^m\to V$. Then the map $V\stackrel{Q}{\to }{R}^d$ provides the requisite universal AND construction. We make this recipe precise in the following section

Construction Details

In this section we use slightly different notation to the rest of the post, dropping overarrows for vectors, and we drop the distinction between features and f-vectors.

Let $V=R^r$ be as above. There is a finite-dimensional space of quadratic functions on $R^r$, with basis $q_{ij}=x_i{x}_j$ of size $r^2$ (such that we can write every quadratic function as a linear combination of these basis functions); alternatively, we can write $q_{ij}\left(v\right)=(v\cdot e_i)(v\cdot e_j),$ for $e_i,e_j$ the basis vectors. We note that this space is spanned by a set of functions which are squares of linear functions of $\left\{x_i\right\}$:

$$\begin{array}{cc}{L}_i^{\left(1\right)}(x_1,\dots ,x_r)& =x_i\\ L_{i,j}^{\left(2\right)}(x_1,\dots ,x_r)& =x_i+x_j\\ L_{i,j}^{\left(3\right)}(x_1,\dots ,x_r)& =x_i-x_j\end{array}$$

The squares of these functions are a valid basis for the space of quadratic functions on $R^r$ since $q_{ii}=(L_i^{\left(1\right)}{)}^2$ and for $i\ne j,$ we have $q_{ij}=\frac{(L_{i,j}^{\left(2\right)}{)}^2-(L_{i,j}^{\left(3\right)}{)}^2}{4}.$ There are $m$ distinct functions of type $\left(1\right)$, and $(\frac{m}{2})$ functions each of type $\left(2\right)$ and $\left(3\right)$, for a total of $r^2$ basis functions as before. Thus there exists a single-layer quadratic-activation neural net $Q:x\mapsto y$ from $R^r\to R^{r^2}$ such that any quadratic function on $R^r$ is realizable as a "linear read-off", i.e., given by composing $Q$ with a linear function $R^{r^2}\to R.$ In particular, we have linear "read-off" functions ${\Lambda }_{ij}:R^{r^2}\to R$ such that $L_{ij}\left(Q\right(x\left)\right)=q_{ij}\left(x\right).$

Now suppose that $f_1,\dots ,f_m$ is a collection of f-vectors which are $ϵ$-almost-orthogonal, i.e., such that $|f_i|=1$ for any $i$ and $|f_i\cdot f_j|<ϵ\forall i<j\le m.$ Note that (for fixed $ϵ<1$), there exist such collections with exponential (in $r$) number of vectors $m$. We can define a new collection of symmetric bilinear functions (i.e., functions in two vectors $v,w\in R^n$ which are linear in each input independently and symmetric to switching $v,w$), ${\varphi }_{i,j},$ for a pair of (not necessarily distinct) indices $0<i\le j\le m,$ defined by ${\varphi }_{i,j}\left(v\right)=(v\cdot f_i)(v\cdot f_j)$ (this is a product of two linear functions, hence quadratic). We will use the following result:

Proposition 1 Suppose ${\varphi }_{i,j}$ is as above and $0<i^{\prime }\le j^{\prime }<m$ is another pair of (not necessarily distinct) indices associated to feature vectors $v_i,v_j.$ Then

$${\varphi }_{i,j}(v_{i^{\prime }},v_{j^{\prime }})\{\begin{array}{cc}=1,& i=i^{\prime }\text{and}j=j^{\prime }\\ \in (-ϵ,ϵ),& (i,j)\ne (i^{\prime },j^{\prime })\\ \in (-{ϵ}^2,{ϵ}^2),& \{i,j\}\cap \{i^{\prime },j^{\prime }\}=\varnothing \text{(i.e., no indices in common)}\end{array}$$

This proposition follows immediately from the definition of ${\varphi }_{k,\ell }$ and the almost orthogonality property. $\square$

Now define the single-valued quadratic function ${\varphi }_{i,j}^{\text{single}}\left(v\right):=\frac{1}{2}{\varphi }_{i,j}(v,v),$ by applying the bilinear form to two copies of the same vector and dividing by 2. Then the proposition above implies that, for two pairs of distinct indices $0<i<j\le m$ and $0<i^{\prime }<j^{\prime }\le m$ we have the following behavior on the sum of two features (the superpositional analog of a two-hot vector):

$${\varphi }_{i,j}^{\text{single}}(v_{i^{\prime }}+v_{j^{\prime }})=\frac{{\varphi }_{i,j}(v_{i^{\prime }},v_{i^{\prime }})+2{\varphi }_{i,j}(v_{i^{\prime }},v_{j^{\prime }})+{\varphi }_{i,j}(v_{j^{\prime }},v_{j^{\prime }})}{2}={\varphi }_{i,j}(v_{i^{\prime }},v_{j^{\prime }})+O\left(ϵ\right).$$

The first formula follows from bilinearity (which is equivalent to the statement that the two entries in ${\varphi }_{i,j}$ behave distributively) and the last formula follows from the proposition since we assumed $(i,j)$ are distinct indices, hence cannot match up with a pair of identical indices $(i^{\prime },i^{\prime })$ or $(j^{\prime },j^{\prime }).$ Moreover, $O\left(ϵ\right)$ term in the formula above is bounded in absolute value by $\frac{2ϵ}{2}=ϵ$.