Two-tower retrieval and embedding-based retrieval

The big idea

Matrix factorization gives every user and item a vector, and a recommendation is the item with the highest dot product. With 5,000 items you can score them all for every request. With 50 million videos, listings or products, and a budget of a few tens of milliseconds per request, you cannot run even a cheap model on every item, let alone the large model that knows the user's last ten clicks, the time of day and the item's price.

So large recommenders work like a hiring process. A quick screen of every CV keeps a few hundred plausible applicants; interviews, which are expensive, are only spent on those; and the final offer goes to a handful. The screen must not miss good people (high recall), while the interviews must order the shortlist well (high precision). In a recommender the screen is called retrievalThe first stage of a recommender: a fast model that narrows the full catalogue to a few hundred or thousand plausible items for one request. Also called candidate generation., and the standard tool for it is a two-tower model served from a nearest-neighbour index.

How it works

Step 1: The funnel

The YouTube paper in the case study below describes exactly this shape: candidate generationAnother name for retrieval: producing the candidate set that later stages rank. narrows millions of videos to hundreds, a separate ranking network scores those hundreds with a much richer set of features, and the ranker's top few dozen are shown. Splitting the work this way has two other benefits. Several retrieval sources can feed the same ranker, and ranking can be retrained often without touching the expensive index.

The two stages also have different metrics. Retrieval is judged by recall@K: of the items the user actually engaged with next, what share appeared among the K candidates? Ranking is judged by how well it orders the shortlist, with metrics such as NDCG from the next lesson, Learning to rank .

Step 2: Two towers

A two-tower modelA retrieval model with two separate networks: one turns user or query features into a vector, the other turns item features into a vector in the same space, and the score is their dot product. has a user tower and an item tower. The user tower reads anything known about the request (user ID, recent clicks, search words, location, device) and outputs a vector $u$. The item tower reads anything known about the item (ID, category, text, price band, engagement counts) and outputs a vector $v$ of the same length. The score is

or the cosine of the two after normalizing them. This is matrix factorization with features: if each tower is just an ID lookup table you get back the $p_u$ and $q_i$ of Collaborative filtering and matrix factorization . The towers can be deep networks, and because the item tower reads content features, a brand-new item gets a sensible vector before anyone has clicked it.

That separation is the whole point, and also the main limitation. Because user and item only meet in the final dot product, the model cannot use cross featuresFeatures that depend on a specific user and item together, such as how many times this user clicked this seller before; available to rankers but not to two-tower retrieval. such as "this user bought from this seller before". Rankers can. That is one reason retrieval and ranking stay separate models.

Step 3: Training with negatives

Training asks the model, for each logged (user, item) pair, to score the item the user engaged with above the alternatives. With the full catalogue as alternatives this is a softmax over millions of classes:

The denominator is too expensive to compute at every step, so you replace the catalogue with a sample of negativesItems used during training as examples of what the user did not choose, so the model learns to score them below the positive item., as word2vec did for words ( Word embeddings ):

• Sampled softmax. Draw a few thousand negatives per example from a background distribution and correct for the sampling. The YouTube paper reports several thousand negatives per example, more than 100 times faster than a full softmax.
• In-batch negatives. Take a batch of B positive pairs; for each user, the other B - 1 items in the batch serve as negatives. They are free, since their vectors are already computed. But items appear in batches in proportion to their popularity, so popular items get pushed down as negatives far too often. Yi and colleagues at Google correct each logit by subtracting the log of the item's sampling probability, $s(u,v_j)-\log p_j$, the "logQ correction".
• Hard negatives. Random items are usually easy to tell apart from the positive: a beach house is obviously wrong for a ski trip. Airbnb found that random negatives made the problem too easy, and instead used listings the guest had seen, or even wishlisted, but did not book.

Step 4: Approximate nearest-neighbour serving

At request time you need the K item vectors with the largest dot product with $u$. Scanning all of them is exact, but its cost grows with the catalogue. An approximate nearest-neighbour (ANN) indexAn index that finds most of the nearest vectors to a query while scanning only a small part of the collection, trading a little recall for large speedups. answers in a fraction of the time and accepts missing a few true neighbours. Two families dominate:

