Anisotropy: The Narrow Cone Problem

Chapter 1: Measuring It

Chapter 0 established that something is wrong. Before we can fix it we need a number — ideally several, because "anisotropy" turns out to be two different geometric facts wearing one name, and the fixes in later chapters attack them separately.

We will build two measurements. The first is the one everybody uses and the one every paper reports. The second is the one that actually tells you which fix to reach for.

Measurement one: average pairwise cosine

The obvious thing. Take n embeddings, compute the cosine between every pair, average.

Ethayarajh calls this the anisotropy baseline, and the name is exactly right: it is what a similarity of "unrelated" looks like in this space. If two random words score A, then a pair scoring A tells you nothing. Only the excess over A is information.

Computing A the naive way is O(n²d), which for a million embeddings in 768 dimensions is 7.7×10¹⁴ multiplications — not happening. But you never need to. There is an exact identity that reduces it to a single mean.

The identity that makes it free

Let v̂_i = v_i/‖v_i‖ be the unit-normalised embeddings, so that cos(v_i, v_j) = v̂_i · v̂_j. Let μ̂ be their mean:

Now expand the squared norm of that mean. This is the whole trick:

The right-hand side is the average of all n² pairwise cosines, including the n self-pairs, each of which equals exactly 1. Peel those off:

Solve for A:

Worked example: five vectors, by hand

Five unit vectors in the plane, at angles 10°, 25°, 40°, 5°, 30°. A cone 35 degrees wide — a caricature of a real embedding cloud, but the arithmetic is checkable.

The slow way. Ten pairs. For unit vectors in the plane the cosine of the pair is just the cosine of the angle difference, so:

all ten pairwise cosines(10,25) -> cos 15 = 0.9659    (25,40) -> cos 15 = 0.9659
(10,40) -> cos 30 = 0.8660    (25, 5) -> cos 20 = 0.9397
(10, 5) -> cos  5 = 0.9962    (25,30) -> cos  5 = 0.9962
(10,30) -> cos 20 = 0.9397    (40, 5) -> cos 35 = 0.8192
                              (40,30) -> cos 10 = 0.9848
                              ( 5,30) -> cos 25 = 0.9063

sum = 9.3799        A = 9.3799 / 10 = 0.93799

The identity way. Sum the five unit vectors component by component:

Two methods, ten pairwise cosines versus one sum, and they agree to five decimal places. Also note ‖μ̂‖ = 0.9749 — the mean of five unit vectors has length 0.975. If those five vectors were pointing in genuinely unrelated directions, the mean of five unit vectors would have length around 1/√5 = 0.447. The cloud's centre of mass has barely moved off the sphere. That is the narrow cone, quantified.

The floor: why zero is not the target

An obvious idea now presents itself: subtract the mean, and A becomes 0. Isotropy achieved. Except it does not, and the reason is a constraint you cannot escape.

Suppose we centre the vectors so that ∑_i u_i = 0. Then the same expansion, run backwards:

The off-diagonal dot products are forced to sum to a negative number. If all the vectors happen to have the same length, dividing through by n(n−1)‖u‖² gives:

The isotropic baseline — what "unrelated" should look like

What should A be for a genuinely isotropic cloud? Zero, by symmetry — but the more useful question is how tightly it concentrates around zero, because that sets the scale of "surprising."

Take v fixed and u uniform on the unit sphere in R^d. Rotate coordinates so v is the first basis vector; then cos(u,v) = u₁, the first coordinate of a uniformly random unit vector. By symmetry every coordinate has the same second moment, and they must sum to ‖u‖² = 1:

One line, no integrals. For d = 768 that is 1/27.71 = 0.0361. So in an honest 768-dimensional space, two random items score 0.00 ± 0.04, and 0.2 would already be a striking similarity.

This table is also why the problem got worse as models got bigger. The isotropic baseline shrinks as 1/√d, so the same absolute drift in the mean becomes a more extreme statistical anomaly, and the useful signal has to fight harder against it.

How many vectors do you need before you believe A?

Every measurement in this lesson is estimated from a sample, so the first question a careful engineer asks is how noisy the estimate is. For A the answer is surprising, and it is worth deriving because it tells you where to stop worrying.

Suppose the truth is isotropy, and you draw n unit vectors. What is the spread of your measured A? Start from the identity: A is the average of the n(n−1) off-diagonal terms û_i·û_j. Each individual term has mean 0 and variance 1/d, which we just derived. If they were independent, the variance of their average would be (1/d)/(n(n−1)) — but they are not independent, because each vector appears in n−1 of the terms.

The dependence turns out to help rather than hurt, and the reason is a small structural fact. Fix one vector û_i and average over its partner: E[ û_i·û_j ] over random û_j is û_i·E[û_j] = û_i·0 = 0, whatever û_i is. No single vector can bias the average in a predictable direction. In the language of U-statistics this is a degenerate kernel, and it means the leading 1/n term of the variance vanishes entirely, leaving:

Measurement two: the spectrum

Average cosine measures one thing: how far the cloud's centre has drifted from the origin. It is blind to a completely different pathology — a cloud that is centred on the origin but is shaped like a pencil rather than a ball.

To see that you need second moments. Define two matrices. The second-moment matrix is the average outer product of the raw vectors:

The covariance matrix is the same thing computed after subtracting the mean μ = (1/n)∑v_i:

And they are related by one line of algebra that is the conceptual key to this whole chapter:

Expand the definition of S, add and subtract μ, and the cross terms vanish because ∑(v_i−μ) = 0 by construction. The mean contributes a rank-one spike to the second-moment spectrum: one enormous eigenvalue of size ‖μ‖² in the direction μ̂, and nothing else.

Reading a spectrum: explained variance and effective rank

Eigendecompose Σ = UΛU^T and sort the eigenvalues λ₁ ≥ λ₂ ≥ … ≥ λ_d ≥ 0. Each λ_k is the variance of the cloud along direction u_k. Normalise them into a probability distribution:

Two summary numbers follow. Top-k explained variance is just p₁ + … + p_k — the fraction of the cloud's spread living in the first k directions. And effective rank, the number of dimensions the cloud genuinely uses, is the exponential of the entropy of that distribution:

Why the exponential of the entropy? Because it has exactly the property you want: if k directions share the variance equally and the rest are zero, then p is uniform over k, H = log k, and erank = k on the nose. If one direction has everything, H = 0 and erank = 1. It interpolates smoothly between those, and it does not care about the units of λ.

Worked example. A three-dimensional cloud whose covariance eigenvalues are λ = (0.90, 0.06, 0.04) — already normalised, so p = λ.

Term by term. ln 0.90 = −0.10536, so the first term is 0.90 × 0.10536 = 0.09483. ln 0.06 = −2.81341, giving 0.06 × 2.81341 = 0.16880. ln 0.04 = −3.21888, giving 0.04 × 3.21888 = 0.12876. Sum:

Three dimensions available; 1.48 dimensions used. That single number is the most compact statement of "this cloud is a pencil" available, and you can compute it for a 768-dimensional embedding matrix in about two seconds.

Second worked example: the effective rank of a power law

Real embedding spectra almost never look like the toy above. They look like power laws: λ_k ∝ k^−α, a long slow decay with no obvious cliff. Working one out by hand builds the intuition for what α you are looking at when you plot your own.

Take five eigenvalues following λ_k = 1/k, that is (1, 1/2, 1/3, 1/4, 1/5). Their sum is the fifth harmonic number:

Normalise, then take each term of the entropy sum. The arithmetic is five divisions and five logarithms:

entropy of a 1/k spectrum, term by termk   lambda_k    p_k = lambda_k/2.283333    ln p_k      -p_k ln p_k
1   1.000000    0.437956                  -0.825636    0.361593
2   0.500000    0.218978                  -1.518784    0.332580
3   0.333333    0.145985                  -1.924249    0.280912
4   0.250000    0.109489                  -2.211931    0.242182
5   0.200000    0.087591                  -2.435074    0.213291
                                                       ----------
                                          H         =  1.430558

erank = e^1.430558 = 4.181    out of 5 available

Five dimensions, 4.18 used. A 1/k spectrum is remarkably healthy — the top direction holds only 43.8% of the variance and the effective rank is 84% of the ambient dimension. That is worth internalising, because a plot of a 1/k spectrum on a linear axis looks alarming: the first bar is five times the last. Effective rank is the number that tells you not to panic.

Push the exponent and the picture changes fast. Computing the same quantity over 768 eigenvalues at several exponents:

Notice how violently effective rank responds between α = 1 and α = 2 — it falls thirtyfold for a doubling of the exponent. Fit a straight line to your log-log spectrum plot and the slope you read off is almost the entire diagnosis. And the earlier report in Chapter 9 with top-1 variance 0.41 and effective rank 12.4 sits at roughly α = 1.6, which is the regime where the space is genuinely being wasted.

The null spectrum: finite samples are anisotropic even when the truth is not

Here is the trap that mirrors the one for A, and it runs in the opposite direction. A is estimated to five decimals from a thousand vectors. The spectrum is not, and if you forget this you will diagnose a healthy space as diseased.

Take n samples of genuinely isotropic d-dimensional data, so the true covariance is exactly the identity and every eigenvalue is exactly 1. Compute the sample covariance and eigendecompose it. You do not get eigenvalues of 1. You get a spread, and random matrix theory tells you exactly how wide it is: with γ = √(d/n), the sample eigenvalues fill the interval

This is the Marchenko–Pastur law, and it is the single most useful null hypothesis in this whole subject. Work an example. At d = 768 and n = 5,000, γ = √0.1536 = 0.3919, so the predicted top eigenvalue is (1.3919)² = 1.937 and the bottom is (0.6081)² = 0.370. A simulation on actual Gaussian noise returns 1.934 and 0.367.

The effective rank has the same bias, in a milder form: that same 768-dimensional isotropic noise at n = 5,000 reports an effective rank of 711 rather than 768, and at n = 2,000 it reports 633. So effective rank is downward biased by finite sampling — a healthy space always looks slightly less healthy than it is. Reassuringly, the bias is small compared to the effect we are hunting: a real embedding cloud reporting 12.4 out of 768 is nowhere near explainable by sampling.

Press pencil only and then mean-centred. The average pairwise cosine reads close to zero — the cloud passes the standard anisotropy test with flying colours — and the effective rank is still about 2 out of 12. That configuration is exactly the trap: a paper that reports only average cosine would call this space isotropic, and every retrieval query run against it would still be dominated by one direction.

Two more measures you will meet in the literature

The partition-function ratio. Arora et al. (2016) showed that if word vectors are isotropic then the partition function Z(c) = ∑_w exp(c^Tv_w) is approximately constant in the unit vector c. Mu & Viswanath turn this into a diagnostic: compute Z(c) at each eigenvector of the embedding matrix and report

A value of 1 means the cloud looks the same from every direction; values well below 1 mean it does not. This is a more sensitive instrument than average cosine because the exponential weights the tails — a few far-out points in one direction move Z a great deal. Mu & Viswanath report that their post-processing pushes this ratio substantially closer to 1 for both word2vec and GloVe.

IsoScore. Rudman et al. (2022) catalogue what is wrong with all of the above: average cosine conflates the offset with the shape; the partition-function ratio is sensitive to the number of samples; neither scales linearly with the fraction of dimensions actually used. Their replacement compares the covariance's eigenvalue vector against the all-ones vector and normalises so that "uses k of d dimensions" scores k/d. You do not need it to follow this lesson, but you should know that "average cosine" is a proxy, not a definition, and that a paper reporting only that number has told you about the offset and almost nothing about the shape.

Computing all of it in one streaming pass

Everything in this chapter can be computed without ever holding your embedding matrix in memory, which matters the moment your corpus is larger than your laptop. Two accumulators do the whole job, and their sizes do not depend on n at all.

the whole diagnostic panel, streamingimport numpy as np

class GeometryAccumulator:
    def __init__(self, d):
        self.n     = 0
        self.sum_u = np.zeros(d)          # (d,)   sum of UNIT vectors  -> gives A
        self.sum_v = np.zeros(d)          # (d,)   sum of RAW vectors   -> gives mu
        self.gram  = np.zeros((d, d))     # (d,d)  sum of v v^T          -> gives Sigma

    def add(self, V):                    # V: (B, d) — any batch size
        U = V / np.linalg.norm(V, axis=1, keepdims=True)
        self.n     += V.shape[0]
        self.sum_u += U.sum(0)
        self.sum_v += V.sum(0)
        self.gram  += V.T @ V             # the only O(B d^2) step

    def report(self):
        n      = self.n
        mu_hat = self.sum_u / n
        offset = float(np.linalg.norm(mu_hat))
        A      = (n * offset**2 - 1) / (n - 1)
        mu     = self.sum_v / n
        # the identity that makes streaming work: Sigma = E[vv^T] - mu mu^T
        cov    = self.gram / n - np.outer(mu, mu)
        lam    = np.linalg.eigvalsh(cov)[::-1]
        p      = lam / lam.sum()
        return dict(A=A, offset=offset, top1=p[0],
                    erank=float(np.exp(-(p * np.log(p + 1e-12)).sum())))

The line that makes this possible is cov = gram/n - outer(mu, mu), which is S = Σ + μμ^T rearranged. You accumulate the uncentred second moment — which needs no knowledge of the mean and therefore no second pass — and subtract the mean's contribution at the end.

What to report, so that someone else can use it

A short checklist, because "the space is anisotropic" is not a reportable measurement and this chapter has now produced five things that are.

What this buys us

We now have a language. "The space is anisotropic" splits into "the mean is at ‖μ‖/mean‖v‖ = 0.85 of a typical vector length" and "the top principal component holds 41% of the centred variance, and the effective rank is 12 out of 768." Those two sentences point at two different fixes, and both are diagnosable in one pass over your data.

The next question is the interesting one. Nobody designed this. No loss function contains a term that says "please collapse into a cone." So why does every large model do it?

Chapter 2: Why It Happens

Nobody wrote a loss term that says "collapse into a cone." No architecture diagram contains a box labelled "common direction." And yet every large pretrained model does it, reliably, across architectures and training corpora and decades of hardware. That reliability is a clue: it means the cause is structural, not incidental.

There are three mechanisms. The first is the one that surprises people, so we start there.

Mechanism one: the loss cannot see it

Start with the last operation in a language model. You have a hidden state h ∈ R^d for the current position and an unembedding matrix E whose rows e_w are one output embedding per vocabulary item. The score for token w is a dot product, and the probabilities come from a softmax:

Now perform the following vandalism. Pick any vector c ∈ R^d you like, and add it to every single row of E:

What happens to the logits?

The added term h · c does not depend on w. It is the same number for every token in the vocabulary. And softmax is invariant to adding a constant to all its inputs — that is the same shift-invariance you exploit numerically when you subtract the max before exponentiating:

Worked example: 3 tokens, 2 dimensions, and a catastrophe

Concrete numbers, small enough to check on paper. Hidden state h = (1, 2) and three output embeddings:

Logits: h·e_a = 1, h·e_b = 2, h·e_c = −1. Softmax them — subtract the max, 2, then exponentiate:

softmax beforeshifted   = (1-2, 2-2, -1-2) = (-1, 0, -3)
exps      = (0.367879, 1.000000, 0.049787)
sum       = 1.417666
p         = (0.259501, 0.705385, 0.035119)

Now the vandalism: c = (3, 3), added to all three rows.

softmax afterlogits    = (1*4+2*3, 1*3+2*4, 1*2+2*3) = (10, 11, 8)
                                       h.c = 1*3 + 2*3 = 9  -> every logit shifted by exactly 9
shifted   = (10-11, 11-11, 8-11) = (-1, 0, -3)
p         = (0.259501, 0.705385, 0.035119)      # identical to seven decimals

Identical predictions. Now measure the geometry. Before, the three embeddings were mutually orthogonal or opposite: cos(a,b) = 0, cos(b,c) = 0, cos(a,c) = −1, so A = −1/3 = −0.3333 — exactly the −1/(n−1) floor from Chapter 1, because those three vectors happened to be centred.

After, compute each cosine longhand. Norms first: ‖e′_a‖ = √(16+9) = 5, ‖e′_b‖ = √(9+16) = 5, ‖e′_c‖ = √(4+9) = √13 = 3.60555.

Where the wandering comes from

A perfectly flat direction does not have to drift; a ball on a level table stays put. So what actually moves it? Three candidates, and the honest answer is that the first two are weaker than most write-ups suggest, which sharpens the case for the third.

Initialisation. The starting E has a mean of roughly zero, but "roughly" means order 1/√|V| times the typical row norm. For a 50,000-token vocabulary that is 1/224 = 0.45% of a row norm — small, nonzero, and with no force to correct it. Note this is a starting offset, not a drift; it stays about that size unless something moves it.

Weight decay. This one is worth doing properly, because the exact answer is the opposite of what you might guess. Sum the cross-entropy gradient over the whole vocabulary at a single training position:

