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: An exploration of GPT-2's embedding weights, published by Adam Scherlis on December 13, 2022 on The AI Alignment Forum.
I wrote this doc in December 2021, while working at Redwood Research. It summarizes a handful of observations about GPT-2's weights -- mostly the embedding matrix, but also the LayerNorm gain parameters -- that I found while doing some open-ended investigation of the model. I wanted to see how much I could learn by studying just those parameters, without looking at the attention layers, MLP layers, or activations.
This is still mostly unedited. Feedback and questions are very welcome.
Embedding matrix glossary
Latent space – 768-channel-dim space that activations between blocks (and in skip connections) live in.
Token space – 50,257-token-dim space of one-hot token vectors
[Token] embedding matrix – W of shape [768, 50257] (this is the transpose of `model.lm_head.weight`)
Embedding vector – Any 768-dimensional column of W
Embedding space – Latent space
Preferred channel basis from LayerNorm
This section involves more math than the rest; you can skip it without missing much. "Preferred basis" means the same thing as "privileged basis" as defined by Anthropic here.
Most of GPT-2’s architecture – the attention layers, fully-connected layers and residual connections – don’t impose a preferred basis on the latent space.
(The skip connections give a canonical map between the latent spaces in different layers, so that it makes sense to talk about “the” latent space globally.)
The LayerNorm layers actually do privilege a basis, though, for two reasons.
LayerNorm(x) is defined as follows, where w and b are learned parameters:
yi=∑j(Idij−1i1j)xj(subtract mean; 1 is an all-ones vector)
zi=yi/√∑jyjyj (normalize variance)
LayerNorm(x)i=wizi−bi (apply bias and gain parameters)
The first line picks out part of a preferred basis – a preferred vector (and the subspace orthogonal to it), namely the all-ones vector 1.
The second line doesn’t pick out a basis, because the stdev of the components of y is just its L2 norm and normalizing a vector is independent of (orthonormal) basis. (It does impose a preferred inner product, but we had that already: GPT-2’s unembedding matrix is the transpose of the embedding matrix, and W^T W is a matrix of inner products.)
The third line is the interesting one. The gain vector w is multiplied elementwise with z, which breaks basis-independence. (You could also write this as D_ij z_j where D is a diagonal matrix, D_ii = w_i.) This makes the standard basis a preferred basis, in a pretty straightforward way: the network can dial up or down the size of different standard-basis components independently (and can’t do this for the components with respect to any other basis).
This provides some justification for breaking things out by latent-space dimension in most of this post.
Token/Position embedding rivalry
Here’s an interesting phenomenon I don’t have a complete explanation for.
This is a chunk of the token embedding matrix:
(X: token dims, Y: channel dims)
The horizontal stripes are dimensions of the latent space where the token embeddings are systematically larger or smaller than average. This gets normalized away inside residual blocks, but is preserved by skip connections.
We can also make histograms from this matrix, by taking either a row or a column and plotting a histogram of the numbers in it. We can overlay a few of these histograms, in different colors, to show a few columns or a few rows at the same time.
Overlaid histograms of the first 30 columns (tokens, left) and the first 30 rows (channel dims, right):
Note that many of the channel dims follow roughly the same distribution, while a few have smaller variance and larger means.
This shows up as a heavy left-tail of the distribution of stdevs, and heavy tails in both direct...
