Incidental polysemanticity

post by Victor Lecomte (victor-lecomte), Kushal Thaman (Kushal_Thaman), tmychow, Rylan Schaeffer · 2023-11-15T04:00:00.000Z · LW · GW · 7 comments

Contents

  Summary
  Intuition
  Outline
  Setup
    Model
    Possible solutions
    Loss and dynamics
  The winning neuron takes it all
    Sparsity force
    How fast does it sparsify?
    Numerical simulations
  Interference arbiters collisions between features
    How strong is the interference?
    Benign and malign collisions
    Experiments
  Discussion and future work
    Implications for mechanistic interpretability
    A more realistic toy model
    Gaps in the theory
  Author contributions
None
7 comments

This is a preliminary research report; we are still building on initial work and would appreciate any feedback.

Summary

Polysemantic neurons (neurons that activate for a set of unrelated features) have been seen as a significant obstacle towards interpretability of task-optimized deep networks,^[1] with implications for AI safety.

The classic origin story of polysemanticity is that the data contains more "features" than there are neurons, such that learning to solve a task forces the network to allocate multiple unrelated features to the same neuron, threatening our ability to understand the network's internal processing.

In this work, we present a second and non-mutually exclusive origin story of polysemanticity. We show that polysemanticity can arise incidentally, even when there are ample neurons to represent all features in the data, using a combination of theory and experiments. This second type of polysemanticity occurs because random initialization can, by chance alone, initially assign multiple features to the same neuron, and the training dynamics then strengthen such overlap. Due to its origin, we term this incidental polysemanticity.

Intuition

The reason why neural networks can learn anything despite starting out with completely random weights is that, just by random chance, some neurons will happen to be very slightly correlated^[2] with some useful feature, and this correlation gets amplified by gradient descent until the feature is accurately represented. If in addition to this there is some incentive for activations to be sparse, then the feature will tend to be represented by a single neuron as opposed to a linear combination of neurons: this is a winner-take-all dynamic.^[3] When a winner-take-all dynamic is present, then by default, the neuron that is initially most correlated with the feature will be the neuron that wins out and represents the feature when training completes.

Therefore, if at the start of training, one neuron happens to be the most correlated neuron with two unrelated features (say dogs and airplanes), then this might^[4] continue being the case throughout the learning process, and that neuron will ultimately end up taking full responsibility for representing both features. We call this phenomenon incidental polysemanticity. Here "incidental" refers to the fact that this phenomenon is contingent on the random initializations of the weights and the dynamics of training, rather than being necessary in order to achieve low loss (and in fact, in some circumstances, incidental polysemanticity might cause the neural network to get stuck in a local optimum).

How often should we expect this to happen? Suppose that we have $n$ useful features to represent and $m\ge n$ neurons to represent them with (so that it is technically possible for each feature to be represented by a different neuron). By symmetry, the probability that the $i^{th}$ and $j^{th}$ feature "collide", in the sense of being initially most correlated with the same neuron, is exactly $1/m$. And there are $(\frac{n}{2})=n(n-1)/2$ pairs of features, so on average we should expect $$\underset{\text{number of pairs}(i,j)}{\underset{}{(\frac{n}{2})}}\times \underset{\text{probability of}(i,j)\text{colliding}}{\underset{}{\frac{1}{m}}}=\frac{n(n-1)}{2m}=\Theta \left(\frac{n^2}{m}\right)$$ collisions^[5] overall. In particular, this means that

• if $m\le O\left(n\right)$ (i.e. the number of neurons is at most a constant factor bigger than the number of features), then $\Omega \left(n^2/n\right)=\Omega \left(n\right)$ collisions will happen: a constant fraction of all neurons will be polysemantic;
• as long as $m$ is significantly smaller than $n^2$, we should expect several collisions to happen.

Our experiments in a toy model show that this is precisely what happens, and a constant fraction of these collisions result in polysemantic neurons, despite the fact that there would be enough neurons to avoid polysemanticity entirely.

Outline

In the rest of this post, we

• set up a toy model for incidental polysemanticity;
• study its winner-take-all dynamic in detail;
• explore what happens over training when features collide, and confirm experimentally that we get as many polysemantic neurons as we expect;
• discuss implications for mechanistic interpretability as well as the limitations of this work, and suggest interesting future work.

Setup