Exactly zero, at every position, for every model, always — because the softmax probabilities sum to one by construction. Cross-entropy therefore applies no net force whatsoever to the mean of the output embedding table. That is the gauge invariance of mechanism one, restated as a gradient statement rather than a loss statement, and it is a stronger version of the same fact.

Now add decoupled weight decay, which subtracts ηλe_w from every row each step. Averaging over the vocabulary, the mean m = (1/|V|)∑_we_w obeys m ← (1 − ηλ)m with no forcing term at all, so it decays geometrically. With a learning rate of 10⁻⁴, decay 0.1, over 10⁶ steps, the shrink factor is (1 − 10⁻⁵)^1000000 = e⁻¹⁰ = 0.000045. Weight decay does not merely help — if it is applied to the embedding table, it annihilates the mean.

Adam on a flat direction. Adam divides each coordinate's step by the running root-mean-square of its own gradients. On a direction where the true gradient is zero and only noise remains, that division normalises the noise to unit scale — so the parameter takes a step of about the full learning rate every time, in a random direction. Over T steps that is a random walk of size η√T. With η = 10⁻⁴ and T = 10⁶, that is 10⁻⁴ × 1000 = 0.1, or roughly a tenth of a typical row norm, accumulated from pure noise. Plain SGD on the same direction would move ησ√T, which is smaller by the factor σ — the actual gradient magnitude, which on a flat direction is tiny. Adam is the optimiser that most aggressively converts a flat direction into a wandering one, precisely because its whole purpose is to make small-gradient directions move as fast as large-gradient ones.

Asymmetric gradients in the softmax. This is the big one, and it deserves its own mechanism.

Mechanism two: rare tokens are pushed, never pulled

Gao et al. (2019) gave this a name — the representation degeneration problem — and a derivation. Start with the gradient of the cross-entropy loss for one training position with target token t. The derivative with respect to a logit is the single most useful fact in all of classification:

And since ℓ_w = h · e_w, the chain rule gives the gradient with respect to the output embedding itself:

Read the sign. If w is the target, the coefficient p_w − 1 is negative, so gradient descent moves e_w toward h — "look more like the context that predicted you." If w is not the target, the coefficient p_w is positive, so e_w moves away from h.

Now consider a genuinely rare token — a rare surname, an obscure unit symbol, a fragment of a URL. Over a training run of, say, 10⁶ gradient steps with a batch of 10⁵ positions, this token might be the target a few hundred times and a non-target 10¹¹ times. Its update is, to an excellent approximation, entirely made of the "push away" term:

where h̄ is the average hidden state over the corpus. Every rare token feels the same push, along the same direction, for the same reason. They are all being shoved away from the corpus's average context — and "away from h̄" is one specific direction, not many.

Worked example: two rare tokens meet in the corner

Two dimensions again. Take the average hidden direction to be h̄ = (0.8, 0.6), a unit vector. Suppose the accumulated drift magnitude works out to 1.0 — say average p̄_w = 0.001, learning rate η = 10⁻³, and N = 10⁶ effective updates, giving 10⁻³ × 10⁻³ × 10⁶ = 1.0.

Both tokens started at random small values. Token 1 started at (0.05, −0.03); token 2 at (−0.04, 0.06). After drift:

Their cosine, computed longhand:

Before training, those two vectors had cosine (0.05)(−0.04) + (−0.03)(0.06) = −0.0038 over norms 0.0583 × 0.0721 = 0.0042, which is −0.90 — nearly opposite. After training on data in which neither ever appears, they are at 0.9919. The model has learned that two completely unrelated tokens are nearly the same thing, and it learned it from the absence of evidence.

Mechanism two attacks the shape, not the mean — worked

The previous example moved two rare tokens and stopped. But we proved above that the total gradient over the vocabulary is exactly zero, which means every unit of "push" applied to a rare token is balanced somewhere by a unit of "pull" applied to a frequent one. Let us follow the whole balance sheet, because the result is not what the phrase "everything drifts into a cone" suggests.

Three tokens in two dimensions, with the average hidden direction h̄ = (0.8, 0.6) again. Token f is frequent: it is the target often, so its accumulated update is a pull toward h̄, say of size 2. Tokens 1 and 2 are rare: each is pushed away by size 1. Two pushes of 1 balance one pull of 2, which is the zero-sum constraint made concrete. Give them initialisations that sum to zero so we can see the effect cleanly:

the whole vocabulary, before and afterinit:   e_f = (-0.01, -0.03)    e_1 = (0.05, -0.03)    e_2 = (-0.04, 0.06)
        sum = (0.00, 0.00)                                # centred at the start

drift:  e_f += 2*(0.8, 0.6) = (1.6, 1.2)                  # frequent: PULLED toward h-bar
        e_1 -=   (0.8, 0.6)                               # rare: PUSHED away
        e_2 -=   (0.8, 0.6)                               # rare: PUSHED away

after:  e_f = ( 1.59,  1.17)    e_1 = (-0.75, -0.63)    e_2 = (-0.84, -0.54)
        sum = (0.00, 0.00)                                # STILL exactly centred

The mean has not moved by a hair, and it never could: the drifts were constructed to cancel, because the gradient made them cancel. Now measure. Norms first: ‖e_f‖ = √(2.5281 + 1.3689) = √3.8970 = 1.974082, ‖e₁‖ = 0.979490, ‖e₂‖ = 0.998599.

Look at what the standard diagnostic says. A is negative, sitting near the centring floor, and a dashboard reporting only average pairwise cosine would call this space perfectly healthy. Now compute the second-moment matrix instead:

the spectrum tells a different storyS = (1/3) * sum_w e_w e_w^T

S11 = (1.59^2 + 0.75^2 + 0.84^2)/3 = (2.5281 + 0.5625 + 0.7056)/3 = 1.2654
S22 = (1.17^2 + 0.63^2 + 0.54^2)/3 = (1.3689 + 0.3969 + 0.2916)/3 = 0.6858
S12 = (1.59*1.17 + 0.75*0.63 + 0.84*0.54)/3 = (1.8603 + 0.4725 + 0.4536)/3 = 0.9288

trace = 1.9512    det = 1.2654*0.6858 - 0.9288^2 = 0.867811 - 0.862669 = 0.005142

eigenvalues = (1.9512 +/- sqrt(1.9512^2 - 4*0.005142))/2 = (1.9512 +/- 1.945923)/2
            = 1.948561  and  0.002639

top-1 share  = 1.948561 / 1.9512 = 99.86%
effective rank = 1.010  out of 2

for comparison, the INITIALISATION had eigenvalues 0.002817 and 0.000383,
a top-1 share of 88% and no dominant scale at all.

One caveat to state, since we just leaned on it: A landed at −0.335 rather than the exact centring floor −1/(n−1) = −0.5 for n = 3. The reason is the norms — 1.974, 0.979, 0.999 — are wildly unequal, and the floor identity requires equal norms. The dot products still obey the exact constraint: −1.9296 − 1.9674 + 0.9702 = −2.9268, which is exactly −(‖e_f‖² + ‖e₁‖² + ‖e₂‖²)/2 = −(3.8970 + 0.9594 + 0.9972)/2 = −2.9268. The identity is about dot products; only the cosine version needs equal lengths.

The empirical fingerprints

If mechanism two is real, it makes falsifiable predictions, and all four papers report the corresponding measurements.

Mechanism three: the residual stream accumulates

Mechanisms one and two are both about the output embedding table. But Ethayarajh measured anisotropy in the hidden states themselves, at every layer, and found it there too. So there must be a third cause living inside the network.

A transformer's hidden state is a residual stream: layer ℓ produces h^(ℓ) = h^(ℓ−1) + δ^(ℓ), where δ is what the attention and MLP blocks wrote this layer. Nothing ever replaces the stream; everything only adds to it.

Now suppose each layer's contribution decomposes into a part that is roughly the same for every input — a bias, a "default" direction, whatever the attention sink or the always-on features write — plus a part that genuinely depends on the token. Call them c and g_i, both of unit scale, with the g_i independent across different inputs i:

Here is the asymmetry that does all the work. The shared parts add coherently: L copies of c give a vector of length L. The individual parts add incoherently — they are independent random directions, so their sum is a random walk of length √L, not L. Write h_i ≈ L·c + √L·g_i with g_i a fresh unit random vector. The cosine between two different inputs:

because in high dimensions independent unit vectors have dot products of order 1/√d, which vanish. So:

This is a toy model, and it should be read as one: it assumes every layer writes the same shared vector at the same scale, which is not literally true, and it ignores the layer norms that partly rescale the stream. But it reproduces the reported curve's shape from one structural fact — that shared components add linearly and idiosyncratic ones add as a square root — and that fact is not going away. Anisotropy is what residual addition does to any signal with a shared component.

Where does the toy break? In two named places, and knowing them is what keeps it a tool rather than a slogan. First, it assumes the shared contribution c has the same magnitude at every layer, which is false — real transformers write much larger updates in some layers than others, and the curve of anisotropy against depth is correspondingly bumpier than L/(L+1). Second, it ignores the LayerNorm between blocks, which rescales the stream to a fixed norm and therefore partially resets the accumulation. Both corrections push the predicted curve down, which is consistent with real measurements sitting below the toy at intermediate depths and catching up at the top. What survives both corrections is the structural asymmetry: nothing about a LayerNorm or a variable write magnitude changes the fact that coherent terms add as L and incoherent ones as √L.

What LayerNorm does and does not fix

A reasonable objection: modern transformers run a LayerNorm right before the unembedding. Does that not remove the common component?

It removes exactly one direction of it. LayerNorm subtracts the mean across features of the hidden vector, which is the same as projecting out the all-ones vector 1 = (1,1,…,1)/√d. After LayerNorm, h · 1 = 0. So the gauge freedom in the direction c = 1 is genuinely eliminated: shifting every output embedding along the all-ones direction now has no effect, because h has no component there to pick it up.

Three dimensions, three lines, and you can see both halves of that claim.

what LayerNorm does and does not removeh = (3, 1, 2)
mean across features = (3 + 1 + 2)/3 = 2
h_centred = (1, -1, 0)                    # this is the "sub-mean" half of LayerNorm
check: h_centred . (1,1,1) = 1 - 1 + 0 = 0     # the all-ones component is gone

now shift every output embedding by c = (1, 1, 1):
  delta logit for token w = h_centred . c = 0       # NOTHING happens. Gauge killed.

but shift by c = (1, -1, 0) instead:
  delta logit = h_centred . c = 1*1 + (-1)*(-1) + 0*0 = 2
  ...which is again the SAME shift for every token w, since it does not
  depend on w — so softmax still cannot see it, and the gauge survives.

# LayerNorm closed 1 of the d gauge directions. 767 remain open at d = 768.

But c can be any of the other d − 1 directions, and for those the invariance is untouched. LayerNorm removes one degree of gauge freedom out of 768. It also rescales h to a fixed norm, which fixes nothing about direction. Empirically, models with final LayerNorm are still severely anisotropic — Ethayarajh's GPT-2 measurements are taken on a model that has one.

Anisotropy is not rank collapse — keep them apart

One terminological clean-up, because these two get used interchangeably and they are different failures with different fixes.

The distinguishing measurement takes one line: count how many eigenvalues exceed a small tolerance. An anisotropic 768-dimensional cloud has 768 nonzero eigenvalues with an effective rank of 12; a rank-collapsed one has 12 nonzero eigenvalues and 756 exact zeros. Effective rank alone cannot tell them apart — which is a genuine limitation of that statistic and the reason to report the raw spectrum alongside it.

How big is the gauge group, exactly?

Shifting is not the only thing the loss cannot see. It is worth counting the whole family, because the count tells you precisely how much of your embedding geometry was ever determined by training — and the answer is "less than you think, but the part that matters is small and specific."

Three transformations leave every logit h · e_w unchanged:

It also explains something otherwise puzzling about Chapter 6. Whitening's solution W = UΛ^−1/2 is famously non-unique — WR works for any orthogonal R. That is not a defect of the derivation; it is the rotation row of the table above showing up again. The ambiguity is real and it is harmless, for exactly the same reason.

A fourth candidate: the attention sink writes a constant

Mechanism three assumed each layer contributes "a part that is roughly the same for every input." That was left abstract. In modern decoder-only transformers there is a concrete, well-documented source of it.

Attention heads in these models place a large, near-constant fraction of their attention mass on the first token of the sequence — the beginning-of-sequence token, or whatever occupies position zero. The phenomenon is usually called an attention sink: the head has to distribute a probability distribution that sums to one, and when it has nothing in particular to attend to, dumping the mass somewhere harmless is the cheapest option available.

Now trace what that does to the residual stream. If a head puts weight a on the sink token, its output at every position includes a · v_sink, where v_sink is the sink's value vector. The same vector. At every position. In every sentence. Sum it over H heads and L layers:

Put plausible numbers on it: 12 layers, 12 heads, an average sink weight of 0.3, and a value-vector norm comparable to the per-head output norm. That is 12 × 12 × 0.3 = 43 units of the same vector added into the stream, against idiosyncratic contributions that add as a square root. This is mechanism three's c, with a name, a location in the architecture, and a measurable weight.

The three mechanisms, side by side

Notice that all three point at the same corrective action from different directions: find the directions carrying the shared component and neutralise them. Which is the algorithm Chapter 3 makes precise, in three lines of numpy.

Could you train an isotropic model on purpose?

The best test of a causal story is whether it tells you how to prevent the effect. Take each mechanism in turn and ask what intervention would kill it, and what that intervention would cost. The answers are instructive, and two of them are things people actually do.

Which table is anisotropic — input or output?

One loose end worth tying, because it is a question people get wrong and it has a clean answer. A language model has (at least conceptually) two embedding tables: the input lookup that turns a token id into a vector, and the output unembedding that turns a hidden state into logits. Which one does mechanism two attack?

The output table, unambiguously. Every gradient in the derivation above — the p_w − 1[w = t] coefficient, the push away from h̄ — is a gradient with respect to e_w, a row of the unembedding. The input table receives a completely different signal: a token's input embedding is updated only when that token actually appears, and always in the direction that improves prediction at that position. There is no "push away from every context you did not appear in", because an input embedding that is never looked up receives no gradient at all.

Chapter 3: All-but-the-Top, Derived

In 2018 Jiaqi Mu and Pramod Viswanath published a paper whose entire method fits in three lines of numpy and whose title is also its algorithm. It works on every static word embedding anybody had trained, it needs no training, no labels, and no hyperparameter search worth the name, and it improves results on word similarity, categorisation, analogy, and downstream sentence tasks.

Here it is, in full.

That is the whole paper's method. The interesting part is why each of the three steps is there, and the answer is Chapter 1's two-pathology decomposition, applied literally.

Why two operations and not one

Recall that the second-moment matrix decomposes as S = Σ + μμ^T. Step 1 removes the μμ^T rank-one spike — the offset pathology. Steps 2 and 3 attack Σ itself — the shape pathology. They are different operations because they fix different diseases, and doing only one of them leaves the other in place.

You can see the necessity of both directly from Chapter 0's expansion. We had:

Subtracting μ kills the leading 1 and the ‖μ‖² compression — you go from measuring cos(μ+x, μ+y) to measuring cos(x, y), and the whole [0.99, 1.00] sliver expands to the full [−1, 1]. But the p and q terms, the ones that corrupt the ranking, live in the components of x and y along μ̂. They survive centring. Only removing that direction — which is step 3, if μ̂ happens to align with a top principal component, which it usually does — kills them.

What removing a direction actually does

Step 3 is a projection. Removing one unit direction u from a vector v is:

P = I − uu^T is an orthogonal projector: it maps R^d onto the (d−1)-dimensional subspace perpendicular to u, and it is idempotent (PP = P — once a direction is gone, removing it again does nothing) and symmetric. Removing D directions at once is P = I − UU^T where U stacks the D orthonormal components.

The consequence for similarity is exact and worth stating precisely: after projection, cos(Pv, Pw) is the cosine of v and w as measured inside the surviving subspace only. Everything the deleted directions had to say about either vector — including whatever genuine information they carried — is gone, irreversibly. That is why D matters.

How large should D be?

Mu and Viswanath's rule of thumb, from their experiments across dimensions and corpora: D ≈ d/100. For 300-dimensional GloVe vectors, D = 3. For 768 dimensions, D = 7 or 8.

The rule has a sensible shape. Too small and you leave the frequency directions in place, so nothing improves. Too large and you start deleting real semantic axes — the components after the first handful are no longer dominated by a shared artefact, and each one you remove takes real content with it. The paper reports the characteristic inverted-U: performance rises, plateaus, then falls as D grows past the useful range.

Worked example: four words in three dimensions, by hand

Everything above is now going to happen in full arithmetic. Three dimensions, four words, every number computed.

Our four words are cat, dog, car, truck. The intended semantics are obvious: cat should be nearest to dog, car nearest to truck. The embeddings are:

The first coordinate is the shared component — every word has a lot of it, and how much varies with frequency, exactly as Chapter 2's mechanisms predict. Coordinates two and three carry the meaning: cat and dog have positive second and third components, car and truck negative ones.

