Embeddings and vector search from first principles

Learn in your browser · Chapter 01 of 06

Search compares lists of numbers, not sentences, so every vector search starts by turning text into one.

A vector is an ordered list of numbers, such as (4, 3). Each position is a dimension. Chapter 3's vectors have 132 dimensions, but the arithmetic is the same as for two.

In the RAG workbook, an embedding was a vector made by a trained neural network, an embedding model, so that similar meanings land close together. This workbook downloads no model: it uses word counts, TF-IDF and a two-dimensional meaning space labelled by hand.

Bag of words. Put every word in the collection in a fixed order, the vocabulary, and count each one in a text. Word order is lost: the textbook notes that 'Mary is quicker than John' and 'John is quicker than Mary' get the same vector.

vocabulary:              books  fines  late  online  renew
Late books, late fines.    1      1      2      0       0
Renew books online.        1      0      0      1       1

Learn in your browser · Chapter 02 of 06

Search finds the stored vectors closest to the question's vector. Three measures are common, and they can disagree.

In our meaning space the first number means about money, the second about time. 'What is the charge for a late book?' is q = (4, 3). A (3, 4) is a short note on late returns. B (8, 6) points exactly the same way as q but is twice as long, like a longer text's count vector. C (0, 10) is the opening hours.

dot(a, b)      = a1 x b1 + a2 x b2
length(a)      = sqrt(a1 x a1 + a2 x a2)
cosine(a, b)   = dot(a, b) / (length(a) x length(b))
distance(a, b) = length(a - b)

Cosine similarity keeps only the angle. For Euclidean distance, smaller is closer. length(q) = sqrt(16 + 9) = 5.

A: dot = 3x4 + 4x3  = 24   cosine = 24 / (5 x 5)  = 0.96   distance = sqrt(1 + 1)   = 1.414
B: dot = 8x4 + 6x3  = 50   cosine = 50 / (10 x 5) = 1.00   distance = sqrt(16 + 9)  = 5.000
C: dot = 0x4 + 10x3 = 30   cosine = 30 / (10 x 5) = 0.60   distance = sqrt(16 + 49) = 8.062

Scroll sideways to see every column.

Normalising divides a vector by its length, giving a unit vector. For unit vectors the dot product is the cosine, as the textbook shows, and squared distance is 2 - 2 x cosine, so all three measures agree. For A: unit A = (0.6, 0.8), unit q = (0.8, 0.6), squared distance 0.04 + 0.04 = 0.08 = 2 - 2 x 0.96. Normalise stored vectors and queries alike.

Learn in your browser · Chapter 03 of 06

Now search sixteen invented library help notes: one TF-IDF vector per note, every note scored, the best k returned.

Tested with Python 3.12.10 on Windows 11; these pins need Python 3.12 or later. In a new folder, run the lines below. They create a virtual environment (a private set of packages) and install pinned versions. On macOS and Linux, type python3 for python and .venv/bin/python for .venv\Scripts\python; those lines were not run.

python -m venv .venv
.venv\Scripts\python -m pip install numpy==2.5.3 scikit-learn==1.9.1

TF-IDF multiplies term frequency (how often a word appears in the text) by inverse document frequency (idf), which is high for rare words. The textbook's idf is log(N / df), where df counts the documents containing the word. scikit-learn's default, below, is ln((1 + n) / (1 + df)) + 1, and each row is then scaled to length 1.

n = 16 notes
'bicycle' is in 1 note:   idf = ln(17 / 2) + 1 = 3.1401
'the' is in 6 notes:      idf = ln(17 / 7) + 1 = 1.8873

Save this as library.py in that folder and run it with .venv\Scripts\python library.py.

# library.py: brute-force TF-IDF search over 16 invented help notes
import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer

DOCS = [
    "Opening hours: open Monday to Saturday, 9am to 6pm. Closed on Sundays.",
    "Summer opening hours: open Monday to Saturday, 9am to 5pm. Closed on Sundays.",
    "Joining is free for anyone who lives, works or studies in the town.",
    "Members can borrow up to 12 books at a time for three weeks.",
    "You can renew books twice, online or by phone, unless they are reserved.",
    "There are no fines for late books, but you cannot borrow more until they are back.",
    "If you lose a book, you pay the cost of replacing it.",
    "Public computers can be booked for one hour a day with your library card.",
    "Free wifi covers the whole building. The network name is on the notice board.",
    "Printing costs 10p per page in black and white and 50p per page in colour.",
    "Children under 12 must be with an adult at all times.",
    "Rhyme time for babies and toddlers runs every Wednesday at 10am.",
    "Study room 7 seats six people and can be booked at the front desk.",
    "There is no car park. Bicycle racks are by the side entrance.",
    "A lift serves every floor, and large print books are on the ground floor.",
    "Assistance dogs are welcome. Other pets are not allowed.",
]