Model

We consider a model similar to the ReLU-output model in Toy Models of Superposition. It is an autoencoder with $n$ features (inputs/outputs) which

• has weight tying between the encoder and the decoder (let $W\in R^{n\times m}$ be those weights),
• uses a single hidden layer of size $m$ with $l_1$ regularization of parameter $\lambda$ on the activations,
• has a ReLU on the output layer,
• has no biases anywhere,
• is trained with the $n$ standard basis vectors as data (so that the "features" are just individual input coordinates): that is, the input/output data pairs are $(e_i,e_i)$ for $i\in \left[n\right]$, where $e_i\in R^n$ is the $i^{th}$ basis vector.

The output is computed as $ReLU\left(W{W}^{T}x\right)$:

The main difference compared to the model from Toy Models of Superposition is the $l_1$ regularization. The role of the $l_1$ regularization is to push for sparsity in the activations and therefore induce a winner-take-all dynamic. We picked this model because it makes incidental polysemanticity particularly easy to demonstrate and study, but we do think the story it tells is representative (see the "Discussion and future work" section for more on this).

We make the following assumptions on parameter values:

Assumption	Reason
the weights $W_{ik}$ are initialized to i.i.d. normals of mean $0$ and standard deviation $\Theta \left(1/\sqrt{m}\right)$	so that the encodings $W_i\in R^m$ start out with constant length
$m\ge n$	just to make it clear that polysemanticity was not necessary in this case
$\lambda \le 1/\sqrt{m}$	so that the $l_1$ regularization doesn't kill all weights immediately

Possible solutions

Let $W_i\in R^m$ be the $i^{th}$ row of $W$. It tells us how the $i^{th}$ feature is encoded in the hidden layer. When the input is $e_i$, the output of the model can then be written as $$\left(ReLU\right(W_1\cdot W_i),\dots ,ReLU(W_n\cdot W_i\left)\right),$$ so for this to be equal to $e_i$ we need $\parallel W_i{\parallel }^2=1$^[6] and $W_i\cdot W_j\le 0$ for $j\ne i$.

Letting $f_k\in R^m$ denote the $k^{th}$ basis vector in $R^m$. There are both monosemantic and polysemantic solutions that satisfy these conditions:

• One solution is to simply let $W_i:=f_i$: the $i^{th}$ hidden neuron represents the $i^{th}$ feature, and there is no polysemanticity.
• But we could also have solutions where two features share the same neuron, with opposite signs. For example, for each $i\in \left[n/2\right]$, we could let $W_{2i-1}:=f_i$ and $W_{2i}:=-f_i$. This satisfies the conditions because $W_{2i-1}\cdot W_{2i}=f_i\cdot (-f_i)=-1\le 0$.
• In general, we can have a mixture of these where each neuron represents either $0$, $1$ or $2$ features, in an arbitrary order.

Loss and dynamics

Let us consider total squared error loss, which can be decomposed as $$L=\sum _i({(1-\parallel W_i{\parallel }^2)}^2+\sum _{j\ne i}ReLU(W_i\cdot W_j{)}^2+\lambda \parallel W_i{\parallel }_1).$$ The training dynamics are $$\frac{d{W}_i}{dt}:=-\frac{\partial L}{\partial W_i}=4\underset{\text{feature benefit}}{\underset{}{(1-\parallel W_i{\parallel }^2)W_i}}-4\underset{\text{interference}}{\underset{}{\sum _{j\ne i}ReLU(W_i\cdot W_j)W_j}}-\underset{\text{regularization}}{\underset{}{\lambda sign\left(W_i\right)}},$$ where $t$ is the training time, which you can roughly think of as the number of training steps. For simplicity, we'll ignore the constants $4$ going forward.^[7]

It can be decomposed into three intuitive "forces" acting on the encodings $W_i$:

• "feature benefit": encodings want to have unit length;
• "interference": different encodings avoid pointing in similar directions;
• "regularization": encodings want to have small $l_1$-norm (which pushes all nonzero weights towards zero with equal strength).

The winning neuron takes it all

See our working notes (in particular, Feature benefit vs regularization) for a more formal treatment.

Sparsity force

For a moment, let's ignore the interference force, and figure out how (and how fast) regularization will push towards sparsity in some encoding $W_i$. Since we're only looking at feature benefit and regularization, the other encodings $W_j$ have no influence at all on what happens in $W_i$.