Before: the six raw cosines. Longhand for the first one, then the table.

One arithmetic habit worth adopting right here: carry the unrounded norm product into the division. If you round each norm to five decimals first and then multiply, you get 8.74192 and a cosine of 0.87966 — wrong in the fifth decimal. That sounds pedantic until you notice that the whole failure this chapter is about turns on a margin of 0.0026. When the signal lives in the fourth decimal, your rounding lives in the fifth.

Step 1: the mean. Add the four vectors component by component and divide by four.

The mean points straight down the shared axis, with a norm of 3.0 against typical vector norms of 1.8 to 4.8. That single number — ‖μ‖ comparable to a typical ‖v‖ — is the offset pathology in one measurement.

Step 1 continued: centre.

Step 2: the scatter matrix. PCA needs the covariance; we will build the scatter matrix S = ∑_i ṽ_iṽ_i^T, which is four times the covariance and has the same eigenvectors. Six entries to compute, since it is symmetric.

building the 3x3 scatter matrix by handS11 = 1.6^2 + (-1.6)^2 + (-1.6)^2 + 1.6^2                  = 4 x 2.56  = 10.24
S22 = 1.0^2 + 0.8^2 + (-0.9)^2 + (-0.9)^2                  = 1.00+0.64+0.81+0.81 = 3.26
S33 = 0.5^2 + 0.9^2 + (-0.6)^2 + (-0.8)^2                  = 0.25+0.81+0.36+0.64 = 2.06

S12 = (1.6)(1.0) + (-1.6)(0.8) + (-1.6)(-0.9) + (1.6)(-0.9)
    = 1.60 - 1.28 + 1.44 - 1.44                            = 0.32
S13 = (1.6)(0.5) + (-1.6)(0.9) + (-1.6)(-0.6) + (1.6)(-0.8)
    = 0.80 - 1.44 + 0.96 - 1.28                            = -0.96
S23 = (1.0)(0.5) + (0.8)(0.9) + (-0.9)(-0.6) + (-0.9)(-0.8)
    = 0.50 + 0.72 + 0.54 + 0.72                            = 2.48

Step 2 continued: power iteration. We need the top eigenvector. The simplest algorithm that works: start with a guess, multiply by S, normalise, repeat. Each multiplication amplifies the dominant direction relative to the others by the ratio of their eigenvalues, so it converges geometrically.

Start with the first basis vector, since S₁₁ is obviously the largest diagonal entry.

power iteration, three roundsu0 = (1, 0, 0)

S u0 = (10.24, 0.32, -0.96)
       norm = sqrt(104.8576 + 0.1024 + 0.9216) = sqrt(105.8816) = 10.28988
u1   = (0.995150, 0.031098, -0.093295)

S u1 = (10.28985, 0.188456, -1.070409)
       norm = sqrt(105.88100 + 0.035516 + 1.145775) = sqrt(107.06229) = 10.34709
u2   = (0.994469, 0.018213, -0.103450)

S u2 = (10.288503, 0.121048, -1.122629)
       norm = sqrt(105.85327 + 0.014653 + 1.260296) = sqrt(107.12822) = 10.35028
u3   = (0.994031, 0.011695, -0.108464)

... two more rounds and it stops moving:
u*   = (0.9936, 0.0060, -0.1129)      lambda1 = 10.3513

Sanity check the eigenvector: 0.9936² + 0.0060² + 0.1129² = 0.98724 + 0.00004 + 0.01275 = 1.00003. Unit length, to rounding. And the eigenvalue tells us how bad the shape pathology is:

For the arithmetic below we use the rounded u₁ = (0.994, 0.006, −0.113).

Step 3a: project each centred vector onto u₁.

the four projection coefficientsp_cat   = (1.6)(0.994) + (1.0)(0.006) + (0.5)(-0.113)
        =  1.5904 + 0.0060 - 0.0565 =  1.5399
p_dog   = (-1.6)(0.994) + (0.8)(0.006) + (0.9)(-0.113)
        = -1.5904 + 0.0048 - 0.1017 = -1.6873
p_car   = (-1.6)(0.994) + (-0.9)(0.006) + (-0.6)(-0.113)
        = -1.5904 - 0.0054 + 0.0678 = -1.5280
p_truck = (1.6)(0.994) + (-0.9)(0.006) + (-0.8)(-0.113)
        =  1.5904 - 0.0054 + 0.0904 =  1.6754

check: 1.5399 - 1.6873 - 1.5280 + 1.6754 = 0.0000   # must sum to zero: the data is centred

Step 3b: subtract the projections. v′ = ṽ − p·u₁.

removing the top componentcat':   (1.6, 1.0, 0.5)     - 1.5399 x (0.994, 0.006, -0.113)
      = (1.6, 1.0, 0.5)     - (1.530661,  0.009239, -0.174009)
      = ( 0.069339,  0.990761,  0.674009)

dog':   (-1.6, 0.8, 0.9)    + 1.6873 x (0.994, 0.006, -0.113)
      = ( 0.077176,  0.810124,  0.709335)

car':   (-1.6, -0.9, -0.6)  + 1.5280 x (0.994, 0.006, -0.113)
      = (-0.081168, -0.890832, -0.772664)

truck': (1.6, -0.9, -0.8)   - 1.6754 x (0.994, 0.006, -0.113)
      = (-0.065348, -0.910052, -0.610680)

check columns sum to zero:
  x:  0.069339 + 0.077176 - 0.081168 - 0.065348 = -0.000001
  y:  0.990761 + 0.810124 - 0.890832 - 0.910052 =  0.000001
  z:  0.674009 + 0.709335 - 0.772664 - 0.610680 =  0.000000

After: the six cosines again. Norms first: ‖cat′‖ = 1.200293, ‖dog′‖ = 1.079543, ‖car′‖ = 1.182023, ‖truck′‖ = 1.097905. Then the one done longhand:

Three things to notice, in increasing order of importance.

One: cat's nearest neighbour is now dog, at +0.99253, with truck at −0.99998. The retrieval is fixed, and not narrowly — the margin went from −0.0026 to +1.99.

Two: the average pairwise cosine is −0.33332, which is −1/3 = −1/(n−1) for n = 4 to within four decimal places. Be precise about why it is not exact. Chapter 1's identity says a centred set of vectors satisfies ∑_i≠j u_i·u_j = −∑_i‖u_i‖² exactly — that is an identity about dot products, and it holds here to machine precision. Turning it into a statement about the average cosine requires dividing each term by ‖u_i‖‖u_j‖, and only if all four norms are identical does that division factor out cleanly. Ours are 1.200, 1.080, 1.182 and 1.098 — close but not equal — so we land at −0.333323 instead of −0.333333. Chapter 1 stated the equal-norm condition; this is what violating it slightly costs you.

Three — and this is the honest one: every single cosine after the transform has magnitude above 0.99. That is not a triumph, it is a symptom of over-correction. We had 3 dimensions, we removed the mean (costing one degree of freedom) and one principal component (costing another), so the four points now live in a 2-dimensional residual space and happen to lie almost along a line in it. The ordering is right, but the geometry has been flattened.

What the transform did to the spectrum

Cosines are what you consume, but the spectrum is what you should watch, because it is the thing Chapter 1 told us average cosine cannot see. Track it through all three stages of the algorithm on the same four words.

Three separate lessons live in that table, and they are all worth more than the cosine numbers.

Centring alone did most of the spectral work. Top-1 share fell from 89.7% to 66.5% and effective rank rose from 1.41 to 1.91, from one subtraction. That is the rank-one spike μμ^T being removed from S = Σ + μμ^T, exactly as Chapter 1 predicted — and it is why "just centre" is the single highest-value line in this entire lesson.

Average cosine went blind after step 1. It reads about −1/3 both after centring and after the full transform, because centring forces that value and nothing afterwards can change it. Every bit of the difference between rows two and three is invisible to A. If you were tuning D by watching average cosine, you would see a flat line and conclude D does not matter.

Removing the top component made the remaining spectrum worse, not better. Top-1 share went from 66.5% to 99.6% and effective rank collapsed from 1.91 to 1.02. This is not a bug and it is not a contradiction — it is arithmetic. Deleting λ₁ leaves λ₂ as the new largest, and shares are computed against a smaller total, so whatever survives is automatically a bigger fraction of it. In three dimensions with a lopsided tail, that is a violent effect.

Choosing D without a rule of thumb

D ≈ d/100 is a starting point, not an answer. The procedure that actually finds the right D takes about twenty minutes and needs no labels for the first two steps.

Step B is the one people skip and the one that pays. It converts "how many components are contaminated?" from a hyperparameter into a measurement, and it gives you the artefact evidence Chapter 7 will demand before you are allowed to delete anything.

Run the sim with subtract mean: off and D = 1. You will find the first principal component of the uncentred data is essentially the mean direction, so removing it does most of the centring for you — which is a common source of confusion when people report that "PCA removal alone" works. It works because it accidentally did step 1 with its first component. Doing them separately is cheaper and clearer, and it frees the first component to attack the shape rather than the offset.

The implementation, and the part everyone gets wrong

all-but-the-top, fitted and applied correctlyimport numpy as np

def abtt_fit(V, D):
    """V: (n, d) matrix of embeddings from your CORPUS. Returns a transform."""
    mu = V.mean(axis=0)                          # (d,)
    X  = V - mu                                     # (n, d) centred
    # right singular vectors of X == eigenvectors of X^T X == PCs of the covariance
    _, _, Vt = np.linalg.svd(X, full_matrices=False)
    U = Vt[:D]                                      # (D, d), orthonormal rows
    return mu, U

def abtt_apply(V, mu, U):
    X = V - mu                                      # same mu as fit time
    return X - (X @ U.T) @ U                        # (n,D) @ (D,d) -> (n,d)

The shapes are the documentation. X @ U.T is (n, D) — one projection coefficient per vector per removed component. Multiplying that by U (D, d) reconstructs the part of each vector that lives in the deleted subspace, which is exactly what we subtract.

Two details in abtt_fit are not stylistic. The SVD is taken on the centred X, not on V, because the right singular vectors of X are the eigenvectors of X^TX, which is the covariance up to a factor of n — run it on the uncentred V and you get the eigenvectors of the second moment S instead, whose top component is essentially μ and which therefore wastes your first removal on redoing step 1. And full_matrices=False matters for cost, not correctness: with it, the SVD returns min(n,d) singular vectors instead of building a full d×d orthogonal matrix you were going to slice down to D rows anyway.

Shipping it when you cannot re-index

The re-index requirement is the operational objection that stops most of these projects, and it has a workaround worth knowing.

ABTT is linear, and linear transforms can be moved to the other side of a dot product. If your index stores raw vectors v and you cannot rewrite fifty million of them, apply the transform to the query twice instead:

The projector's symmetry and idempotence mean transforming both sides is identical to transforming one side twice — and since P is idempotent, that is identical to transforming the query once and leaving the documents alone. So a plain dot-product index over raw vectors, queried with Pq, returns exactly the ABTT-transformed dot products, for free, with no re-indexing at all.

What the paper actually reported

Mu and Viswanath ran this on word2vec and GloVe across a wide battery: word similarity (RG65, WordSim-353, MEN, SimLex-999 and others), concept categorisation, word analogy, semantic textual similarity, and supervised sentence-classification tasks with the embeddings frozen. The post-processed vectors improved on essentially all of them, at zero training cost. They also verified the mechanism directly: the top principal components of the raw vectors correlate with word frequency, and the isotropy measure I(V) — the partition-function ratio from Chapter 1 — moves substantially closer to 1 after processing.

Two details worth carrying:

And one warning the paper is explicit about: the vectors coming out of ABTT are not drop-in replacements in a system that assumed the old geometry. Everything downstream — thresholds, cached similarities, approximate-nearest-neighbour indexes — must be rebuilt.

Does it compose with itself?

A question worth settling because the answer is instructive. If removing one component helps, does running the whole algorithm twice help twice?

Applied to the same data with the same fitted U, no — and this is exactly the idempotence of the projector from earlier in this chapter. P = I − UU^T satisfies PP = P, so the second application changes nothing at all. Running abtt_apply twice is a no-op, and if your pipeline accidentally does it, you will see zero difference and correctly conclude nothing is wrong.

Applied with a refitted U, yes, and this is the composition the paper points at. After removing D components, re-run the fit: the covariance of the output has a new top eigenvector, and if the variance is still concentrated it may be worth removing too. The reason it is not simply the same as choosing a larger D up front is that the intermediate step usually includes a dimensionality reduction — ABTT, then PCA down to k, then ABTT again — and the reduction re-concentrates variance into the surviving directions, which can manufacture a new dominant direction that did not exist before. Whitening-k in Chapter 6 is the closed-form version of the same idea, doing the flatten and the reduce in one step instead of alternating.

Chapter 4: The Layer-Wise Anatomy

Everything so far assumed each word has one vector. Contextual models broke that assumption, and in doing so they made anisotropy dramatically worse — while also making it much harder to reason about, because now the same word has thousands of vectors and you have to decide what you are even comparing.

Kawin Ethayarajh's 2019 paper is the measurement that sorted this out. It asks a deceptively simple question — how contextual are contextualised word representations, really? — and to answer it honestly he had to first establish that you cannot read a cosine similarity off a contextual model without correcting for the layer's anisotropy. That correction is the part of the paper this lesson needs most, but the findings that come out of it are worth the whole chapter.

Four measurements, defined exactly

Fix a model and a layer ℓ. Write f_ℓ(s, i) for the layer-ℓ representation of the token at position i of sentence s.

1. The anisotropy baseline. Sample two words uniformly at random from the corpus, each in its own context, and take the expected cosine. This is A from Chapter 1, computed on contextual vectors:

2. Self-similarity. Take a single word w that occurs in n different sentences, and average the cosine over all pairs of its occurrences:

That is structurally identical to the anisotropy baseline — average pairwise cosine — just restricted to occurrences of one word. Which means Chapter 1's identity applies: SelfSim_ℓ(w) = (n‖μ̂_w‖² − 1)/(n−1), where μ̂_w is the mean of the unit-normalised occurrence vectors. A word whose occurrences all point the same way has a long mean and a self-similarity near 1; a word whose meaning is entirely determined by context has a short mean and a low one.

3. Intra-sentence similarity. Within one sentence, take the sentence vector to be the mean of its word vectors, then average each word's cosine to it:

4. Maximum explainable variance. Collect a word's occurrence vectors into a matrix and run PCA. MEV(w) is the fraction of variance the first principal component explains — the answer to "how much of this word's contextual behaviour could a single static vector capture?"

Note the different shapes of these four, because it is easy to conflate them. Measurement 1 averages over all words in all contexts — one number per layer. Measurements 2 and 4 are per-word, computed over that word's occurrences — they answer "how stable is this word." Measurement 3 is per-sentence, computed over the tokens inside it — it answers "how much do the words in a sentence converge on one representation." Same operation, three different populations, and the population is what gives each number its meaning. A common error is to compute intra-sentence similarity, find it high, and report it as evidence of anisotropy — it is not, until you subtract the baseline, because words in a sentence and words drawn at random both inherit the same cone.

The correction that makes all of it meaningful

Here is the methodological hinge. Every one of measurements 2, 3, 4 is a cosine similarity, and Chapter 0 showed that cosines in an anisotropic space are inflated by an amount that has nothing to do with the quantity you meant to measure. So each must be reported relative to the layer's own baseline:

Worked example: the same word, two layers. Take the word run, appearing in four sentences, and measure at two depths of the same model.

layer 12 — the last layerthe six pairwise cosines among "run"'s four occurrence vectors:
  0.991  0.988  0.983  0.990  0.986  0.984
  sum = 5.922      SelfSim = 5.922 / 6 = 0.987

anisotropy baseline for this layer  = 0.985
adjusted self-similarity            = 0.987 - 0.985 = 0.002

layer 1 — near the bottomthe six pairwise cosines among the same four occurrences:
  0.81  0.77  0.79  0.76  0.80  0.75
  sum = 4.68       SelfSim = 4.68 / 6 = 0.780

anisotropy baseline for this layer  = 0.250
adjusted self-similarity            = 0.780 - 0.250 = 0.530

This is the reason you should be suspicious of any paper, blog post, or dashboard that reports a raw cosine between contextual embeddings without saying what the baseline was. A cosine of 0.85 could be an extraordinary match or a below-average one, and the number alone cannot tell you which.

The correction reorders words — it is not a rescaling

A tempting misreading: "subtracting a per-layer constant is a monotone transform, so it moves all the numbers down together and the ranking is safe." That is true within a layer and false the moment you compare across layers — which is exactly what anyone selecting a layer, or building a per-layer feature, is doing. Here is the flip, in five rows of arithmetic.

five words, raw self-similarity measured at whichever layer each was measuredword      layer   raw SelfSim   layer baseline   adjusted
------------------------------------------------------------------
mortgage    11        0.962          0.955          +0.007
crane        6        0.804          0.410          +0.394
the         12        0.978          0.981          -0.003
photosynth.  6        0.851          0.410          +0.441
bat          9        0.905          0.760          +0.145