vectorizer = TfidfVectorizer()           # defaults: lowercase, unit-length rows
X = vectorizer.fit_transform(DOCS)       # one row per document

def search(query, k=3):
    q = vectorizer.transform([query])
    scores = (X @ q.T).toarray().ravel() # dot product = cosine here
    top = np.argsort(-scores, kind="stable")[:k]
    return [(f"D{i + 1}", round(float(scores[i]), 2)) for i in top if scores[i] > 0]

if __name__ == "__main__":
    print(X.shape)
    for query in ["Can I renew a book online?", "Where can I leave my bike?"]:
        print(query, search(query))

(16, 132)
Can I renew a book online? [('D5', 0.42), ('D7', 0.16), ('D4', 0.09)]
Where can I leave my bike? [('D4', 0.24), ('D13', 0.22), ('D5', 0.22)]

This is brute-force, or exact, nearest-neighbour search: score everything, sort, keep the top k. Its cost grows in step with the collection. kind="stable" makes ties repeatable, and zero scores (no shared word) are dropped.

Renewal works. Bike fails: D14 says 'Bicycle', shares no word with the question and scores 0, while the top three share only 'can'. A neural embedding model is trained so that synonyms like these can land close, but we ran none, so we show no score.

Learn in your browser · Chapter 04 of 06

Two questions prove nothing. Score the search on labelled questions, and keep that set fixed while you change things.

A labelled set pairs each question with its gold documents, the ones that answer it.

• Recall@k: the share of a question's gold documents in the top k, averaged over questions. The textbook defines recall as the fraction of relevant documents retrieved. With one gold document per question, it is the RAG workbook's hit rate.
• Mean reciprocal rank (MRR): 1 divided by the rank of the first correct result, or 0 if none, averaged over questions. The TREC-8 question answering track used it.

Save this as evaluate.py beside library.py and run it.

# evaluate.py: recall@k and mean reciprocal rank on labelled questions
from library import search

LABELLED = {                              # question: documents that answer it
    "Can I renew a book online?": {"D5"},
    "How many books can I borrow?": {"D4"},
    "Is the library open on Sunday?": {"D1", "D2"},
    "What does it cost to print in colour?": {"D10"},
    "Are pets allowed?": {"D16"},
    "Do you charge for printing?": {"D10"},
    "Where can I leave my bike?": {"D14"},
    "Do I need to pay to join?": {"D3"},
}

firsts, r1, r3 = [], 0, 0
for question, gold in LABELLED.items():
    ranked = [doc for doc, score in search(question, k=16)]
    r1 += len(gold & set(ranked[:1])) / len(gold)
    r3 += len(gold & set(ranked[:3])) / len(gold)
    hits = [n for n, doc in enumerate(ranked, 1) if doc in gold]
    firsts.append(hits[0] if hits else None)
n = len(LABELLED)
mrr = sum(1 / f for f in firsts if f) / n     # a miss adds 0
print("first hits:", firsts)
print(f"recall@1 {r1 / n:.3f}  recall@3 {r3 / n:.3f}  MRR {mrr:.3f}")

first hits: [1, 1, 2, 2, 1, 3, None, None]
recall@1 0.375  recall@3 0.750  MRR 0.542

MRR is (1 + 1 + 0.5 + 0.5 + 1 + 0.333 + 0 + 0) / 8 = 0.542. The Sunday question has two gold notes, at ranks 2 and 3, so it scores 1 for recall@3. The misses are the honest failure case: 'bike' against 'Bicycle', 'join' against 'Joining'. The scikit-learn documentation notes that a bag of words ignores word derivations.

Learn in your browser · Chapter 05 of 06

Vector search fails in a few predictable ways. Reproduce them, then ask which a neural embedding model would change.

Save this as failures.py beside library.py and run it.

# failures.py: numbers, near-duplicates, negation, out of domain, chunk size
from library import vectorizer, search, DOCS

def cos(x, y):
    a, b = vectorizer.transform([x, y])
    return round(float((a @ b.T).toarray()[0, 0]), 2)

print(vectorizer.build_analyzer()("Study room 7 seats six people"))
print(search("When are you open on Saturday?"))
print(cos("Other pets are allowed.", "Other pets are not allowed."))
print(search("Do you sell train tickets?"))
print(cos("colour printing", DOCS[9]), cos("colour printing", " ".join([DOCS[0], DOCS[7], DOCS[9]])))

