# Polysemanticity and Capacity in Neural Networks

Adam Scherlis1, Kshitij Sachan1, Adam S. Jermyn2, Joe Benton, Buck Shlegeris1

1Redwood Research
2Flatiron Institute

###### Abstract

Individual neurons in neural networks often represent a mixture of unrelated
features. This phenomenon, called polysemanticity, can make interpreting
neural networks more difficult and so we aim to understand its causes. We
propose doing so through the lens of feature _capacity_ , which is the
fractional dimension each feature consumes in the embedding space. We show
that in a toy model the optimal capacity allocation tends to monosemantically
represent the most important features, polysemantically represent less
important features (in proportion to their impact on the loss), and entirely
ignore the least important features. Polysemanticity is more prevalent when
the inputs have higher kurtosis or sparsity and more prevalent in some
architectures than others. Given an optimal allocation of capacity, we go on
to study the geometry of the embedding space. We find a block-semi-orthogonal
structure, with differing block sizes in different models, highlighting the
impact of model architecture on the interpretability of its neurons.

##  1 Introduction

Individual neurons in neural networks often represent multiple unrelated
features in the input [OMS17, OCS+20]. This phenomenon is known as
polysemanticity, and makes it more difficult to interpret neural networks
[OCS+20]. While “feature” is a somewhat fuzzy concept [EHO+22b], there are at
least some cases where we “know it when we see it”. For example, when the
input features are independent random variables that do not interact in the
data-generating process, neurons that represent combinations of these input
features can be confidently called polysemantic. In this work we explore how
loss functions incentivize polysemanticity in this setting, and the structure
of the learned solutions.

Fittingly, there are multiple ways that polysemanticity can manifest. Here we
focus on one form that seems particularly fundamental, namely superposition
[EHO+22b]. Suppose we have a linear layer that embeds features which then pass
through a layer with a nonlinear activation function. The feature embedding
vectors might not be orthogonal, in which case multiple neurons (nonlinear
units) are involved in representing each feature. When there are at least as
many features as neurons this means that some neurons represent multiple
features, and so are polysemantic (Figure 1, right). There are other causes of
polysemanticity, e.g. feature embedding vectors could be rotated relative to
the neuron basis (Figure 1, left), but we do not study these in this work.

Figure 1: Feature embedding vectors are shown in two dimensions. The neuron
basis corresponds to the coordinate axes. Left: rotated embeddings. Right:
non-orthogonal embeddings. In both cases the result is polysemanticity because
each neuron receives some input when either feature is present.

Here we build on the work of [EHO+22b], who studied polysemanticity in the
context of toy models of autoencoders. They found that models can support both
monosemantic and polysemantic neurons, that polysemantic neurons can perform
certain kinds of computations, and that the embedding vectors of features
often formed repeating motifs of a few features symmetrically embedded in a
low-dimensional subspace. Moreover, in their models they found distinct
“phases” where superposition was either significant or completely absent.
Sparser inputs resulted in more superposition. Features with similar
importance were more likely to be in superposition. This reflects an abundance
of unexpected structure, and gives new handles on the phenomenon of
polysemanticity.

We study these phenomena through the lens of _capacity_ , or the fraction of
an embedding dimension allocated to each feature (Section 2, also termed
“dimensionality” by [EHO+22b]). This ranges from 0-1 for each feature, and the
total capacity across all features is bounded by the dimension of the
embedding space. Because the model has a limited number of neurons and so a
limited number of embedding dimensions, there is a trade-off between
representing different features. We find that the capacity constraint on
individual features (0-1) means that many features are either ignored
altogether (not embedded) or else allocated a full dimension orthogonal to all
the other features in the embedding space, depending on the relative
importance of each feature to the loss. Features are represented
polysemantically only when the marginal loss reduction of assigning more
capacity to each is equal (Figure 2). This neatly explains the sharp “pinning”
of features to either 0 or 1 capacity noted by [EHO+22b], and gives us a
framework for understanding the circumstances under which features are
represented polysemantically.

Figure 2: The marginal loss reduction $-\partial L/\partial C_{i}$ is shown
for several features as a function of feature capacity in our toy model.
Circles represent the optimal capacity allocation for a particular total
embedding dimension. Colors vary to make individual curves more
distinguishable.

To explore capacity allocation in a concrete model, we instantiate our theory
for a one-layer model with quadratic activations (Section 3). Our model
differs from the Anthropic toy model in that ours uses a different activation
function to make the math more tractable, and, more importantly, ours is
focused on polysemantic computation rather than data compression. We contrast
these toy models in Figure 3.

Figure 3: Comparison between the Anthropic toy model of [EHO+22b] (left) and
our toy model (right). The key difference is that the Anthropic model studies
the compression and recovery of high-dimensional vectors, while ours examines
how a smaller number of polysemantic neurons can simulate the computation done
by a larger number of monosemantic ones. Figure kindly provided by Chris Olah.

For our toy model we can analytically determine the capacity allocation as a
function of feature sparsity and importance (i.e. weight in the loss), and so
construct a “phase diagram” (Figure 4). While the details of our phase diagram
differ from those of [EHO+22b], reflecting our different toy model, there are
three qualitative features that are in good agreement. First, when a feature
is much more important than the rest, it is always represented fully with its
own embedding dimension. Second, when a feature is much less important than
the rest, it is ignored entirely. Finally, in a sparsity-dependent
intermediate region features are partially represented, sharing embedding
dimensions. In addition, this confirms our theoretical expectation that
capacity is allocated according to how much each feature matters to the loss
(a mixture of importance and sparsity) and that it is often allocated to fully
ignore some features while fully representing others. We supplement this with
empirical results for a variety of activation functions showing that the phase
diagram predicts the behavior of a broad family of 2-layer models.

Figure 4: Upper: Analytical and empirical phase diagrams for our toy model
with 6 features and 3 neurons. In both panels one feature has a different
importance from the rest, and colors show the resulting capacity allocation
for that feature as a function of sparsity and relative importance. Lower:
Plots of marginal loss reduction $\partial L/\partial C_{i}$ as a function of
feature capacity for each labelled point in the analytical phase diagram. The
blue curve represents the feature with varied importance and the black one
represents the constant important feature. Black dots are optimal allocations
of capacity.

We then turn to study the geometry of the embedding space (Section 4). When
embedding matrices fully utilize the available capacity we call them
“efficient”. We find that every efficient embedding matrix has a block-semi-
orthogonal structure, with features partitioned into different blocks. When
multiple features in a block are present they interfere with each other,
causing spurious correlations in the output and hence greater loss. Features
do not, however, interfere across blocks.

Figure 5: Left: An embedding matrix with two blocks. Center: The relationship
between features and (principal-component-aligned) neurons for this matrix.
Right: Embedding vector geometry for this matrix.

The blocks in efficient matrices correspond to the polytope structure
[EHO+22b] found, with small blocks corresponding to features embedded as
regular polytopes and large blocks corresponding to less-ordered structures.
Large- and small-block arrangements come with different advantages. With large
blocks there is significant freedom to allocate capacity across features,
whereas with small blocks there is the additional constraint that the capacity
of each block be an integer and that the block capacities add up to the total
capacity. On the other hand, with small blocks the lengths of embedding
vectors can be chosen more freely because blocks can be scaled independently
of each other without affecting the capacity allocation.

In our quadratic model the embedding matrices in our toy model always have one
large block, which is correspondingly less structured. We expect that
differences in architecture can lead to different sizes of blocks, which could
provide a way to control the extent of polysemanticity in models, alongside
other approaches such as changing the activation function [EHO+22a].

##  2 Capacity and Superposition

###  2.1 Definitions

Suppose we have a model composed of stacks of linear layers with nonlinear
activation functions. In each layer, the model applies a linear transform to
the input vector $x$ to produce an embedding vector $e$, and then performs an
element-wise non-linear calculation on those embeddings to produce the non-
linear activation vector $h$. For instance, we might have

$\displaystyle e$ $\displaystyle\equiv W\cdot x$ (1) $\displaystyle h$
$\displaystyle\equiv\mathrm{ReLU}(e)$ (2)

with $W\in\mathbb{R}^{d\times p}$, $x\in\mathbb{R}^{p}$,
$e,h\in\mathbb{R}^{d}$. We associate each dimension of the input vector $x$
with a feature, and we call each dimension of the non-linear layer a neuron.

For simplicity, in the rest of this paper we work with a one-layer model, but
our capacity definition should be valid for any layer in a multi-layer model.

When a model represents a feature in the input space, it is convenient to
think that it expends some capacity to do so. Our intuition here is that as we
ask a model to represent more and more features we eventually exhaust its
ability to do so, resulting in features interfering. We study the
superposition phenomena by asking the question: "How do models allocate
limited representation capacity to input features?" In what follows we assume
that each input feature is assigned a unique dimension in the input space
(e.g. feature $i$ is input dimension $i$), and capacity we define below.

Let $W_{i}\in\mathbb{R}^{d}$ be the embedding vector for feature $i$. The
capacity allocated to feature i is

$\displaystyle C_{i}=\frac{(W_{i}\cdot W_{i})^{2}}{\sum_{j}(W_{i}\cdot
W_{j})^{2}}$ (3)