RAW ranking (most self-similar first):
   the 0.978 > mortgage 0.962 > bat 0.905 > photosynthesis 0.851 > crane 0.804

ADJUSTED ranking:
   photosynthesis 0.441 > crane 0.394 > bat 0.145 > mortgage 0.007 > the -0.003

# the ordering is very nearly REVERSED. Every conclusion you would draw
# from the raw column about which words carry stable identity is wrong.

The reversal is not a coincidence, it is structural. Raw self-similarity is dominated by the layer's baseline, and the baseline rises with depth, so the raw column is mostly ranking which layer you measured at. The word with the most stable identity in this table, photosynthesis, has one sense and appears in one kind of context; it scores fifth from the top on the raw number purely because it was measured at layer 6.

Mean-pooling multiplies the anisotropy — derived and checked

Nearly every sentence-embedding pipeline averages token vectors. That operation has a large, predictable, and almost never mentioned effect on the geometry, and Chapter 2's coherent-versus-incoherent argument gives it to us in three lines.

Model each token vector as a shared part plus an idiosyncratic part, h_i = c + g_i, with ‖c‖ = C, ‖g_i‖ = G, and the g_i independent across tokens. Average m of them for a sentence:

The shared part comes through untouched — averaging m copies of c gives c. The idiosyncratic part is an average of m independent random vectors, so its norm is G/√m, not G. Pooling shrinks the signal and leaves the noise floor exactly where it was. Feed that into the same expansion Chapter 0 used and the anisotropy baseline for pooled vectors is:

Take a model whose single-token residual is half the size of its shared component, G/C = 0.5, and read off what pooling does. The right column is a simulation on 768-dimensional vectors built to exactly that recipe, 2,000 sentences each:

Two practical consequences follow immediately. First, if your chunks have variable length, your anisotropy is variable too, and long chunks get a systematic similarity bonus to everything — which is precisely the "corr(u1, token_count) = 0.78" line you will meet in Chapter 9's diagnostic report. Second, centring helps here more than almost anywhere else, because c is exactly what mean-pooling preserved perfectly and μ is exactly an estimate of c.

Finding one: it gets worse with depth

Across ELMo, BERT and GPT-2, the anisotropy baseline rises monotonically through the network. The input layer is close to isotropic — that layer is a lookup table, non-contextual, and it never went through the residual accumulation of Chapter 2's mechanism three. Every layer above it is anisotropic, and the last layer is dramatically so.

GPT-2 is the extreme case: in its last layer, two words sampled uniformly at random have an average cosine similarity of approximately 0.99. Effectively, every representation in that layer points the same way.

Switch to self-similarity (raw) and look at GPT-2's curve. It dips through the middle of the network and then climbs back to nearly 1.0 at the top — which, read naively, says the last layer has recovered a stable word identity. Now switch to adjusted. The climb disappears entirely; the curve keeps falling. The apparent recovery was the anisotropy baseline rising underneath it.

Finding two: upper layers are more context-specific

With the correction applied, adjusted self-similarity decreases monotonically with depth in all three models. The higher you go, the less a word's representation is determined by which word it is and the more by where it sits.

That is what "contextualised" was supposed to mean, and it is satisfying that the measurement confirms it. But the same measurement gives a much less obvious result one step further in.

Finding three: stopwords are the most contextual words

Which words have the lowest adjusted self-similarity — the most context-dependent representations? Not the famous polysemes like bank or bat. Ethayarajh found that among the least self-similar words are stopwords: the, of, and, to.

The reason, once you see it, is inevitable. A word's representation is built out of its own identity plus its context. Bank carries a lot of its own identity — a handful of distinct senses, each with a stable signature. The carries almost none. It appears in every context, contributes nearly nothing semantically, and so its representation is close to a pure readout of whatever surrounds it. A word with no meaning of its own has nothing to be self-similar about.

Finding four: contextual is not "static plus noise"

The most consequential negative result in the paper. If contextual representations were just a static vector with contextual jitter added, then a word's first principal component would explain most of the variance in its occurrences. It does not: on average, less than 5% of the variance in a word's contextualised representations is explained by its first principal component.

Get a feel for how small that is. If a word occurs 100 times, the centred occurrence matrix has at most 99 non-zero eigenvalues. Perfectly uniform variance across all of them would give each component 1/99 = 1.01%. A measured 4% is four times uniform — a real dominant direction exists — but it is nowhere near the 60–90% you would need to say "this word has an embedding, plus noise."

And yet: Ethayarajh also showed that if you do take that first principal component from lower layers and use it as a static embedding, it outperforms GloVe and FastText on many static-embedding benchmarks. Both facts are true at once. The static component is a small fraction of what is going on, and it is still a very good static embedding.

There is one arithmetic subtlety to keep straight when you compute MEV yourself, and it is the same finite-sample effect Chapter 1 warned about. A word with m occurrences has a centred occurrence matrix of rank at most m − 1, so the uniform baseline for its top component is 1/(m − 1), not 1/d. For m = 10 occurrences the uniform floor is 11%, so a measured MEV of 12% is nothing. For m = 1,000 the floor is 0.1%, and the same 12% is enormous. Always compare MEV against 1/(m − 1) for that specific word, and never pool words with very different occurrence counts into one average without doing so — otherwise your "average MEV" is mostly a measurement of how often each word appears.

Finding five: the three models are contextual in different ways

Intra-sentence similarity separates the architectures, and the differences are not cosmetic.

Three models, three geometries, all called "contextual embeddings" in the same sentence of a hundred papers. If you swap one for another in a pipeline and your similarity thresholds do not move, you have not been measuring what you thought.

Is the anisotropy doing something useful?

It is worth resisting the assumption that anisotropy is purely pathological. Ethayarajh raises the possibility that it is a byproduct of context-specificity — that a model which must make representations depend heavily on context may unavoidably produce a geometry like this, and that the two are linked rather than one being a defect on top of the other.

There is a plausible mechanism. If upper layers are computing something closer to "what should be predicted next given everything" than "what does this token mean," then all positions in all sentences are computing answers to the same kind of question, and answers to the same question live in the same region. The cone might be the model's answer space, not a bug in its coordinate system.

This matters practically, because it predicts that aggressive isotropy fixes on the last layer will destroy something. Which is exactly what Chapter 7 finds.

Worked example: self-similarity two ways

Because self-similarity is an average pairwise cosine, Chapter 1's identity computes it in one pass instead of n². Worth verifying once on numbers you can check, because on a real corpus a word can occur a hundred thousand times and the quadratic route is not available.

Take a word with four occurrences whose (unit-normalised) representations sit at angles 20°, 25°, 65° and 70° in a two-dimensional slice — two occurrences in one sense, two in another.

route one — all six pairs(20,25) -> cos  5 = 0.99619      (25,65) -> cos 40 = 0.76604
(20,65) -> cos 45 = 0.70711      (25,70) -> cos 45 = 0.70711
(20,70) -> cos 50 = 0.64279      (65,70) -> cos  5 = 0.99619

sum = 4.81543          SelfSim = 4.81543 / 6 = 0.802572

route two — one mean, one dot productangle   cos        sin
 20     0.93969    0.34202
 25     0.90631    0.42262
 65     0.42262    0.90631
 70     0.34202    0.93969
sum     2.61064    2.61064
mu_hat  0.65266    0.65266

||mu_hat||^2 = 0.65266^2 + 0.65266^2 = 0.425965 + 0.425965 = 0.851930

SelfSim = (n * ||mu_hat||^2 - 1) / (n - 1)
        = (4 * 0.851930 - 1) / 3 = 2.407720 / 3 = 0.802573

Six cosines versus one mean and one dot product, agreeing to six decimals. On a hundred thousand occurrences the first route is 5×10⁹ cosines and the second is still one mean and one dot product.

Computing your own baseline — the code

The baseline is layer-specific and corpus-specific, so the published numbers are not yours. Fortunately it is the cheapest measurement in this lesson.

layer-wise anisotropy baseline for any HF modelimport torch, numpy as np

def layer_baselines(model, tok, sentences):
    """Returns one anisotropy baseline per layer, using the Chapter 1 identity."""
    sums, counts = None, 0
    for s in sentences:
        enc = tok(s, return_tensors="pt", truncation=True)
        out = model(**enc, output_hidden_states=True).hidden_states   # tuple of (1, T, d)
        if sums is None:
            sums = [torch.zeros(h.shape[-1]) for h in out]
        for L, h in enumerate(out):
            v = h[0]                                             # (T, d)
            v = v / v.norm(dim=-1, keepdim=True)                    # unit vectors
            sums[L] += v.sum(0).detach()
        counts += out[0].shape[1]
    return [float((counts * (m / counts).pow(2).sum() - 1) / (counts - 1))
            for m in sums]

Accumulate one running sum of unit vectors per layer, then square its length. One pass over a few thousand sentences, and you have the number that makes every future cosine on that model interpretable. Store it beside the model, and never report a similarity without it again.

The static-embedding recipe, in five lines

Ethayarajh's fourth finding is also a small free win, and it takes about an hour to build. Take the occurrences of a word in a corpus, collect their layer-ℓ representations, and take the first principal component:

a static embedding table distilled from a contextual modelocc = collect_occurrences(word, corpus, layer=5)      # (m, d) — a LOW layer, not the last
X   = occ - occ.mean(0)
u1  = np.linalg.svd(X, full_matrices=False)[2][0]        # (d,) the dominant direction
# fix the sign so the component points with the majority of occurrences
static[word] = u1 * np.sign((X @ u1).sum())

Two details carry the result. Use a low layer, because upper layers are more context-specific and less self-similar, so their first component captures less. And fix the sign, because a principal component is only defined up to sign and a library will hand you either one — a sign flip turns a word into its own antonym as far as cosine is concerned, and this bug is silent.

So which layer should you actually pool from?

Everything in this chapter converges on one decision that every sentence-embedding pipeline has to make and that most make by default. The anatomy gives us the reasoning, and the reasoning contradicts the default.

Three forces are in tension as you move up the network.

Two of the three say do not use the last layer, and the third says do not use the first. That is the whole argument for the middle, and it is why the standard strong baseline in this literature is first-last averaging — pool the token vectors of layer 0 and layer L and average the two — rather than taking the final layer alone.

It is worth seeing why that particular combination works, because it is not obvious. Layer 0 is the input embedding table: non-contextual, but also barely anisotropic, since it has never been through the residual accumulation of Chapter 2's mechanism three. The last layer is highly contextual and highly anisotropic. Averaging them is a crude but effective trade: you get some contextualisation, and you dilute the cone with a layer that does not have one.

the pooling decisions, ranked by what they cost you1. WHICH LAYER    — biggest effect, costs one line
     last layer only ......... worst geometry, most specialised
     first-last average ...... the standard strong baseline
     a middle layer .......... often best; sweep 6..9 of 12 on your own eval
     layer 0 only ............ isotropic but non-contextual — a static embedding

2. WHICH TOKENS   — second biggest
     [CLS] ................... only meaningful if something TRAINED it. In a plain
                               MLM checkpoint the CLS vector was trained for next-
                               sentence prediction or nothing at all
     mean over tokens ........ the default. See Chapter 4's pooling derivation:
                               it multiplies the anisotropy by the token count
     mean over content words . lower m, so a wider usable band — but stopwords are
                               the MOST context-saturated tokens (finding three),
                               so you are discarding sentence-level signal
     max over tokens ......... rarely better; unstable under paraphrase

3. NORMALISE?     — no effect on cosine whatsoever (Chapter 0). Do it for storage
                    hygiene if you like, but do not expect a quality change.

A note on subwords, which quietly breaks all of this

Every definition in this chapter said "the representation of the word at position i." Real tokenizers do not give you words. They give you subword pieces, and a rare word — exactly the kind whose geometry you most want to inspect — is the one most likely to be split into several.

That creates two distinct problems and they need different answers.

What the baseline costs to compute, and how often

One practical objection to "always report the anisotropy-adjusted number" is that it sounds like extra infrastructure. It is not, and the arithmetic is worth doing so you can say so with confidence.

What to do with this, today

Chapter 5: BERT-flow — Learning the Map

All-but-the-Top asks: which directions are bad? Then it deletes them. It is surgery.

Li et al. in 2020 asked a different question: what distribution do I wish my embeddings had? Then they learned an invertible function that turns the one they have into the one they want. It is not surgery, it is a change of coordinates — nothing is deleted, everything is rearranged.

The reframing is worth dwelling on, because it changes what counts as a solution. Under the surgical view, the fix is a projection and you must decide how many dimensions to sacrifice. Under the distributional view, the fix is a bijection and you sacrifice nothing — the transform is reversible in principle, and the only question is whether the target distribution is a good one.

The three complaints

The paper's diagnosis of BERT's sentence-embedding space has three parts, and only the first is the one we have been discussing.

Complaint 3 is the one that motivates the whole method, and it is easy to miss. Suppose you embed two sentences and they land on either side of a hole. The straight-line distance between them is small; the amount of trained, meaningful space between them is zero. Any similarity you compute across a hole is an extrapolation the model never had reason to get right. A distribution with no holes is not a cosmetic preference — it is the condition under which distance means something everywhere.

Why a standard Gaussian is the right target

They pick N(0, I) — the standard multivariate normal. Three properties, each answering one complaint.

There is also a practical reason: we know exactly how to build invertible maps to a Gaussian, because the machine-learning community spent five years doing it for image generation. The tool is called a normalising flow.

Normalising flows from zero

A flow is an invertible function f between two spaces, used to move a probability distribution from one to the other. The single fact you need is the change-of-variables formula, and it is easier to derive in one dimension than to look up.

Suppose z is a random variable with density p_Z, and u = f(z) for an invertible, differentiable f. Probability mass is conserved: the chance of landing in a tiny interval around u must equal the chance of landing in the corresponding interval around z.

The derivative factor is a bookkeeping correction: if f stretches a region, the density there must thin out to compensate. In d dimensions, "how much does f stretch volume" is the absolute determinant of the Jacobian:

Check it on a case you already know. Let z ~ N(0,1) and f(z) = 2z + 3, so u should be N(3, 4). Then f⁻¹(u) = (u−3)/2 and the derivative is 1/2. The formula predicts:

At u = 3: φ(0) = 1/√(2π) = 0.398942, halved gives 0.199471. And the true N(3,4) density at its mean is 1/(σ√(2π)) = 1/(2 × 2.506628) = 0.199471. At u = 5: φ(1) = 0.241971, halved gives 0.120985; the true value is 0.199471 × e^−0.5 = 0.199471 × 0.606531 = 0.120985. Exact both times.

The objective

Now the setup. Let u = BERT(sentence) be a frozen sentence embedding, drawn from whatever unknown distribution BERT induces. We want an invertible f_φ mapping a standard Gaussian z into that distribution, so that f_φ⁻¹(u) is the Gaussianised version of our embedding. Train φ by maximum likelihood:

Two terms, and they are in tension — which is the whole reason the method does not collapse.

The first term wants every mapped point to land where the Gaussian's density is high, that is, near the origin. On its own, it would be maximised by the constant map f⁻¹(u) = 0 for all u. Everything to one point, infinite likelihood, zero information. That is representation collapse, arriving by a different road than in contrastive learning but landing in the same ditch.

The second term forbids it. If f⁻¹ squashes volume, its Jacobian determinant goes to zero and log|det| goes to −∞. The log-determinant is an explicit, unavoidable anti-collapse penalty, and it is not a heuristic — it is required by the change-of-variables formula for the first term to be a valid density at all.

Note what is not in that objective: labels. No sentence pairs, no similarity annotations, no NLI data. You need only a pile of sentences from the domain you care about, which means you can fit the flow on your own unlabelled corpus. BERT stays frozen throughout — only φ is trained.

Coupling layers: invertible networks that are actually cheap

The catch in the objective is det(∂f⁻¹/∂u). For a general 768×768 Jacobian, that determinant costs O(d³) per example and the Jacobian itself costs a backward pass per output. Unusable.

The Glow family, which BERT-flow uses, solves this with additive coupling layers. Split the vector in half, u = [u_a; u_b], and define:

where m is any neural network at all — it does not need to be invertible, monotone, or anything else. Two consequences fall out immediately.

It inverts in closed form. Given v, recover u_a = v_a, then compute m(v_a) — the same network, forward — and set u_b = v_b − m(v_a). One forward pass, no solving.

Its log-determinant is zero. Write the Jacobian in block form:

Lower block-triangular with identity blocks on the diagonal, so its determinant is the product of the diagonal blocks' determinants: 1 × 1 = 1. The log-determinant term is exactly 0, free of charge, for any m whatsoever.

Worked example in four dimensions. Let u = (1, 2, 3, 4), split as u_a = (1, 2), u_b = (3, 4), and let m(a₁, a₂) = (0.5a₁ + a₂, a₁ − 0.5a₂).