• IVF (inverted file). Cluster the item vectors offline with k-means. At query time, find the nprobe closest cluster centres and scan only the items in those clusters. More clusters probed means higher recall and more work.
• HNSW (hierarchical navigable small world). Link each vector to its near neighbours in a layered graph and search greedily from a coarse top layer down. It usually gives excellent recall per unit of time, at the cost of more memory and slower updates.

Retrieval quality is then two numbers multiplied: how good the towers are, and how much recall the index gives up. Measure both. The full mechanics of ANN indexes are covered in Vector databases and ANN indexes .

Going deeper

Code: two towers with in-batch negatives

import numpy as np

# Synthetic catalogue: 5,000 items with 12 content features, users with 8 profile features.
rng = np.random.default_rng(0)
n_items, n_users, d = 5000, 3000, 16
X = rng.normal(0, 1, (n_items, 12))                         # item features
Z = rng.normal(0, 1, (n_users, 8))                          # user features
A = rng.normal(0, 0.25, (8, 12))                            # hidden true taste
pop = 1.5 * rng.normal(0, 1, n_items)                       # some items are far more popular
logit = Z @ A @ X.T + pop
P = np.exp(logit - logit.max(1, keepdims=True)); P /= P.sum(1, keepdims=True)
clicks = np.array([[u, rng.choice(n_items, p=P[u])] for u in range(n_users) for _ in range(12)])
rng.shuffle(clicks)
train, test = clicks[:30000], clicks[30000:]
q = np.bincount(train[:, 1], minlength=n_items) / len(train)   # how often each item lands in a batch

def train_towers(logq, epochs=8, B=256, lr=0.05):
    r = np.random.default_rng(1)
    Wu, Wi = r.normal(0, .1, (8, d)), r.normal(0, .1, (12, d))
    E = np.zeros((n_items, d))                              # item-ID embedding inside the item tower
    for _ in range(epochs):
        for s in range(0, len(train), B):
            u, i = train[s:s + B, 0], train[s:s + B, 1]
            U, V = Z[u] @ Wu, X[i] @ Wi + E[i]                  # user tower, item tower
            S = U @ V.T - (np.log(q[i]) if logq else 0)          # in-batch: other rows are negatives
            G = np.exp(S - S.max(1, keepdims=True)); G /= G.sum(1, keepdims=True)
            G[np.arange(len(u)), np.arange(len(u))] -= 1; G /= len(u)   # softmax cross-entropy gradient
            gV = G.T @ U
            Wu -= lr * Z[u].T @ (G @ V); Wi -= lr * X[i].T @ gV; np.add.at(E, i, -lr * 5 * gV)
    return Z @ Wu, X @ Wi + E

def recall(Uemb, Vemb, K=100):
    top = np.argsort(-(Uemb[test[:, 0]] @ Vemb.T), axis=1)[:, :K]
    return np.mean([i in t for i, t in zip(test[:, 1], top)])

popular = np.argsort(-q)[:100]
print(f"random 100 items         recall@100 {100 / n_items:.3f}")
print(f"100 most popular items   recall@100 {np.mean(np.isin(test[:, 1], popular)):.3f}")
for logq in (False, True):
    U, V = train_towers(logq)
    print(f"two towers, logQ {str(logq):5s}   recall@100 {recall(U, V):.3f}")

random 100 items         recall@100 0.020
100 most popular items   recall@100 0.500
two towers, logQ False   recall@100 0.634
two towers, logQ True    recall@100 0.696

5,000 synthetic items with 12 content features and 3,000 users with 8 profile features; each user clicks 12 items drawn from a hidden taste function plus a strong popularity term. Both towers are linear (the item tower also has a small per-item ID vector) and train with in-batch negatives, batches of 256, by hand-written gradients.

• Popularity is a strong baseline: the 100 most popular items contain 50 percent of the held-out clicks. Random items manage 2 percent.
• Two towers beat it: recall@100 of 0.634 with plain in-batch softmax.
• The logQ correction helps: subtracting $\log p_j$ from each logit raises recall@100 to 0.696, because popular items are no longer over-punished for showing up in every batch as negatives.

