Two-tower retrieval

Retrieval is a different job

A production recommender is usually two parts. Retrieval narrows a catalogue of millions to a few hundred candidates. Ranking orders those candidates with an expensive model. The two-tower model is the standard retrieval design, because it turns the narrowing into a nearest-neighbour search.

One tower encodes the user and the context. The other encodes the item. Each ends in a vector, and the score is a dot product or a cosine. The towers never see each other’s features. That limit is what lets the item vectors be computed ahead of time. It is also the model’s main weakness.

The split is the point. Item vectors are built offline and loaded into an index. The user tower runs at request time and produces one vector. A query is then a nearest-neighbour lookup, not a forward pass over every item. Without that split, retrieval at catalogue scale is not possible.

Training with in-batch negatives

The loss is a softmax over the scores of one positive item and many negatives. In-batch negatives use the other items in the same batch as negatives for every user. That is cheap and works well when the batch is large. The scores are scaled by a temperature, which sets how sharply the model separates the positive.

There is a subtle failure. An item that appears twice in a batch is a false negative for the other user who touched it. Remove duplicates within the batch, or mask the diagonal correctly, so the model is not punished for a correct match. With small batches, add negatives sampled from the catalogue to make up.

Hard negatives sharpen the boundary. A random item is easy to reject. An item the user almost clicked teaches more. Hard negatives are usually mined from a previous model’s high-scoring misses. Then they are mixed with random negatives, so the model does not overfit to a narrow set of mistakes.

A two-tower model in TensorFlow Recommenders

import tensorflow as tf
import tensorflow_recommenders as tfrs

class TwoTower(tfrs.models.Model):
    def __init__(self, user_tower, item_tower):
        super().__init__()
        self.user_tower = user_tower
        self.item_tower = item_tower
        self.task = tfrs.tasks.Retrieval()

    def compute_loss(self, features, training=False):
        user = self.user_tower(features["user_id"])
        item = self.item_tower(features["item_id"])
        return self.task(query_embeddings=user, candidate_embeddings=item)

The index and its freshness

At serving time the item tower’s output lives in an approximate nearest-neighbour index (a fast index that may miss a few true matches). FAISS, HNSW, ScaNN and managed vector stores are examples. The index trades recall (the share of true matches it returns) for speed, and that recall is a product metric. An index that returns 90 percent of the true top hundred is fine for retrieval. It is not fine if the missing ten percent are the new releases.

Freshness is the harder problem. A new item has no vector until the item tower runs. A changed price or title changes the features the tower reads. Re-embed on a schedule, and keep a path for items that arrive between rebuilds. The usual answer is a small exact index for recent items beside the large approximate one.

Judge retrieval on its own terms. Measure recall at k against the full catalogue, coverage of the item catalogue, and the share of requests with fewer than k candidates. A retrieval model that improves ranking metrics while returning the same head items has not improved retrieval. It has moved the popularity bias one stage earlier.

Feature choices and leakage

The user tower can read history, demographics and context. The item tower can read text, image vectors, price and category. Neither can read the other side’s raw features at serving time, because the item vector is computed before the request exists. A cross feature between user and item belongs in the ranking model, not in retrieval.

Leakage in retrieval is usually a feature that is only known after the interaction. A click count updated in real time can do it. So can a rating average that includes the test window, or an item’s total watch time. Each makes the towers look excellent offline and useless online. Build every feature as of the request time. Test the pipeline with a temporal split.

Cold items need a content path. If the item tower only has an id vector, a new item gets a random vector and will never be retrieved. Text and image features give it a sensible starting point. A small exploration budget gives it the impressions that turn the prior into evidence.