Two-Tower Retrieval with In-Batch Negatives

Train a two-tower retrieval model, the kind that picks a few hundred candidates out of millions of items at the first stage of a recommender.

A user tower maps user features to a vector, and an item tower maps item features to a vector. The score of a pair is the cosine similarity of the two vectors. Training uses in-batch negatives: in a batch of B true (user, item) pairs, the other B - 1 items act as negatives for each user.

u = normalise(user_vecs), v = normalise(item_vecs)   # (B, d) each
logits = u @ v^T / temperature                       # (B, B)
loss = cross_entropy(logits, labels = [0, 1, ..., B-1])

Write two functions:

• in_batch_softmax_loss(user_vecs, item_vecs, temperature): the loss above.
• train_two_tower(user_feats, item_feats, pairs): trains both towers and returns them as (user_tower, item_tower). Each one is a callable that maps an (n, features) tensor to an (n, d) tensor.

Function signature

def in_batch_softmax_loss(user_vecs, item_vecs, temperature):

def train_two_tower(user_feats, item_feats, pairs):

Example

Input

user_vecs = [[1, 0], [0, 1]]
item_vecs = [[2, 0], [0, 3]]
temperature = 1.0

Output

0.3133

Reasoning

After normalisation both sets are [[1, 0], [0, 1]], so the logits are the identity matrix. Each row has a logit of 1 for its own item and 0 for the other. The cross entropy of each row is -log(e / (e + 1)) = 0.3133, and so is the mean.

Constraints and notes

• user_feats is (n_users, 16), item_feats is (n_items, 16) and pairs is (N, 2) of (user id, item id). Look up the features of a batch by indexing with the ids.
• The tests score every item for each held-out pair and need recall@10 of at least 0.25: the true item must be in the user's top 10 for at least a quarter of the pairs. A popularity ranking gets about 0.03.
• Small MLP towers, a temperature near 0.1, Adam and about ten epochs work. Return the towers in eval mode.