We can think of $C_{i}$ as “the fraction of a dimension” allocated to feature
$i$ ([EHO+22b])111We can also interpret $C_{i}$ as the squared correlation
coefficient between $x_{i}$ and $(W^{T}Wx)_{i}$ – see Appendix E.. The
numerator measures the size of the embedding and the denominator tracks the
interference from other features. By this definition, $C_{i}$ is bounded
between 0 and 1. In the case $W_{i}=0$, where this expression is undefined, we
set $C_{i}=0$.222This is the limit of the expression from almost all
directions, so in practice $W_{i}\approx 0$ implies $C_{i}\approx 0$.

We define the total model capacity to be $C=\sum_{i}C_{i}$ (in a multi-layer
model, this would be a single layer-pair’s capacity). This is bounded between
1 and the embedding dimension $D$ (see Appendix F for a proof of the upper
bound).

Figure 6: Example capacity allocations for different embeddings.

Note that a set of capacities does not uniquely specify a weight matrix. For
example, capacity is invariant to the overall scaling and rotation of $W$. In
what follows it will be useful to have a full parameterization of $W$ that
includes $C_{i}$, so we define $S$ to be a set of additional parameters that
uniquely specify a weight matrix $W$ given its capacities
$C_{1},\ldots,C_{N}$. We can then parameterize the loss using
$(C_{1},\ldots,C_{N},S)$ rather than $W$.

###  2.2 Loss Minimization

We are interested in how loss minimization allocates capacity among different
features. Because the capacity of each feature lies in $[0,1]$ and there is
also a constraint on the total capacity of a model, this is a constrained
optimization problem:

$\displaystyle\min_{C_{1:n},S}$ $\displaystyle L(C_{1:n},S)$ s.t.
$\displaystyle 0\leq C_{i}\leq 1$ $\displaystyle 1\leq\sum_{i}C_{i}\leq D$

In brief, we minimize over $S$ first to find the $S^{*}(C_{i})$ minimizing the
loss333There are multiple ways to define $S$, which give different $(\partial
L/\partial C_{i})$ holding ”$S$” constant. Minimizing over $S$ removes this
ambiguity.. We then minimize over $C_{i}$ Intuitively capacity should always
be a good thing, so we are interested in situations where the loss function
$L$ is monotonically decreasing in $C_{i}$. In this case, the bound on the
total capacity $\sum_{i}C_{i}\leq D$ will be saturated, and we can replace
this constraint with a Lagrange multiplier:

$\displaystyle\min_{C_{1:n},\lambda}\quad$ $\displaystyle
L(C_{1:n})+\lambda\left(\sum_{i}C_{i}-D\right)$ (4) s.t. $\displaystyle 0\leq
C_{i}\leq 1$

Because we have independent constraints on each $C_{i}$, for each $i$ the
minimum occurs either at a boundary ($C_{i}=0$ or $1$) or in between at a
point where

$\displaystyle\frac{\partial}{\partial
C_{i}}\left(L(C_{1:n})+\lambda\left(\sum_{i}C_{i}-D\right)\right)=0$ (5)

The derivative in the three cases will obey

$\displaystyle\frac{\partial L(C_{1:n})}{\partial C_{i}}$
$\displaystyle\geq-\lambda\quad(C_{i}=0)$ (6) $\displaystyle\frac{\partial
L(C_{1:n})}{\partial C_{i}}$ $\displaystyle=-\lambda\quad(0<C_{i}<1)$ (7)
$\displaystyle\frac{\partial L(C_{1:n})}{\partial C_{i}}$
$\displaystyle\leq-\lambda\quad(C_{i}=1)$ (8)

In other words, at the minimum of the loss, capacity is allocated to each
feature until the marginal value (decrease in loss) from additional capacity
is at a constant threshold – unless the feature is so important that the
marginal value is above threshold even when fully represented, or so
unimportant that the marginal value is below threshold even when the feature
is ignored.

Note that none of this derivation depended on the precise definition of
capacity. We only had to assume that (1) allocating more capacity to a feature
never makes the loss greater, (2) all capacity allocations satisfying the
feature-level and total constraints are realizable by some embedding matrix.
Thus we could have used other definitions, and indeed in Appendix G we explore
one such alternative.

We can visualize this in terms of the graph of $\partial L/\partial C_{i}$ as
a function of $C_{i}$, holding $C_{j}$ constant for $j\neq i$. For instance,
in Figure 2 there are diminishing returns: $\partial^{2}L/\partial
C_{i}^{2}>0$, so that each additional unit of capacity makes the next one less
useful. If $\partial L/\partial C_{i}$ crosses $-\lambda$, the value of
$C_{i}$ where it crosses will be the optimal allocation. Features only get 0
or 1 capacity if the entire $\partial L/\partial C_{i}$ curve is above or
below $-\lambda$ (respectively).

By contrast, in the case of increasing returns ($\partial^{2}L/\partial
C_{i}^{2}<0$), every feature will have $C_{i}=0$ or $C_{i}=1$ and
superposition will not occur at all. As we show below for several different
toy models, diminishing returns occur for inputs of high sparsity or kurtosis,
and accelerating return for inputs of low sparsity or kurtosis.

##  3 Quadratic model

We now instantiate our theory for a two-layer model with quadratic
activations. This choice of activation function allows us to study this model
analytically in some detail.

We define our model in Section 3.1. In Section 3.2 we analyze the loss
function for the case of independent random input features. We are able to
write the expected loss entirely in terms of the embedding lengths and the
variance and kurtosis of feature strengths.

In Section 3.3 we demonstrate a surprising feature of this quadratic model:
the loss function can be written in terms of the feature capacity and the
embedding lengths. This allows us to analytically explore the marginal benefit
of allocating more capacity to each feature.

We then go on to solve for the optimum allocation of capacity across features
(Section 3.4) and relate this to our results from Section 2).

###  3.1 The model

We study a regression task, with ground-truth data $x,y(x)$ and modeled data
$x,\tilde{y}(x)$. The ground truth data is generated by an input distribution
over $x$ and a set of “importance" coefficients $v_{i}$:

$\displaystyle y$ $\displaystyle=\sum_{i}v_{i}x_{i}^{2}.$ (9)

We think of $v_{i}$ as the importance of feature $i$, and highlight that $y$
here is a scalar. The model $\tilde{y}(x)$ is parameterized as

$\displaystyle\tilde{y}$ $\displaystyle=\sum_{i}(W_{i}\cdot x)^{2}+b$ (10)

where $b$ is a bias and $W$ is a rectangular embedding matrix $W$. Like $y$,
$\tilde{y}$ is a scalar. Note that the ground-truth and modeled data are both
of the form $x^{T}Ax$ for a matrix $A$. For the ground truth, this is a full-
rank diagonal444In general, we could consider a ground truth with an arbitrary
full-rank $A$ instead of a diagonal matrix. Forcing $A$ to be diagonal is then
equivalent to choosing to always write $x$ in the principal-component (SVD)
basis of $A$. This gives a slightly stronger motivation for our choice to
interpret the components of $x$ as separate “features”: we are really choosing
features based on the principal axes of $A$. matrix; for our model it is
limited in rank by the shape of $W$. Therefore, perfect loss is not possible
in general.

Our loss function is squared error

$\displaystyle L(x)=(\tilde{y}(x)-y(x))^{2}.$ (11)

###  3.2 Interpreting the loss function

We now assume that $\mathbb{E}[x_{i}]=0$ and
$\text{Cov}[x_{i},x_{j}]=\delta_{ij}$, but do not impose any additional
restrictions on the distribution of $\\{x_{i}\\}$. With this, we find
(Appendix A)

$\mathbb{E}[L]=(\mathbb{E}[x_{i}^{4}]-1)\sum_{i}(||W_{i}||^{2}-v_{i})^{2}+2\sum_{i\neq
j}(W_{i}\cdot W_{j})^{2}$ (12)

The loss is a weighted sum of two components. The first term captures true
correlations in the data, because
$v_{i}\propto\text{Cov}[\tilde{y},x_{i}x_{j}]$, while the second represents
“hallucinated correlations” which are not present in the data ($W_{i}\cdot
W_{j}\propto\text{Cov}[\tilde{y},x_{i}x_{j}]$) but which the model produces
anyway (Appendix B). When the fourth moment is large (for example due to
sparsity, see Appendix D) the hallucinated term becomes less important and we
get polysemanticity/superposition.

More generally, when a model is incapable of representing the precise
correlations in the data there is often a tradeoff, where the model can
improve its representation of the true correlations at the cost of
incorporating hallucinated ones. This will often be beneficial to the overall
loss, but can end up spreading the task of representing a single feature
across multiple neurons. If the model has a limited number of neurons this can
result in polysemanticity.

###  3.3 Capacity

We would like to write the expected loss in terms of the capacity of each
feature. Recalling the definition of capacity, the expected loss can be
written as

$\displaystyle\mathcal{L}(\vec{C},\vec{n})=(k-1)\sum_{i}(n_{i}-v_{i})^{2}-2\sum_{i}n_{i}^{2}+2\sum_{i}\frac{n_{i}^{2}}{C_{i}},$
(13)