Code: an IVF index from scratch

import numpy as np

# 100,000 item embeddings (32 numbers each) and 200 user embeddings, as two towers would produce.
rng = np.random.default_rng(0)
n, d = 100_000, 32
centres = rng.normal(0, 1, (300, d))
items = centres[rng.integers(0, 300, n)] + 1.0 * rng.normal(0, 1, (n, d))
users = centres[rng.integers(0, 300, 200)] + 1.0 * rng.normal(0, 1, (200, d))

def top10(q, cand):                                   # exact nearest neighbours by Euclidean distance
    dist = ((items[cand] - q) ** 2).sum(1)
    return cand[np.argsort(dist)[:10]]

exact = [top10(q, np.arange(n)) for q in users]       # brute force: scan every item

# IVF index: k-means into 256 lists offline; at query time scan only the nprobe closest lists.
k = 256
cent = items[rng.choice(n, k, replace=False)]
for _ in range(10):
    d2 = (items ** 2).sum(1)[:, None] - 2 * items @ cent.T + (cent ** 2).sum(1)[None]
    assign = d2.argmin(1)
    cent = np.array([items[assign == c].mean(0) if (assign == c).any() else cent[c] for c in range(k)])
lists = [np.flatnonzero(assign == c) for c in range(k)]
print(f"{n:,} items in {k} lists, sizes {min(map(len, lists))} to {max(map(len, lists))}")

for nprobe in (1, 4, 16, 64):
    hit, scanned = 0, 0
    for q, truth in zip(users, exact):
        near = np.argsort(((cent - q) ** 2).sum(1))[:nprobe]
        cand = np.concatenate([lists[c] for c in near])
        scanned += len(cand)
        hit += len(np.intersect1d(top10(q, cand), truth))
    print(f"nprobe {nprobe:3d}: scans {scanned / len(users) / n:6.1%} of items, recall@10 {hit / (10 * len(users)):.3f}")

100,000 items in 256 lists, sizes 128 to 956
nprobe   1: scans   0.4% of items, recall@10 0.786
nprobe   4: scans   1.7% of items, recall@10 0.882
nprobe  16: scans   6.7% of items, recall@10 0.961
nprobe  64: scans  26.6% of items, recall@10 0.997

100,000 item vectors of 32 numbers, clustered into 256 lists. Brute force gives the true top 10 for 200 queries; IVF scans only the closest lists.

• Probing 1 list scans 0.4 percent of the items and finds 78.6 percent of the true top 10.
• Probing 16 lists scans 6.7 percent and finds 96.1 percent.
• Probing 64 lists scans 26.6 percent and finds 99.7 percent. Pick nprobe from a curve like this one, for your latency budget.
• The list sizes range from 128 to 956: uneven clusters mean uneven query cost, a point that matters in the Airbnb case below.

Practical details

• Normalize the vectors. With raw dot products, items with large vectors score high for everyone. The Yi et al. paper reports that normalizing embeddings, with a temperature on the logits, helps. Airbnb, on the other hand, chose Euclidean distance because it gave much more balanced IVF clusters than dot product; they suspect magnitude matters because many of their features are historical counts.
• Freshness. If item vectors are rebuilt daily, a listing created an hour ago is not in the index. Systems either insert new vectors in near real time or run a separate retrieval source for new items.
• Several sources. Production systems merge candidates from several retrievers (two towers, item-to-item co-occurrence, popular in your region, followed creators) and let the ranker sort them out.

Trade-offs

Common mistakes

• Judging retrieval by ranking metrics. Retrieval's job is recall@K; ordering is the ranker's job.
• Training on in-batch negatives without the logQ correction, which quietly buries popular items.
• Using only random negatives: the model learns to separate the obvious and fails on near misses.
• Testing towers with brute-force search, then serving with an ANN index tuned for speed: the recall lost in the index never shows up offline.
• Letting the user tower and the item index drift: a new user-tower version must be paired with item vectors from the same training run.

In the wild

Check yourself

Recap