one additive coupling layer, forward and backforward:
  m(1, 2) = (0.5*1 + 2,  1 - 0.5*2) = (2.5, 0.0)
  v_a = (1, 2)
  v_b = (3, 4) + (2.5, 0.0) = (5.5, 4.0)
  v   = (1, 2, 5.5, 4.0)

inverse (uses m in the FORWARD direction only):
  u_a = v_a = (1, 2)
  m(1, 2) = (2.5, 0.0)
  u_b = (5.5, 4.0) - (2.5, 0.0) = (3.0, 4.0)
  u   = (1, 2, 3, 4)          # exact round trip

jacobian dv/du =
  [ 1     0    0  0 ]
  [ 0     1    0  0 ]
  [ 0.5   1.0  1  0 ]         # row from dm_1/du_a
  [ 1.0  -0.5  0  1 ]         # row from dm_2/du_a
  lower triangular, unit diagonal  ->  det = 1  ->  log|det| = 0

One coupling layer leaves half the coordinates untouched, so you stack many and permute the coordinates between them. After enough layers every coordinate has been transformed as a function of every other, the composite is still exactly invertible, and the total log-determinant is still zero.

Worked example: the collapse, and the term that stops it

The claim that the log-determinant prevents collapse is easy to assert and easy to verify, because the simplest possible flow makes the whole thing a one-variable calculus problem.

Let the flow be a single scalar rescaling: f⁻¹(u) = su for some s > 0, in d dimensions. Take one data point at radius r, so ‖u‖ = r. The two terms are:

So the objective, dropping the constant, is:

Without the second term, dL/ds = −sr², which is negative for every s > 0. Gradient ascent drives s to 0 and every data point to the origin. The "likelihood" grows without bound, and the flow has learned to throw the data away.

With the second term, differentiate and set to zero:

Read the answer: the optimum maps a point of radius r to radius sr = √d. And √d is exactly the typical radius of a standard d-dimensional Gaussian, since E‖z‖² = d. The flow has done the correct thing — matched the data's scale to the target's scale — and it did so because the log-determinant made shrinking expensive.

Put numbers on it. Take d = 768 and a typical sentence-embedding norm r = 12:

the objective at three values of soptimum:  s = sqrt(768)/12 = 27.7128/12 = 2.3094
   L = -(2.3094^2 * 144)/2 + 768 * ln(2.3094)
     = -(5.3333 * 144)/2 + 768 * 0.83698
     = -384.00 + 642.80 = +258.80

no change:  s = 1
   L = -(1 * 144)/2 + 768 * ln(1) = -72.00 + 0 = -72.00

collapsing:  s = 0.1
   L = -(0.01 * 144)/2 + 768 * ln(0.1) = -0.72 - 1768.39 = -1769.11

# drop the log-det term and the third row becomes the BEST of the three (-0.72),
# with s -> 0 better still. That is the collapse, in one column of numbers.

Why "no holes" is a theorem, not a hope

Complaint 3 was that BERT's distribution has low-density gaps in which distance is meaningless. Why does mapping to a Gaussian fix that, rather than merely rearranging where the gaps are?

Because the Gaussian density is log-concave: log p is a concave function, since log p(z) = −‖z‖²/2 + const and −‖z‖² is concave. And log-concavity has exactly the property you want, in one line. For any two points x, y and any λ in [0,1]:

The density along the straight line between any two points never dips below the lower of the two endpoints. There is no interior valley, anywhere, ever. If both of your sentences landed in trained, high-density regions, then every point between them is at least as well-supported as the worse of the two — so interpolation and distance are meaningful along the whole segment.

Is a Gaussian not itself a thin shell?

An objection that should occur to you, because it occurs to everyone who has met concentration of measure: in high dimensions a standard Gaussian does not look like a filled ball at all. Its samples nearly all sit on a thin spherical shell. Have we just swapped one degenerate geometry for another?

Get the numbers first. For z ~ N(0, I_d), the squared norm ‖z‖² is a chi-squared variable with d degrees of freedom, so it has mean d and variance 2d. Propagate that through the square root with a first-order expansion — ‖z‖ = √d · √(1 + δ) ≈ √d (1 + δ/2) with δ = (‖z‖² − d)/d — and the variance of the norm is (1/4)·(2d/d) = 1/2:

At d = 768: every sample has norm 27.71 ± 0.71, a shell of relative thickness 2.6%. So yes, it is a shell.

But a shell is not a cone, and the distinction is the whole point. The shell is centred on the origin and covers every direction equally. Two independent samples from it have cosine 0 ± 1/√d, which is the isotropic ideal from Chapter 1 — measured, not assumed. A cone is a small patch of a shell. Concentration of measure controls the radial coordinate; anisotropy is entirely about the angular ones, and the Gaussian is perfect in exactly the coordinates that matter.

What no invertible map can do

Before we admire the method, we should fence in what it can possibly achieve, because the fence is exactly where Chapter 6's result lands.

A bijection cannot separate points that coincide. If f is invertible and BERT put two semantically different sentences at the same vector, then f maps them to the same vector too — that is what "invertible" means, read in reverse. No flow, however deep, can recover a distinction the encoder never made. The method's ceiling is the information already in the frozen embedding.

The objective never mentions a pair. Look again at what is being maximised: the likelihood of individual embeddings under a density. There is no term containing two sentences, no similarity label, no notion of "these should be close." Any improvement in semantic similarity is a side effect of reshaping the marginal distribution. That is a remarkable thing for it to achieve, and it is also a hard ceiling — a method with no pairwise signal cannot repair a pairwise ordering error that the encoder made deliberately.

The stack, and how it is trained

How many layers do you need? Count what one layer can reach. A single additive coupling makes each of the d/2 coordinates in u_b a function of the d/2 coordinates in u_a, and leaves u_a alone entirely. So after one layer, half the coordinates have not been transformed at all. Add a permutation and a second coupling and the previously untouched half is now transformed as a function of the previously transformed half — which means every output coordinate depends on every input coordinate after two couplings, and the composite has a full Jacobian rather than a block-triangular one.

Two is the minimum for full mixing; it is not enough for expressiveness, because the composite of two additive couplings is still a fairly constrained family. In practice BERT-flow-style stacks use on the order of 8 to 16, and the diminishing return is easy to reason about: each additional layer adds capacity but also adds an inference-time forward pass on the query path, and the whole method's justification is that BERT is frozen and the fix should be cheap. If you find yourself needing 64 coupling layers, you are training a second encoder and should consider whether Chapter 7's answer — retrain the first one — is not simply better value.

One coupling layer transforms half the coordinates. The Glow recipe interleaves it with a permutation so that over several layers every coordinate is transformed as a function of every other:

the flow, in the shape it is actually builtclass Coupling(nn.Module):
    def __init__(self, d, hidden=1024):
        super().__init__()
        self.half = d // 2
        self.m = nn.Sequential(nn.Linear(self.half, hidden), nn.ReLU(),
                               nn.Linear(hidden, d - self.half))

    def forward(self, u):                 # u: (B, d)
        a, b = u[:, :self.half], u[:, self.half:]
        return torch.cat([a, b + self.m(a)], dim=-1)   # log|det| = 0

    def inverse(self, v):
        a, b = v[:, :self.half], v[:, self.half:]
        return torch.cat([a, b - self.m(a)], dim=-1)   # same m, forward only

# flow = [Coupling, Permute, Coupling, Permute, ...] x K
# total log|det| = 0 for additive couplings; permutations contribute 0 too
# so the objective reduces to  max  E[ log p_Z( flow.inverse(u) ) ]
# which for a standard Gaussian is  min  E[ ||flow.inverse(u)||^2 / 2 ]
#
# WAIT. That is the collapse objective again — and it does not collapse only
# because an additive coupling CANNOT shrink volume. The constraint is
# architectural rather than a penalty: det = 1 is enforced by construction,
# so the map is volume-preserving and s -> 0 is not in the hypothesis class.

The permutation between couplings deserves a word, because Glow's contribution was to make it learnable. A fixed permutation — reverse the coordinates, say — works and has determinant ±1, so log|det| is 0 and it is free. Glow replaces it with an invertible 1×1 convolution, which in this setting is just a learned d×d invertible matrix applied between couplings. That is strictly more expressive, and its log-determinant is not zero — it is log|det W|, computed once per layer from an LU factorisation, which is cheap because it does not depend on the batch. So a Glow stack's total log-determinant is the sum of those matrix terms and nothing else, since the additive couplings contribute zero. If you build the stack with fixed permutations instead, the total is exactly zero and the objective genuinely reduces to minimising ‖f⁻¹(u)‖².

That comment is the detail most write-ups skip, and it is the whole trick. With additive couplings the log-determinant is identically zero, so it disappears from the loss — but it disappears because the architecture has already made collapse impossible, not because it was safe to ignore. If you swap in an affine coupling (v_b = u_b·exp(g(u_a)) + m(u_a)), the determinant is no longer 1, the log-det term comes back into the loss, and forgetting it will collapse your flow within a few hundred steps.

Worked example: the affine coupling, and the bug you will write. Same four-dimensional vector as before, u = (1, 2, 3, 4) split into u_a = (1,2) and u_b = (3,4). Keep m(a₁,a₂) = (0.5a₁ + a₂, a₁ − 0.5a₂) and add a scale network g(a₁,a₂) = (0.5a₁ − 0.2a₂, 0.3a₁ − 0.1a₂).

affine coupling — the determinant is no longer freeg(1, 2) = (0.5*1 - 0.2*2,  0.3*1 - 0.1*2) = (0.1, 0.1)
exp(g)  = (1.105171, 1.105171)
m(1, 2) = (2.5, 0.0)

v_a = (1, 2)
v_b = (3, 4) * (1.105171, 1.105171) + (2.5, 0.0)
    = (3.315513, 4.420684) + (2.5, 0.0) = (5.815513, 4.420684)

jacobian: the b-block diagonal is exp(g), so
  log|det| = g_1 + g_2 = 0.1 + 0.1 = 0.2          # NOT zero any more
  det      = e^0.2 = 1.221403                     # this block EXPANDS volume by 22%

inverse, still one forward pass:
  u_b = (v_b - m(v_a)) / exp(g(v_a))
      = (5.815513 - 2.5, 4.420684 - 0.0) / 1.105171
      = (3.315513, 4.420684) / 1.105171 = (3.000000, 4.000000)   # exact

Now watch the failure. Suppose you copied the additive-coupling training loop, where log|det| was identically zero and could be dropped, and you kept that shortcut. Your loss is then just ‖f⁻¹(u)‖²/2, and the fastest way to reduce it is for g to become large and negative, because the inverse divides by exp(g).

what the missing term costs, over 8 coupling layers at d = 768suppose g drifts to -1 on every coordinate of every layer:

  scale per layer  = e^-1 = 0.367879
  after 8 layers   = e^-8 = 0.000335              # every embedding is now ~1/3000 its size
  so ||z||^2 / 2   shrinks by e^-16 = 1.1e-7      # the logged loss goes to ~0 — it looks GREAT

  the log-det you dropped:
  per layer  = sum over 384 scaled coords of (-1) = -384
  8 layers   = -3072

  so the TRUE objective went DOWN by 3072 nats while your logged loss went down too.
  the flow has learned the identity-with-a-shrink and destroyed all your geometry,
  and the training curve looks perfect the whole way.

Why insist on invertibility at all?

A fair challenge: if the goal is "make the distribution look Gaussian", why not fit a Gaussian mixture to the embeddings and map each point through its own component, or train any old network with a distribution-matching loss? Why pay for the constraint that the map be a bijection?

Three reasons, in ascending order of importance.

It makes the objective computable. Maximum likelihood needs a density, and change of variables only gives you a density if the map is invertible and differentiable. Drop invertibility and there is no likelihood to maximise; you are back to adversarial or moment-matching objectives, which are far harder to train and offer no guarantee about what happened to individual points.

It guarantees nothing is destroyed. A non-invertible map can improve every distributional statistic you measure by simply collapsing the distinctions you were not measuring. That is the failure mode of "make it look Gaussian" as a loss: mapping everything to fresh Gaussian noise satisfies it perfectly and is useless. Invertibility rules that out by construction, in the same structural way that additive couplings rule out volume collapse. Note the pattern — twice now, the method has been made safe by an architectural constraint rather than a penalty term.

It makes the result auditable. Because f is invertible you can push a Gaussianised vector back and check you recover the original embedding to machine precision. That single round-trip assertion catches essentially every implementation bug in the transform, and no non-invertible method offers an equivalent.

The results, and the number that stung

Evaluated on the standard semantic textual similarity suite — STS-B and STS 2012 through 2016 — measured by Spearman correlation with human similarity ratings. Using the comparison table later compiled in the SimCSE paper, which puts all these methods on identical footing:

What it costs you

And then the awkward question, which the next chapter answers with a single line of linear algebra: given that the target is a Gaussian, and a Gaussian is fully specified by its mean and covariance, how much of that ten-point gain actually required a nonlinear map?

Chapter 6: Whitening — the Closed Form

Chapter 5 ended on an uncomfortable observation. BERT-flow's target is a standard Gaussian. A Gaussian is completely determined by two things: its mean and its covariance. So if the only defects in your embedding distribution are that its mean is wrong and its covariance is wrong, the minimal transform that fixes them is affine — a translation and a matrix multiply. No network, no training, no gradient steps.

Su et al. (2021) tried it. It works about as well as the flow, it takes one call to a linear-algebra routine, and it can reduce your dimension by a factor of three on the way. This chapter derives it completely and then does a 2×2 example in full arithmetic.

Deriving the transform

We want a map from raw embeddings x to new vectors z such that the transformed cloud has zero mean and identity covariance. Write it as an affine map with a matrix W (using the row-vector convention, so vectors are rows and matrices act on the right — this matches the paper and matches numpy):

The mean is handled: subtracting μ makes the mean of z zero for any W. Now impose the covariance condition. Since z is a linear function of the centred vectors, its covariance is:

and we want that to equal the identity:

Now solve for W. Σ is a covariance matrix, so it is symmetric and positive semidefinite, which means the spectral theorem hands us an eigendecomposition with an orthogonal U and a non-negative diagonal Λ:

Guess W = UΛ^−1/2, where Λ^−1/2 is the diagonal matrix of 1/√λ_k. Substitute and watch it collapse:

One subtlety worth naming: the solution is not unique. If W works then so does WR for any orthogonal R, because R^T(W^TΣW)R = R^TIR = I. There is a whole rotational family of valid whitening matrices. For our purposes this does not matter at all, because cosine similarity is invariant under a shared rotation — all members of the family give identical similarities. If you ever need a canonical choice, the symmetric one W = Σ^−1/2 = UΛ^−1/2U^T is the one that moves each point the least.

The geometry: three moves

Whitening versus All-but-the-Top

Both methods reshape the covariance spectrum. Put them side by side and the relationship is exact:

So ABTT is the crude special case: instead of asking "how much should this direction count?", it answers "zero" or "unchanged." Whitening answers "exactly one, for everything." That sounds strictly better, and Chapter 7 explains the circumstances under which it is not.

Worked example: five sentences, two dimensions, everything by hand

Five sentence embeddings. Two dimensions, so we can draw them and so the eigendecomposition is a quadratic formula rather than a library call.

The symptom first. Their pairwise cosines, all ten of them:

before whitening — the ten cosines|s1|=12.80625  |s2|=10.63015  |s3|=8.48528  |s4|=5.00000  |s5|=2.23607

s1.s2 = 136 / (12.80625 x 10.63015) = 136 / 136.1345 = 0.999012
s1.s3 = 108 / (12.80625 x  8.48528) = 108 / 108.6647 = 0.993884
s1.s4 =  64 / (12.80625 x  5.00000) =  64 /  64.0312 = 0.999512
s1.s5 =  28 / (12.80625 x  2.23607) =  28 /  28.6357 = 0.977800
s2.s3 =  90 / (10.63015 x  8.48528) =  90 /  90.1998 = 0.997784
s2.s4 =  53 / (10.63015 x  5.00000) =  53 /  53.1507 = 0.997164
s2.s5 =  23 / (10.63015 x  2.23607) =  23 /  23.7697 = 0.967622
s3.s4 =  42 / ( 8.48528 x  5.00000) =  42 /  42.4264 = 0.989949
s3.s5 =  18 / ( 8.48528 x  2.23607) =  18 /  18.9737 = 0.948681
s4.s5 =  11 / ( 5.00000 x  2.23607) =  11 /  11.1803 = 0.983870

sum = 9.855278        A = 0.985528

Average pairwise cosine 0.9855, in a two-dimensional space, from five points chosen without any conspiracy. This is the whole disease reproduced in a toy: the cloud simply does not surround the origin.

Step 1: the mean.

Centred vectors: (4,3), (2,2), (0,1), (−2,−2), (−4,−4). They sum to (0,0), as they must.

Step 2: the covariance. Three numbers, each a sum of five products divided by five.