where $\mathcal{L}\equiv\mathbb{E}[L]$ is the expected loss,
$k\equiv\mathbb{E}[x_{i}^{4}]$ is the kurtosis, and $n_{i}\equiv||W_{i}||^{2}$
is the squared embedding length. This form is surprising because it says that
the loss is composed of (1) a term that depends only on the embedding lengths
and (2) a sum of terms that each depend on the capacity of just one feature.
In particular, what is surprising is that the effects of interference between
features are entirely captured by the feature capacity, so the second term can
be interpreted as the loss due to interference.

Perhaps more striking, the partial derivative of loss with respect to capacity
is just

$\frac{\partial\mathcal{L}}{\partial C_{i}}=-2\frac{n_{i}^{2}}{C_{i}^{2}}.$
(14)

That is, the marginal benefit of capacity for feature $i$ depends only on the
embedding length and capacity of that feature and not at all on any other
aspect of the model. In what follows we make extensive use of this property to
find the optimal capacity allocation.

###  3.4 Optimal capacity allocation

We now minimize the loss to find (Appendix C):

$\displaystyle C_{i}$
$\displaystyle=\max\left(\min\left(\frac{k-1}{k-3}\frac{v_{i}}{\lambda}-\frac{2}{k-3},1\right),0\right)$
(15) $\displaystyle n_{i}$ $\displaystyle=\begin{cases}0&\quad C_{i}=0\\\
\frac{k-1}{k-3}v_{i}-\frac{2\lambda}{k-3}&\quad 0<C_{i}<1\\\ v_{i}&\quad
C_{i}=1\end{cases}$ (16)

Here $\lambda$ is a Lagrange multiplier determined by the condition that
$\sum_{i}C_{i}=D$. In Appendix C we also verify that this solution is indeed
feasible (i.e. there is an embedding matrix producing the above $C_{i}$ and
$n_{i}$).

For kurtosis $k>3$, features will be ignored if

$\displaystyle\frac{k-1}{2}v_{i}<\lambda$ (17)

and fully represented if

$\displaystyle v_{i}>\lambda$ (18)

If $k\leq 3$, features will always either be ignored or fully represented
because the loss is concave-down (or flat) in $C_{i}$. Hence the $D$ features
with highest $v_{i}$ will be represented.

###  3.5 Phase diagrams

We now calculate phase diagrams for capacity allocation.

####  3.5.1 Phase diagram boundaries

First, we solve for $\lambda$ by inserting the set of $v_{i}$ into equation
(56):

$C_{i}=\max\left(\min\left(\frac{k-1}{k-3}\frac{v_{i}}{\lambda}-\frac{2}{k-3},1\right),0\right).$

We then sum over $i$ and set this equal to $D$, allowing us to solve for
$\lambda$.

Features will be polysemantic if (equations 17 and 18)

$v_{i}<\lambda<\frac{k-1}{2}v_{i}$

or equivalently

$\frac{2}{k-1}\lambda<v_{i}<\lambda$

Note that for $k<3$, this is impossible, so features will never be partially
represented (polysemantic).

####  3.5.2 Most-importances-equal case

We now consider the case with $N$ features, $D<N$ dimensions, and

$\displaystyle v_{i}$ $\displaystyle=V$ (19) $\displaystyle v_{i}$
$\displaystyle=1$ $\displaystyle i\neq 1$ (20)

When all features are partially represented we can use the constraint
$\sum_{i}C_{i}=D$ to solve for $\lambda$ and find

$\lambda=\frac{(k-1)(N-1+V)}{(k-3)D+2N}.$

Feature #1 therefore has capacity

$C_{1}=\frac{(k-3)D+2N}{(k-3)(N-1+V)}V-\frac{2}{k-3}.$

As a sanity check, this is $D/N$ when $V=1$.

Feature #1 becomes fully represented $(C_{1}=1)$ when

$V\geq\frac{(k-1)(N-1)}{(k-3)(D-1)+2(N-1)}$

When $k$ is large, this approaches $(N-1)/(D-1)$. When $V\geq(N-1)/(D-1)$,
feature #1 is fully represented regardless of $k$.

Feature #1 is ignored $(C_{1}=0)$ when

$V\leq\frac{2(N-1)}{(k-3)D+2(N-1)}$

When $k$ is large, this goes to zero like $1/k$.

For $k=3$, both bounds are at $V=1$ and the feature instantly changes from
being ignored to fully represented when $V$ crosses 1.

###  3.6 Phase Diagram Intuition

Figure 4 shows the analytic phase diagram and the corresponding empirical
results. The lower panel shows the marginal loss benefit ($-\partial
L/\partial C_{i}$) as a function of feature capacity for each labelled point
in the analytical phase diagram.

We first focus on relative feature importance. In our model an increase in
relative feature importance corresponds to shifting $\partial L/\partial
C_{i}$ up by a constant factor. Going from point D to point F, the blue
feature increases in relative importance, and the corresponding marginal loss
curve is shifted up. At point D the blue feature is too unimportant to be
represented, at point E the blue feature is similar enough in importance to be
allocated a fractional dimension with equal marginal loss benefit to the black
reference feature, and at point F the blue feature is so much more important
than the black feature that it consumes all available capacity.

We next turn to sparsity (kurtosis). As sparsity increases, the marginal loss
curves go from convex (A) to flat (C), and eventually to concave (E). Thus,
there are diminishing returns to representing a feature when it is sparse and
increasing returns when it is dense. When features are dense, the curves are
concave, so the feature with larger marginal loss benefit is purely
represented and the other feature is ignored. This explains the jump from
ignored to monosemantic at the top of the phase diagram between points A and
B. When features are sparse, we see a smoother transition from ignored to
polysemantic to monosemantic representations.

At the critical point (labelled C in the figure), the marginal loss benefit
curves are flat and equal in scale, so _any_ allocation of capacity that makes
full use of the available embedding dimensions is optimal. Thus, any given
feature could be ignored, represented monosemantically, or represented
polysemantically.

The shapes of the $\partial L/\partial C_{i}$ curves are specific to our toy
model, but it is generally true that the boundaries between phases are
determined by the marginal loss benefit of capacity. Note that this is true
for _any_ notion of capacity satisfying the properties discussed in Section
2.2.

Finally, we note that our phase diagram bears a striking resemblance to that
of a physical system with a second-order phase transition: there is a line of
discontinuous “first order” change (above the critical point), giving way to a
continuous transition at higher sparsity. The resemblance is stronger if we
think of $\partial L/\partial C_{i}$ as our “order parameter" instead of
$C_{i}$ itself: $\partial L/\partial C_{i}$ is constant along one side of the
discontinuity, and changes as a function of sparsity along the other side,
with the discontinuity going to zero at the critical point. $\partial
L/\partial C_{i}$ is continuous everywhere else. That said, there are some
notable differences from typical second-order transitions; most importantly,
$\partial L/\partial C_{i}$ is not a smooth analytic function of importance
and sparsity, but rather a piecewise-analytic function whose derivatives are
discontinuous at the edges of the superposition phase. It is also unclear
whether there is any meaningful analog of critical-point phenomena such as
critical exponents.

###  3.7 Other models and nonlinearities

We now compare the numerical phase diagram for our toy model to that of
Anthropic’s autoencoder model [EHO+22b]. The key difference between these
models is that we sum the output of the nonlinear layer before computing the
loss whereas in Anthropic’s model the loss is computed by summing the squared
differences on the output layer. This means that theirs is tasked with data
compression (i.e. recovering an encoded vector) while ours is tasked with a
nonlinear computation.

We compare these models across three nonlinearities: quadratic, ReLU, and
GeLU. In our model when we vary the nonlinearity we do so both in the ground
truth and in the activations. We further consider two cases: one with six
features embedded in three dimensions (“6 in 3”) and one with six features
embedded in five dimensions (“6 in 5”).

The top row of Figure 7 shows the results for our toy model. The same general
shape is the same in all cases, as predicted from our theoretical results: a
sharp 0-1 transition at mild sparsity (low kurtosis), and a smoother
transition with a region of superposition for more extreme sparsity (higher
kurtosis).

The superposition phase of our model also seems to be smoother, with capacity
changing steadily from 0 to 1 (as predicted by our analytic results). In
contrast, the Anthropic model has a large region where the capacity used by a
feature is “stuck” close to 1/2. This presumably corresponds to an antipodal
embedding geometry; or, in the terminology of Section 4, a separate block of
the embedding matrix containing two features in a one-dimensional subspace. If
the neuron basis is rotated appropriately, this becomes a single neuron that
represents two features, with neither the neuron nor the features overlapping
with anything else. By contrast, the superposed features in our model
unavoidably project onto many neurons regardless of how the neuron basis is
rotated.

Interestingly, some aspects of the phase diagram seem to depend on the
nonlinearity and/or model. Our model has a narrower superposed region than
Anthropic’s model. For the “6 in 3” case with significantly more features than
neurons, Anthropic’s model shows features in superposition even when their
importance is orders of magnitude larger than any other feature, which does
not happen for our model (or for Anthropic’s when there is only one more
feature than neurons).

Figure 7: Empirical phase diagrams are shown for our toy model, Anthropic’s
toy model with 6 features in 3 dimensions, and Anthropic’s toy model with 6
features in 5 dimensions. For each model we study three different
nonlinearities. Our toy model looks similar for 6 features in 3 dimensions and
6 features in 5 dimensions for all nonlinearities, so we only display plots of
5 features in 3 dimensions. The contour at $c=0.5$ has been replaced with two
contours at $c=0.47,0.53$ to give more detail for the large $c\approx 1/2$
regions, which correspond to the antipode polytope from [EHO+22b].

