Two-tower models

Problem

Matrix factorization learns one vector per user and one per item from the interaction matrix, and scores a pair with a dot product. That is fast, but it knows nothing beyond IDs: a new user or item has no row, and the user's recent history, the time of day, an item's text or category cannot be used. A model that reads all of those features and scores a (user, item) pair jointly, with a ranking network, is far more accurate, and far too slow to run against the whole catalogue: at 2,000,000 multiply-adds per pair, 10,000,000 items cost 20,000,000,000,000 multiply-adds per request.

The way out is to give up on modelling the pair jointly for a first, cheap stage: score with a dot product between a query vector and an item vector, where each vector is a rich function of its own side's features. Then item vectors are fixed between model updates, and finding the best ones is nearest-neighbour search.

Intuition

Two networks, two "towers", map different kinds of input into one shared embedding space. The query tower reads who is asking and in what context: user ID, the items they recently watched, their language, the time. The item tower reads what an item is: its ID, title, category, age. Training pulls each observed (query, item) pair together and pushes the query away from other items, so that geometric closeness comes to mean "this person would engage with this item".

The only place the two sides meet is the final dot product. That restriction is what makes serving cheap, and it is also the model's main weakness: no feature of the query can change how an item's features are read. So two-tower models do retrieval, narrowing millions to hundreds, and a heavier ranker that may cross every feature does the final ordering of those few.

Mechanics

Separate query and item encoders

Each tower is an ordinary network. Categorical features enter through nn.Embedding, a lookup table of shape [num_embeddings, embedding_dim] initialized from a standard normal. Variable-length sets of IDs, such as a watch history, enter through nn.EmbeddingBag, which looks up every ID in a bag and reduces them with mode="mean" (the default), "sum" or "max" without materializing the per-ID vectors. The pooled and dense features are concatenated and passed through an MLP, and the output is usually L2-normalized with F.normalize, so the dot product is a cosine similarity in [−1, 1].

The score is s(q, v) = q · v / τ. With normalized vectors the temperature τ sets how sharp the softmax over items can get; values around 0.05 are common. The towers need not share anything; when both sides are the same kind of object (two sentences, a question and a passage) they often share weights, as in DSSM or DPR's dual encoder.

In-batch negatives and the sampled softmax

Training data is usually implicit: a log of (query, item) pairs that happened, with no record of what the user disliked. The natural objective treats retrieval as classification over the whole catalogue, a softmax over every item with the observed one as the label:

L = −log [ exp s(q, v⁺) / Σ_{j ∈ catalogue} exp s(q, v_j) ]

The denominator runs over millions of items, each needing a pass through the item tower. A sampled softmax replaces it with the positive plus a sample of negatives. The cheapest sample is free: in-batch negatives. A batch of B pairs gives B query vectors and B item vectors; the B × B matrix of scores has the positives on its diagonal, and every off-diagonal entry is another query's positive, used as a negative. The loss is F.cross_entropy on that matrix with targets 0, 1, …, B − 1. Each query gets B − 1 negatives for the cost of one matrix multiply, so larger batches give a better approximation of the full softmax. This is the InfoNCE loss of contrastive learning.

Popularity bias and the logQ correction

In-batch negatives are not a uniform sample: an item appears in the batch as often as it appears in the log, so popular items are used as negatives far more than rare ones and the model learns to under-score them. Yi et al. correct the sampled softmax by subtracting the log of each item's sampling probability from its logit:

s^c(q, v_j) = s(q, v_j) − log p_j

where p_j is the probability that item j appears in a batch, estimated from its frequency in the training stream. Rare items get a large boost and popular items a small one. This is the importance-sampling correction long used for sampled softmax (Jean et al.), which removes the sampling bias from the estimate of the full softmax. The correction is applied only in training; at serving time the raw score s is used. A related problem is the false negative: if the same item is the positive for two queries in a batch, each query's positive also appears as its own negative, so duplicates are usually masked out.

Hard negatives

Most random negatives are easy: their score is already far below the positive, their softmax probability is close to zero, and since the gradient of the loss with respect to a negative's logit is exactly that probability, they teach almost nothing. Hard negatives are items the model currently scores high that are not the positive: same category, retrieved by the current index, or retrieved by a lexical method like BM25. DPR added one BM25 hard negative per question to its in-batch negatives; ANCE mines them from an index of the model being trained, refreshed asynchronously. The risk is that many "hard negatives" are unlabeled positives, so the hardest ones are often skipped, and training on only hard negatives tends to collapse; mixing them with easy in-batch negatives is the usual recipe.

Serving with a vector index

After training, run the item tower over the whole catalogue once and load the vectors into a vector index. A request runs the query tower once, then asks the index for the top k items by inner product. Normalizing both sides makes maximum inner product equal to maximum cosine, which every index supports. Exact search over N items costs N · d multiply-adds; IVF, HNSW and product quantization trade a little recall for orders of magnitude less work and memory. Item vectors are recomputed when the item tower is retrained or an item's features change; new items need only a forward pass, not retraining, which is how two-tower models soften the cold-start problem of ID-only factorization.

Retrieval, then ranking