the 2x2 covariance, entry by entrySxx = (16 + 4 + 0 + 4 + 16) / 5 = 40 / 5 = 8.0
Syy = ( 9 + 4 + 1 + 4 + 16) / 5 = 34 / 5 = 6.8
Sxy = (12 + 4 + 0 + 4 + 16) / 5 = 36 / 5 = 7.2

Sigma = [ 8.0  7.2 ]
        [ 7.2  6.8 ]

correlation = 7.2 / sqrt(8.0 x 6.8) = 7.2 / 7.37564 = 0.97619    # a very thin cigar

Step 3: eigenvalues. For a 2×2 symmetric matrix the characteristic polynomial is λ² − (trace)λ + (det) = 0.

Check: λ₁λ₂ = 14.62496 × 0.17504 = 2.5600 = det. Correct.

And look at the ratio: λ₁/λ₂ = 83.55. In two dimensions the cloud is 83 times more spread along one axis than the other. That is the shape pathology, in a number.

Step 4: eigenvectors. For λ₁, solve (Σ − λ₁I)u = 0 using the first row:

Take u = (1, 0.920133), whose length is √(1 + 0.846645) = √1.846645 = 1.358913. Normalise:

u₂ is u₁ rotated by 90 degrees, which is forced in two dimensions by orthogonality. Check u₁: 0.541522 + 0.458477 = 0.999999.

Step 5: assemble W. The scale factors are the reciprocal square roots:

W = UΛ^−1/2 scales each column of U by the corresponding factor:

building Wcolumn 1 = u1 x 0.261489 = ( 0.735882, 0.677109) x 0.261489 = ( 0.192425, 0.177056)
column 2 = u2 x 2.390162 = (-0.677109, 0.735882) x 2.390162 = (-1.618401, 1.758878)

W = [  0.192425  -1.618401 ]
    [  0.177056   1.758878 ]

Verify W^TΣW = I — one entry longhand, because this is the kind of thing that is worth doing once in your life.

checking the top-left entrySigma . col1 = ( 8.0*0.192425 + 7.2*0.177056 ,  7.2*0.192425 + 6.8*0.177056 )
             = ( 1.539398 + 1.274806 , 1.385459 + 1.203983 )
             = ( 2.814204 , 2.589442 )

col1 . that  = 0.192425*2.814204 + 0.177056*2.589442
             = 0.541523 + 0.458477
             = 1.000000          # exactly 1, to six decimals

off-diagonal:
Sigma . col2 = ( 8.0*(-1.618401) + 7.2*1.758878 , 7.2*(-1.618401) + 6.8*1.758878 )
             = ( -12.947210 + 12.663920 , -11.652489 + 11.960369 )
             = ( -0.283290 , 0.307880 )
col1 . that  = 0.192425*(-0.283290) + 0.177056*0.307880 = -0.054512 + 0.054512 = 0.000000

Step 6: transform the five points. Applying z = (x−μ)W is the same as computing each coordinate as a projection divided by a standard deviation: z_k = (u_k · (x−μ)) / √λ_k. Longhand for s₁:

transforming s1 = (10, 8), centred to (4, 3)u1 . (4,3) =  0.735882*4 + 0.677109*3 =  2.943529 + 2.031328 =  4.974857
u2 . (4,3) = -0.677109*4 + 0.735882*3 = -2.708438 + 2.207647 = -0.500791

z1 = ( 4.974857 / 3.824259 , -0.500791 / 0.418382 ) = ( 1.30087 , -1.19697 )

Step 7: verify the output covariance really is the identity. This is the payoff, so compute it rather than trusting the algebra.

covariance of the whitened pointsSxx = (1.30087^2 + 0.73896^2 + 0.17706^2 + 0.73896^2 + 1.47792^2) / 5
    = (1.692262 + 0.546062 + 0.031350 + 0.546062 + 2.184247) / 5
    = 4.999983 / 5 = 0.999997

Syy = (1.19697^2 + 0.28095^2 + 1.75888^2 + 0.28095^2 + 0.56191^2) / 5
    = (1.432737 + 0.078933 + 3.093659 + 0.078933 + 0.315743) / 5
    = 5.000005 / 5 = 1.000001

Sxy = (1.30087*(-1.19697) + 0.73896*0.28095 + 0.17706*1.75888
       + (-0.73896)*(-0.28095) + (-1.47792)*(-0.56191)) / 5
    = (-1.557105 + 0.207612 + 0.311408 + 0.207612 + 0.830449) / 5
    = -0.000024 / 5 = -0.000005

Identity covariance, to five decimal places, done by hand. The cigar is a ball.

Step 8: the new cosines. All ten again, on the whitened points:

Three of those after-values are exactly ±1, and the reason is worth naming rather than glossing. Look at the centred column: s₂ → (2,2), s₄ → (−2,−2), s₅ → (−4,−4). Those three lie on a single line through the origin — s₅ = 2s₄ = −2s₂. Whitening is a linear map applied after centring, and a linear map sends a line through the origin to a line through the origin, so those three points must still be scalar multiples of each other afterwards, and their cosines must be exactly ±1. The raw column does not show this because those cosines were computed on the uncentred vectors (8,7), (4,3), (2,1), which are not collinear at all. Centring created the exact ±1; whitening merely could not destroy it. And the recurring ±0.44721 is 1/√5 — the sort of clean surd a five-point hand-built example produces. Do not expect either pattern at d = 768.

Step from 1 to 2 — the rotation — and watch the average cosine readout not move at all. Rotation is an orthogonal transform; it preserves every dot product and every norm, so it cannot change a single cosine. It is there purely so that step 3 can be a diagonal scaling. Skipping it and scaling in the original coordinates would not whiten anything, because the cloud's axes are not the coordinate axes.

Whitening-k: the free lunch

Keep only the first k columns of U. Then W is d×k, the output has k dimensions, and it is still exactly white in those k directions.

Su et al. report that on 768-dimensional BERT embeddings, keeping k = 256 typically matches or beats full whitening — while cutting storage and retrieval cost by three. Getting better results from throwing information away deserves an explanation, and there is a good one.

Whitening multiplies direction k by 1/√λ_k. The smallest eigenvalues get the largest amplification. And the smallest eigenvalues are exactly the directions where your covariance estimate is least reliable — a direction with tiny true variance is dominated by sampling noise, and whitening then blows that noise up to the same scale as your signal.

Quantify it on the worked example. The amplification factors were 0.261 for the big direction and 2.390 for the small one — a relative reamplification of 2.390/0.261 = 9.14×, which is exactly √(λ₁/λ₂) = √83.55. In a 768-dimensional embedding cloud the condition number routinely runs into the thousands, so the ratio is 30× to 100×. You are handing a megaphone to your noisiest measurement.

The implementation

whitening, fitted and appliedimport numpy as np

def whiten_fit(V, k=None, alpha=0.0):
    """V: (n, d) corpus embeddings.  Returns (mu, W) with W of shape (d, k)."""
    mu  = V.mean(axis=0, keepdims=True)              # (1, d)
    X   = V - mu
    cov = (X.T @ X) / X.shape[0]                     # (d, d)
    if alpha:                                        # shrinkage toward a scaled identity
        cov = (1 - alpha) * cov + alpha * np.trace(cov) / cov.shape[0] * np.eye(cov.shape[0])
    lam, U = np.linalg.eigh(cov)                     # eigh: ascending, orthonormal U
    lam, U = lam[::-1], U[:, ::-1]                  # descending
    lam = np.maximum(lam, 1e-12)                     # guard the numerical zeros
    W = U / np.sqrt(lam)                             # broadcasts over columns == U @ diag(lam**-0.5)
    return mu, (W[:, :k] if k else W)

def whiten_apply(V, mu, W):
    Z = (V - mu) @ W                                 # (n, d) @ (d, k) -> (n, k)
    return Z / np.linalg.norm(Z, axis=1, keepdims=True)

Three implementation notes that matter more than they look.

Use eigh, not eig. A covariance matrix is symmetric; eigh exploits that, returns real eigenvalues in ascending order, and guarantees an orthonormal U. eig returns complex types and a non-orthogonal basis, and the whole derivation above assumed U^TU = I.

Clamp the eigenvalues. If n < d, the covariance is rank-deficient by construction — you cannot estimate d(d+1)/2 covariance entries from fewer than d points. Some λ_k will be zero to machine precision, and 1/√0 is an inf that will silently poison every embedding. Either clamp, or use whitening-k with k < n, or apply shrinkage. Ideally all three.

Fit once, apply everywhere. As with ABTT: μ and W are fitted parameters of your retrieval system. Store them, version them, and apply the identical transform to queries and documents alike. Refitting on new data means re-indexing everything.

Whitening when you have fewer vectors than dimensions

The clamp in the code above hides a structural fact worth stating plainly: if n < d, the covariance is guaranteed to be singular, and no amount of numerical care changes that. The centred matrix X is n×d, so its rank is at most n − 1, and X^TX therefore has at least d − n + 1 eigenvalues that are exactly zero in exact arithmetic and around 10⁻¹⁶ in floating point. Whitening asks you to divide by their square roots.

Do the arithmetic on what that means. A zero eigenvalue that lands at 10⁻¹⁶ because of rounding gets amplified by 1/√(10⁻¹⁶) = 10⁸. Your embedding now has a coordinate a hundred million times larger than every real one, and it is pure floating-point noise. Every cosine you compute afterwards is a cosine between two noise vectors.

There is also a genuine efficiency win available in this regime, and it is the same trick that makes PCA on images tractable. Do not build the d×d matrix at all — build the n×n one:

the dual formulation — never form a d x d matrix# X is (n, d) with n << d.  X^T X is (d, d) and mostly zeros in its spectrum.
# X X^T is (n, n) and has the SAME nonzero eigenvalues.

G = X @ X.T                       # (n, n)   — for n = 2000, d = 4096 this is 16x smaller
lam, W = np.linalg.eigh(G)        # n eigenvalues, n eigenvectors
lam, W = lam[::-1], W[:, ::-1]

# recover the right singular vectors (the PCs in d-space) from the left ones
U = X.T @ W / np.sqrt(lam + 1e-12)      # (d, n), orthonormal columns

# the identity that makes this work:
#   X = P S Q^T  ->  X X^T = P S^2 P^T   and   X^T X = Q S^2 Q^T
#   so Q = X^T P / S, and the nonzero spectra are identical.

The other reason to reduce k: it is free money

The quality argument for whitening-k is the noise argument above. The engineering argument is independent of it and often the one that actually gets the change approved.

what 768 -> 256 actually buys, for 10 million vectorsstorage, float32:
  768 dims : 10e6 * 768 * 4 bytes = 30.7 GB
  256 dims : 10e6 * 256 * 4 bytes = 10.2 GB     -> 20.5 GB saved

brute-force scan of the whole index, per query:
  768 dims : 10e6 * 768  = 7.68e9 multiply-adds
  256 dims : 10e6 * 256  = 2.56e9 multiply-adds  -> 3x faster, exactly

memory bandwidth, which is what actually limits a scan:
  the same 3x, and bandwidth is the binding constraint on almost every
  vector search that is not already using a quantised index.

and the quality change, per Su et al.: zero to slightly positive.

What the paper reported

Su et al. found that whitening matches BERT-flow's semantic-textual-similarity results almost exactly — 66.28 versus 66.55 average Spearman for BERT-base in the compiled comparison — while requiring no training whatsoever, and that whitening-256 preserved or improved that result with a threefold reduction in storage and retrieval cost.

Chapter 7: When the Fixes Hurt

Four chapters of increasingly elegant machinery, all of it real, all of it published, all of it reproducible. This chapter is the counterweight, because every one of these methods has a regime where applying it makes your system worse, and two of those regimes are common enough that you will hit them.

Then, at the end, the reason none of this is the state of the art any more.

Failure one: the direction you deleted was the answer

The diagnostics from Chapter 1 detect concentration of variance. They cannot detect whether that concentration is an artefact or the signal. Here is a case where every measurement screams "anisotropic" and the correct action is to do nothing.

Take a corpus with two well-separated topics. Model it in three dimensions: cluster A centred at (+2, 0, 0), cluster B at (−2, 0, 0), each with isotropic spread σ = 0.5 in every direction, equal numbers of each.

the covariance of two separated clustersglobal mean = (0, 0, 0)                      # the clusters cancel

Sigma_xx = 2^2 + 0.5^2 = 4.25                 # separation^2 + within-cluster variance
Sigma_yy = Sigma_zz = 0.5^2 = 0.25

eigenvalues: (4.25, 0.25, 0.25)      trace = 4.75
top component explains 4.25 / 4.75 = 89.5% of the variance

p = (0.8947, 0.0526, 0.0526)
H = 0.8947*0.11125 + 2*(0.0526*2.94504) = 0.09954 + 0.30982 = 0.40935
effective rank = e^0.40935 = 1.506   out of 3

Every number says the same thing: 89.5% of the variance in one direction, effective rank 1.5 out of 3. By the criteria of Chapters 1 and 3, this is a textbook narrow cone and you should remove the top principal component.

Do it, and the two clusters land on top of each other. The x coordinate was the topic label. A classifier that was at 100% is now at chance. You have deleted the only thing the embedding knew.

This is also the mechanism behind a real finding in the literature. Cai et al. (ICLR 2021) showed that contextual embedding spaces are globally anisotropic but locally isotropic — inside individual clusters, the geometry looks fine. The global cone is substantially a story about a handful of well-separated clusters, not a single uniform drift. Rajaee and Pilehvar followed with cluster-based isotropy enhancement: subtract the cluster means and remove the dominant directions within each cluster, so the between-cluster structure survives.

When the top component is artefact and signal at once

The advice "check what u₁ correlates with" implies a clean verdict. Often there is not one, because the same direction can be both — and this case is common enough to deserve its own answer.

The archetype: your top component correlates 0.7 with document length and separates your two document types, because one type happens to be longer. Delete it and you lose the type distinction; keep it and every long document floats to the top of every result list. Neither of the two prescriptions is right.

Failure two: the transform is fitted, and your queries are not from that distribution

μ and W are estimated from a sample. That makes them parameters, and parameters can be wrong for the data they are applied to.

Concretely: you fit whitening on your document corpus — product descriptions, say — and at serve time you apply it to user queries, which are short, keyword-shaped, and grammatically nothing like a product description. Their true mean is somewhere else.

Write the query mean as μ_q = μ_c + δ. Subtracting the corpus mean leaves every query carrying the same residual offset δ. If typical centred vectors have norm 1 and ‖δ‖ = 0.3, then by Chapter 1's identity the queries have an induced anisotropy among themselves of:

That is not catastrophic on its own. What is worse is the asymmetry: the documents are properly centred and the queries are not, so every query gets a systematic bonus toward documents that happen to align with δ. The same handful of documents drifts up every result list, for a reason that has nothing to do with any query. That failure looks like a relevance bug and is a preprocessing bug.

Failure two and a half: your corpus is a mixture, and the mean belongs to nobody

Failure two assumed two populations — queries and documents — and one fitted transform. The nastier version has two populations inside the corpus you fitted on, in unequal proportions, and it is the common case: 90% product descriptions and 10% support-ticket text, or 95% English and 5% everything else.

The arithmetic is short and the conclusion is unpleasant. Let the two sub-populations be centred at a and b with ‖a − b‖ = 1, in proportions 0.9 and 0.1, with within-population residuals of typical norm 1. The fitted global mean is

Subtract it. Population A's vectors now carry a residual offset of a − μ = 0.1(a − b), of norm 0.1 — negligible. Population B's carry b − μ = 0.9(b − a), of norm 0.9 — nearly as large as their own spread. Apply Chapter 1's identity to population B alone:

The fix is the one Rajaee and Pilehvar arrived at from a different direction: cluster first, then centre within each cluster. In this framing it is obvious — a single μ can only be correct for a single population, and you had two. Practical recipe: run k-means with a modest k (say 8 to 32) on a sample, store one μ_k per cluster along with the cluster assignment model, and subtract the mean of whichever cluster a vector lands in. It costs one nearest-centroid lookup per embedding and it is the cheapest fix for the most common failure.

Failure three: the fix is invisible on your actual task

Isotropy calibration is measured with geometric metrics and justified with geometric arguments, and the connection to downstream quality is weaker than the literature's enthusiasm suggests. Follow-up work has repeatedly found that making a space more isotropic by these measures does not reliably improve fine-tuned task performance, and sometimes degrades it.

There is a coherent reason. If you are going to fine-tune a task head on top of these embeddings, the head can learn any linear reweighting of directions it likes — including the whitening transform itself. Applying it beforehand gives the head nothing it could not have learned, and costs it the ability not to, in the directions where the raw scaling was correct.

Failure four: it may not even be about the vocabulary

Chapter 2's mechanism two — rare tokens pushed into a common direction by the softmax — is a satisfying story that predicts frequency-aligned principal components, which are observed. But Godey et al. (2024) found the same cone geometry inside the self-attention blocks of transformers trained without any vocabulary softmax at all, and on non-language data.