Assuming $\parallel W_i\parallel <1$, each weight $W_{ik}$ is

• pushed up with strength $(1-\parallel W_i{\parallel }^2)W_{ik}$ by the feature benefit force;
• pushed down with strength $\lambda sign\left(W_{ik}\right)$ by the regularization.

Crucially, the upwards push is relative to how large $W_{ik}$ is, while the downwards push is absolute. This means that weights whose absolute value is above some threshold $\theta$ will grow, while those below the threshold will shrink, creating a "rich get richer and poor get poorer" dynamic that will push for sparsity. This threshold is given by $$(1-\parallel W_i{\parallel }^2)W_{ik}=\lambda sign\left(W_i\right)⟺|W_{ik}|=\frac{\lambda }{1-\parallel W_i{\parallel }^2}=:\theta ,$$ so we have $$\begin{array}{cc}\frac{d|W_{ik}|}{dt}& =\underset{\text{feature benefit}}{\underset{}{(1-\parallel W_i{\parallel }^2)|W_{ik}|}}-\underset{\text{regularization}}{\underset{}{\lambda f1[W_{ik}\ne 0]}}& & =\{\begin{array}{cc}\underset{\text{constant in}k}{\underset{}{(1-\parallel W_i{\parallel }^2)}}\underset{\text{distance from threshold}}{\underset{}{(|W_{ik}|-\theta )}}& \text{if}{W}_{ik}\ne 0\\ 0& \text{otherwise.}\end{array}\\ \text{(1)}\end{array}$$

We call this combination of feature benefit and regularization force the sparsity force. It uniformly stretches the gaps between (the absolute values of) different nonzero weights.

Note that the threshold $\theta$ is not fixed: we will see that as $W_i$ gets sparser, $\parallel W_i{\parallel }^2$ will get closer to $1$, which increases the threshold and allows it to get rid of larger and larger entries, until only one is left. But how fast will this go?

How fast does it sparsify?

The next two subsections are not critical for understanding the overall message; feel free to skip directly to the section titled "Interference arbiters collisions between features" if you're happy with just accepting the fact that $W_i$ will progressively sparsify over some predictable length of training time.

In order to track how fast $W_i$ sparsifies, we will look at its $l_1$ norm $$\parallel W_i{\parallel }_1=\sum _k|W_{ik}|$$ as a proxy for how many nonzero coordinates are left. Indeed, we will have $\parallel W_i\parallel \approx 1$ throughout, so if $W_i$ has $m^{\prime }$ nonzero values at any point in time, their typical value will be $\pm 1/\sqrt{m^{\prime }}$, which means $\parallel W_i{\parallel }_1\approx m^{\prime }\frac{1}{\sqrt{m^{\prime }}}=\sqrt{m^{\prime }}$.

Since the sparsity force is proportional to $1-\parallel W_i{\parallel }^2$, we need to get a sense of what values $\parallel W_i\parallel$ will take over time. As it turns out, $\parallel W_i\parallel$ changes relatively slowly, so we can get useful information by assuming the derivative $\frac{d\parallel W_i{\parallel }^2}{dt}$ is $0$: $$\begin{array}{cc}0& \approx \frac{d\parallel W_i{\parallel }^2}{dt}\\ & =2\frac{d{W}_i}{dt}\cdot W_i\\ & =2(\underset{\text{from feature benefit}}{\underset{}{(1-\parallel W_i{\parallel }^2)\parallel W_i{\parallel }^2}}-\underset{\text{from regularization}}{\underset{}{\lambda \parallel W_i{\parallel }_1}}),\end{array}$$ which means $$1-\parallel W_i{\parallel }^2\approx \frac{\lambda \parallel W_i{\parallel }_1}{\parallel W_i{\parallel }^2}.$$ Plugging this back into $$\frac{d\parallel W_i{\parallel }_1}{dt}=\sum _k\frac{d|W_{ik}|}{dt}$$ and using reasonable assumptions about the initial distribution of $W_i$ (see our working notes for details), we can prove that $\parallel W_i{\parallel }_1$ will decrease as $1/\lambda t$ with training time $t$: $$\parallel W_i\left(t\right){\parallel }_1=\{\begin{array}{cc}\Theta \left(\sqrt{m}\right)& t\le \frac{1}{\lambda \sqrt{m}}\\ \Theta \left(1/\lambda t\right)& \frac{1}{\lambda \sqrt{m}}\le t\le \frac{1}{\lambda }\\ \Theta \left(1\right)& t\ge \frac{1}{\lambda }.\end{array}$$ Correspondingly, if we approximate the number $m^{\prime }$ of nonzero cooordinates as $\parallel W_i{\parallel }_1^2$, it will start out at $m$, decrease as $1/(\lambda t{)}^2$, then reach $1$ at training time $t=\Theta \left(1/\lambda \right)$.

Numerical simulations

We compared our theoretical predictions for $\parallel W_i{\parallel }_1$ and $m^{\prime }$ (if the constants hidden in $\Theta (\cdot )$ are assumed to be $1$) to their actual values over training time when the interference force is turned off. The specific values of parameters are $m:={10}^5$ and $\lambda :={10}^{-5}$, and the standard deviation of the $W_{ik}$'s was $\frac{0.9}{\sqrt{m}}$.

Code is available here.

Interference arbiters collisions between features

What happens when you bring the interference force into this picture? In this section, we argue informally that the interference is initially weak if $m\ge n$, and only becomes significant later on in training, in cases where two of the encodings $W_i$ and $W_j$ have a coordinate $k$ such that $W_{ik}$ and $W_{jk}$ are both large and have the same sign—when that's the case, the larger of the two wins out.

How strong is the interference?

First, observe that in the expression for the interference force on $W_i$ $$-\sum _{j\ne i}ReLU(W_i\cdot W_j)W_j,$$ each $W_j$ contributes only if the angle it forms with $W_i$ is less than ${90}^{\circ }$. So the force will mostly be in the same direction as $W_i$, but opposite. That means that we can get a good grasp on its strength by measuring its component in the direction of $W_i$, which we can do by taking an inner product with $W_i$.

We have $$\begin{array}{cc}\left(\sum _{j\ne i}ReLU\right(W_i\cdot W_j\left)W_j\right)\cdot W_i& =\sum _{j\ne i}ReLU(W_i\cdot W_j)(W_i\cdot W_j)\\ & =\sum _{j\ne i}ReLU(W_i\cdot W_j{)}^2.\end{array}$$ Initially, each encoding is a vector of $m$ i.i.d. normals of mean $0$ and standard deviation $\Theta \left(1/\sqrt{m}\right)$, so the distribution of the inner products $W_i\cdot W_j$ is symmetric around $0$ and also has standard deviation $\Theta \left(1/\sqrt{m}\right)$. This means that $ReLU(W_i\cdot W_j{)}^2$ has mean $\Theta \left(1/m\right)$, and thus the sum has mean $\Theta \left(n/m\right)$. As long as $m\ge n$, this is dominated by the feature benefit force: indeed, the same computation for the feature benefit gives $$\left((1-\parallel W_i{\parallel }^2)W_i\right)\cdot W_i=(1-\parallel W_i{\parallel }^2)\parallel W_i{\parallel }^2=\Theta \left(1\right)$$ as long as $\Omega \left(1\right)\le \parallel W_i{\parallel }^2\le 1-\Omega \left(1\right)$.

Moreover, over time, the positive inner products $W_i\cdot W_j>0$ will tend to decrease exponentially. This is because the interference force on $W_i$ includes the term $-ReLU(W_i\cdot W_j)W_j$ and the interference force on $W_j$ includes the term $-ReLU(W_i\cdot W_j)W_i$. Together, they affect $W_i\cdot W_j$ as $$\begin{array}{cc}(-ReLU(W_i\cdot W_j\left)W_j\right)\cdot W_j+(-ReLU(W_i\cdot W_j\left)W_i\right)\cdot W_i& =-(W_i\cdot W_j)(\parallel W_i{\parallel }^2+\parallel W_j{\parallel }^2)\\ & =-\Theta (W_i\cdot W_j)\end{array}$$ as long as $\parallel W_i{\parallel }^2,\parallel W_j{\parallel }^2=\Theta \left(1\right)$, which is definitely the case at the start and will continue to hold true throughout training.

Benign and malign collisions

On the other hand, the interference between two encodings $W_i$ and $W_j$ starts to matter significantly when it affects one coordinate much more strongly than the others (rather than affecting all coordinates proportionally, like the feature benefit force does). This is the case when $W_i$ and $W_j$ share only one nonzero coordinate: a single $k$ such that $W_{ik},W_{jk}\ne 0$. Indeed, when that's the case, the interference force $-ReLU(W_i\cdot W_j)W_j$

• only affects the coordinates of $W_i$ that are nonzero in $j$,
• and will probably not be strong enough counter the $l_1$-regularization and revive coordinates of $W_i$ that are currently zero,

so only $W_{ik}$ can be affected by this force.

When this happens, there are two cases:

• If $W_{ik}$ and $W_{jk}$ have opposite signs, we have $W_i\cdot W_j=W_{ik}{W}_{jk}<0$, so nothing actually happens, since the ReLU clips this to $0$. Let's call this a benign collision.
• If $W_{ik}$ and $W_{jk}$ have the same sign, we have $W_i\cdot W_j=W_{ik}{W}_{jk}>0$, and both weights will be under pressure to shrink, with strength $-W_{ik}{W}_{jk}^2$ and $-W_{ik}^2{W}_{jk}$ respectively. Depending on their relative size, one or both of them will quickly drop to $0$, thus putting the $k^{th}$ neuron out of the running in terms of representing the corresponding features. Let's call this a malign collision.

Polysemanticity will happen when the largest^[8] coordinates in encodings $W_i$ and $W_j$ get into a benign collision. This happens with probability $$\underset{\text{largest weight in}{W}_i\text{is also largest in}{W}_j}{\underset{}{\frac{1}{m}}}\times \underset{\text{they have opposite signs}}{\underset{}{\frac{1}{2}}}=\frac{1}{2m},$$ so we should expect roughly $$(\frac{n}{2})\times \frac{1}{2m}\sim \frac{n^2}{4m}$$ polysemantic neurons by the end.

Experiments

Training the model we described on $n:=256$ and $m$ ranging from $256$ to $4096$ shows that this trend of $\Theta \left(\frac{n^2}{m}\right)$ does hold, and the constant $\frac{1}{4}$ seems to be fairly accurate as well.

Discussion and future work

Implications for mechanistic interpretability

The fact that there are two completely different ways for polysemanticity to occur could have important consequences on how to deal with it.

To our knowledge, polysemanticity has mostly been studied in settings where the encoding space has no privileged basis: the space can be arbitrarily rotated without changing the dynamics, and in particular the corresponding layer doesn't have non-linearities or any regularization other than $l_2$. In this setting, the features can be represented arbitrarily in the encoding space, and we usually observe interference (non-orthogonal encodings) only when there are more features than dimensions.

On the other hand, the incidental polysemanticity we have demonstrated here is inherently tied to the canonical basis, contingent on the random initialization and dynamics, and happens even when there are significantly more dimensions available than features.

This means that some tools that work against one type of polysemanticity might not work against the other. For example:

• Tools that make assumptions about the linear structure of the encodings might not work as well when non-linearities are present.
• A costly but technically feasible way to get rid of superposition when there is no privileged basis is to just increase the number of neurons so that it matches the number of features. It is much less realistic to do away with incidental polysemanticity in this way, since we saw that it can happen until the number of neurons is roughly equal to the number of features squared.
• On the other hand, since incidental polysemanticity is contingent on the random initializations and the dynamics of training, it could be solved by nudging the trajectory of learning in various ways, without necessarily changing anything about the neural architecture.
  • As an example, here is one possible way one might get rid of incidental polysemanticity in a neuron that currently represents two features $i$ and $j$. Duplicate that neuron, divide its outgoing weights by $2$ (so that this doesn't affect downstream layers), add a small amount of noise to the incoming weights of each copy, then run gradient descent for a few more steps. One might hope that this will cause the copies to diverge away from each other, with one of the copies eventually taking full ownership of feature $i$ while the other copy takes full ownership of feature $j$.

In addition, it would be interesting to find ways to distinguish incidental polysemanticity from necessary polysemanticity.

• Can we distinguish them based only on the final, trained state of the model, or do we need to know more about what happened during training?
• Is "most" of the polysemanticity in real-world neural networks necessary or incidental? How does this depend on the architecture and the data?

A more realistic toy model

The setup we studied is simplistic in several ways. Some of these ways are without loss of (much) generality, such as the fact that encoding and decoding matrices are tied together,^[9] or the fact that the input features are basis vectors.^[10]

But there are also some choices that we made for simplicity which might be more significant, and which it would be nice to investigate. In particular:

• We introduced a winner-takes-all dynamic using $l_1$ regularization. This made the study of sparsification quite nice, but $l_1$ regularization is not very commonly used, and is not the main cause of privileged bases in modern neural networks. It would be interesting to study winner-takes-all dynamics that occur for different reasons (such as nonlinear activation functions or layer normalization) and see if the dynamics of sparsification are similar.
• In our setup, polysemanticity is benign because one feature can be represented with a positive activation and the other one with a negative activation, and this is disambiguated by the ReLU in the output layer. Polysemanticity can be benign for other reasons: for example, it could be that two features never occur in the same context and that the context can be used to disambiguate between them. We give a possible toy model for this scenario in a working note.

Gaps in the theory

We were able to give strong theoretical guarantees for the sparsification process by considering how the feature benefit force and regularization interact when interference is ignored, but we haven't yet been able to make confident theoretical claims about how the three forces interact together.

In particular:

• It would be nice to theoretically predict the likelihood that the largest coordinate in $W_i$ remains largest throughout training. This would require a better understanding of exactly how impactful malign collisions are when one of the weights involved is significantly larger than the other, and how much this affects the race between the largest weights in an encoding.
• Perhaps an interesting case to study is the limit $\lambda \to 0$ (very slow regularization). This would make the interference work much faster than the regularization, and from the perspective of studying sparsification, this could allow us to assume that the inner products $W_i\cdot W_j$ are nonpositive throughout training.

Author contributions

• Ideation: Victor initially proposed the project, and it took shape in discussions with Rylan.
• Experiments: Kushal and Trevor ran the experiments which confirmed that incidental polysemanticity occurs and measured the number of polysemantic neurons. They were mentored by Victor and Rylan.
• Theory: Victor drove the theory work and implemented the numerical simulations of sparsification, with significant help from Kushal.
• Writing: This post was primarily written by Victor, with feedback from Kushal and Rylan.

---

1. see e.g. the "Polysemantic Neurons" section in Zoom In: An Introduction to Circuits ↩︎

2. When we say a neuron is correlated with a feature, what we more formally mean is that the neuron's activation is correlated with whether the feature is present in the input (where the correlation is taken over the data points). But the former is easier to say. ↩︎

3. Analogous phenomena are known under other names, such as "privileged basis". ↩︎

4. depending mostly on the specifics of the neural architecture and the data (but also on the random initializations of the weights) ↩︎

5. Here, we define a "collision" as the event that two features $i$ and $j$ collide. So for example there is a three-way collision between $i$, $j$ and $k$, that would count as three collisions between $i$ and $j$, $i$ and $k$, and $j$ and $k$. ↩︎

6. We use $\parallel \cdot \parallel$ to denote Euclidean length ($l_2$ norm), and $\parallel \cdot {\parallel }_1$ to denote Manhattan length ($l_1$ norm). ↩︎

7. It's equivalent to making $\lambda$ four times larger and making training time four times slower. ↩︎

8. This would not necessarily be the largest weight at initialization, since there might be significant collisions with other encodings, but the largest weight at initialization is still the most likely to win the race all things considered. ↩︎

9. We're referring to the fact that the encoding matrix $W^T$ is forced to be the transpose of the decoding matrix $W$. This assumption makes sense because even if they were kept independent and initialized to different values, they would naturally acquire similar values over time because of the learning dynamics. Indeed, the $i^{th}$ column of the encoding matrix and the $i^{th}$ row of the decoding matrix "reinforce each other" through the feature benefit force until they have an inner product of $1$, and ao as long as they start out small or if there is some weight decay, they would end up almost identical by the end of training. ↩︎

10. If the input features are not the canonical basis vectors but are still orthogonal (and the outputs are still basis vectors), then we could apply a fixed linear transformation to the encoding matrix and recover the same training dynamics. And in general it makes sense to consider orthogonal input features, because when the features themselves are not orthogonal (or at least approximately orthogonal), the question of what polysemanticity even is becomes more murky. ↩︎

7 comments

Comments sorted by top scores.