Covington et al.'s YouTube system split recommendation into candidate generation, which narrows millions of videos to hundreds, and ranking, which scores those hundreds with many more features and a different objective. Most large systems keep that funnel, often with several retrievers (two-tower, co-occurrence, popular in your region) merged before a ranker and a final re-ranking stage for diversity and business rules. Each stage is judged differently: retrieval by recall@k (is the item the user wanted among the candidates at all?), ranking by NDCG or similar on the order.

Worked example

Every number below comes from tested helpers in ml_models_docs.py (tt_bag_mean, tt_normalize, tt_in_batch, tt_full_softmax, tt_hard_negative, tt_serving). The towers are tiny and hand-set so each step can be checked; the serving numbers are an invented, labelled configuration, not a measured system.

Implementation

Both towers, the in-batch loss with the logQ correction and duplicate masking, and top-k retrieval. Item frequencies log_p are estimated from the training stream and passed in.

import torch
import torch.nn as nn
import torch.nn.functional as F


class QueryTower(nn.Module):
    def __init__(self, num_users, num_items, dim=64):
        super().__init__()
        self.user = nn.Embedding(num_users, dim)
        self.history = nn.EmbeddingBag(num_items, dim, mode="mean")
        self.mlp = nn.Sequential(nn.Linear(2 * dim, 128), nn.ReLU(), nn.Linear(128, dim))

    def forward(self, user_ids, history_ids, offsets):
        x = torch.cat([self.user(user_ids), self.history(history_ids, offsets)], dim=-1)
        return F.normalize(self.mlp(x), dim=-1)


class ItemTower(nn.Module):
    def __init__(self, num_items, num_categories, dim=64):
        super().__init__()
        self.item = nn.Embedding(num_items, dim)
        self.category = nn.Embedding(num_categories, dim)
        self.mlp = nn.Sequential(nn.Linear(2 * dim, 128), nn.ReLU(), nn.Linear(128, dim))

    def forward(self, item_ids, category_ids):
        x = torch.cat([self.item(item_ids), self.category(category_ids)], dim=-1)
        return F.normalize(self.mlp(x), dim=-1)


def in_batch_loss(q, v, item_ids, log_p, tau=0.05):
    logits = q @ v.T / tau - log_p[item_ids]
    same = item_ids[:, None] == item_ids[None, :]
    diagonal = torch.eye(len(q), dtype=torch.bool, device=q.device)
    logits = logits.masked_fill(same & ~diagonal, float("-inf"))
    return F.cross_entropy(logits, torch.arange(len(q), device=q.device))


@torch.no_grad()
def retrieve(query_tower, item_vectors, user_ids, history_ids, offsets, k=1000):
    q = query_tower(user_ids, history_ids, offsets)
    return (q @ item_vectors.T).topk(k, dim=-1)

retrieve is exact search; at catalogue scale item_vectors lives in an approximate index instead (vector indexes). log_p is subtracted only in the training loss. For tables with millions of rows, nn.Embedding(..., sparse=True) with a sparse-aware optimizer avoids updating every row on every step.

Tradeoffs

• Separable scoring is what makes millisecond retrieval over millions of items possible, and it caps accuracy: no query–item interaction before the dot product. Late-interaction models (ColBERT) and rankers recover some of that at higher cost.
• In-batch negatives are free but biased toward popular items and limited by batch size; logQ correction fixes the bias, and gradient caching or cross-device batches raise the count (see gradient accumulation).
• Hard negatives speed learning and sharpen fine distinctions, and add false negatives and the cost of mining from a stale index.
• Feature-based towers handle new users and items better than ID-only factorization, but ID embeddings still dominate for items with history, and their tables dominate memory.
• Approximate indexes trade recall for latency; the index must be rebuilt or updated whenever the item tower changes, and query and item vectors must come from compatible model versions.
• Training on logged interactions learns what the previous system chose to show, a feedback loop that retrieval of fresh or rare items needs exploration to break.

Connections

References

• Covington, Adams and Sargin. Deep Neural Networks for YouTube Recommendations. RecSys 2016. doi:10.1145/2959100.2959190
• Yi et al. Sampling-Bias-Corrected Neural Modeling for Large Corpus Item Recommendations. RecSys 2019. doi:10.1145/3298689.3346996
• Huang et al. Learning Deep Structured Semantic Models for Web Search using Clickthrough Data (DSSM). CIKM 2013. doi:10.1145/2505515.2505665
• Jean et al. On Using Very Large Target Vocabulary for Neural Machine Translation (sampled softmax). ACL 2015. arXiv:1412.2007
• van den Oord, Li and Vinyals. Representation Learning with Contrastive Predictive Coding (InfoNCE). 2018. arXiv:1807.03748
• Karpukhin et al. Dense Passage Retrieval for Open-Domain Question Answering. EMNLP 2020. arXiv:2004.04906
• Xiong et al. Approximate Nearest Neighbor Negative Contrastive Learning for Dense Text Retrieval (ANCE). ICLR 2021. arXiv:2007.00808
• Johnson, Douze and Jégou. Billion-scale similarity search with GPUs (Faiss). 2017. arXiv:1702.08734
• Khattab and Zaharia. ColBERT: Efficient and Effective Passage Search via Contextualized Late Interaction over BERT. SIGIR 2020. arXiv:2004.12832
• PyTorch v2.15.0: Embedding, EmbeddingBag, normalize, cross_entropy.