##  4 Geometry of efficient matrices

In this section we study the geometric of _efficient_ embedding matrices,
namely those which saturate the total capacity bound $\sum_{i}C_{i}\leq D$
with the capacity defined in equation (3). In particular, we exhaustively
enumerate the forms these matrices can take and comment on the implications
for the optimal capacity allocation in models.

###  4.1 Diagonal (monosemantic)

The simplest way to make $W$ efficient is to put $D$ of the vectors orthogonal
to each other with arbitrary lengths and set the rest to zero. This makes $W$
a diagonal matrix padded with zeros.

In this case, the capacities $C_{i}$ are either 0 or 1. The norms-squared
$n_{i}$ are zero (if $C_{i}=0$) or nonzero but arbitrary (if $C_{i}=1$).

###  4.2 Semiorthogonal (“everything bagel”)

At the other extreme, we can choose $W=\sqrt{\lambda}R$ where $R$ is
semiorthogonal, meaning that $RR^{T}=I$. (This means the rows of $W$ are
orthonormal, but the columns aren’t.)

In this case, we have $WW^{T}=\lambda I$, so

$\displaystyle(W^{T}W)^{2}$ $\displaystyle=\lambda W^{T}W$ (21)

and the capacity simplifies to

$\displaystyle C_{i}=\frac{1}{\lambda}[W^{T}W]_{ii}=\frac{n_{i}}{\lambda}$
(22)

So the capacities and norms of vectors are directly related.

We now confirm that this is efficient:

$\displaystyle\sum_{i}C_{i}$
$\displaystyle=\frac{1}{\lambda}\text{Tr}(W^{T}W)=\frac{1}{\lambda}\text{Tr}(WW^{T})=\frac{1}{\lambda}\lambda
D=D$ (23)

###### Theorem 1

Such a matrix exists for _any_ $\vec{C}$ that sums to $D$, so this gives us a
way of dividing up capacity efficiently however we like.

The proof is shown in Appendix H.

The singular values of a semiorthogonal matrix are all equal (and vice versa),
so the condition of semiorthogonality is equivalent to requiring embedding
vectors to be distributed in all directions in a balanced, isotropic
way.555The exact sense of “balanced” is a little tricky: the vectors don’t
have to be centered around the origin! Instead, think about a dataset whose
principal components all have the same variance. This gives some geometric
intuition for what the angles between vectors are doing: shorter vectors, for
features with less capacity, will be spaced closer together in order to keep
the overall arrangement balanced.

###  4.3 Mixed (tegum product)

We can combine these approaches by splitting our embedding space into
orthogonal subspaces of dimension $D_{k}$ and using a semiorthogonal matrix to
embed $N_{k}$ of our vectors into the $k$th subspace.

To do this, make $W$ block-diagonal with blocks $\lambda_{k}R_{k}$, where
$R_{k}$ is a $D_{k}\times N_{k}$ semiorthogonal matrix. Here $\sum_{k}D_{k}=D$
and $\sum_{k}N_{k}=N$. For an example with two blocks, see Figure 5.

Within each block, the dimensions of a subspace are divided up as before:

$\displaystyle n_{i}$ $\displaystyle=\lambda_{k}C_{i}\quad\text{if $i\in$
block }k$ (24) $\displaystyle\sum_{i\in\text{ block }k}C_{i}$
$\displaystyle=D_{k}$ (25)

Note that the independent scalars $\lambda_{k}$ let us scale different
subspaces arbitrarily. (This is the only thing that makes this more general
than the previous case.)

###  4.4 General

So far we have chosen subspaces aligned with the standard basis on
$\mathbb{R}^{D}$ (i.e. the neurons), but we can rotate our set of embedding
vectors however we like without changing $\vec{C}$ or $\vec{n}$. Similarly, we
have assigned dimensions of $\mathbb{R}^{N}$ (i.e. features) to subspaces
sequentially, but we can do it in any order.

In other words: a sufficient condition for efficiency is that $W=QBP$ where
$Q$ is an orthogonal matrix, $P$ is a permutation matrix, and $B$ is a
(rectangular) block-diagonal matrix whose blocks are semiorthogonal matrices
$R_{k}$. In terms of the SVD decomposition $W=QSR$ , an equivalent condition
is that rows of $R$ corresponding to different singular values $S_{ii}$ are
never nonzero in the same column.

###### Theorem 2

This is also a necessary condition; all efficient matrices are of this form.

The proof is shown in Appendix H.

###  4.5 Interpolating between regimes

Consider some block of an efficient matrix (factor of a tegum product), and
some vector $\vec{w}_{i}$ in that block. If the norm-squared $n_{i}$ of that
vector is increased, its capacity $C_{i}$ increases in order to keep $c\propto
n$ within the block. At the same time, the vector becomes more orthogonal to
the rest of the block.

If the capacity reaches $C_{i}=1$, the vector becomes fully orthogonal and
separates into a new block. Once that happens, $n_{i}$ can be varied
arbitrarily without affecting $C_{i}$.

On the other hand, if the capacity decreases to $C_{i}=0$, then $n_{i}=0$ and
the vector becomes zero; at that point the feature is ignored.

###  4.6 Application to toy models

Empirically, Anthropic found that (when importances are equal) their embedding
matrices often factor into orthogonal subspaces containing simple geometric
structures [EHO+22b]. This is equivalent to the statement that their embedding
matrices have a large number of small blocks. (This also explains some
features of the phase diagram for their model, such as a large region near
$c=1/2$; see above.) In contrast, the embedding matrices in our model have one
large superposed block (except for ignored or fully-represented features).

We can understand this phenomenon by thinking about the constraints on
$\vec{C}$ and $\vec{n}$ in the large-block and small-block regime. When there
is one large block, capacity can be divided up arbitrarily; however, the
lengths of embedding vectors are determined entirely by the capacities (up to
an overall scalar). For our model, this is compatible with good loss, because
minimizing loss gives the relationship $C_{i}\propto n_{i}$ anyway (see
above). However, for Anthropic’s model, this may be more constraining of
embedding vector norms than the loss function “wants".

In contrast, a collection of small blocks imposes constraints on capacity
allocation, because each block must separately add up to the correct capacity.
Once capacity is allocated, however, the lengths of vectors can be chosen much
more freely by scaling blocks independently of each other. This doesn’t
provide any benefit for our toy model, but seems to be useful to Anthropic’s.

This qualitative difference matters because it affects the interpretability of
the network: when there are many small blocks, it is possible to choose an
orthogonal basis of neurons such that each neuron is only polysemantic between
a small number of features. On the other hand, large blocks imply that neurons
will be polysemantic between many features at once.

##  5 Conclusions

We have studied polysemanticity through the lens of feature _capacity_ , or
the fraction of an embedding dimension allocated to each feature. Treating
capacity allocation as a constrained optimization problem, we find that many
features are either ignored (not embedded) or else allocated a full embedding
dimension orthogonal to all other features, depending on their relative
importance to the loss. Features are represented polysemantically only when
the marginal loss reduction of assigning more capacity to each is equal and
significant (Figure 2). This neatly explains the sharp “pinning” of features
to either 0 or 1 capacity noted by [EHO+22b].

To explore capacity allocation in a concrete model, we investigated a 2-layer
model with quadratic activations (Section 3) and constructed a phase diagram
of capacity allocation. We found good qualitative agreement to the phase
diagram of [EHO+22b], and confirmed our theoretical predictions regarding
optimal capacity allocation.

Finally, we studied the geometry of the embedding space (Section 4), and
characterized embedding matrices which make full use of the available
dimensions. We found that efficient embedding matrices are block-
semiorthogonal. These blocks correspond to the polytope structure [EHO+22b]
found, with small blocks corresponding to features embedded as regular
polytopes and large blocks corresponding to less-ordered structures. Large-
and small-block arrangements come with different advantages. With large blocks
there is significant freedom to allocate capacity across features, whereas
with small blocks there is the additional constraint that the capacity of each
block be an integer and that the block capacities add up to the total
capacity. On the other hand, with small blocks the lengths of embedding
vectors can be chosen more freely because blocks can be scaled independently
of each other without affecting the capacity allocation.

In our quadratic model the embedding matrices in our toy model always have one
large block, while [EHO+22b] found a range of block sizes depending on the
properties of the training distribution. This suggests that differences in
architecture can lead to different sizes of blocks, which could provide a way
to control the extent of polysemanticity in models, alongside other approaches
such as changing the activation function [EHO+22a].

## Author Contributions

AS derived the geometry of efficient matrices, proved the upper bound on the
model capacity, and proposed an alternate definition of capacity. KS performed
the numerical calculations in this work. ASJ, KS, and BS developed the
capacity allocation and constrained optimization framework. AS, KS, and BS
developed the quadratic model, and derived the expected loss, and related this
to the input kurtosis and sparsity. AS, KS, ASJ, and BS derived the analytic
optima for the quadratic model. JB proved the existence of feasible matrices
for optimal capacity allocations and proved that single-block semiorthogonal
matrices can be found for any capacity allocation. JS, KS, ASJ, and BS
provided feedback on drafts and the organization of this manuscript. AS, ASJ,
and KS contributed to writing this manuscript.

