What is a two-tower model, and how do you train one?

In this article

1. Before you start
2. Why interviewers ask this
3. Two-tower model from first principles
4. Worked example
5. Trade-offs and edge cases
6. How to answer it in an interview
7. Common mistakes
8. Follow-up questions
9. Sources

Before you start

This page is one step of the embeddings and vector search learning path.

Why interviewers ask this

The two-tower model is a common first stage in large recommendation systems, for example the YouTube retrieval model of Yi et al. (2019), so a design round on one often reaches it early. Interviewers use it to see whether a candidate knows why the first stage must be cheap, what the training loss looks like, and what goes wrong when training batches are built from popular items.

A weak answer says "two networks produce embeddings and we compare them." A strong answer says why the towers cannot see each other until the final score, writes the loss as a softmax over the batch, and names the failure that comes from sampling negatives from the batch. A Staff-level answer adds the serving picture: precompute item vectors, search them with an approximate nearest-neighbour index, then rerank a short list with a slower model.

Two-tower model from first principles

Two-tower model

A two-tower model is a pair of neural networks, a query tower and an item tower, that each turn their input into a vector of the same size, scoring a pair by the dot product of the two vectors, with no interaction between the sides before that score.

In-batch negative

An in-batch negative is an item that belongs to another example in the same training batch and is reused as a wrong answer for the current query, so one batch of B pairs gives each query B minus one negatives at no extra encoding cost.

InfoNCE loss

The InfoNCE loss is the cross-entropy of a softmax over the scores of one query against its true item and its negatives, so minimising it pushes the true pair's score above every negative's.

In plain words

Think of a library where every book has a call number and every reader's request is also converted to a call number by a fixed rule. Matching a request to books is then just comparing numbers. The librarian converts every book once, writes the numbers on a shelf list, and converts only the request when it arrives. The analogy breaks in one place. In a two-tower model both conversion rules are learned together, from examples of requests and the books that satisfied them.

In one precise sentence, a two-tower model learns a function f for queries and a function g for items so that the dot product of f(query) and g(item) is high for matching pairs and low for the rest.

Huang et al. (2013) mapped queries and documents into a shared low-dimensional space. They used it for web search and trained on clicks. Covington et al. (2016) described the two-stage design used for YouTube: candidate generation, then ranking.

Step 1: Two separate towers encode the two sides

The query tower reads whatever describes the request, such as the text of a search or a user's history, and outputs one vector. The item tower reads whatever describes the item, such as its title and category, and outputs a vector of the same size, with a size such as 128 numbers (illustrative). The two towers have separate weights because queries and items are different kinds of data.

Input: raw features of one side. Output: one vector. Neither tower sees the other side, and that restriction is the point. You now know the model's shape: two independent encoders whose only meeting point comes in the next step. A model that shares weights between the sides is a variant called a siamese model, a kind of bi-encoder compared below.

Step 2: A dot product scores a query and an item

The score of a pair is the dot product of the query vector and the item vector. A high score means the model considers the item relevant to the query. Cosine similarity is a common alternative that first rescales each vector to length one.

Input: two vectors. Output: one number. This is the only place the query and the item interact, which is cheap and has a major consequence: the item vector does not depend on the query, so it can be computed once and reused for every query. A model that reads the query and the item together, a cross-encoder, can match more finely but cannot be precomputed; the cross-encoder and reranker page covers that trade.

Step 3: In-batch negatives and the popularity correction train the towers

Training uses pairs of a query and an item it matched, such as a search and the result that was clicked. For a batch of B such pairs, every query is scored against all B items: its own item is the positive, and the other B minus one are negatives. A softmax over those B scores and a cross-entropy against the true item is the InfoNCE loss from van den Oord et al. (2018). A temperature, a number that divides the scores before the softmax, sets how sharply close competitors are punished.

The catch is sampling bias. Items in a batch are drawn in proportion to how often they were clicked, so popular items appear as negatives far more often, and the model learns to push them down. Yi et al. (2019) correct it by estimating each item's frequency from streaming data and adjusting its score; the usual form subtracts the logarithm of the item's sampling probability from its logit. At probabilities 0.01 and 0.0001 (illustrative), the popular item's logit drops 4.6052 more.¹ Candidates usually leave this term out.

Step 4: Item vectors are precomputed and searched

