Summary & Motivation

This post is a continuation and clarification of Circuits in Superposition: Compressing many small neural networks into one. That post presented a sketch of a general mathematical framework for compressing different circuits into a network in superposition. On closer inspection, some of it turned out to be wrong, though. The error propagation calculations for networks with multiple layers were incorrect. With the framework used in that post, the errors blow up too much over multiple layers. 

This post presents a slightly changed construction that fixes those problems, and improves on the original in some other ways as well.^[1] 

By computation in superposition we mean that a network represents features in superposition and performs more computations with them than it has neurons, across multiple layers. Having better models of this is important for understanding how and even if networks use superposition, which in turn is important for mech-interp in general. 

Performing computation in superposition over multiple layers introduces additional noise compared to just storing features in superposition^[2]. This restricts the amount and type of computation that can be implemented in a network of a given size, because the noise needs to be reduced or suppressed to stay smaller than the signal.

Takeaways

Our setup in this post (see the Section Construction for details) is as follows:

• We have $T$ small circuits, each of which can be described as a $d$-dimensional multilayer perceptron (MLP) with $L$ layers.
• We have one large $D$-dimensional MLP with $L$ layers, where $D>d$ , but $D<Td$. So we can't just dedicate $d$ neurons in the large MLP to each circuit.
• We embed all $T$ circuits into the large network, such that the network approximately implements the computations of every circuit, conditional on no more than $z<<T$ circuits being used on any given forward pass.

The number of circuits we can fit in scales linearly with the number of network parameters

Similar to the previous post, we end up concluding that the total number of parameters in the circuits must be smaller than the number of parameters in the large network. This result makes a lot of intuitive sense, since the parameters determine the maximum amount of information a network can possibly store.^[3]

More specifically, we find that the term $\sqrt{\frac{zT{d}^2}{D^2}}$ needs to be smaller than $1$ for the errors on the computations of individual circuits to stay smaller than the signal. 

Here $T$ is the total number of circuits, $d$ is the width of each small circuit, $D$ is the layer width of the large network and $z$ is the number of circuits active on a given forward pass. 

This gives us an approximate upper bound on the maximum number of $d$-dimensional circuits we can fit into a network with $D$ neurons per layer:

If you only remember one formula from this post, let it be that one.

This is much smaller than the number of $d$-dimensional features we can store in superposition in a layer of width $D$, if we don't intend to use them for calculations within the network. That number is^[4] 

So, while storage capacity scales exponentially with $D$, capacity for computation only scales quadratically.

Each circuit will only use a small subset of neurons in the larger network

For this construction to work, each circuit can only be using a small fraction of the large network's neurons. 

I, Linda, expect this to be true more generally. I think basically any construction that achieves computation in superposition, across multiple layers, in the sense we mean here will (approximately) have this property. My reasons for thinking this are pretty entangled with details of the error propagation math, so I've relegated them to the Discussion section. 

Implications for experiments on computation in superposition

The leading source of error in this construction is signals from the active circuits (used on the forward pass) bleeding into currently inactive circuits that shouldn't be doing anything. This bleed-over then enters back into the active circuits as noise in the next layer. 

This means that the biggest errors don't appear until layer 2^[5]. This is important to keep in mind for experimental investigations of computation in superposition, e.g. when training toy models. If your network only has one computational layer, then it doesn't have to implement a way to reduce this noise.^[6] 

Reality really does have a surprising amount of detail

To make sure the math was actually really correct this time around, Linda coded up a little model implementing some circuits in superposition by hand. 

Naturally, while doing this, she found that there were still a bunch of fiddly details left to figure out how to make circuits in superposition actually work in real life, even on a pretty simple example, because the math makes a bunch of vague assumptions about the circuits that turn out to be important when you actually get down to making things work in practice.

The math presented in this post won't deal with those fiddly details. It is intended to be a relatively simple high level description of a general framework. E.g, we assume that individual circuits have some level of noise robustness around values representing 'inactive' for that circuit, without worrying about how it's achieved.