## Acknowledgments

The Flatiron Institute is supported by the Simons Foundation. We are grateful
to Jacob Steinhardt, David Lindner, Oliver Balfour, Jeff Wu, Ryan Greenblatt,
and Martin Wattenberg for helpful comments on this manuscript, as well as to
Chris Olah and Nicholas Schiefer for discussions about this work. We thank
Chris Olah for producing Figure 3.

##  Appendix A Quadratic Model Expected Loss

Our loss function is squared error

$\displaystyle L(x)=(\tilde{y}(x)-y(x))^{2}.$ (26)

We assume that $\mathbb{E}[x_{i}]=0$ and
$\text{Cov}[x_{i},x_{j}]=\delta_{ij}$, but do not impose any additional
restrictions on the distribution of $\\{x_{i}\\}$. With this, we can write the
expected loss as

$\displaystyle\mathbb{E}[L]$ $\displaystyle=\mathbb{E}[(x^{T}Dx+b)^{2}]$ (27)
$\displaystyle=\mathbb{E}\left[\left(\sum_{i,j}D_{ij}x_{i}x_{j}+b\right)^{2}\right]$
(28)
$\displaystyle=\sum_{i,j,k,l}D_{ij}D_{kl}\mathbb{E}\left[x_{i}x_{j}x_{k}x_{l}\right]+2b\sum_{i,j}D_{ij}\mathbb{E}[x_{i}x_{j}]+b^{2},$
(29)

where $D=W^{T}W-v$ and $v$ is the diagonal matrix formed of $\\{v_{i}\\}$.

Using the moment-cumulants formula we consider every partition of four terms
((1234),(123,4),(12,34),(12,3,4),(1,2,3,4))). Because the $x$’s are
independent with mean 0 and variance 1, the only partition terms that are non-
zero are 4 and 2|2, and we find

$\displaystyle\mathbb{E}[L]$
$\displaystyle=\mathbb{E}[x_{i}^{4}]\sum_{i}D_{ii}^{2}+\sum_{i\neq
j}D_{ii}D_{jj}+2\sum_{i\neq j}D_{ij}^{2}+2b\sum_{i}D_{ii}^{2}+b^{2}$ (30)
$\displaystyle=(\mathbb{E}[x_{i}^{4}]-1)\sum_{i}D_{ii}^{2}+\sum_{i,j}D_{ii}D_{jj}+2\sum_{i\neq
j}D_{ij}^{2}+2b\sum_{i}D_{ii}^{2}+b^{2}$ (31)
$\displaystyle=(\mathbb{E}[x_{i}^{4}]-1)\sum_{i}D_{ii}^{2}+\left(\sum_{i}D_{ii}\right)^{2}+2\sum_{i\neq
j}D_{ij}^{2}+2b\sum_{i}D_{ii}^{2}+b^{2}$ (32)
$\displaystyle=(\mathbb{E}[x_{i}^{4}]-1)\sum_{i}D_{ii}^{2}+\left(b+\sum_{i}D_{ii}\right)^{2}+2\sum_{i\neq
j}D_{ij}^{2}.$ (33)

Inspecting the second term we see that at optimum $b=-\sum_{i}D_{ii}$. Making
this assignment we find

$\displaystyle\mathbb{E}[L]$
$\displaystyle=(\mathbb{E}[x_{i}^{4}]-1)\sum_{i}D_{ii}^{2}+2\sum_{i\neq
j}D_{ij}^{2},$ (34)

Substituting in for $D$ gives us

$\mathbb{E}[L]=(\mathbb{E}[x_{i}^{4}]-1)\sum_{i}(||W_{i}||^{2}-v_{i})^{2}+2\sum_{i\neq
j}(W_{i}\cdot W_{j})^{2}$ (35)

##  Appendix B Covariances and Correlations

Here we derive the relationship the terms in the quadratic model loss and
true/hallucinated correlations:

$\displaystyle\text{Cov}[y,x_{i}^{2}]$
$\displaystyle=\text{Cov}\left[\sum_{j}v_{j}x_{j}^{2},x_{i}^{2}\right]$ (36)
$\displaystyle=v_{i}\text{Cov}[x_{i}^{2},x_{i}^{2}]$ (37)
$\displaystyle=v_{i}(\mathbb{E}[x_{i}^{4}]-\mathbb{E}[x_{i}^{2}]^{2})$ (38)
$\displaystyle=v_{i}(\mathbb{E}[x_{i}^{4}]-1)$ (39)
$\displaystyle\text{Cov}[\tilde{y},x_{i}^{2}]$
$\displaystyle=\text{Cov}\left[\sum_{j,k}(W^{T}W)_{jk}x_{j}x_{k},x_{i}^{2}\right]$
(40) $\displaystyle=(W^{T}W)_{ii}\text{Cov}[x_{i}^{2},x_{i}^{2}]$ (41)
$\displaystyle=||W_{i}||^{2}(\mathbb{E}[x_{i}^{4}]-1)$ (42)
$\displaystyle\text{Cov}[\tilde{y},x_{i}x_{j}]$
$\displaystyle=\text{Cov}\left[\sum_{j,k}(W^{T}W)_{jk}x_{j}x_{k},x_{i}x_{j}\right]$
$\displaystyle(i\neq j)$ (43)
$\displaystyle=2(W^{T}W)_{ij}\text{Cov}[x_{i}x_{j},x_{i}x_{j}]$ (44)
$\displaystyle=2(W_{i}\cdot W_{j})$ (45)
$\displaystyle\text{Cov}[y,x_{i}x_{j}]$ $\displaystyle=0$ $\displaystyle(i\neq
j)$ (46)

Here the true correlations in the model are described by the term
$\text{Cov}[\tilde{y},x_{i}^{2}]$, as these are what appear in the ground
truth ($\text{Cov}[y,x_{i}^{2}]$). By contrast, the hallucinated correlations
are those of the form $\text{Cov}[\tilde{y},x_{i}x_{j}],i\neq j$, which are
zero in the ground truth but are generally non-zero in the quadratic toy
model.

##  Appendix C Quadratic Model Loss Minmization

In minimizing the loss we have to consider three constraints:

  1. 1.

$1\leq\sum_{i}C_{i}\leq D$

  2. 2.

$0\leq C_{i}\leq 1$

  3. 3.

$(\vec{C},\vec{n})$ is feasible

Here feasibility means that there is actually a set of embedding directions we
can choose such that $W$ realizes the capacity allocation $\vec{C}$ and
embedding lengths $\vec{n}$. Our strategy here will be to “ask forgiveness
rather than permission": We minimize $\mathcal{L}$, ignoring the feasibility
constraint, and check for feasibility after the fact.

First we hold $\vec{C}$ fixed and optimize over $\vec{n}$:

$\displaystyle\frac{\partial\mathcal{L}}{\partial{n_{i}}}$
$\displaystyle=2(k-1)(n_{i}-v_{i})-4n_{i}+4\frac{n_{i}}{C_{i}}$ (48)
$\displaystyle 0$
$\displaystyle=2(k-1)(n_{i}^{*}-v_{i})-4n_{i}^{*}+4\frac{n_{i}^{*}}{C_{i}}$
(49) $\displaystyle n_{i}^{*}((k-3)+2/C_{i})$ $\displaystyle=(k-1)v_{i}$ (50)
$\displaystyle n_{i}^{*}$ $\displaystyle=\frac{(k-1)v_{i}}{(k-3)+2/C_{i}}$
(51)

Note that $\lim_{C_{i}\to 0}n_{i}^{*}=0$ and $\lim_{C_{i}\to
0}(n_{i}^{*}{}^{2}/C_{i})=0$, so the final two terms of the loss are zero for
non-represented features.

Next we differentiate with respect to capacity to find

$\displaystyle\frac{\partial\mathcal{L}(\vec{C},\vec{n}^{*}(\vec{C}))}{\partial
C_{i}}=\frac{\partial\mathcal{L}}{\partial
C_{i}}+\sum_{j}\frac{\partial\mathcal{L}}{\partial n_{j}}\frac{\partial
n_{j}^{*}}{\partial C_{i}}$ (52)

The second term vanishes666This is because we optimized over $\\{n_{i}\\}$
without any constraints, so $\vec{n}^{*}$ is a local optimum for any
$\vec{C}$., so we only have the first term

$\displaystyle\frac{\partial\mathcal{L}}{\partial
C_{i}}=-2\frac{(n_{i}^{*}(C_{i}))^{2}}{C_{i}^{2}}$ (53)

This is non-positive, so our loss is monotonically decreasing in $C_{i}$. That
means that we will saturate our first constraint, $\sum_{i}C_{i}=D$.

Next let’s solve for $C_{i}$ and $n_{i}$:

$\displaystyle\lambda C_{i}$ $\displaystyle=\frac{(k-1)v_{i}}{(k-3)+2/C_{i}}$
(54) $\displaystyle\frac{v_{i}}{\lambda}$
$\displaystyle=C_{i}\frac{k-3}{k-1}+\frac{2}{k-1}$ (55) $\displaystyle C_{i}$
$\displaystyle=\max\left(\min\left(\frac{k-1}{k-3}\frac{v_{i}}{\lambda}-\frac{2}{k-3},1\right),0\right)$
(56)