After training, the item tower runs once over the whole catalogue and the vectors go into an approximate nearest-neighbour index; approximate means it checks a small part of the catalogue and may miss a few true top items in exchange for speed. At request time only the query tower runs, and the index returns the items whose vectors have the largest dot product with the query vector. Libraries such as Faiss, whose graphics processing unit (GPU) search (Johnson et al., 2017) handles billions of vectors, are built for this search.

Input: one query vector. Output: the top few hundred item ids. Candidate generation of this kind is followed by a slower ranking model on the short list, the two-stage design of Covington et al. (2016). You now have the whole system: two towers, a dot product, a batch loss with a correction, and an index.

1 · Two towers2 · Dot product3 · Batch loss4 · IndexReplay

Step-by-step description

1. Query and item are encoded by separate networks.
2. The two vectors meet in one dot product.
3. Other items in the batch act as negatives.
4. Item vectors are precomputed and searched.

• A query tower and an item tower are two separate networks that each output one vector.
• A dot product of the two vectors is the relevance score.
• A softmax over the other items in the batch trains both towers together.
• Item vectors are computed ahead of time and stored in a nearest-neighbour index.

Worked example

A batch of four pairs shows the loss by hand. The score matrix below is illustrative: entry (i, j) is the dot product of query i with item j, and the diagonal entries are the true pairs. Row 1 is (4.0, 1.0, 0.0, 1.0), row 2 is (0.5, 3.0, 1.0, 0.0), row 3 is (1.0, 0.0, 2.0, 1.5) and row 4 is (0.0, 1.0, 1.0, 3.5).¹

With a temperature of 1, the softmax-cross-entropy of each row is 0.1114, 0.2368, 0.7466 and 0.1776, a mean loss of 0.3181; with a temperature of 0.5 the mean falls to 0.1166, because dividing by a smaller number widens the gaps between scores.¹ Row 3 is the hardest: its true item scores 2.0 but item 4 scores 1.5, so a near miss costs 0.7466. Take row 3 at temperature 1: the exponentials of 1.0, 0, 2.0 and 1.5 are 2.7183, 1, 7.3891 and 4.4817, and they sum to 15.5891; the true item's probability is 7.3891 / 15.5891 = 0.4740, and its loss is -ln 0.4740 = 0.7466.¹

The equal-score baseline is ln B, the natural logarithm of the batch size B: with equal scores each of the B items gets probability 1/B, and the loss −ln(1/B) is ln B.¹ It is also the ceiling of the InfoNCE bound on mutual information, which van den Oord et al. (2018) write as ln B minus the loss and which tightens as B grows; more negatives also approximate a softmax over the whole catalogue more closely. Chen et al. (2020) report that contrastive image models gain from larger batches.

Text description of the diagram

• The equal-score baseline is ln B, the ceiling of the InfoNCE bound on mutual information.
• That baseline is 1.3863 at a batch of 4 and 8.3178 at a batch of 4,096.

For serving, brute-force search over the 10,000,000-item index costs 1,280,000,000 multiply-adds per query, which is why an approximate index is used instead.¹

Trade-offs and edge cases

A two-tower model trades accuracy for speed. Because the towers meet only at a dot product, it cannot model fine interactions between a query and an item, so it is used to cut millions of items to hundreds and a slower model reranks them. Karpukhin et al. (2020) trained a dual encoder for passage retrieval and report that it beat a strong keyword baseline by 9% to 19% absolute in top-20 accuracy across several question-answering datasets.

The same idea has a shared-weights variant. A siamese model uses one network for both sides, which suits two inputs of the same kind, such as two sentences; Reimers and Gurevych (2019) report a siamese network that cuts the cost of finding the most similar pair in a set of 10,000 sentences from about 65 hours to about 5 seconds against feeding every pair through a Bidirectional Encoder Representations from Transformers (BERT) model. A two-tower model with separate weights suits inputs of different kinds, such as a user history and a product.

Text description of the diagram

• A two-tower model uses different weights for query and item.
• A siamese model shares one set of weights between both sides.
• A cross-encoder reads query and item together, so item vectors cannot be precomputed.

Edge cases: a popular item can appear twice in one batch and then counts as a negative for its own duplicate (a false negative), so mask duplicates of the positive; the batch loss needs a large batch, which strains memory (the follow-up below covers fixes); in-batch negatives are mostly easy ones, so systems add hard negatives mined from earlier retrieval results; and an index of precomputed vectors goes stale when the item tower is retrained, so every item must be re-encoded and the index rebuilt. The machine learning system design guide places this stage inside a full design, and the transformer overview covers the encoders many towers are built from.

