Bezen

ML System Design · Module 13 of 47 · 13 concepts · about 16 min

Embeddings and Candidate Retrieval

How a system picks a few hundred promising items out of millions in milliseconds: embeddings, the two-tower model, training with negatives, approximate nearest-neighbour search, other candidate sources, filtering, index freshness, cold start, and how retrieval is measured.

At work: The first stage of every large recommender, search engine and ad system. If retrieval misses the right item, no ranking model can recover it.

Why retrieval comes first

What it is

A good ranking model is too slow to run on every item. So large systems work in stages:

  • retrieval (candidate generation): a fast, rough method picks a few hundred likely items from the whole catalogue;
  • ranking: the accurate, slower model scores only those.

Retrieval is judged by whether the good items make it into the shortlist (recall); ranking by whether they end up at the top.

Why it matters

The arithmetic forces it.

  • Scoring a whole catalogue with the ranking model takes seconds per request.
  • Scoring a shortlist takes milliseconds.
  • The two stages trade a little accuracy (retrieval can miss things) for a thousandfold saving.

In the real world: YouTube, Pinterest, Amazon and Spotify all describe two-stage (or more) systems: retrieval narrows millions of items to hundreds, ranking orders them.

Visual example

Computed

The ranking model costs about 0.08 ms per item (module 4: 40 ms for 500). Only a shortlist fits in the budget.

Time to rank, per request
items rankedtime
Whole city catalogue (50,000)4,000 ms
Retrieved shortlist (500)40 ms

Four seconds per home-screen open is out of the question; forty milliseconds fits.

Spending all the design effort on the ranker. If retrieval's shortlist does not contain what the user wants, a perfect ranker only orders the wrong items well.

Embeddings for retrieval

What it is

Retrieval usually works with embeddings: users and items become lists of numbers (vectors) in the same space, placed so that:

  • a user lands close to items they are likely to want;
  • similar items land close to each other.

'Close' is measured by dot product or cosine similarity. Retrieval then means: find the item vectors nearest to this user's vector.

Why it matters

Turning preference into distance makes retrieval a geometry problem, which can be solved very fast.

  • Item vectors are computed in advance and stored in an index.
  • At request time only the user's vector is computed, then the index finds its nearest items.

In the real world: Kirana Kart represents each customer and each product as a 64-number vector; the products nearest to a customer's vector become their home-screen candidates.

Comparing vectors from different models. A user vector from this week's model and item vectors from last week's are in different spaces; their distances mean nothing. Retrain and re-index together.

Practise it: Design a Recommendation System

The two-tower model

What it is

A two-tower model learns user and item embeddings together:

  • a user tower: a neural network that turns user features (history, location, time) into a vector;
  • an item tower: a network that turns item features (ID, category, text, image) into a vector of the same size;
  • the score is the dot product of the two vectors;
  • training pushes the score up for pairs that interacted (bought) and down for pairs that did not.

Why it matters

The two towers never look at each other's inputs, and that is the point.

  • Item vectors can be computed once for the whole catalogue and indexed.
  • Only the user tower runs per request.
  • The cost is that it cannot model fine interactions between a specific user feature and a specific item feature; the ranker, which sees both together, does that.

In the real world: Two-tower retrieval is described by YouTube, Google Play, Pinterest and many others; it is the standard first stage of modern recommenders.

import torch.nn as nn

class Tower(nn.Module):
    def __init__(self, n_in, dim=64):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(n_in, 256), nn.ReLU(), nn.Linear(256, dim))
    def forward(self, x):
        return nn.functional.normalize(self.net(x), dim=-1)

user_tower, item_tower = Tower(n_user_feats), Tower(n_item_feats)

def in_batch_loss(user_x, item_x, temperature=0.05):
    u, v = user_tower(user_x), item_tower(item_x)      # (B, 64) each
    logits = u @ v.T / temperature                      # every user against every item in the batch
    labels = torch.arange(len(u))                       # the matching item is on the diagonal
    return nn.functional.cross_entropy(logits, labels)