So, in actual practice, the details of this construction may need some adjustment depending on how exactly individual circuits implement their noise robustness, and whether any of them are doing similar things.^[7]

A post with the hand coded model and the fiddly details should be coming out "${\text{Soon}}^{\text{TM}}$".

Construction

The construction in this post has some significant differences from the previous one. 

To simplify things, the networks here don't have a residual stream, they're just deep MLPs.^[8] We have one large MLP with $L$ layers, neuron dimension $D$, activation vectors $A^l\in R^D$, and weight matrices $W^l\in R^{D\times D}$.

We also have $T$ circuits, indexed by $0,\dots ,T-1$, each described by a small MLP with $L$ layers, neuron dimension $d$, activations vectors $a_t^l\in R^d$, and weight matrices $w_t^l\in R^{d\times d}$.  

Our goal is to figure out a construction for the weight matrices $W^l$, which embeds the circuits into the network, such that the outputs of each circuit can be read-out from the final output of the network $A^L$ with linear projections, up to some small error.^[9]

Assumptions

For this construction to work as intended, we need to assume that:

1. Only  $z\ll \frac{D}{d}$ circuits can be active on any given forward pass.
2. Small circuits are robust to noise when inactive. I.e. a small deviation to the activation value of an inactive circuit applied in layer $l$ will not change the activation value of that circuit in layer $l+1$.^[10]
3. If a circuit is inactive, all of its neurons have activation value zero. I.e. $a_t^l=0$  if circuit $t$ is inactive.
4. The entries of the weight matrices $w_t^l$ for different circuits in the same layer are uncorrelated with each other.

Assumption 1 is just the standard sparsity condition for superposition. 

Assumption 2 is necessary, but if it is not true for some of the circuits we want to implement, we can make it true by modifying them slightly, in a way that doesn't change their functionality.^[11] How this works will not be covered in this post though.

Assumptions 3 and 4 are not actually necessary for something similar to this construction to work, but without them the construction becomes more complicated. The details of this are also beyond the the scope of this post.

Embedding the circuits into the network

The important takeaways from this section are Equations (1.11) and (1.13)-(1.14), which we will make use of in the Error Calculation section. If you're happy with these and don't think they require further motivation, then you don't need to read the rest of this section.

Remember that the circuit weights $w_t^l$ and their activation vectors $a_t^l$ are handed to us and we can't change them, but we are free to choose the weights $W^l$ of the network to be whatever we like. We also assume that we get to choose how to linearly embed the input vectors of the circuits $a_t^0$ into the input vector of the network $A^0$ at the start.

To help with the embedding we will introduce:

• Embedding matrices $E_t^l\in R^{D\times d}$ for each circuit $t$ in each layer $l\ge 1$
• Unembedding matrices $U_t^l\in R^{d\times D}$ for each circuit $t$ in each layer $l\ge 1$.

Our goal is to calculate Equation (1.2) using the network, which (due to our choice of $U_t^l$ and $E_t^l$, see next section) can be re-expressed as

We approximate this as

and

If we combine Equations (1.4) and (1.5) while pretending^[12] they are exact relations, we get

If we combine that with Equation (1.1) the network weights $W^l$ for $l\ge 2$ are now defined via the embedding and unembedding matrices as

You might notice that this leaves $W^1$ undefined, and that there are no embedding and unembedding matrices for $l=0$. That's because layer zero is a bit special.

Layer 0

There are no embedding and unembedding matrices for layer 0, because we can just skip ahead and use our free choice of how to linearly embed $a_t^0$ into $A_t^0$ to implement the first matrix multiplications in each circuit $w_t^1$ without any interference errors. 

We choose

which gives us:

I.e, (1.4) is exactly true in the first layer. Down in the Error Calculation section, this will have consequences for which layer each error term first shows up in.

Maybe you think this is sort of cheating. Perhaps it is, but a model can train to cheat like this as well. That's part of the point we want to make in this post: Having more than one layer makes a difference. From layer 2 onward, this kind of thing doesn't work anymore. We no longer gets to assume that every feature always starts embedded exactly where we want it. We have to actually compute on the features in a way that leaves every intermediary result embedded such that later layers can effectively compute on them in turn. 