where we have introduced the clipping at 0 and 1 to make this also apply to
the capacities on the boundary. With $C_{i}$ we then obtain

$\displaystyle n_{i}=\begin{cases}0&\quad C_{i}=0\\\
\frac{k-1}{k-3}v_{i}-\frac{2\lambda}{k-3}&\quad 0<C_{i}<1\\\ v_{i}&\quad
C_{i}=1\end{cases}$ (57)

####  C.0.1 Checking feasibility

For our solution $(\vec{C},\vec{n})$ to be feasible, there must exist a
corresponding matrix $W$.

As shown in Appendix H, for any $\vec{C}$ with $\sum_{i}C_{i}=D\leq N$ and
$0\leq C_{i}\leq 1$, there is a semiorthogonal matrix $W$ of shape $(N,D)$
with $C_{i}(W)=C_{i}$.

We next need to show that $\vec{n}$ is compatible with $\vec{C}$. The
condition $n_{i}\propto C_{i}$ holds for any semiorthogonal matrix, and we can
multiply each $W_{i}$ by a scalar to get the constant of proportionality
right.

Therefore, this solution is feasible!

##  Appendix D Kurtosis and Sparsity

###  D.1 Sparsity increases fourth moment

Suppose that, with probability $p$, $x$ takes the value of a new random
variable $z$ and otherwise it is 0 (i.e. $x=sz$, $s\sim B(p)$). We set the
mean $\mathbb{E}[z]=0$ and variance $\text{Var}[z]=1/p$ so that
$\mathbb{E}[x]=0,\text{Var}[x]=1$. As $p$ decreases (i.e. sparsity increases),
$\mathbb{E}[x^{4}]$ increases:

$\displaystyle\mathbb{E}[x^{4}]=\mathbb{E}[s^{4}z^{4}]=\mathbb{E}[s^{4}]\mathbb{E}[z^{4}]=p\cdot\text{Kurt}[z]\text{Var}[z]^{2}=\text{Kurt}[z]/p$
(58)

###  D.2 Uniform Random Variable Moments

We want to choose $a$ such that for $x\sim U(-a,a)\cdot B(p)$,
$\text{Var}[x]=1$. Choose $a=\sqrt{\frac{3}{p}}$:

$\displaystyle\text{Var}[x]$
$\displaystyle=\mathbb{E}[s^{2}]\mathbb{E}[z^{2}]=p\cdot\frac{\left(2\sqrt{3/p}\right)^{2}}{12}=1$
(59) $\displaystyle\text{Kurt}[x]$
$\displaystyle=\frac{\mathbb{E}[x^{4}]}{\mathbb{E}[x^{2}]^{2}}=\frac{9}{5}\cdot\frac{1}{p}$
(60)

##  Appendix E Interpretation as correlation coefficient

The (population) Pearson correlation coefficient $\rho^{2}$, also called
“fraction of variance explained”, is defined as the ratio between covariance
squared and the product of variances:

$\displaystyle\rho^{2}(X,Y)=\frac{(k-1)(X,Y)^{2}}{\kappa(X,X)\kappa(Y,Y)}$
(61)

(The more-familiar $r^{2}$ is an estimator for $\rho^{2}$ computed from a
sample.)

Let’s assume our input features $X_{i}$ are linearly uncorrelated and have
mean zero and variance one:

$\displaystyle(k-1)(X_{i})$ $\displaystyle=0$ (62)
$\displaystyle(k-1)(X_{i},X_{j})$ $\displaystyle=\delta_{ij}$ (63)

We’ll embed them with embedding vectors $w_{i}=W\hat{e}_{i}=W_{:,i}$ and
immediately unembed with $W^{T}$. Then $\rho^{2}$ between the $i$th input and
output is (using Einstein notation for everything except $i$, which is never
summed)

$\displaystyle\rho^{2}(X_{i},[W^{T}WX]_{i})$
$\displaystyle=\frac{(k-1)(X_{i},W*{ik}W_{kj}X_{j})^{2}}{(k-1)(X_{i},X_{i})(k-1)(W_{ik}W_{kj}X_{j},W_{ik^{\prime}}W_{k^{\prime}j^{\prime}}X_{j^{\prime}})}$
(64)
$\displaystyle=\frac{(W_{ik}W_{kj}(k-1)(X_{i},X_{j}))^{2}}{W_{ik}W_{kj}W_{ik^{\prime}}W_{k^{\prime}j^{\prime}}(k-1)(X_{i},X_{i})(k-1)(X_{j},X_{j^{\prime}})}$
(65)
$\displaystyle=\frac{(W_{ik}W_{kj}\delta_{ij})^{2}}{W_{ik}W_{kj}W_{ik^{\prime}}W_{k^{\prime}j^{\prime}}\delta_{jj^{\prime}}}$
(66)
$\displaystyle=\frac{(W_{ik}W_{ki})^{2}}{(W_{ik}W_{kj})(W_{ik^{\prime}}W_{k^{\prime}j})}$
(67) $\displaystyle=\frac{(w_{i}\cdot w_{i})^{2}}{\sum_{j}(w_{i}\cdot
w_{j})^{2}}$ (68) $\displaystyle=C_{i}$ (69)

where we’ve written the sum over $j$ explicitly for the sake of clarity.

Note that $W^{T}$ is not always the optimal way to reconstruct $X_{i}$ from
$WX$. We haven’t checked the math, but it seems to be the case that using the
optimal unembedding matrix instead of $W^{T}$ gives the alternate definition
of capacity in Appendix G.

##  Appendix F Proof of capacity constraint

Suppose we have $N$ features and $D$ neurons, with $N\geq D$.

Let our embedding matrix be $W_{ai}$, with shape $[D,N]$. The embedding
vectors in $\mathbb{R}^{D}$ are column vectors $\vec{w}_{i}:=W_{:,i}$ with
components $[w_{i}]_{a}=W_{ai}$.

Define

$\displaystyle V_{ai}$
$\displaystyle:=\frac{W_{ai}}{\sqrt{[W^{T}WW^{T}W]_{ii}}}$ (70)
$\displaystyle\vec{v}_{i}:=V_{:,i}$
$\displaystyle=\frac{\vec{w}_{i}}{\sqrt{\sum_{j}(\vec{w}_{i}\cdot\vec{w}_{j})^{2}}}$
(71)

and

$\displaystyle U$ $\displaystyle:=V^{T}W$ (72) $\displaystyle U_{ij}$
$\displaystyle=\vec{v}_{i}\cdot\vec{w}_{j}$ (73)

Note that

$\displaystyle U_{ii}^{2}$
$\displaystyle=\frac{(\vec{w}_{i}\cdot\vec{w}_{i})^{2}}{\sum_{j}(\vec{w}_{i}\cdot\vec{w}_{j})^{2}}=C_{i}$
(74)

Also,

$\displaystyle\sum_{j=1}^{N}U_{ij}^{2}=1$ (75)

so the row vectors $\hat{u}_{i}:=U_{i,:}\in\mathbb{R}^{N}$ are unit vectors.

Let $D^{\prime}:=\text{rank}(U)\leq D$, so that the vectors $\hat{u}_{i}$ span
a subspace of dimension $D^{\prime}$.

Let the $D^{\prime}$ unit vectors $\hat{y}_{b}\in\mathbb{R}^{N}$ with
components $[y_{b}]_{i}=:Y_{bi}$ be an orthonormal basis for this subspace.

Let $X_{ib}$ be the components of $\hat{u}_{i}$ in the $\hat{y}_{b}$ basis, so
that

$\displaystyle U$ $\displaystyle=XY$ (76) $\displaystyle\hat{u}_{i}$
$\displaystyle=\sum_{b=1}^{D^{\prime}}X_{ib}\hat{y}_{b}$ (77)

$\hat{u}_{i}$ is a unit vector, so its components in an orthonormal basis obey

$\displaystyle\sum_{b=1}^{D^{\prime}}X_{ib}^{2}$ $\displaystyle=1$ (78)

Therefore, the row vectors $\hat{x}_{i}:=X_{i,:}\in\mathbb{R}^{D^{\prime}}$
are unit vectors.

We can also define vectors $\vec{y}_{i}:=Y_{:,i}$; note that these are the
columns of $Y$ and $\hat{y}_{b}$ are the rows.

Note that $U_{ij}=\hat{x}_{i}\cdot\vec{y}_{i}$, so we have

$\displaystyle C_{i}=U_{ii}^{2}$
$\displaystyle=(\hat{x}_{i}\cdot\vec{y}_{i})^{2}$ (79)
$\displaystyle\leq(\hat{x}_{i}\cdot\hat{x}_{i})(\vec{y}_{i}\cdot\vec{y}_{i})\text{\quad
by Cauchy-Schwarz}$ (80) $\displaystyle=(\vec{y}_{i}\cdot\vec{y}_{i})$ (81)
$\displaystyle=\sum_{b=1}^{D^{\prime}}Y_{bi}^{2}$ (82)

Finally, summing over $i$,

$\displaystyle\sum_{i=1}^{N}C_{i}$
$\displaystyle\leq\sum_{i=1}^{N}\sum_{b=1}^{D^{\prime}}Y_{bi}^{2}$ (83)
$\displaystyle=\sum_{b=1}^{D^{\prime}}\sum_{i=1}^{N}Y_{bi}^{2}$ (84)
$\displaystyle=\sum_{b=1}^{D^{\prime}}|\hat{y}_{b}|^{2}$ (85)
$\displaystyle=D^{\prime}\leq D$ (86)

