Dot products and cosine similarity: one vector against a hundred thousand in a single line

An embedding is a row of floats

Ask an embedding model for a vector and you get back a list of numbers — 384 for a small open model, 1,536 or 3,072 for the hosted ones. What the numbers mean is the business of the machine learning course; here they are a (768,) float32 array, and a corpus of a hundred thousand documents is a (100000, 768) matrix. Everything below is arithmetic on that matrix.

q = np.asarray(embed("how do I reset my password"), dtype=np.float32)   # (768,)
M = np.load("embeddings.npy")                                            # (100000, 768)

np.save/np.load write the raw array with a small header — the fastest way to persist one, and much smaller than JSON.

Dot product

The dot product of two vectors multiplies them elementwise and sums:

a @ b                    # same as np.dot(a, b), same as (a * b).sum()

@ is the matrix-multiplication operator, and for a matrix against a vector it computes the dot product of every row with the vector at once:

M @ q                    # (100000, 768) @ (768,) → (100000,)

One line, one compiled loop, one hundred thousand dot products. Compare:

[np.dot(row, q) for row in M]     # ~400 ms
M @ q                             # ~3 ms

Same numbers, same arithmetic, different plumbing. The loop pays Python overhead per row; the matrix product hands the whole block to a compiled routine that reads memory in order and uses the CPU's vector units — the previous lessons in one measurement.

Cosine similarity

The dot product grows with the lengths of the vectors as well as their alignment. Cosine similarity removes the length:

def cosine(a, b):
    return (a @ b) / (np.linalg.norm(a) * np.linalg.norm(b))

It is the cosine of the angle between them: 1 for the same direction, 0 for perpendicular, −1 for opposite. Embedding models are used with cosine because a document's vector length carries little meaning while its direction carries the content.

Doing the division per query is wasteful. Normalise each row of M to unit length once, when you build it:

norms = np.linalg.norm(M, axis=1, keepdims=True)    # (100000, 1)
M_unit = M / norms                                  # every row now has length 1

Then every query is:

q_unit = q / np.linalg.norm(q)
sims = M_unit @ q_unit                              # (100000,) of cosines

A dot product of two unit vectors is their cosine. The keepdims=True is the broadcasting lesson paying off: (100000, 768) / (100000, 1) divides each row by its own norm.

Guard the zero vector: a document that embedded to all zeros has norm 0 and produces nan everywhere. norms[norms == 0] = 1 before dividing, or filter such rows out and log them.

Top k

k = 5
top = np.argpartition(sims, -k)[-k:]         # indices of the k largest, unordered
top = top[np.argsort(sims[top])[::-1]]       # sort just those k
for i in top:
    print(f"{sims[i]:.3f}  {docs[i][:60]}")

argpartition finds the k largest in linear time without sorting all hundred thousand; the second line orders only the five. For a hundred thousand rows argsort on everything is fine too — it is a few milliseconds — but the habit matters once the corpus is tens of millions.

Many queries at once

A batch of queries is a matrix too, and the matrix-matrix product gives every pairwise similarity:

Q_unit = Q / np.linalg.norm(Q, axis=1, keepdims=True)   # (50, 768)
S = Q_unit @ M_unit.T                                   # (50, 100000)

Row i of S is the similarity of query i to every document. The .T is a view, so nothing is copied. Fifty queries cost little more than one, because the compiled routine is doing the same memory pass with more arithmetic per element loaded.

Where brute force stops

The arithmetic is rows × dimensions multiply-adds per query: 100,000 × 768 ≈ 77 million, a few milliseconds on one core. At ten million rows it is 7.7 billion, around a second, and the matrix is 30 GB in float32 — more than most machines hold. That is the point where an approximate index (FAISS, HNSW via hnswlib, or a vector database) earns its complexity by trading a little recall for a thousand-fold speed-up. Below a million rows, M_unit @ q in NumPy is simpler, exact, and fast enough, and that covers nearly every project a person builds alone.

A cheap intermediate step: store M_unit as float16 to halve memory, and upcast the query result. The precision lesson said float16 keeps three figures; for ranking similarities that is usually enough, and it is a measurable decision rather than a guess.

Try this now

Make a random (100000, 768) float32 matrix with default_rng, normalise its rows, and time M_unit @ q against the Python loop. Then plant a known vector at row 4,321, query with a slightly noisy copy of it, and confirm that row comes out on top.