The geometry
of the embedding

1 The discovery that motivated the project

In the previous course, we tried semantic search with RoBERTa's embeddings. It failed in a specific way: every similarity came out above +0.99. Sentences about completely different subjects looked identical.

At the time we diagnosed it as anisotropy and moved on. Here the question became a different one: is this a defect of RoBERTa, or does it happen with any model? We measured it on ours.

Why training does this

The cause is in the gradient of the cross-entropy. For each position, one token is the target and receives positive gradient; the other 16,391 receive negative gradient — "it wasn't you".

A rare token is the target almost never and a "non-target" millions of times. Its net push is systematically negative, in the direction opposite to the mean of the hidden states. And since all rare tokens get the same push, they are piled up in the same region — a cone.

2 Four instruments, four blind spots

A single measurement does not describe the health of a space. We use four, and the point of this section is that each one is blind to a failure another one sees.

Anisotropy — mean cosine between pairs

Measures the offset from the origin: do the vectors all point the same way? It is the most direct measure of the cone. Blind to: rank collapse, redundancy between dimensions.

Decorrelation — mean off-diagonal |correlation|

Do the dimensions carry different information, or repeat the same? It is what Barlow Twins and VICReg's covariance term attack. Blind to: offset from the origin.

Effective rank — exp(H(σ / Σσ))

How many of the 256 dimensions the space really uses. Entropy of the singular values, exponentiated: goes from 1 (everything on a line) to 256 (full use). Blind to: collapse to a point — because the measure centers the data first, and a single point plus noise becomes "full rank" once centered.

SigReg — distance to an isotropic Gaussian

Projects onto random directions and compares each projection with N(0,1). By Cramér–Wold, if all the 1-D projections are standard Gaussian, the whole distribution is isotropic Gaussian. It is the most comprehensive measure — and the hardest to interpret on its own.

The proof of the blind spots

We built five synthetic distributions and measured all four:

3 A measurement error that almost slipped through

The first version of our SigReg started like this:

def sigreg(Z, …):
    Z = Z - Z.mean(0, keepdim=True)     # "normalize"
    Z = Z / (Z.std() + 1e-6)

It looked reasonable. The synthetic test showed otherwise:

caso                      sigreg
gaussiana isotropica      0.0037
cone (anisotropica)       0.0029   <- MENOR que a gaussiana!

The cone — mean cosine +0.99, the worst pathology in the table — scored better than the perfect space.

The cause: our anisotropy is an offset from the origin. The synthetic cone is small noise around a point far from zero. Subtracting the mean makes that point vanish, and what is left is isotropic noise. Centering erased exactly what we wanted to measure.

The paper's SigReg does not center, and now we understand why: the target is N(0, I), and zero mean is part of the target, not of the preparation. With the line removed:

cone (anisotropica)       2.7687   <- 440x o piso de ruido

4 The simple remedy, and what it costs

We add to the loss a term that penalizes the mean cosine between rows of the table, sampling 512 per step:

def reg_cosseno(W, n=512):
    i = torch.randint(W.shape[0], (n,))
    E = F.normalize(W[i], dim=-1)
    G = E @ E.T
    return (G.sum() - G.diagonal().sum()) / (n * (n - 1))

7.6× less anisotropy for 0.17% of quality — a difference within the oscillation validation already has between evaluations. And as a bonus, the rank goes up from 232 to 245.

That is why the regularizer is on by default in train.py, with the evidence written into the code itself. Turning it off requires an explicit --reg-cos 0.

A hypothesis the data did not confirm

Tokens never seen in training end up with a norm 3.15× larger than the trained ones (1.52 against 0.48). With weight tying, the table is also the output layer — so a large-norm vector should produce a large logit, and junk tokens would show up in the predictions.

We tested on six prompts:

The data did not confirm the hypothesis. The norm grows, but the direction learned to be opposite to the real hidden states — the model pushed these tokens below random. It is a good sign about the training, and a reminder: a large norm does not imply a large logit when the direction is wrong.

5 Why this matters for the final goal

Three reasons, in order of concreteness:

1. Sentence embeddings. If the model is going to produce vectors for search or clustering, a cone with cosine 0.38 makes similarity nearly useless. It was literally what broke semantic search in the previous course.
2. The compass (the first version of the rudder, which only predicted the direction of the next state) depends on it. The heading term at generation adds β · cos(E[v], g) to the logits. In an anisotropic space, every token has a similar cosine with every target — and the compass points nowhere. Chapter 05 shows this measured.
3. Transferred knowledge. A better-conditioned space is a better starting point for any fine-tuning later.

With the diagnosis made and the simple remedy working, chapter 04 — once the experiment queue closes — walks through the lineage
that computer vision built for the same problem, and tests each link on ours.