Welcome to The Nonlinear Library, where we use Text-to-Speech software to convert the best writing from the Rationalist and EA communities into audio. This is: Toy Models and Tegum Products, published by Adam Jermyn on November 4, 2022 on The AI Alignment Forum.
(Thanks to Adam Scherlis, Kshitij Sachan, Buck Shlegeris, Chris Olah, and Nicholas Schiefer for conversations that informed this post.)
Anthropic recently published a paper on toy models of superposition [Elhage+22]. One of the most striking results is that, when features are sparse, feature embeddings can become divided into disjoint subspaces with just a few vectors per subspace. This type of decomposition is known as a tegum product.
This post aims to give some intuition for why tegum products are a natural feature of the minima of certain loss functions.
Setup
Task
Suppose we’ve got d embedding dimensions and n>d features. We want to embed features into dimensions in a way that minimizes overlap between their embedding vectors. In the simplest case this could be because we’re building an autoencoder and want to compress the features into a low-dimensional space.
One approach is to encode all of the features in all of the dimensions. With this approach there is some interference between every pair of features (i.e. no pair is embedded in a fully orthogonal way), but we have a lot of degrees of freedom that we can use to minimize this interference.
Another approach is to split the d dimensions into k orthogonal subspaces of d/k dimensions. This has the advantage of making most pairs of vectors exactly orthogonal, but at the cost that some vectors are packed more closely together. In the limit where k=1 this reduces to the first approach.
Our aim is to figure out the k that minimizes the loss on this task.
Loss
Suppose our loss has the following properties:
L=∑i≠jℓ(cosθij). That is, the loss decomposes into a sum of terms involving the cosine similarities of feature vectors, and all features are equally important.
ℓ(0)=0. The loss vanishes for orthogonal vectors.
dℓ/dcosθ>0. The loss is greater the more the vectors overlap.
Using these properties, we find that the loss is roughly
where ϵ is the typical cosine similarity between vectors in a subspace.
Loss-Minimizing Subspaces
The Johnson-Lindenstrauss lemma says that we can pack m nearly-orthogonal vectors into D dimensions, with mutual angles satisfying
where
and ϵ0 is a constant. Setting m=n/k and D=d/k gives
Assuming we pick our vectors optimally to saturate the Johnson-Lindenstrauss bound, we can substitute this for ϵ in the loss and differentiate with respect to k to find
There are three possible cases: either the minimum occurs at k=d (the greatest value it can take), or at k=1 (the smallest value it can take) or at some point in between where dL/dk vanishes.
The derivative vanishes if
which gives
where
When α≥2 there is no place where the derivative vanishes, and the optimum is k=1. Otherwise there is an optimum at
so long as this is less than d. If it reaches d, the optimum sticks to k=d.
Interpretation
We can think of α as the sensitivity of the loss to interference. Specifically, which moment of the interference distribution do we care about?
When α is large, we care more about decreasing higher moments, and in the limit of infinite α what matters is just the maximum interference between vectors. Hence when α is large we want to have fewer subspaces, each with more vectors but smaller cosine similarities.
By contrast, when α is small, we care more about decreasing smaller moments, and in the limit as α0 what matters is the fraction of vectors that interfere at all. Hence when α is small we want to have more subspaces, each with fewer vectors but larger cosine similarities.
So tegum products are preferred when we can tolerate larger “peak” interference and want fewer instances of interference, whereas a single large subspace is preferred when we can tolerate lots of instances of interference and want to mi...