##  Appendix G An alternative definition of capacity

Instead of applying equation (2) directly to $W$, we can take the (compact)
SVD $W=QSR$ and apply it to $R$,

$\displaystyle\tilde{C}_{i}(W):=\frac{([R^{T}R]_{ii})^{2}}{[(R^{T}R)^{2}]_{ii}}=[R^{T}R]_{ii}$
(87)

An equivalent prescription is to first rotate (by $Q^{-1}$) and then scale (by
$S^{-1}$) the embedding space $\mathbb{R}^{D}$ so that the embedding vectors
are isotropically distributed, then take $\tilde{C}_{i}$ to be the norm-
squared of the new embedding vector $S^{-1}Q^{-1}\vec{w}_{i}$.

This is equivalent to our previous definition when $W$ is “$c$-efficient”,

$\displaystyle\sum_{i=1}^{N}C_{i}=D\implies\tilde{C}_{i}=C_{i}$ (88)

but relaxes the condition for “$\tilde{c}$-efficiency” dramatically, to the
point that generic matrices are efficient:

$\displaystyle\sum_{i=1}^{N}\tilde{C}_{i}$ $\displaystyle=\text{rank}(W)$ (89)
$\displaystyle W\text{ full rank}$
$\displaystyle\implies\sum_{i=1}^{N}\tilde{C}_{i}=D$ (90)

On the other hand, $\tilde{C}_{i}(W)$ is a somewhat more discontinuous
function of $W$ than $C_{i}(W)$ is.

Our results for the quadratic toy model are essentially unchanged; $W$ is
efficient at the minimum of $\mathbb{\mathbb{missing}}{E}[\mathcal{L}]$, so
these definitions coincide (although we’re not sure their derivatives do). In
more realistic models, it remains to be seen which version is more useful.

This new definition also sheds some light on “$c$-efficiency”: a matrix is
efficient if its embedding vectors are isotropically distributed within each
block. In the case of only one block, a matrix is efficient if identity
covariance for a random $\vec{x}$ implies identity covariance for $W\vec{x}$.

##  Appendix H Feasibility Proofs

##### Set-up:

We have vectors $w_{1},\dots,w_{N}\in\mathbb{R}^{D}$ spanning
$\mathbb{R}^{D}$, where $D\leq N$, arranged into an $D\times N$ matrix

$W=\begin{pmatrix}|&&|\\\ w_{1}&\dots&w_{N}\\\ |&&|\end{pmatrix}.$

Our aim is to show that

$\sum_{i=1}^{N}C_{i}\leq D$

and to determine the conditions under which equality holds.

##### Solution:

We begin by finding a maximal decomposition of $\mathbb{R}^{D}$ into
orthogonal subspaces $V_{1}\oplus\dots\oplus V_{r}$ such that each $w_{i}$
lies in some $V_{s}$ and partition the $w_{i}$ accordingly, discarding any
index with $w_{i}=0$. It then suffices to examine the problem separately on
each subspace $V_{s}$.

Therefore, we assume that $w_{1},\dots,w_{N}$ are such that we cannot
decompose our problem further. In particular, this implies that $w_{i}\neq 0$
for all $i$ and we cannot partition the indices into two subsets $S,T$ such
that $w_{i}\cdot w_{j}=0$ for all $i\in S,j\in T$.

###### Lemma 1

Let $\\{e_{1},\dots,e_{N}\\}$ be the standard basis of $\mathbb{R}^{N}$ and
let $V\subseteq\mathbb{R}^{N}$ be a linear subspace of dimension $D$, such
that $e_{i}\not\in V^{\perp}$ for all $i$. Then for any arbitrary unit vectors
$v_{1},\dots,v_{N}\in V$, we have

$\sum_{i=1}^{N}(v_{i}\cdot e_{i})^{2}\leq D$

with equality if and only if $v_{i}$ is parallel to
$\operatorname{proj}(e_{i},V)$, where $\operatorname{proj}(u,W)$ denotes the
orthogonal projection of a vector $u$ onto the subspace $W$.

_Proof:_ Define $v_{i}^{\ast}\in V$ to be a unit vector in the direction
$\operatorname{proj}(e_{i},V)$, and note that this exists by the assumption
$e_{i}\not\in V^{\perp}$. Then we have $|v_{i}\cdot
e_{i}|\leq|v^{\ast}_{i}\cdot e_{i}|$ for all other unit vectors $v\in V$, with
equality if and only if $v_{i}$ is parallel to $v^{\ast}_{i}$. So it suffices
to show that

$\sum_{i=1}^{N}(v_{i}^{\ast}\cdot e_{i})^{2}=D.$

Let $b_{1},\dots,b_{D}$ be an orthonormal basis for $V$, and let
$b_{D+1},\dots,b_{N}$ be an orthonormal basis for $V^{\perp}$. Then

$v_{i}^{\ast}=\lambda_{i}(b_{1}\cdot e_{i})b_{1}+\dots+(b_{D}\cdot
e_{i})b_{D}$

where the normalising constant $\lambda_{i}=(\sum_{j=1}^{D}(b_{j}\cdot
e_{i})^{2})^{-1/2}$ is non-zero as $v^{\ast}_{i}$ is non-zero (and we may wlog
assume $\lambda_{i}>0$). Then

$(v_{i}^{\ast}\cdot e_{i})^{2}=\lambda_{i}^{2}\left(\sum_{j=1}^{D}(b_{j}\cdot
e_{i})^{2}\right)^{2}=\sum_{j=1}^{D}(b_{j}\cdot e_{i})^{2}$

from which the lemma follows.

Define the matrix $A=(a_{ij})_{i,j=1,\dots,N}$ via $a_{ij}=w_{i}\cdot w_{j}$,
and denote by $a_{i}$ the vector $(a_{i1},\dots,a_{in})^{T}\in\mathbb{R}^{N}$,
so that

$A=\begin{pmatrix}|&&|\\\ a_{1}&\dots&a_{N}\\\ |&&|\end{pmatrix}$

and the capacities can be expressed as $C_{i}=a_{ii}^{2}/||a_{i}||^{2}$. Note
that $a_{ii}=||w_{i}||^{2}\neq 0$ and so non of the $a_{i}$ are trivial. Also,
note that $A=W^{T}W$, so the rank of $A$ is the same as the rank of $W$, and
thus equal to $D$.

Now let $v_{i}=a_{i}/||a_{i}||$, so that $v_{i}$ is a unit vector. Since the
rank of $A$ is $D$, the $a_{i}$ span a space $V\subseteq\mathbb{R}^{N}$ of
dimension $D$. We can write

$v_{i}\cdot e_{i}=a_{ii}/||a_{i}||,$

from which it also follows that $e_{i}\not\in V^{\perp}$ (otherwise we would
have $a_{ii}=0$ which is not possible). Therefore we may apply Lemma 1 to get

$\sum_{i=1}^{N}C_{i}\leq D$

with equality if and only if $a_{i}$ is parallel to
$\operatorname{proj}(e_{i},V)$.

Assuming that equality holds, using the proof of Lemma 1, we can write

$a_{i}=\mu_{i}\sum_{k=1}^{D}(b_{k}\cdot e_{i})b_{k},\hskip
14.22636pta_{ij}=\mu_{i}\sum_{k=1}^{D}(b_{k}\cdot e_{i})(b_{k}\cdot e_{j})$

for some constant $\mu_{i}$. Since $a_{ii}=a_{i}\cdot e_{i}>0$, we see that
$\mu_{i}>0$.

Now, if $\mu_{i}\neq\mu_{j}$ then

$\mu_{j}a_{ij}=\mu_{j}\mu_{i}\sum_{k=1}^{D}(b_{k}\cdot e_{i})(b_{k}\cdot
e_{j})=\mu_{i}a_{ji}$

from which we conclude by the symmetry of $a_{ij}$ that $a_{ij}=w_{i}\cdot
w_{j}=0$. Hence if we have any $\mu_{i}\neq\mu_{j}$ we can perform a non-
trivial decomposition into orthogonal subspaces $V_{1}\oplus\dots\oplus V_{r}$
satisfying the conditions at the start of the solution. Since we assumed this
was not possible, we must in fact have $\mu:=\mu_{i}$ is constant for all $i$.

Consider the map

$\Lambda:w_{i}\mapsto\mu^{1/2}\sum_{k=1}^{D}(b_{k}\cdot e_{i})b_{k}.$

If we let $T\subseteq\mathbb{R}^{N}$ be the linear subspace spanned by
$w_{1},\dots,w_{N}$, then $\dim T=m$ and $\Lambda$ maps $T$ to $V$. We can
also check that the orthogonality of the bases $\\{e_{i}\\}_{i}$ and
$\\{b_{i}\\}_{i}$ implies that $\Lambda$ preserves the inner products of the
$w_{i}$. It follows that $\Lambda$ is invertible and orthogonal.

In addition, we have that

$b_{j}^{T}\Lambda We_{i}=b_{j}^{T}\Lambda
w_{i}=\mu^{1/2}\sum_{k=1}^{D}(b_{k}\cdot e_{i})(b_{k}\cdot
b_{j})=\mu^{1/2}(b_{j}\cdot e_{i})$