Feeding the item tower features that change by the minute (current stock, current price). Item vectors are computed in advance, so they will be stale; keep fast-changing facts for filtering and ranking.

Negatives for retrieval training

What it is

A retrieval model learns as much from what users did not choose as from what they did. Sources of negatives:

  • random items from the catalogue: easy, so the model learns little from them after a while;
  • in-batch negatives: the other items in the same training batch. Cheap and effective, but biased towards popular items, which appear in more batches;
  • hard negatives: items that are close but wrong (shown and skipped; similar category but not bought). Teach fine distinctions.

Why it matters

The choice of negatives decides what the model can tell apart.

  • Only easy negatives: the model separates 'milk' from 'lawnmowers' but not 'full-cream milk' from 'toned milk'.
  • In-batch negatives punish popular items too often; a sampling-bias correction (subtracting the log of each item's frequency from its score during training) offsets this.

In the real world: Google's paper on sampling-bias-corrected two-tower models showed that correcting in-batch negatives for item popularity noticeably improved retrieval for less popular items.

Using every item shown and not clicked as a hard negative. Many were never looked at (position bias, module 7), so they are not really rejected.

Approximate nearest-neighbour search

What it is

Finding the exact nearest vectors among millions is slow. Approximate nearest-neighbour (ANN) search finds almost all of them, very fast. Common methods:

  • HNSW: a graph of vectors, searched by hopping towards the query. Fast and accurate; uses a lot of memory.
  • IVF: cluster the vectors; search only the few clusters nearest the query.
  • Product quantisation (PQ): compress each vector to a few bytes; compare compressed forms.

Libraries and databases: FAISS, ScaNN, hnswlib, and vector databases such as Milvus, Qdrant, Weaviate and Pinecone.

Why it matters

ANN trades a little recall for enormous speed.

  • Typical settings return 95–99% of the true nearest neighbours in a millisecond or two.
  • Each method has knobs (clusters searched, graph connections) that trade recall, speed and memory; tune them against a measured recall target.

In the real world: Spotify's Annoy and Meta's FAISS were both built for this exact job; FAISS is now a standard component of retrieval systems.

Visual example

ANN

General tendencies; tune on your own data.

Approximate nearest-neighbour methods compared
methodspeedrecallmemorynotes
Exact (brute force)Slow at scale100%Full vectorsFine for a few hundred thousand on GPU
HNSWVery fastVery highHighThe common default when memory allows
IVFFastHigh; tunableFull vectorsSearch more clusters for more recall
IVF + PQFastSomewhat lowerVery lowFor hundreds of millions to billions

Not measuring the ANN index's own recall. A badly tuned index can silently drop a fifth of the true neighbours, and every downstream metric suffers without an obvious cause.

Practise it: Design an Image Search System

Index memory and compression

What it is

An index's memory is roughly items × dimensions × bytes per number, plus the method's overhead:

  • 128 floating-point numbers take 512 bytes per item;
  • product quantisation can compress that to 32 or 64 bytes, with some loss of accuracy;
  • fewer dimensions, or 16-bit numbers, also shrink it.

Why it matters

Memory decides how many machines the index needs.

  • A few million items fit comfortably on one machine at full precision.
  • A billion do not, so they are compressed and sharded (split across machines, each searching its part, results merged).

In the real world: Billion-scale image and video search systems combine compression with sharding: each machine holds a compressed slice of the index, and a query is sent to all of them.

Visual example

Computed

Vector memory for 128-number embeddings: full 512 bytes per item, or 64 bytes compressed.

Vectors only; index structures add more
catalogueitemsfull precisioncompressed (64 B)
Kirana Kart, one city50,00025.6 MB3.2 MB
A large marketplace10,000,0005.1 GB640 MB
A video platform1,000,000,000512 GB64 GB

Choosing an embedding size by habit. Going from 64 to 256 dimensions quadruples memory; measure whether recall actually improves.

Many candidate sources

What it is

Production systems rarely rely on one retrieval method. They merge several candidate sources:

  • embedding retrieval: nearest items to the user's vector;
  • item-to-item: items often bought with what is in the basket (co-occurrence);
  • history: the user's own usual purchases ('buy again');
  • popularity: trending in this city or store now;
  • business: promotions, new arrivals.

Each contributes a quota; duplicates are removed; the ranker scores the union.

Why it matters

Each source catches what others miss.

  • Embeddings generalise but can miss obvious repeat purchases.
  • Popularity covers new users with no history.
  • Item-to-item captures 'pasta goes with sauce'.
  • Quotas keep any one source from crowding out the rest.

In the real world: Kirana Kart's home screen draws about 800 candidates from embeddings, 600 from the customer's history, 400 from item-to-item on the basket, and 200 from what is trending in their store; module 14 shows how those 2,000 are cut down.

Visual example

Sources

Several cheap sources, merged, then one accurate ranker.

Candidates for one home-screen request

  1. four sources, in parallel
    Embeddings~800Nearest to the customer
    Buy again~600Their own history
    Goes with basket~400Bought together
    Trending~200Popular in the store
  2. Merge and filter~2,000Duplicates and out-of-stock removed
  3. Rank~500 → 20Light ranker, then heavy ranker

Letting one source dominate the shortlist. If 'buy again' fills 1,800 of 2,000 slots, the user is never shown anything new, and the system cannot learn about other products.

Filtering: before or after the search

What it is

Most retrieval must respect hard conditions: in stock at the customer's store, allowed in their region, suitable for their age. Two places to apply them:

  • Pre-filter: restrict the search to allowed items first, then find nearest neighbours among them.
  • Post-filter: find nearest neighbours in the whole index, then drop disallowed ones.

Why it matters

Each has a failure mode.

  • Post-filtering is simple but can leave too few results when most neighbours are filtered out ('in stock at a small store').
  • Pre-filtering always returns enough, but some ANN indexes slow down badly under restrictive filters.
  • Common fixes: fetch extra before post-filtering, keep a separate index per store or region, or use a vector database with efficient filtered search.

In the real world: Kirana Kart builds a separate small index per dark store each night from the products it stocks, so every search is already limited to available items.

Post-filtering with a strict filter. Ask for 500 neighbours, discard the 470 not stocked at this store, and the ranker gets 30 candidates.

Practise it: Improve a Low-Recall Model, Design a Search Ranking System

Index freshness

What it is

Items change: new products arrive, others are discontinued, and the embedding model is retrained. Keeping the index current:

  • full rebuild on a schedule (nightly), after retraining;
  • incremental updates: add new items' vectors as they arrive, mark removed ones deleted;
  • a small delta index for items added since the last rebuild, searched alongside the main one.

Why it matters

A stale index cannot show new items.

  • In fast-moving catalogues (news, fresh listings) items must be searchable within minutes.
  • After retraining, every vector must be recomputed and swapped in at once, so old and new vectors are never mixed.

In the real world: News and marketplace platforms commonly pair a large index rebuilt daily with a small, frequently updated one for the newest items.

Updating the embedding model but only re-embedding new items. The index then mixes two incompatible vector spaces.

Practise it: Design an Image Search System, Design a Personalized News Feed

Cold start

What it is

Cold start is predicting for users or items with little or no history:

  • new items: use content features (title, category, image embedding) in the item tower, so a new product gets a sensible vector immediately; reserve some exposure for new items to gather data (exploration);
  • new users: start with popularity in their area, ask a few preferences, use context (location, time, device), and update quickly from the first clicks.

Why it matters

Every system meets cold start every day.

  • New products are often the ones the business most wants to sell.
  • A new user's first session decides whether they come back.
  • A model that relies only on ID embeddings is helpless for both.

In the real world: Kirana Kart guarantees each new product a small share of impressions for its first two weeks, so the model can learn whether customers want it.

Ignoring new items in evaluation. Offline tests on past data favour items with history, so a model that never recommends anything new can look excellent.

Practise it: Design a Demand Forecasting Platform, Design a Real-Time Anomaly Detection System, Design a Recommendation System, Design a Model Serving Platform, Design a Personalized News Feed, Design a Search Ranking System

Measuring retrieval

What it is

Retrieval is measured by recall at k: of the items the user actually engaged with, how many were in the top k candidates?

  • k is the shortlist size (e.g. 500);
  • measure on a later time period than training;
  • measure per source and overall, and on slices (new users, new items).

Why it matters

Recall is the right measure because the ranker can only work with what retrieval returns.

  • Precision matters little at this stage: the ranker removes the poor candidates.
  • Recall at k against k shows how big the shortlist must be.

In the real world: Kirana Kart tracks recall at 500 for each candidate source weekly; a drop in the embedding source's recall after a retrain blocks that model's release.

Judging retrieval only by the final online result. When the whole system improves or worsens, you cannot tell which stage did it without recall measured separately.

Interview questions: retrieval

What it is

The questions interviewers ask about retrieval, from the first screening call to the senior round. Cover the answer, say yours aloud, then compare.

Why it matters

Reading an idea and explaining it are different skills.

  • Each answer below is the two or three sentences an interviewer is listening for.
  • The reasoning behind each one is in the cards above. If an answer surprises you, go back to that card.

In the real world: Recommendation and search design interviews spend a large share of time here. 'How do you get from millions of items to the hundred you show?' is the pivotal question.

Q: Why do recommenders use multiple stages?

The accurate ranking model is too slow to score millions of items, so a fast retrieval stage picks a few hundred candidates and the ranker scores only those. Retrieval is judged on recall, ranking on order.

Q: What is a two-tower model?

Two networks that map user features and item features into the same vector space; the score is their dot product. Item vectors are precomputed and indexed, so only the user tower runs per request.

Q: How do you choose negatives for training retrieval?

Random, in-batch and hard negatives. In-batch negatives are efficient but favour popular items, so correct for sampling frequency; hard negatives teach fine distinctions.

Q: What is approximate nearest-neighbour search?

Index structures such as HNSW, IVF and product quantisation that find almost all nearest vectors in milliseconds, trading a little recall for huge speed. Measure the index's own recall.

Q: Pre-filter or post-filter?

Post-filtering is simple but can leave too few results under strict filters; pre-filtering always returns enough but can be slow in some indexes. Per-segment indexes or fetching extra are common fixes.

Q: How do you handle cold start?

Content features in the item tower for new items, guaranteed exploration traffic, and popularity, context and quick updates for new users.

Proposing to score every item with a deep model per request.

Practise it: Design a Recommendation System, Design a Search Ranking System

Exercise: build two-tower retrieval with an ANN index

What it is

A short build that uses this module's ideas. Write it yourself before looking anything up; the checks at the end tell you whether it works.

Why it matters

An idea you have only read fades in a week.

  • Building it once, even in fifty lines, is what makes it yours.
  • Interviewers can tell the difference between someone who has read about a loop and someone who has debugged one.

In the real world: This is the core of most recommender take-home tasks; interviewers often ask how you measured the index's recall.

Your task. Train a small two-tower model on public interaction data, index the item vectors and measure recall.

Steps

  1. Use MovieLens or a public e-commerce interaction dataset; split interactions by time.
  2. Build a popularity baseline and measure recall at 10, 100 and 500 on the later period.
  3. Train a two-tower model with in-batch negatives (or use a matrix-factorisation model if simpler).
  4. Index item vectors with FAISS, first exact and then HNSW or IVF; measure the ANN index's recall against exact search.
  5. Measure the model's recall at 10, 100 and 500, overall and for users with fewer than five interactions.

You are done when

  • The two-tower model beats popularity on recall at 100.
  • You know the ANN index's recall against exact search, and its query time.
  • You can say how the model does for users with little history.

Stretch. Add item content features (genre, title embedding) to the item tower and measure recall on items that appear only in the test period.

Practise it: Design a Recommendation System