Recommendation is the machine learning workload most enterprises already have and least often modernise. The collaborative-filtering job that was written in 2019 still runs, still retrains weekly, and nobody wants to touch it. Meanwhile the architecture that powers every large-scale recommender in production today — a two-tower retrieval model feeding a heavier ranker — is a couple of hundred lines of PyTorch and trains on a single GPU for most catalogues.
This tutorial builds that system end to end: the retrieval tower, the loss that makes it work, the negatives that make it good, the ranker on top, the metrics that actually predict online behaviour, and the serving path. Everything here is plain PyTorch plus an approximate-nearest-neighbour index.
The architecture, in one paragraph
A production recommender is two stages. Retrieval narrows a catalogue of millions to a few hundred candidates in single-digit milliseconds; it must be cheap, so the user and the item are encoded independently into the same embedding space and matched by dot product, which lets you precompute every item vector and use an ANN index. Ranking scores those few hundred candidates with a much richer model that sees user-item cross features, recency, price, context — anything, because it only runs on a few hundred rows. Almost every recommendation quality problem is really a retrieval problem: the ranker can only reorder what retrieval gave it.
1. Features and vocabularies
Start with the smallest useful feature set and resist the urge to add more before the baseline works.
- User tower: user ID embedding, recent interaction history (a sequence of item IDs), coarse context (country, device, hour bucket).
- Item tower: item ID embedding, category, brand, a text embedding of the title if you have one, and a few numeric features (age of item, log popularity).
Hash or index every categorical into a contiguous vocabulary and persist the mapping alongside the model. A recommender that cannot reproduce its vocabularies cannot reproduce its predictions; store the mapping as an artefact, versioned with the checkpoint.
import torch
import torch.nn as nn
import torch.nn.functional as F
EMB = 128
class UserTower(nn.Module):
def __init__(self, n_users, n_items, n_countries, hist_len=32):
super().__init__()
self.user_emb = nn.Embedding(n_users, EMB)
self.item_emb = nn.Embedding(n_items, EMB, padding_idx=0) # shared history embedding
self.country_emb = nn.Embedding(n_countries, 16)
self.mlp = nn.Sequential(
nn.Linear(EMB * 2 + 16, 256), nn.ReLU(),
nn.Linear(256, EMB),
)
def forward(self, user_id, history, country):
# history: (B, hist_len) of item ids, 0 = pad
mask = (history != 0).float().unsqueeze(-1)
hist = (self.item_emb(history) * mask).sum(1) / mask.sum(1).clamp(min=1.0)
x = torch.cat([self.user_emb(user_id), hist, self.country_emb(country)], dim=-1)
return F.normalize(self.mlp(x), dim=-1)
class ItemTower(nn.Module):
def __init__(self, n_items, n_categories, n_numeric=4):
super().__init__()
self.item_emb = nn.Embedding(n_items, EMB, padding_idx=0)
self.cat_emb = nn.Embedding(n_categories, 32)
self.mlp = nn.Sequential(
nn.Linear(EMB + 32 + n_numeric, 256), nn.ReLU(),
nn.Linear(256, EMB),
)
def forward(self, item_id, category, numeric):
x = torch.cat([self.item_emb(item_id), self.cat_emb(category), numeric], dim=-1)
return F.normalize(self.mlp(x), dim=-1)
Two choices in that code matter. The history pooling shares the item embedding table with the item tower, which ties the two spaces together and speeds convergence considerably. And both outputs are L2-normalised, so the dot product is a cosine similarity — which keeps the logit scale stable and matches what your ANN index will compute at serving time.
2. The loss: in-batch softmax with logQ correction
The standard retrieval objective treats every other item in the batch as a negative. It is efficient — one matrix multiply gives you B x B logits — but it has a well-known bias: popular items appear as in-batch negatives far more often than rare ones, so the model over-penalises them and under-recommends the things users actually click. The fix is the logQ correction: subtract the log sampling probability of each candidate from its logit.
class TwoTower(nn.Module):
def __init__(self, user_tower, item_tower, temperature=0.05):
super().__init__()
self.user_tower, self.item_tower = user_tower, item_tower
self.temperature = temperature
def forward(self, user_batch, item_batch, log_q):
u = self.user_tower(**user_batch) # (B, EMB)
v = self.item_tower(**item_batch) # (B, EMB)
logits = (u @ v.T) / self.temperature # (B, B)
logits = logits - log_q.unsqueeze(0) # logQ correction
labels = torch.arange(u.size(0), device=u.device)
return F.cross_entropy(logits, labels)
log_q is the log of each in-batch item's empirical sampling frequency, estimated from your training stream (a simple running count over item impressions per epoch is enough). Temperature around 0.05 is a reasonable start for normalised embeddings; it is worth a small sweep because it interacts strongly with batch size.
Batch size is a model hyperparameter here, not just a throughput knob. More in-batch negatives means a better-calibrated softmax, so push the batch as large as memory allows — 8k to 32k rows is normal for retrieval towers, which is achievable on one GPU because the towers are small. Gradient checkpointing is rarely needed; BF16 autocast is.
3. Hard negatives, where the quality comes from
In-batch negatives are random, and random negatives are easy: after a few epochs the model separates "running shoes" from "garden furniture" perfectly and still cannot tell two similar shoes apart. Mixed negative sampling fixes this. Every N epochs, index the current item embeddings, retrieve the top candidates for a sample of users, drop the true positives, and keep the remainder as hard negatives appended to each row's candidate set:
# extra_negatives: (B, K) item ids mined from the current index
neg = self.item_tower(**extra_negative_batch) # (B*K, EMB)
neg = neg.view(B, K, EMB)
hard_logits = torch.einsum("be,bke->bk", u, neg) / self.temperature
logits = torch.cat([logits, hard_logits], dim=1) # (B, B+K)
loss = F.cross_entropy(logits, labels)
Keep the ratio modest — K of 4 to 16 — and keep the random in-batch negatives alongside. All-hard training collapses: the model learns to discriminate near-duplicates and loses the coarse structure that makes retrieval robust. Re-mine periodically rather than every step; a stale index from the previous epoch works fine and costs almost nothing.
4. Offline metrics that predict online behaviour
Evaluate on a time-based split, never a random one. Train on interactions up to day T, evaluate on days T+1 to T+7. A random split leaks the future into training and will flatter every model you build by a wide margin.
Compute, over the full catalogue rather than a sampled subset:
- Recall@k (k = 100, 500) for the retrieval stage — the only number that bounds everything downstream.
- NDCG@10 for the ranking stage.
- Catalogue coverage: the fraction of distinct items appearing in anyone's top 100. A model with good recall and 3% coverage is a popularity baseline in disguise.
- Cold-start recall: recall restricted to items first seen in the last 14 days. This is where ID-only models fail and content features earn their place.
@torch.no_grad()
def recall_at_k(model, users, positives, item_matrix, k=100):
u = model.user_tower(**users) # (N, EMB)
scores = u @ item_matrix.T # (N, n_items)
topk = scores.topk(k, dim=1).indices
hits = (topk == positives.unsqueeze(1)).any(dim=1).float()
return hits.mean().item()
Always report the popularity baseline (recommend the globally most popular items) next to your model. It is embarrassing how often a first two-tower model fails to beat it, and much better to learn that offline.
5. The ranker
Retrieval hands over a few hundred candidates; the ranker can be as expensive as your latency budget allows. A gradient-boosted tree over hand-built features is a perfectly respectable baseline and often hard to beat. If you want it in PyTorch — worth it when you have sequence features or want to share embeddings with retrieval — a straightforward DLRM-style model works: embed the categoricals, take explicit pairwise feature interactions, concatenate with the dense features, and pass through an MLP.
Train it on logged impressions with clicks as positives and non-clicked impressions as negatives, and train it on the output distribution of your retrieval model, not on random items. A ranker trained on random negatives sees a completely different candidate distribution at serving time and will be miscalibrated in exactly the region that matters.
Use a listwise or pairwise loss if your product surfaces a ranked list; use plain binary cross-entropy if you need calibrated click probabilities for downstream business logic (blending with margin, for instance). You usually cannot have both from one head — train two heads if you need both.
6. Serving
The serving path follows directly from the architecture.
- Batch-encode the catalogue with the item tower after each training run and write the vectors to an ANN index — FAISS, ScaNN, or a vector database if you already run one. For catalogues under a few million items an
IndexHNSWFlaton CPU is fast enough and much simpler to operate than a GPU index. - Serve the user tower online. It is small;
torch.compileplus BF16 puts a single forward under a millisecond on CPU. Export it withtorch.exportif you want the Python-free path. - Query, filter, rank. ANN search for top 500, apply business filters (in stock, not already purchased, regional availability), rank with the heavier model, apply diversity or exploration rules last.
- Log everything you served, in order, with the model version. Impression logs are the training data for the next model; if you do not log the candidates you showed and their positions, you cannot correct for position bias later, and you will not get that data back.
The part teams underestimate is embedding drift between towers. The user tower is served live; the item index is rebuilt on a schedule. If you deploy a new user tower against an item index built by the previous model, the spaces do not match and quality collapses silently — recall metrics look fine offline, engagement drops online. Version the towers and the index together, and refuse to serve a mismatched pair at startup.
What we see go wrong
Random train/test splits. Covered above; still the single most common error and it invalidates every comparison built on it.
No popularity baseline. Without it there is no way to know whether the model learned anything beyond "show what everyone clicks".
Training on clicks only. Clicks measure what was shown as much as what was wanted. Include the position in the ranker's features at training time and set it to a constant at inference, or weight examples by inverse propensity — otherwise you are modelling your own UI.
Weekly full retrains and nothing else. Catalogues change daily. Incremental training on recent interactions with a periodic full rebuild keeps cold-start recall from decaying between runs.
No offline-to-online correlation check. Track, over several A/B tests, whether offline recall improvements actually moved the online metric. If they never do, your offline evaluation is measuring the wrong thing and every future decision made with it is a coin flip.
Where to go next
Once the two-stage baseline works, the upgrades in rough order of payoff are: sequence modelling in the user tower (a small transformer over the interaction history rather than mean pooling), content embeddings for cold-start items, multi-task ranking heads (click, add-to-cart, purchase, dwell) and only then the generative retrieval approaches built on semantic item IDs. Do them one at a time, with the same time-based evaluation, or you will not know which one helped.
We build and rescue systems like this for a living — see PyTorch for recommender systems and deep learning model development. If your recommender has stopped improving, or you are still evaluating on a random split, get in touch and we will look at it with you.