['study', 'room', 'seats', 'six', 'people']
[('D1', 0.37), ('D2', 0.35), ('D6', 0.23)]
0.88
[('D7', 0.47), ('D5', 0.24), ('D6', 0.21)]
0.32 0.22

Scroll sideways to see every column.

Chunk size and overlap. In vector terms, mixing topics dilutes a chunk's score, and overlap copies text into neighbouring chunks, creating near-duplicates. Measure both.

A stale index. An edited note keeps its old vector until you re-embed it. TF-IDF's idf depends on the whole collection, so one new note changes every word's weight, and new words stay unknown until you refit and rebuild. With a neural model, re-embed changed chunks, and everything if the model changes: each model makes its own space, so vectors from different models should not be compared. Store each vector with its chunk id, a text hash, the model version and a date.

Learn in your browser · Chapter 06 of 06

Brute force is exact but grows with the collection. See how high dimensions behave, what approximate indexes trade away, and why vectors need protecting.

High dimensions behave oddly. dims.py compares the farthest and nearest of 1,000 random points from a random query:

# dims.py: farthest over nearest distance for random points
import numpy as np

rng = np.random.default_rng(0)
for d in [2, 10, 100, 1000]:
    dist = np.linalg.norm(rng.random((1000, d)) - rng.random(d), axis=1)
    print(f"{d:>4} dimensions: {dist.max() / dist.min():.2f}")

   2 dimensions: 135.62
  10 dimensions: 4.48
 100 dimensions: 1.39
1000 dimensions: 1.11

In 2 dimensions the farthest point is over 135 times as far as the nearest; in 1,000, only 1.11 times. Beyer and colleagues showed that, under broad conditions, nearest and farthest distances converge as dimensions grow, with the effect appearing from as few as 10 to 15 dimensions. They stress that high-dimensional search is not therefore meaningless, and text vectors are not uniform random points. Distrust 2-D intuition, and compare any index with a plain linear scan.

Approximate nearest-neighbour search scores only a promising part of the collection and may miss true neighbours. The textbook's cluster pruning picks sqrt(N) random documents as leaders, attaches every document to its nearest leader, then scores only the groups of the b leaders nearest the query. pruning.py tests it on 10,000 synthetic unit vectors around 50 hidden topics, counting scores, not seconds.

# pruning.py: cluster pruning (leaders and followers) against brute force
import numpy as np

rng = np.random.default_rng(0)
centres = rng.normal(size=(50, 32))                  # 50 hidden topics
def sample(n):
    x = centres[rng.integers(50, size=n)] + 0.5 * rng.normal(size=(n, 32))
    return x / np.linalg.norm(x, axis=1, keepdims=True)

X, queries = sample(10_000), sample(200)
leaders = rng.choice(10_000, size=100, replace=False)   # sqrt(10,000)
owner = np.argmax(X @ X[leaders].T, axis=1)             # nearest leader

for b in [1, 2, 5, 10]:                                 # leaders searched
    recall, work = [], []
    for q in queries:
        exact = set(np.argsort(-(X @ q))[:10])          # brute force
        near = np.argsort(-(X[leaders] @ q))[:b]
        cand = np.flatnonzero(np.isin(owner, near))
        found = cand[np.argsort(-(X[cand] @ q))[:10]]
        recall.append(len(exact & set(found)) / 10)
        work.append(100 + len(cand))
    print(f"b={b:>2}: recall@10 {np.mean(recall):.2f}, scores {np.mean(work):,.0f}")

b= 1: recall@10 0.76, scores 255
b= 2: recall@10 0.94, scores 335
b= 5: recall@10 1.00, scores 594
b=10: recall@10 1.00, scores 1,063

Brute force computes 10,000 scores. One group costs 255 on average but finds only 76% of the true top 10; two reach 94% for 335. These topics are tidy; yours may not be.

Graph indexes make the same trade differently. The HNSW paper links each vector to near neighbours in layered graphs; a search moves greedily from an entry point towards the query, and a candidate-list size (ef) sets how much it explores. We have not run one. For any index, measure recall@k against brute force on your data.

Privacy. Vectors are not anonymous. vectorizer.inverse_transform(X[15]) hands back D16's eight distinct words. Neural embeddings are harder to read but not safe: one study recovered 50–70% of input words from popular sentence embeddings, and another recovered 92% of 32-token texts exactly and, in a separate test, full names from clinical notes.

Under the UK GDPR, as the ICO explains, personal data relates to someone identifiable directly or indirectly, and pseudonymised data is still personal data. A working rule, not legal advice: treat vectors of personal text as personal data, with the same access control, retention and deletion. The ICO notes this guidance is under review after the Data (Use and Access) Act.