Constructing the Embedding and Unembedding matrices

In this subsection, we first lay out some requirements for the embedding matrices $E_t^l$ and unembedding matrices $U_t^l$, then describe an algorithm for constructing them to fulfil those requirements.

The main takeaway here is that we can construct $E_t^l$ and $U_t^l$ such that they have the properties described in the final subsection. If you are happy to just trust us on this, then you don't need to read the rest of this section.

Remember that each neuron in the large network will be used in the embeddings of multiple small circuit neurons. This is inevitable because the total numbers of neurons across all small circuits is larger than the number of neurons in the large network, $dT>D$, as per assumption 1.

Requirements

1. Neurons from the same circuit should be embedded into non-overlapping sets of neurons in the large network. We want this because neurons in the same circuit will tend to coactivate a lot, so reducing interference between them is especially important.
2. Neurons from different small circuits should be embedded into sets of neurons in the large network that have at most overlap one. So, no pair of circuit neurons shares more than one network neuron. This ensures that there is a hard bound on how bad errors from interference between different circuits can get.
3.  $E_t^l$ and $U_t^l$ should only contain non-negative values. This ensures that the embedding does not change the sign of any circuit neuron pre-activations, so that the ReLUs still work correctly.
4. The embedding should distribute circuit neurons approximately evenly over the network neurons. Otherwise we're just wasting space.

Our construction will satisfy these requirements and redundantly embed each circuit neuron into a different subset of $S>1$ network neurons. We call S the embedding redundancy.

Generally speaking, larger values of $S$ are good, because they make each circuit more robust to worst-case errors from interference with other circuits. However, our requirement 2 creates an upper bound on $S$, because the more network neurons we assign to each circuit neuron, the harder it becomes to ensure that no two circuit neurons share more than one network neuron.^[13]

Our allocation scheme is based on prime number factorisation. We promise it's not as complicated as that description may make it sound. The details of how it works are not important for understanding the rest of this post though, so feel free to skip it if you're not interested. 

Step 1 

First, we split the network neurons (for each layer) into $d$ sets of $\frac{D}{d}$ neurons^[14]. The first neuron of each circuit will be embedded in the first of these sets, the second neuron of each circuit will be embedded into the second set and so on, with the $d$-th neuron of each circuit embedded into the $d$-th set of network neurons. 

So next, we need to describe how to embed each set of $T$ circuit neurons into a set of $\frac{D}{d}$ network neurons.

Step 2 

The following algorithm allocates $S>1$ out of  $\frac{D}{d}$ large neurons to each small circuit, while ensuring that each pair of circuits shares at most one such allocation, and that the allocations are fairly evenly spread out over the $\frac{D}{d}$ neurons. 

The allocation algorithm:

Each small circuit neuron is assigned to $S$ large network neurons that are spaced step neurons apart. The set of possible_steps is chosen such that no two different neuron_allocations will overlap more than once. The set of possible_steps is based on prime factors.

First, we define a set of possible_steps^[15] and a function that draws a new instance from that set:

Here, $N$ denotes the natural numbers, $N=\{1,2,3\dots \}$, and $N_0$ denotes the natural numbers including zero, $N_0=\{0,1,2,3\dots \}$.

Next we chose one step from possible_steps. We use this to generate approximately $\frac{D}{Sd}$ non-overlapping neuron_allocations, where each neuron_alocation consists of S subsequent neurons where each one is step neurons away from the previous one.

When we can't generate more non-overlapping neuron_allocations from our first step we draw a new step from possible_steps. Repeat until we have $T$ neuron allocations, i.e, one for each small circuit.

Pseudo code:

next_step = function that generates a new element from possible_steps when called

step = next_step()
shift = 0
start = shift
current_small_circuit = 0
while current_small_circuit < T
	neuron_allocation[current_small_circuit] = [
				start, start + step, ..., start + (S-1)*step]
	start += S*step[17][17][16][16]

	if start + (S-1)*step >= D/d:
		shift += 1
		if shift == step:
			step = next_step()
			shift = 0
		start = shift	
	current_small_circuit+=1

