I wrote this doc in December 2021, while working at Redwood Research. It summarizes a handful of observations about GPT-2-small'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:

$y_{i} = \sum_{j} ({Id}_{i j} - 1_{i} 1_{j}) x_{j}$ (subtract mean; 1 is an all-ones vector)
$z_{i} = y_{i} / \sqrt{\sum_{j} y_{j} y_{j}}$ (normalize variance)
$LayerNorm (x)_{i} = w_{i} z_{i} - b_{i}$ (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.)...

Martin Milan

Martin Milan