Measure anisotropy in an embedding space

CodingIntermediate⏱ 1 hourRAGanisotropywhiteningsimulationdiagnostics

Theory: Embedding Geometry and Similarity →

Problem

Write a diagnostic for embedding anisotropy.

Measure the mean and standard deviation of the cosine between random unrelated pairs and compare with the isotropic predictions $0$ and $1/\sqrt d$ .
Compute the spectrum of the embedding covariance and the stable rank $\|\Sigma\|_F^2/\|\Sigma\|^2$ .
Apply mean-centring and whitening and re-measure.
Report what each step does to the contrast between matched and unmatched pairs.

Implement it, then check the implementation against something you can verify independently.

Write it

Solve it here8 tests · Python runs in your browser

import numpy as np

def cosine_stats(X, pairs=20000, seed=0):
    """Mean and standard deviation of the cosine between random distinct rows.

X has shape (n, d) and its rows are embeddings. Sample `pairs` random
    pairs i != j with a generator seeded with `seed` and return
    (mean, std) of cos(x_i, x_j) as plain floats.

For an isotropic cloud these should come out near 0 and 1/sqrt(d). A mean
    far from 0 is the "cone effect": every embedding shares a direction, so
    everything looks similar to everything else.
    """
    raise NotImplementedError

def stable_rank(X):
    """Stable rank of the covariance of the rows of X.

With Sigma the covariance of the (mean-centred) rows, return
    ||Sigma||_F^2 / ||Sigma||_op^2. It is d when the spectrum is flat and
    O(1) when one direction dominates, so it measures how many directions the
    cloud genuinely uses.
    """
    raise NotImplementedError

def whiten(X):
    """Mean-centre the rows of X and whiten them.

Return an array of the same shape whose column means are 0 and whose
    covariance is the identity to numerical precision.
    """
    raise NotImplementedError

⌘/Ctrl + ↵ to run · limit 60s · draft saved in this browser

8 tests will run against your code. The first run downloads the Python runtime — several megabytes, once per browser.

Solution

Opening this before you have a real attempt on paper is the fastest way to feel like you understand something you cannot reproduce.

Problem	Field	Type	Difficulty	Time
Three similarities agree on the unit sphere	RAG	Proof	Foundation	15 min
Simulate distance contrast run	Classical ML	Coding	Beginner	30 min
A sufficient condition for correct retrieval	RAG	Proof	Intermediate	30 min
Balls into bins, with one choice and with two run	AI Systems	Coding	Intermediate	1 hour
Cache hit rates under Zipf popularity run	AI Systems	Coding	Intermediate	1 hour

Problem

Write it

Solution

Related problems