Why it works:

This will make sure no neuron_allocations share more than one neuron because^[17] 

which means that if two different neuron_allocations with different steps, share one neuron, all their other neurons are guaranteed to be different.

Step 3

We construct the embedding matrix $E_t^l$ from $d$ column vectors, each of which has $S$ non-zero values, based on the allocation from Step 2.

Pseudo code:

E_t = zeros(D, d)

E_t[neuron_allocations[t], 0] = 1
E_t[D/d + neuron_allocations[t], 1] = 1
...
E_t[(d-1) * D/d + neuron_allocations[t], (d-1)] = 1

Step 4

To reduce noise, it's important that circuit neurons don't keep having the same "neighbours" from layer to layer, where "neighbours" are neurons from other different circuits that share a network neuron.^[18] To ensure this, we just re-shuffle the circuit index $t$ at every layer. I.e. for each layer, we perform a random permutation on what set of S network neurons is allocated to each circuit neuron.

Step 5

The unembedding matrix is now just the transpose of the embedding matrix normalised by the embedding redundancy $S$:

Real python code

If you prefer real code over words + pseudo code, then here's a Colab Notebook with the full code for creating the embedding and unembedding matrices.

Properties of $E$ and $U$

From I, II, and (1.10) we get that:

From III and (1.10) we can derive that, for any vector $v\in R^d$:

As a reminder, $S$ here is the number of network neurons each circuit neuron is redundantly embedded into. 

The derivation for this assumes that all network neurons are used by the same number of small circuits. In general this will not be strictly true, but it will be close enough to the truth.

Error calculation

In this section we'll calculate approximately how much error is produced and propagated though each layer. We start by giving a mathematical definition of the total error we want to calculate. Then we split this error term into three parts, depending on the origin of the error. Then we estimate each error term separately; finally, we add them all together.

The main results can be found in the summary subsection at the end.

Defining the error terms

In order to define the errors, we introduce $b_t^l$, which is the linear read-off of $a_t^l$ that the large networks can access as input for the next layer:

Inserting this into Equations (1.6) and (1.9), we get:

The error term ${ϵ}_t^l$ is then defined as the discrepancy between the activations of the small networks $a_t^l$ and their linear read-off in the big network $b_t^l$:

Inserting the definitions for both $a_t^l$ and $b_t^l$, we find that the discrepancy is $0$ at the first layer. 

For later layers, the error is:

We can use Equation (1.11) to make the second term more similar to the first term.

In order to simplify this further, we will break up the expression inside the first ReLU. 

To do this, we first notice that if we are holding pre-activation constant, then a ReLU is just a diagonal matrix, with diagonal values equal to 1 or 0, depending on the sign of the pre-activation. We use this fact to replace the first $ReL{U}_D$ with the diagonal matrix $R^l\in R^{D\times D}$ and the second one with $R_t^l\in R^{D\times D}$:

Note that $R^l$ depends on the input; all other matrices here are the same for all inputs.

Now, we split the expression for the error in Equation (2.8) into three parts:

1) The embedding overlap error ${\mathring{ϵ}}_t^l$ is the part of the error that is due to correct activation in active circuits is spilling over into other circuits, because we are using an overcomplete basis. I.e, it's because the embeddings are not completely orthogonal.

This will turn out to be the leading error term for inactive circuits.

2) The propagation error ${\widetilde{ϵ}}_t^l$ is the part of the error that is caused by propagating errors from the previous layer.

This will turn out to be the leading error term in active circuits from layer 2 onward, and the largest error overall.

3) The ReLU activation error ${\mathring{ϵ}}_t^l$ occurs when the ReLUs used by a circuit activate differently that they would if there were no interference, i.e. no noise and no other circuits active:

This error term will turn out to be basically irrelevant in our construction.

The total error in Equation (2.8) is the sum of these three error components, Equations (2.11)-(2.13):

${\mathring{ϵ}}_t^l$ – The embedding overlap error