If anisotropy shows up in models that do not have the mechanism, the mechanism is not the whole cause. Their conclusion is that anisotropy is closer to an intrinsic property of trained self-attention than a corpus artefact — which, if right, means the honest expectation for post-hoc fixes is "reshapes a symptom" rather than "removes a defect."

The modern answer: build the geometry in during training

Every method in Chapters 3 through 6 is a change of coordinates applied to a frozen encoder. None of them can create a distinction the encoder never made. If two sentences landed at the same point, no invertible map will separate them — that is what invertible means.

Contrastive learning attacks the problem one level up: change what the encoder puts in the space.

SimCSE (Gao, Yao and Chen, 2021) is the canonical instance, and its unsupervised version is almost absurdly simple. Take a sentence. Run it through the encoder twice. Because dropout is on, the two passes produce slightly different vectors — and that is your positive pair. Every other sentence in the batch is a negative. The loss is InfoNCE:

where sim is cosine similarity and τ is a temperature, typically 0.05. The supervised version swaps the dropout trick for natural-language-inference pairs: entailment pairs are positives, contradiction pairs are explicit hard negatives.

Alignment and uniformity — and why the loss contains an anisotropy penalty

Wang and Isola (2020) decomposed what a contrastive objective is actually asking for, into two measurable quantities on the unit hypersphere:

Alignment wants positive pairs close. Uniformity wants everything else spread evenly over the sphere. Both lower is better, and they pull in opposite directions — the same two-term tension as the flow's likelihood-versus-log-determinant in Chapter 5.

That uniformity formula looks arbitrary until you evaluate it, at which point it turns into Chapter 1's A with a constant in front. Do the substitution. For unit vectors, ‖f(x) − f(y)‖² = 2 − 2cos, so:

Now evaluate the expectation in the two extreme cases, with actual numbers.

uniformity, evaluatedPERFECTLY UNIFORM on the sphere in d = 768:
  cos is centred at 0 with sd 1/sqrt(768) = 0.0361  (Chapter 1)
  for a roughly normal cos, E[exp(4 cos)] = exp(16 * 0.0361^2 / 2) = exp(0.01042) = 1.01047
  L_uniform = -4 + 0.01042 = -3.990          # essentially the floor, -4

GPT-2's LAST LAYER, where A = 0.99 and the spread is tiny:
  E[exp(4 cos)] = exp(4 * 0.99) = exp(3.96)
  L_uniform = -4 + 3.96 = -0.040            # essentially the ceiling, 0

and in general, whenever the cosine distribution is tight around A:
  L_uniform = -4 (1 - A)

Now the connection that closes the loop on this whole lesson. SimCSE's authors analyse the asymptotic form of the loss's denominator and show it reduces to minimising the sum of all entries of the Gram matrix of the (unit-norm) embeddings:

Look at that expression next to Chapter 1. It is exactly the double sum we expanded to derive the identity — and we showed it equals n²‖μ̂‖², which is n(n−1)A + n. Minimising it is minimising the average pairwise cosine.

And the results are not close:

Ten points from post-hoc geometry; another ten from training the geometry in. Every strong text embedder shipped since — the E5 family, GTE, BGE, the commercial APIs — is contrastively trained, and none of them ship a whitening step, because they do not need one.

Put the same table on a cost axis, because the ranking changes when you do:

Where the repulsion comes from, in gradients

Chapter 2 diagnosed the missing ingredient precisely: cross-entropy pushes each non-target output embedding away from the hidden state h, but it never pushes two output embeddings away from each other. There is no term in the loss whose gradient separates two representations. Contrastive learning adds exactly that term, and you can see it in one differentiation.

Write the InfoNCE loss for anchor i with similarities s_ij = cos(h_i, h_j⁺):

Differentiate with respect to any similarity. The log-sum-exp derivative is a softmax, so with p_ij = softmax_j(s_ij/τ):

The same shape as Chapter 2's logit gradient — but the objects are different. Here s_ij is a similarity between two representations, so a positive gradient on s_ij means gradient descent actively decreases the similarity between h_i and h_j. That is real, direct repulsion between the embeddings, which the language-modelling loss never had.

Worked example: where the repulsion actually goes. Take a batch of three with τ = 0.05. Anchor 1's cosines against its own positive and two negatives:

easy negatives — almost no gradientcosines   = (0.90, 0.40, 0.35)
logits    = (18.0,  8.0,  7.0)          # divided by tau = 0.05
shift by 18 -> (0, -10, -11)
exps      = (1.000000, 0.0000454, 0.0000167)
sum       = 1.0000621
p         = (0.999938, 0.0000454, 0.0000167)

gradient pushing away from negative 2 = 0.0000454
gradient pushing away from negative 3 = 0.0000167

one hard negative — the gradient relocates entirelycosines   = (0.90, 0.85, 0.35)
logits    = (18.0, 17.0,  7.0)
shift by 18 -> (0, -1, -11)
exps      = (1.000000, 0.367879, 0.0000167)
sum       = 1.367896
p         = (0.731050, 0.268938, 0.0000122)

gradient pushing away from the hard negative = 0.268938

ratio to the easy case: 0.268938 / 0.0000454 = 5,920x

Temperature is the knob that sets how hard "hard" is

That 5,920× ratio was computed at τ = 0.05, and it is worth seeing how violently it depends on that number, because τ is the hyperparameter people copy without understanding and then wonder why their contrastive run either collapses or does nothing.

Recompute both scenarios — the easy batch (0.90, 0.40, 0.35) and the hard batch (0.90, 0.85, 0.35) — across a range of temperatures. The quantity that matters is p_ij, which is the gradient on that negative:

Read the second column down its length and notice it is not monotone. It rises to a peak around τ = 0.2 and then falls. That is the whole story of temperature in one column, and it has a clean explanation: as τ shrinks, the softmax concentrates, which first sharpens the focus onto the hard negative (good) and then concentrates so hard onto the positive that even the hard negative stops receiving anything (bad). The positive's own probability in the last row is 0.993 — the loss is essentially zero and there is nothing left to learn from.

The evaluation that is allowed to decide

Every section of this chapter ends in "measure it on your own data", so it is worth being concrete about what that measurement has to look like. Geometry transforms produce small effects on the metrics that matter, and small effects are exactly where sloppy evaluation goes wrong.

Where the positive pairs come from when you have no labels

The obvious objection to "just train contrastively" is that you do not have labelled similar pairs. Usually you do not need them.

Note the last row against the gradient arithmetic. Supervised SimCSE beats the unsupervised version by roughly five points, and the mechanism is visible in the numbers above: a contradiction pair is a sentence that is lexically close and semantically opposite, which is precisely a hard negative, which is precisely where the gradient concentrates.

The same arithmetic explains why batch size is not a mere throughput parameter here. In-batch negatives means a batch of N gives each anchor N − 1 negatives, and the temperature table showed that only the negatives within about τ of the positive contribute anything. So the useful quantity is not N, it is the expected number of hard negatives in the batch — N times the probability that a random other sentence is confusably similar. If that probability is 1 in 500, a batch of 64 contains a hard negative about 12% of the time and 88% of your steps are learning almost nothing about uniformity. A batch of 512 contains one about 64% of the time. That is the entire reason contrastive training papers report batch sizes in the thousands and reach for gradient caching and cross-device negatives to get there, and it is also why a single well-chosen hard negative per example — the NLI contradiction — is worth more than a tenfold batch increase.

Failure six: correcting something that was already corrected

The last failure is the one that arrives after this lesson has been read by more than one person on the team, and it is entirely self-inflicted.

Modern embedding APIs and open-weights embedding models are contrastively trained. Their uniformity term, as the previous section showed, already minimises the anisotropy baseline directly. If you then apply centring and whitening on top, you are correcting a defect that is not there, and you pay the full price of the transform for it: a fitted artefact, a re-index requirement, and the noise amplification of Chapter 6.

The arithmetic of the harm is straightforward. Chapter 6 showed that whitening amplifies direction k by 1/√λ_k. On a well-trained contrastive embedder the spectrum is already fairly flat by design, so λ_k is close to uniform for the leading directions and small only in the tail — which is exactly the tail that is estimation noise. You have removed no real anisotropy and amplified the noisiest directions in the space.

So when do the post-hoc fixes still earn their keep?

Often, actually. "Contrastive training won" is a statement about what to do if you can train. Most people, most of the time, cannot.

Chapter 8: Anisotropy vs Superposition

Two communities have spent the last several years staring at the same vectors and describing what they see in incompatible language.

The embedding-geometry literature — every paper in this lesson — says: the representations occupy a narrow cone, the space is badly used, and the fix is to reshape the distribution.

The interpretability literature says: the model packs far more features into the space than it has dimensions, giving each a nearly orthogonal direction, and the fix is to recover the dictionary.

Those sound like they cannot both be right. A space full of nearly orthogonal directions is the definition of well-used; a narrow cone is the definition of badly used. This chapter resolves that, and the resolution turns out to be practically useful, not merely tidy.

Superposition, from zero

Start with the problem it solves. Suppose a layer has d = 768 dimensions, and the world it is modelling has far more than 768 things worth tracking — is this text code, is it French, is the subject plural, is the tone sarcastic, is there a date nearby, is the current token inside a quotation. Tens of thousands of such features, plausibly.

If each feature needed its own orthogonal direction, the layer could hold 768 of them and no more. But features are sparse: at any moment almost all of them are off. Sarcasm and Fortran syntax and Japanese honorifics are not simultaneously active in the same token position.

Superposition is the strategy that exploits this: assign each feature a direction, allow the directions to be merely nearly orthogonal, and accept the resulting interference — because with sparse activation, two interfering features are rarely on at once, so the interference rarely costs anything. Elhage et al.'s Toy Models of Superposition demonstrated this happening in networks small enough to analyse completely, and showed it appears exactly when features are sparse and there are more of them than dimensions.

How much capacity does near-orthogonality buy?

This is quantifiable, and the numbers make the trade-off concrete. For two independent uniformly random unit vectors in R^d, Chapter 1 gave us the standard deviation of their cosine as 1/√d, and the distribution concentrates sharply:

So if you scatter N random directions, the expected number of pairs violating an interference budget ε is roughly N²exp(−dε²/2). Setting that to order 1 and solving gives the capacity:

That bound is the Johnson–Lindenstrauss lemma seen from the other side. JL is usually quoted as "you can embed N points into O(log N / ε²) dimensions with distortion ε" — invert it and you get "d dimensions hold exp(dε²) points at distortion ε", which is the same exponential. The interpretability framing and the dimensionality-reduction framing are the same theorem, and the only difference is which variable you solve for.

Two things to read off. First, the growth is savage — going from tolerating 0.2 interference to 0.4 buys ten orders of magnitude. Superposition is not a marginal trick; it is the difference between hundreds and trillions. Second, the top row is a warning: demanding exact orthogonality is worse than demanding almost-orthogonality with a tiny budget, because the exact constraint caps you at d while the loose one is exponential. The interior of that table is where trained networks live.

Worked example: three features in two dimensions

Capacity tables are abstract. The smallest case where superposition actually happens is three features in a two-dimensional space, and every number in it fits on one line, so let us do it completely.

The model needs three directions in R². The best it can do — the arrangement that maximises the smallest angle — is to space them 120 degrees apart:

Every pair has cosine −0.5. Check one: u_A·u_B = (1)(−0.5) + (0)(0.866025) = −0.5. So the interference budget is a whopping ε = 0.5 — ten times what you would tolerate in a capacity calculation. This should be a disaster. It is not, and the reason is the single most elegant idea in the toy-models paper.

The model reads feature B out with ReLU(u_B · x). Run three cases.

case 1 — only A fires, with magnitude 1x = 1 * u_A = (1, 0)

readouts before ReLU:  u_A.x =  1.0
                       u_B.x = (-0.5)(1) + (0.866025)(0)  = -0.5
                       u_C.x = (-0.5)(1) + (-0.866025)(0) = -0.5

after ReLU:            (1.0, 0.0, 0.0)      # EXACT. Interference of -0.5 was clipped to 0.

That is the trick. The interference is negative, and a ReLU discards negative values for free. A budget of ε = 0.5 costs nothing at all when only one feature is on, because the false readings land below zero and never survive the nonlinearity. This is why a network will happily choose a highly non-orthogonal arrangement: with sparse, non-negative activations, negative interference is not interference.

case 2 — A and B both fire, each with magnitude 1x = u_A + u_B = (1 - 0.5, 0 + 0.866025) = (0.5, 0.866025)

readouts:  u_A.x = (1)(0.5) + (0)(0.866025)             =  0.500
           u_B.x = (-0.5)(0.5) + (0.866025)(0.866025)   = -0.25 + 0.75 =  0.500
           u_C.x = (-0.5)(0.5) + (-0.866025)(0.866025)  = -0.25 - 0.75 = -1.000

after ReLU: (0.5, 0.5, 0.0)     # both features DETECTED — but at HALF magnitude

Two active features and the readout is still qualitatively right: A and B are on, C is off, and C's false reading has gone more negative, not less. What breaks is the magnitude — each feature reads 0.5 instead of 1.0, because each has absorbed −0.5 of interference from the other. A downstream consumer that only asks "is this feature present?" is fine. One that uses the magnitude is now wrong by a factor of two.

case 3 — all three fire, each with magnitude 1x = u_A + u_B + u_C = (1 - 0.5 - 0.5, 0 + 0.866025 - 0.866025) = (0, 0)

readouts:  (0, 0, 0)            # TOTAL information loss. Three features on reads
                                # identically to zero features on.

The second reconciliation: non-negativity alone produces a cone

This one is a derivation, and it is the most underappreciated fact in the area.

Suppose a representation's coordinates are all non-negative — because they came through a ReLU, or because they are the codes of a sparse autoencoder, which are non-negative by construction. Then any two such vectors have a non-negative dot product, so their cosine is at least 0. The cone is forced by the sign constraint before any learning happens.

How large is the effect? Take two vectors with independent, identically distributed non-negative coordinates, with mean m and second moment s = E[x²]. The expected dot product factorises because the coordinates are independent:

and the squared norm concentrates around its mean:

so, for large d, the cosine between two completely unrelated such vectors is:

Now plug in the two cases you actually meet.

Post-ReLU Gaussian pre-activations. If the pre-activation is standard normal, then x = max(0, g). Its mean is the half-normal mean, m = 1/√(2π), and its second moment is E[max(0,g)²] = ½E[g²] = ½, since exactly half the mass is kept. Therefore:

Exponentially distributed activations (a reasonable model for sparse positive codes). For Exp(1), m = 1 and s = 2, so cos ≈ 1/2.

This is a genuinely load-bearing correction. Much of the anisotropy literature works on residual-stream or attention outputs, which are sign-symmetric and where zero is the right null. But the moment you measure the geometry of MLP hidden states, sparse-autoencoder codes, or anything after a ReLU or GELU, you must compare against the sign-constrained baseline instead.

Sparsity rescues the sign constraint: the q/2 floor

The 1/π result assumed dense post-ReLU activations — every coordinate is a fresh half-normal, roughly half of them positive. Sparse codes are a different regime, and running the same formula on them gives a result that is genuinely reassuring.

Model a sparse non-negative code: each coordinate is zero with probability 1 − q, and when it fires it is drawn from Exp(1). Compute the two moments the formula needs.

The second moment of an exponential with rate 1 is 2, because Var = 1 and mean = 1, so E[x²] = 1 + 1 = 2. Substituting into cos ≈ m²/s:

The amplification law: a nearly-isotropic dictionary still makes a cone

Now the quantitative version of the reconciliation, which is the most useful thing in this chapter. We claimed the dictionary and the activations can have different geometries. They can — but they are not independent, and the relationship between them has a factor in it that nobody expects.

An activation is a sparse non-negative combination of dictionary directions, x = ∑_i a_iu_i, with the same coefficient statistics as above: mean m, second moment s, each coordinate independent. Take two such activations and compute the expected dot product, using independence of the coefficients:

where ū is the mean of the N dictionary directions. Chapter 1's identity says ‖ū‖² is the dictionary's own anisotropy, A_dict. The squared norm of a single activation is dominated by the diagonal terms, N·s. Assemble the cosine, write K = qN for the expected number of features active at once, and after cancelling:

Read the factor K. The dictionary's anisotropy is multiplied by the number of simultaneously active features before it reaches the activations. Sum many vectors that each lean slightly the same way and the lean adds coherently — the exact same coherent-versus-incoherent argument as Chapter 2's residual stream, now applied across features instead of across layers.

Put numbers on it. Simulate N = 4,096 dictionary directions in d = 512, activation rate q = 0.02 (so K = 82 features active per vector), and tilt the dictionary by adding a small common direction:

Three consequences follow, and they are all actionable.

What a cone does to a sparse autoencoder, in practice

The theory above has a very concrete operational consequence for anyone who trains one of these.