so if we consider $\Lambda W$ as a map from $\mathbb{R}^{D}$ with basis
$\\{e_{i}\\}_{i}$ to $V$ with basis $\\{b_{i}\\}_{i}$ then it has matrix

$\Lambda W=\mu^{1/2}\begin{pmatrix}b_{1}\cdot e_{1}&\dots&b_{1}\cdot e_{N}\\\
\vdots&&\vdots\\\ b_{D}\cdot e_{1}&\dots&b_{D}\cdot e_{N}\end{pmatrix}.$

Using $\sum_{j=1}^{N}(b_{i}\cdot e_{j})^{2}=1$, it follows that $\Lambda
WW^{T}\Lambda^{T}=\mu I_{D}$. Since $\Lambda$ is orthogonal, we deduce that
$WW^{T}=\mu I_{D}$ as required.

##### Showing that every saturating case is achievable:

We want to show that for every tuple of real numbers $(C_{1},\dots,C_{N})$
with $0\leq C_{i}\leq 1$ and $\sum_{i=1}^{N}C_{i}=D$, we can find
$w_{1},\dots,w_{N}\in\mathbb{R}^{D}$ such that

$C_{i}=\frac{(w_{i}\cdot w_{i})^{2}}{\sum_{j=1}^{N}(w_{i}\cdot w_{j})^{2}}.$

A tuple $(C_{1},\dots,C_{N})$ is feasible if there exists an $N\times N$
unitary matrix $U$ such that if we let $\hat{U}$ denote the first $D$ rows of
$U$ and define $u_{1},\dots,u_{N}\in\mathbb{R}^{D}$ by

$\hat{U}=\begin{pmatrix}|&&|\\\ u_{1}&\dots&u_{N}\\\ |&&|\end{pmatrix},$

then $||u_{i}||^{2}=C_{i}$ for each $i=1,\dots,n$. Note that if
$(C_{1},\dots,C_{N})$ is feasible then for the corresponding $u_{i}$, if we
denote their components by $\\{u_{ki}\\}_{k=1}^{D}$, we have

$\displaystyle\frac{(u_{i}\cdot u_{i})^{2}}{\sum_{j=1}^{N}(u_{i}\cdot
u_{j})^{2}}$ $\displaystyle=\frac{(u_{i}\cdot
u_{i})^{2}}{\sum_{j=1}^{N}\sum_{k,l=1}^{D}u_{ik}u_{jk}u_{il}u_{jl}}$
$\displaystyle=\frac{(u_{i}\cdot
u_{i})^{2}}{\sum_{k,l=1}^{D}u_{ki}u_{li}\delta_{kl}}$
$\displaystyle=\frac{(u_{i}\cdot
u_{i})^{2}}{\sum_{k=1}^{D}u_{ki}u_{ki}}=||u_{i}||^{2}=C_{i}$

so we see that $(C_{1},\dots,C_{N})$ is a valid set of capacities by taking
$w_{i}=u_{i}$ for each $i$. Hence it suffices to show that all the relevant
tuples are feasible.

###### Lemma 2

Suppose that $(C_{1},\dots,C_{N})$ is a feasible tuple. Then

  1. 1.

any permutation of $(C_{1},\dots,C_{N})$ is feasible;

  2. 2.

for any $C^{\prime}_{N},C^{\prime}_{N+1}\geq 0$ such that
$C^{\prime}_{N}+C^{\prime}_{N+1}=C_{N}$, the tuple
$(C_{1},\dots,C_{n-1},C^{\prime}_{N},C^{\prime}_{N+1})$ is feasible;

  3. 3.

the tuple $(1-C_{1},\dots,1-C_{N})$ is feasible, with $D$ replaced by $N-D$.

_Proof:_ Let $U$ be an $N\times N$ unitary matrix corresponding to
$(C_{1},\dots,C_{N})$ as above. Then (i) is immediate by permuting the columns
of $U$, and (iii) follows by exchanging the first $D$ and last $N-D$ rows of
$U$, noting that the squared norm of each column is 1.

To prove (ii), consider extending $U$ to an $(N+1)\times(N+1)$ unitary matrix
$U^{\prime}$ by placing $U$ in the upper $n\times n$ quadrant and placing a 1
at the bottom right position. Then write

$U^{\prime}=\begin{pmatrix}\begin{matrix}&&\\\ &U&\\\
&&\end{matrix}&\vline&0\\\ \hline\cr
0&\vline&\begin{matrix}1\end{matrix}\end{pmatrix}=\begin{pmatrix}|&&|&|\\\
u^{\prime}_{1}&\dots&u^{\prime}_{N}&u^{\prime}_{N+1}\\\ |&&|&|\end{pmatrix},$

for vectors $u^{\prime}_{i}\in\mathbb{R}^{N+1}$. Now we may rotate the final
pair of vectors together by an arbitrary angle $\theta$ and the resulting
matrix remains orthogonal:

$U^{\prime}_{\theta}=\begin{pmatrix}&&&\\\
u^{\prime}_{1}&\dots&u^{\prime}_{N}\cos(\theta)+u^{\prime}_{N+1}\sin(\theta)&-u^{\prime}_{N}\sin(\theta)+u^{\prime}_{N+1}\cos(\theta)\\\
&&&\end{pmatrix},$

Restricting to the first $D$ rows of $U^{\prime}_{\theta}$, we see that the
first $n-1$ columns still have squared norms $C_{1},\dots,C_{n-1}$
respectively (inherited from $U$). Therefore, as $\theta$ varies the sums of
the squared norms of the final two columns must be $C_{N}$. Moreover, the
split between the two columns varies continuously between $(C_{N},0)$ and
$(0,C_{N})$ as $\theta$ ranges between $0$ and $\pi/2$. Therefore we can
achieve any split $(C^{\prime}_{N},c^{\prime}_{N+1})$ between the final two
columns, proving part (2).

Finally, we show that every tuple $(C_{1},\dots,C_{N})$ satisfying $0\leq
C_{i}\leq 1$ and $\sum_{i=1}^{N}C_{i}=m$ is feasible. Suppose not, for a
contradiction. Then we may pick such a tuple with minimal $n$. Clearly all
such tuples with $n\leq 1$ are feasible, so we may assume $n\geq 2$.

If there are any $C_{i},C_{j}$ with $C_{i}+C_{j}\leq 1$ then we may replace
$C_{i},C_{j}$ by $C_{i}+C_{j}$ to get another tuple satisfying the given
conditions but with smaller $n$. By our minimality assumption, we conclude
this new tuple must be feasible. But then our original tuple is feasible by
Lemma 2.1 and 2.2.

On the other hand, if there are no $C_{i},C_{j}$ with $C_{i}+C_{j}\leq 1$ then
consider the tuple $(1-C_{1},\dots,1-C_{N})$. We know that we must have
$(1-C_{i})+(1-C_{j})\leq 1$ for any $i,j$. Replacing $(1-C_{i}),(1-C_{j})$ by
$(1-C_{i})+(1-C_{j})$ to get a new tuple and following the argument of the
previous paragraph, we conclude that the new tuple must be feasible, and so
$(1-C_{1},\dots,1-C_{N})$ must be feasible also. But then
$(C_{1},\dots,C_{N})$ is feasible too by Lemma 2.3. We are thus done in all
cases.

## References

  * [EHO+22a] Nelson Elhage, Tristan Hume, Catherine Olsson, Neel Nanda, Tom Henighan, Scott Johnston, Sheer ElShowk, Nicholas Joseph, Nova DasSarma, Ben Mann, Danny Hernandez, Amanda Askell, Kamal Ndousse, Andy Jones, Dawn Drain, Anna Chen, Yuntao Bai, Deep Ganguli, Liane Lovitt, Zac Hatfield-Dodds, Jackson Kernion, Tom Conerly, Shauna Kravec, Stanislav Fort, Saurav Kadavath, Josh Jacobson, Eli Tran-Johnson, Jared Kaplan, Jack Clark, Tom Brown, Sam McCandlish, Dario Amodei, and Christopher Olah.  Softmax linear units.  Transformer Circuits Thread, 2022.  https://transformer-circuits.pub/2022/solu/index.html.
  * [EHO+22b] Nelson Elhage, Tristan Hume, Catherine Olsson, Nicholas Schiefer, Tom Henighan, Shauna Kravec, Zac Hatfield-Dodds, Robert Lasenby, Dawn Drain, Carol Chen, Roger Grosse, Sam McCandlish, Jared Kaplan, Dario Amodei, Martin Wattenberg, and Christopher Olah.  Toy models of superposition.  Transformer Circuits Thread, 2022.  https://transformer-circuits.pub/2022/toy_model/index.html.
  * [OCS+20] Chris Olah, Nick Cammarata, Ludwig Schubert, Gabriel Goh, Michael Petrov, and Shan Carter.  Zoom in: An introduction to circuits.  Distill, 2020.  doi:10.23915/distill.00024.001.  https://distill.pub/2020/circuits/zoom-in.
  * [OMS17] Chris Olah, Alexander Mordvintsev, and Ludwig Schubert.  Feature visualization.  Distill, 2017.  doi:10.23915/distill.00007.  https://distill.pub/2017/feature-visualization.