The embedding overlap error, defined by Equation (2.11), is the error we get from storing the circuit neurons into the network neurons in superposition. 

Calculations:

Looking at Equation (2.11), we first note that we only have to sum over active circuits, since we assumed that $a_u^{l-1}=0$ for inactive circuits. Remember that there are $T$ circuits in total, but only $z\ll T$ are active at a time:

So now we only have to care about the network neurons that are used by active circuits. In general we can't make any assumptions about whether these neurons are active (i.e. $(R^l{)}_{i,i}=1$) or inactive (i.e. $(R^l{)}_{i,i}=0$). We'll therefore go with the most conservative estimate that increases the error the most, which is $R$${}^l\approx I$:

To calculate the mean square of this error, we assume that $w_u^l{a}_u^{l-1}$ for different circuits $u$ are uncorrelated, and then use Equation (1.13):

The sum over $u$ then has $(z-1)$ or $z$ terms depending on whether circuit $t$ is active or inactive:

This gives us the typical size of the embedding overlap error:

${\widetilde{ϵ}}_t^l$ – The propagation error

The propagation error, defined by Equation (2.12) is the largest and most important error overall. This error occurs when we perform computation in superposition over multiple layers, instead of just storing variables in superposition, or performing other types of single-layer superposition operations.

The existence of this error term is related to the fact that we need to embed not just the neurons, but also the weights of the circuits into the network. As opposed to the neuron activations of a circuit, the weights of a circuit don't go away just because the circuit is turned off. This is why we get a factor $\sqrt{T}$ in this term, where $T$ is the total number of circuits, not just the number of active circuits $z\ll T$. This why this error ends up being so large.

Since there are no errors at layer $0$, we get ${\widetilde{ϵ}}_t^1=0$, i.e. the propagation error does not show up until layer 2.

Calculation:

If we were to just conservatively assume that $R^l=I$ for the purpose of this error estimate, as we did with the embedding error, we'd end up with an estimate:

Since $Td>D$, such an error term would quickly overtake the signal. Fortunately, the propagation error is actually much smaller than this, because of how our construction influence $R^l$.

As a reminder, $R^l$ is a diagonal matrix of ReLU activations, defined as:

We will estimate the propagation error in Equation (2.12) by breaking it up into two cases: The error on neurons that are only used by inactive circuits, and the error on neurons that are used by at least one active circuit.

Case 1: Neurons used only by inactive circuits. For neurons $i$ that are only used by inactive circuits:

Our assumption 3 at the start of the Construction section was that $a_u^l=0$ for inactive circuits. Combining this with Equation (1.2), we have:

This is where our crucial assumption 2 from the start of the Construction section comes into play. We required the circuits to be noise robust when inactive, meaning that:

So, assuming that the previous error ${ϵ}_u^{l-1}$ is sufficiently small compared to the noise tolerance of the circuits, we get $(R^l{)}_{i,i}=0$, provided that neuron $i$ is only used by inactive circuits.

Case 2: Neurons that are used by at least one active circuit. Here, we can assume the same conservative estimate on $R^l$ we made use of when calculating the embedding overlap error. I.e. $(R^l{)}_{i,i}=1$ for these neurons.

This means that the propagation error can flow from active to active, inactive to active, and active to inactive circuits. 

There is also a small amount of error propagation from inactive to inactive circuits, whenever the embedding overlap between two inactive circuits also overlaps with an active circuit. But this flow is very suppressed.

To model this, we start with the approximation of $R^l$ we derived in our two cases above:

The minimum in this expression is very annoying to deal with. So we overestimate the error a tiny bit more by using the approximation:

As a reminder, the definition of the propagation error ${\widetilde{ϵ}}_t^l$ was:

Inserting our approximation of $U_t^l{R}^l{E}_u^l$ into this yields:

Using similar calculations as those for the embedding overlap error, we get:

${\ddot{ϵ}}_t^l$ – The ReLU activation error

This error term, defined in Equation (2.13) ends up being negligible. Sometimes it might even reduce the overall error a little.

Calculations:

To help us show this, we introduce ${\Delta }_t$:

With this, the definition of $R^l$, (2.9) becomes:

Remember that the definition of $R_t^l$, is:

 ${\ddot{ϵ}}_t^l$, (2.13) can thus be written as:

There are two combinations of $(R^l{)}_{i,i}$ and $(R_t^l{)}_{i,i}$ that can contribute to ${\ddot{ϵ}}_t^l$. These are $(R^l{)}_{i,i}=1$, $(R_t^l{)}_{i,i}=0$, and $(R^l{)}_{i,i}=0$, $(R_t^l{)}_{i,i}=1$.

Case 1: $(R^l{)}_{i,i}=1,(R_t^l{)}_{i,i}=0$

This will happen if and only if

and in this case the ${\ddot{ϵ}}_t^l$ contribution term will be:

${\Delta }_t$ is the source of the error calculated in the previous sections. Notice that in this case the ReLU activation error contribution ${\left(E_t^l{w}_t^l{a}_t^{l-1}\right)}_i$ is smaller and has opposite sign compared to ${\Delta }_t$. We can therefore safely assume that it will not increase the overall error.

Case 2: $(R^l{)}_{i,i}=0,(R_t^l{)}_{i,i}=1$

This will happen if and only if:

In this case, the ${\ddot{ϵ}}_t^l$ contribution term will be: 

So, the ReLU activation error contribution $-{\left(E_t^l{w}_t^l{a}_t^{l-1}\right)}_i$ is still smaller in magnitude than ${\Delta }_t$ but does have the same sign as ${\Delta }_t$.

However, since $(R^l{)}_{i,i}=0$, ${\Delta }_t$ does not contribute to the total error at all. But in our calculations for the other two error terms, we didn't know the value of $(R^l{)}_{i,i}$, so we included this error term anyway. 

So in this case, the error term coming from the ReLU activation error is also already more than accounted for.

${ϵ}_t^l$ – Adding up all the errors

Let's see how the three error terms add up, layer by layer:

Layer 0

There are no errors here because nothing has happened yet. See equation (2.5):

Layer 1

Since there was no error in the previous layer, we only get the embedding overlap error. From Equations (3.5) and (3.6):

Layer 2

This is the first layer where we get both the embedding overlap error (3.5) - (3.6) and the propagation error (4.10) - (4.11):

The leading term is the last term, i.e. the propagation error flowing in from the inactive circuits in the previous layer: 

The typical noise in any inactive circuit is:

Assuming that the noise in the previous layer is small, the leading term is the embedding overlap error, $\mathring{ϵ}$.

Layer 3

Now both the last terms are of the same size.

For the same reason as last layer, the inactive error is

From here on, it just keeps going like this.

Worst-case errors vs mean square errors

So far, our error calculations have only dealt with mean square errors. However, we also need to briefly address worst-case errors. Those are why we need to have an embedding redundancy $S>1$ in the construction.

For this, we're mainly concerned with the embedding overlap error, because that error term comes from just a few sources (the $z$ active circuits), which means its variance may be high. In contrast, the propagation error comes from adding up many smaller contributions (the approximately $\frac{z{S}^{2}Td}{D}$^[19] circuits that have non-zero embedding overlap error in the previous layer), so we expect it to be well behaved, i.e. well described by the previous mean square error calculations.

The worst-case embedding overlap error happens if some circuit is unlucky enough to be an embedding neighbour to all $z$ active circuits. For active circuits, the maximum number of active neighbours is (z-1) since a circuit can't be its own neighbour. 

From (3.1) and (1.13), we calculate

So, as long as the embedding redundancy $S$ is sufficiently large compared to the number of active circuits $z$, we should be fine. 

Summary:

The main source of error is signal from the active circuits bleeding over into the inactive ones, which then enters back into the active circuits as noise in the next layer. 

The noise in the active circuits accumulate from layer to layer. The noise in the inactive circuits does not accumulate.

At layer $0$, there are no errors, because nothing has happened yet.

At layer $1$ and onward, the leading term for the error on inactive circuits is:

At layer $1$, the leading term for the error on active circuits is: