Using What It Never Learned — Embeddings and Retrieval

Textbook › Part III Deep Learning: what assumption you add

CHAPTER 10

Using What It Never Learned — Embeddings and Retrieval

★ MEMBER17 min

One conclusion of Chapter 9 was that what was never learned does not come out of showing pairs. SFT has one ten-thousandth of pre-training's targets; what it changes is the delivery, not the contents.

So what do you do when the model has to use something that is not in its weights? Two roads: put it into the weights (train again), or hand it over on the spot. This chapter takes the second.

Handing it over means choosing. Choosing means deciding "close" as a number. The cosine similarity left on deposit in Chapter 1 of Part 0 is collected here.

1. "Just hand over everything" does not work

First, confirm that the road without choosing is closed.

Cut documents into 512-token chunks. How many fit in a context of length $T$ ?

context $T$	chunks that fit
8,192	16
32,768	64
131,072	256

With a million chunks on hand, 16 is $1.6\times10^{-5}$ of them. 999,984 are being thrown away.

"Just make the context longer" runs straight into the $T^2$ of Chapter 6. Going from 8k to 128k multiplies the attention table by $16^2 = 256$ while the number of chunks that fit only goes up 16-fold.

Retrieval is not a clever mechanism. It is the question of which 16 out of a million go into the context. Once you are choosing, how you choose *is* the performance.

2. Deciding closeness — cosine, not the raw dot product

A trained model maps both documents and queries to vectors (embeddings). Closeness looks measurable by a dot product, but the raw dot product has a trap.

Measured. Against a query $q$ :

$a$ : close in direction ( $\cos = 0.9929$ )
$b$ : far in direction ( $\cos = 0.5659$ ) but three times as long

The dot products are $q\cdot a = 0.9929$ against $q\cdot b = 1.6977$ . The one pointing the wrong way wins. Length can come from a document being long or wordy — nothing to do with closeness.

So divide the length out. This is the equation from Chapter 1 of Part 0.

\cos\theta \;=\; \frac{x^{\top}y}{\|x\|\,\|y\|}

(1)

Once everything is normalized, Euclidean distance returns the same ranking. With $\|x\|=\|y\|=1$ ,

\|x-y\|^{2} \;=\; 2 - 2\,x^{\top}y

(2)

and measurement agrees (0.0143 and 0.0143 on the $a$ side, 0.8682 and 0.8682 on the $b$ side). The document with the largest cosine is the document with the smallest L2.

FIG 1Drag the query around the clusters of documents. Cosine, dot product and L2 produce different top-5 line-ups — switch to the dot product and cluster C, which is merely longer rather than better aligned, muscles in

A high cosine means "the directions line up under this embedding." It is no guarantee that the answer is in there. What counts as close was decided by the embedding's training, and that training is not measured in this book.

3. In high dimensions, unrelated means nearly orthogonal

Embeddings have hundreds to thousands of dimensions. High dimensions have a property that happens to suit retrieval.

Two random unit vectors, cosine measured, 4,000 pairs per dimension.

$d$	standard deviation	$1/\sqrt{d}$	largest $\lvert\cos\rvert$
2	0.7079	0.7071	1.000
8	0.3525	0.3536	0.941
64	0.1273	0.1250	0.501
512	0.0436	0.0442	0.162
4,096	0.0164	0.0156	0.055

The mean is 0 at every dimension and the spread shrinks as $1/\sqrt{d}$ (the same shape as the $\sigma/\sqrt{n}$ of Chapter 3 of Part 0). At $d = 4096$ , the most lopsided of 4,000 pairs still only reached $\lvert\cos\rvert = 0.055$ .

Unrelated things pile up around 0, so anything with even a little relation stands out. That is the ground under "you can search with embeddings."

The other side shows up at the same time. At $d=4096$ , anything past 0.05 is already unusual — which also means reading the gap between 0.82 and 0.79 as a difference in meaning is risky.

4. How many multiplications is brute force

Hold $N$ documents in $d$ dimensions and dot every one of them against each query: $2Nd$ FLOPs, counting one multiply and one add.

$N$	FLOPs per query	memory (float32)
$10^{6}$	$1.54\times10^{9}$	3.1 GB
$10^{8}$	$1.54\times10^{11}$	307 GB

At $d = 768$ . This is paid on every query. Small next to the $10^{21}$ of training in Chapter 8 — but training happens once and retrieval happens every time.

So a way of not looking at everything is needed.

5. Approximate nearest neighbours — fraction scanned against what is missed

The method is simple. Split the documents into clusters ahead of time (k-means) and only look inside the clusters nearest the query. The number of clusters looked at is $n_{\text{probe}}$ .

FIG 2Left, the distribution of pairwise cosine. The further apart same-cluster (gold) and different-cluster (grey) sit, the more "clumped" the data is. Right, recall@10 against the fraction scanned; the dashed line is picking the same number of documents at random. Raise the spread to blur the clusters and the curve drops a long way — but never onto the dashed line (the figure runs a smaller version of the experiment live)

An experiment. $N = 10{,}000$ , $d = 64$ , 100 clusters, $K = 10$ , 200 queries. The data is scattered around 50 centres; the mean cosine within a cluster is 0.410. As a control, recall was also measured for the same number of documents picked at random.