How to answer it in an interview

The spoken structure: say what the two towers are, say why that allows precomputation, give the loss, then name the popularity correction and the reranker.

Common mistakes

• "The two towers share weights." (wrong) Two-tower usually means separate weights; sharing them is the siamese variant. (right) "Separate towers for different data types, shared weights for same-type inputs."
• "At serving time I run both towers for every candidate." (wrong) That removes the benefit. (right) "Item vectors are precomputed; only the query tower runs per request."
• "In-batch negatives are free and unbiased." (wrong) They are free but sampled in proportion to popularity. (right) "I correct the logits by the log of the sampling probability."
• "A bigger batch always means better retrieval." (wrong) Gains shrink, and a bigger batch holds more copies of popular positives, which count as false negatives unless they are masked; the popularity bias itself does not change with batch size, because the optimal score is proportional to p(item | query) / p(item) at any number of negatives. (right) "Larger batches help with diminishing returns; I mask duplicates, add hard negatives and keep the logQ correction."

Follow-up questions

How do you train a two-tower model?

Collect pairs of a query and an item it matched, encode both sides with their towers, score every query against every item in the batch, and minimise the softmax cross-entropy against the true item. Add a popularity correction to the logits if items are sampled by frequency.

How do you fine-tune embedding models?

Start from a pretrained encoder, then continue training on your own matched pairs with the same in-batch softmax loss, often adding hard negatives mined from a first-stage search. Evaluate with a retrieval metric such as recall at k on a held-out set of queries, and re-encode the whole catalogue afterwards.

How would you increase in-batch negatives for InfoNCE if GPU memory is full?

Use gradient caching: Gao et al. (2021) separate the backward pass through the loss from the backward pass through the encoder, so the encoder processes a subset of the batch at a time and memory stays nearly constant. Another route is a queue of stored vectors from recent batches, as in momentum contrast (He et al., 2019).

What are the trade-offs of the InfoNCE loss?

Three knobs matter. A lower temperature widens the gaps between scores before the softmax; in the worked example the mean loss falls from 0.3181 to 0.1166 as it drops from 1 to 0.5. A larger batch gives more negatives but costs memory. In-batch negatives bring false negatives (duplicates of the positive) and popularity bias, so mask duplicates and keep the sampling correction.

When is a siamese setup better than two towers?

A siamese setup is better when both sides are the same kind of input, such as two questions or two sentences, because shared weights mean one model to train and one vector space by construction. Two towers fit better when the sides differ in type or features.

Why does a two-tower model need a reranker after it?

It scores a query and an item with one dot product, so it misses interactions that a model reading both together can catch. The first stage keeps recall high and cheap; the reranker spends more compute on the short list. The deep learning guide covers the neighbouring training topics.

Sources

1. Learning Deep Structured Semantic Models for Web Search using Clickthrough Data — Microsoft Research (Huang et al., CIKM 2013), 27 Oct 2013
2. Deep Neural Networks for YouTube Recommendations — Google Research (Covington et al., RecSys 2016), 15 Sep 2016
3. Sampling-Bias-Corrected Neural Modeling for Large Corpus Item Recommendations — Google Research (Yi et al., RecSys 2019), 16 Sep 2019
4. Dense Passage Retrieval for Open-Domain Question Answering (arXiv 2004.04906) — arXiv (Karpukhin et al.), 10 Apr 2020
5. Representation Learning with Contrastive Predictive Coding (arXiv 1807.03748) — arXiv (van den Oord et al.), 10 Jul 2018
6. Scaling Deep Contrastive Learning Batch Size under Memory Limited Setup (arXiv 2101.06983) — arXiv (Gao et al.), 18 Jan 2021
7. Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks (arXiv 1908.10084) — arXiv (Reimers and Gurevych), 27 Aug 2019
8. A Simple Framework for Contrastive Learning of Visual Representations (arXiv 2002.05709) — arXiv (Chen et al.), 13 Feb 2020
9. Momentum Contrast for Unsupervised Visual Representation Learning (arXiv 1911.05722) — arXiv (He et al.), 13 Nov 2019
10. Billion-scale similarity search with GPUs (arXiv 1702.08734) — arXiv (Johnson, Douze, Jegou), 28 Feb 2017