If the activation cloud has a large mean offset μ, the autoencoder must reconstruct that offset for every single input. Without a decoder bias, the cheapest way to do that is for a handful of latents to fire on everything and encode the mean between them. Those latents are then uninterpretable by construction — they are not features, they are the offset — and the capacity they consume is gone. Worse, the sparsity penalty is fighting them, so the loss lands in an unhappy compromise where the mean is only partly reconstructed and the reconstruction error is large everywhere.

That is why the standard form subtracts b_dec on the way in and adds it back on the way out, and why initialising b_dec to the empirical mean of the activations is standard practice rather than a refinement. It removes the offset from the problem before the dictionary is asked to represent anything.

Estimate what it saves. If ‖μ‖ is comparable to a typical activation norm — which Chapter 9's diagnostic says is exactly the regime — then roughly half of the squared norm the autoencoder is asked to reconstruct is the offset. With a dictionary of N latents at activation rate q, reconstructing a constant direction with sparse codes needs at least one always-on latent, so you have paid one latent plus a permanently violated sparsity constraint, forever, for a vector you could have stored in d floats.

What interference costs, in numbers

The capacity table above only counted directions. To understand when superposition is actually a good deal, you have to price the interference, and that is a short calculation.

Suppose feature B has direction u_B, and the model reads B out by taking a dot product with u_B. Now feature A fires with magnitude 1 while B is off. The readout for B picks up:

A false reading of ε on a feature that is not there. With one interferer and a detection threshold of, say, 0.5, an ε of 0.3 is harmless.

But k features are typically active at once, and their interference contributions have random signs, so they add as a random walk rather than linearly:

Why the two pictures use different tools

The interference calculation also explains why interpretability work reaches for sparse autoencoders rather than PCA, which is worth spelling out because both produce "directions" and they are not the same kind of object.

That last row is the crux. If a model represents three co-occurring features, PCA is structurally obliged to return a mixture — nothing in its objective rewards separating them, and orthogonality actively forbids returning three directions that are not perpendicular. Anisotropy analysis is built on PCA, so its "top component" is a variance-weighted blend of whatever the model does most. Reading that blend as "a thing the model represents" is a category error, and it is a common one.

What a cone does to cosine-based feature analysis

One last practical consequence, because interpretability work uses cosine similarity constantly — between activations and feature directions, between features across layers, between a probe direction and a candidate concept — and every one of those numbers sits on the same inflated baseline.

Take the standard move of asking "which dictionary feature is this activation most aligned with?" by ranking cos(x, u_i). If the activations sit in a cone around a direction c, then every activation has a large component along c, and the feature whose direction happens to be closest to c wins every single query, regardless of content. You will find one feature that appears to fire on everything and conclude it is a general-purpose feature. It is the mean.

Where the two lenses disagree: the top principal component

Now the practical fork, and it is the same fork as Chapter 7's failure one, stated in the other vocabulary.

Both are legitimate, and which one is right for you is decided by what you plan to do next. If the vectors are going into a cosine similarity, the anisotropy lens is the operative one — an always-on feature is noise for retrieval no matter how meaningful it is, because it is present in every query and every document equally. If the vectors are going into an analysis of what the model computed, the superposition lens is operative, and deleting the top component means deleting evidence.

Where they agree: subtract the mean

Here is the pleasing convergence. A sparse autoencoder, in its standard form, does not encode its input directly. It first subtracts a learned decoder bias:

That subtraction of b_dec is there because the activation cloud has a large mean offset and no useful feature should have to spend capacity re-encoding it. In practice b_dec is often initialised to the empirical mean of the activations — which is, exactly and literally, step 1 of All-but-the-Top.

One thing not to do

Do not whiten before running a sparse autoencoder, or before any dictionary-learning analysis.

Whitening sets every direction's variance to 1. That is exactly the information a dictionary-learning method needs: a feature that fires strongly and often occupies a high-variance direction, and that magnitude is signal about the feature's importance. Whitening erases the distinction between a direction the model uses constantly and one it barely touches, and it does so by construction, since flattening the spectrum is its definition.

The same warning applies in reverse to a subtler case. If you whiten and then interpret the resulting principal axes as "the model's features", you are interpreting an artefact of the transform — after whitening, the covariance is the identity, every orthogonal basis is equally principal, and the axes your library happens to return are arbitrary. There is nothing to interpret.

The formal statement is the rotational ambiguity from Chapter 6, read as a warning instead of a convenience. Whitening's solution is unique only up to W → WR for orthogonal R, and after whitening every direction has variance exactly 1, so no direction is distinguished by the data any more. Rerun the eigendecomposition on a machine with a different BLAS and you will get a different basis, equally valid, with different "features" in it. If your interpretability result changes when the library version changes, that is the tell.

A shared measurement protocol

The two lenses do not have to talk past each other. The amplification law says exactly which quantities each community should report so the other can use them, and it is a short list.

both lenses, one functiondef geometry_report(acts, sae=None):
    """acts: (n, d) layer activations.  sae: an optional trained autoencoder."""
    def aniso(M):
        U = M / np.linalg.norm(M, axis=1, keepdims=True)
        n = M.shape[0]; mu = U.mean(0)
        return float((n * (mu @ mu) - 1) / (n - 1))

    rep = {"A_act": aniso(acts)}

    # the correct null depends on the SIGN structure, not on the model
    if (acts < 0).mean() < 0.01:                    # effectively non-negative
        q = float((acts > 1e-6).mean())
        rep["null"] = q / 2 if q < 0.3 else 1 / np.pi
        rep["K"]    = q * acts.shape[1]           # active features per vector
    else:
        rep["null"] = 0.0                            # sign-symmetric: zero is right

    if sae is not None:
        rep["A_dict"] = aniso(sae.W_dec)              # (N, d) — the FEATURE directions
        rep["K"]      = float((sae.encode(acts) > 0).sum(1).mean())
        # the amplification law, used as a consistency check
        k = rep["K"] * rep["A_dict"] / 2
        rep["A_act_predicted"] = k / (1 + k)
    return rep

A synthesis you can act on

Chapter 9: Diagnose Your Own Embeddings

Everything in this lesson now becomes one script and one decision. The script takes a sample of your actual embeddings and returns five numbers. The decision takes those five numbers and tells you what, if anything, to do.

The script

the whole diagnosisimport numpy as np

def diagnose(V):
    """V: (n, d) embeddings sampled from YOUR live distribution. n should be >> d."""
    n, d = V.shape
    U = V / np.linalg.norm(V, axis=1, keepdims=True)      # unit vectors

    # 1 — the offset: how far has the centre of mass drifted from the origin?
    mu_hat = U.mean(0)
    offset = float(np.linalg.norm(mu_hat))                # in [0, 1]

    # 2 — the anisotropy baseline, EXACTLY, in one pass (the Chapter 1 identity)
    A = (n * offset**2 - 1) / (n - 1)

    # 3 — the shape, from the CENTRED covariance spectrum
    X   = V - V.mean(0)
    sv  = np.linalg.svd(X, compute_uv=False)
    lam = sv**2 / n                                       # covariance eigenvalues
    p   = lam / lam.sum()
    erank = float(np.exp(-(p * np.log(p + 1e-12)).sum()))     # effective rank

    # 4 — the dominant direction itself, so you can ask what it encodes
    u1   = np.linalg.svd(X, full_matrices=False)[2][0]
    proj = X @ u1

    return dict(n=n, d=d,
                offset_ratio   = offset,
                anisotropy_A   = float(A),
                isotropic_sd   = float(1 / np.sqrt(d)),
                top1_var       = float(p[0]),
                top10_var      = float(p[:10].sum()),
                effective_rank = erank,
                erank_frac     = erank / d,
                u1_projection  = proj)

And the fifth number, which does not fit in the function because it needs your metadata — and which is the one that decides whether any fix is appropriate at all:

the artefact test — do this before you transform anythingrep = diagnose(V)
proj = rep["u1_projection"]

# is the dominant direction an artefact?
print("corr with token count :", np.corrcoef(proj, token_counts)[0, 1])
print("corr with doc length  :", np.corrcoef(proj, char_lengths)[0, 1])

# or is it the signal?
for label in set(labels):
    print(label, proj[labels == label].mean())      # do the classes separate along u1?

The whole diagnosis, worked by hand on six vectors

Before you trust that function on twenty thousand embeddings, run it on something you can check with a pencil. Here is a six-vector, three-dimensional set built to have every pathology this lesson describes, with every number computed longhand. Paste it into your test suite and you will never again wonder whether an implementation bug is masquerading as a finding.

Number one — the offset ratio. Average the six unit vectors, coordinate by coordinate.

the mean of the unit vectorsx: 4*(0.948683) + 0.984732 + 0.942809 = 3.794733 + 1.927541 = 5.722274   -> /6 = 0.953712
y: 0.316228 + 0 - 0.316228 + 0 + 0.123091 - 0.235702 = -0.112611        -> /6 = -0.018769
z: same as y, by the symmetry of the construction                        -> /6 = -0.018769

||mu_hat||^2 = 0.953712^2 + 2*(0.018769^2) = 0.909567 + 0.000704 = 0.910271
offset_ratio = sqrt(0.910271) = 0.954082

Number two — the anisotropy baseline. Straight from the Chapter 1 identity, and worth verifying once against the brute-force route so you believe it forever:

Summing all thirty ordered pairwise cosines directly gives 26.769756, and 26.769756/30 = 0.892325. Identical to six decimals, one pass versus thirty dot products.

Number three — the isotropic scale. 1/√d = 1/√3 = 0.577350. So A = 0.892 is 1.5 isotropic standard deviations up — unremarkable at d = 3, which is exactly why the same A would be a screaming result at d = 768, where the same calculation gives 24.7 standard deviations. The dimension is doing all the interpretive work here, and that is the point of reporting both numbers together.

Number four — the centred spectrum. The raw mean is μ = (3, 0, 0), so the centred vectors are (0,1,0), (0,0,1), (0,−1,0), (0,0,−1), (1,0.5,0.5) and (−1,−0.5,−0.5).

the 3x3 covariance, then its spectrumSxx = (0+0+0+0+1+1)/6           = 2.0/6 = 0.333333
Syy = (1+0+1+0+0.25+0.25)/6     = 2.5/6 = 0.416667
Szz = (0+1+0+1+0.25+0.25)/6     = 2.5/6 = 0.416667
Sxy = (0+0+0+0+0.5+0.5)/6       = 1.0/6 = 0.166667
Sxz = (0+0+0+0+0.5+0.5)/6       = 1.0/6 = 0.166667
Syz = (0+0+0+0+0.25+0.25)/6     = 0.5/6 = 0.083333

trace = 0.333333 + 0.416667 + 0.416667 = 1.166667 = 7/6

eigenvalues (exact fractions, which is why this example was chosen):
  lambda = 2/3, 1/3, 1/6   ->  sum = 7/6  # check against the trace
  p      = 4/7, 2/7, 1/7   =  0.571429, 0.285714, 0.142857

entropy, term by term:
  (4/7) * 0.559616 = 0.319781        # -ln(4/7) = 0.559616
  (2/7) * 1.252763 = 0.357932
  (1/7) * 1.945910 = 0.277987
  H = 0.955700       erank = e^0.955700 = 2.6005   out of 3

Number five — the artefact test. The top eigenvector is u₁ = (1,1,1)/√3 = (0.577350, 0.577350, 0.577350). Project the centred vectors onto it and correlate against the metadata column:

is the dominant direction an artefact?projections:  (0,1,0).u1     =  0.577350
              (0,0,1).u1     =  0.577350
              (0,-1,0).u1    = -0.577350
              (0,0,-1).u1    = -0.577350
              (1,0.5,0.5).u1 =  2 * 0.577350 =  1.154701
              (-1,-0.5,-0.5) = -1.154701

lengths:      10, 10, 10, 10, 14, 6          mean 10

covariance   = (1.154701*4 + (-1.154701)*(-4)) / 6 = 9.237604/6 = 1.539601
sd(proj)     = sqrt((4*0.333333 + 2*1.333333)/6) = sqrt(0.666667) = 0.816497
sd(lengths)  = sqrt((16 + 16)/6) = sqrt(5.333333) = 2.309401

correlation  = 1.539601 / (0.816497 * 2.309401) = 1.539601/1.885618 = +0.8165

the golden test — paste this into your repodef test_diagnose_golden():
    V = np.array([[3,1,0], [3,0,1], [3,-1,0], [3,0,-1],
                  [4,0.5,0.5], [2,-0.5,-0.5]], dtype=float)
    r = diagnose(V)
    assert abs(r["offset_ratio"]   - 0.954082) < 1e-5
    assert abs(r["anisotropy_A"]   - 0.892325) < 1e-5
    assert abs(r["top1_var"]       - 4/7)      < 1e-6
    assert abs(r["effective_rank"] - 2.6005)   < 1e-3
    # the identity that must hold for ANY input: A == offset^2 in the large-n limit,
    # and exactly A = (n*offset^2 - 1)/(n - 1) for every n
    n = V.shape[0]
    assert abs(r["anisotropy_A"] - (n*r["offset_ratio"]**2 - 1)/(n-1)) < 1e-12
    # and the sign-free property: the top eigenvector is only defined up to sign,
    # so always test |corr|, never corr
    assert abs(abs(np.corrcoef(r["u1_projection"], [10,10,10,10,14,6])[0,1])
               - 0.816497) < 1e-5

That last assertion is not decoration. A principal component is defined only up to sign, and different library versions, different BLAS backends, and different machines will hand you either one. A test that asserts a signed correlation will pass on your laptop and fail in continuous integration for reasons that will consume an afternoon.

Reading the report

Here is a real-shaped output and how to read every line of it.

a typical decoder-pooled embedding setn              : 20000
d              : 768
offset_ratio   : 0.870
anisotropy_A   : 0.757
isotropic_sd   : 0.036
top1_var       : 0.41
top10_var      : 0.72
effective_rank : 12.4
erank_frac     : 0.016
corr(u1, token_count) : 0.78

The decision tree

The checklist of things that will bite you

Glossary of every bolded term

The cheat sheet — every formula in one place

The numbers worth remembering

Where to go from here

References

1. Mu, J., Viswanath, P. "All-but-the-Top: Simple and Effective Postprocessing for Word Representations." ICLR 2018 — arXiv:1702.01417. Chapter 3.
2. Ethayarajh, K. "How Contextual are Contextualized Word Representations? Comparing the Geometry of BERT, ELMo, and GPT-2 Embeddings." EMNLP 2019 — arXiv:1909.00512. Chapter 4.
3. Li, B., Zhou, H., He, J., Wang, M., Yang, Y., Li, L. "On the Sentence Embeddings from Pre-trained Language Models" (BERT-flow). EMNLP 2020 — arXiv:2011.05864. Chapter 5.
4. Su, J., Cao, J., Liu, W., Ou, Y. "Whitening Sentence Representations for Better Semantics and Faster Retrieval." 2021 — arXiv:2103.15316. Chapter 6.
5. Gao, J., He, D., Tan, X., Qin, T., Wang, L., Liu, T.-Y. "Representation Degeneration Problem in Training Natural Language Generation Models." ICLR 2019 — arXiv:1907.12009. The rare-token mechanism in Chapter 2.
6. Gao, T., Yao, X., Chen, D. "SimCSE: Simple Contrastive Learning of Sentence Embeddings." EMNLP 2021 — arXiv:2104.08821. Chapter 7, and the source of the comparison table used throughout.
7. Wang, T., Isola, P. "Understanding Contrastive Representation Learning through Alignment and Uniformity on the Hypersphere." ICML 2020 — arXiv:2005.10242. Chapter 7.
8. Arora, S., Li, Y., Liang, Y., Ma, T., Risteski, A. "A Latent Variable Model Approach to PMI-based Word Embeddings." TACL 2016 — arXiv:1502.03520. The partition-function isotropy measure in Chapter 1.
9. Cai, X., Huang, J., Bian, Y., Church, K. "Isotropy in the Contextual Embedding Space: Clusters and Manifolds." ICLR 2021. Global anisotropy, local isotropy — Chapter 7.
10. Rajaee, S., Pilehvar, M. T. "A Cluster-based Approach for Improving Isotropy in Contextual Embedding Space." ACL 2021. Cluster-wise correction — Chapter 7.
11. Rudman, W., Gillman, N., Rayne, T., Eickhoff, C. "IsoScore: Measuring the Uniformity of Embedding Space Utilization." Findings of ACL 2022. The critique of average cosine in Chapter 1.
12. Godey, N., de la Clergerie, É., Sagot, B. "Anisotropy Is Inherent to Self-Attention in Transformers." EACL 2024. Anisotropy without a vocabulary softmax — Chapters 2 and 7.
13. Elhage, N. et al. "Toy Models of Superposition." Anthropic, 2022 — transformer-circuits.pub. Chapter 8.
14. Kingma, D. P., Dhariwal, P. "Glow: Generative Flow with Invertible 1×1 Convolutions." NeurIPS 2018 — arXiv:1807.03039. The flow architecture BERT-flow borrows.
15. Reimers, N., Gurevych, I. "Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks." EMNLP 2019 — arXiv:1908.10084. The setting all the STS numbers live in.