Two-tower retrieval models

Setup#

pip install numpy

NumPy only, on purpose. Everything below is fifty lines of arithmetic, and writing the gradients by hand once makes the in-batch-negatives trick permanent in a way that loss.backward() never does. The PyTorch version is at the end.

A two-tower model with in-batch negatives#

The user tower is a lookup: one learned vector per person. The item tower is a linear map from tags to a vector, so an item the model has never seen still gets a position.

import numpy as np

USERS = ["Aarav", "Bhavna", "Chetan", "Divya", "Esha", "Farhan"]
FEATURES = ["drama", "horror", "comedy", "sports", "action", "fantasy"]
ITEMS = {                     # every item is described by tags, never by an id
    "Lagaan":    [1, 0, 0, 1, 0, 0],
    "Dangal":    [1, 0, 0, 1, 0, 0],
    "3 Idiots":  [1, 0, 1, 0, 0, 0],
    "Sholay":    [1, 0, 0, 0, 1, 0],
    "Tumbbad":   [0, 1, 0, 0, 0, 1],
    "Stree":     [0, 1, 1, 0, 0, 0],
    "Bhoot":     [0, 1, 0, 0, 1, 0],
    "Naya-Film": [1, 0, 0, 1, 0, 0],      # added after training, never seen by the model
}
TITLES = list(ITEMS)
F = np.array([ITEMS[t] for t in TITLES], dtype=float)

# (user, item) pairs the model may learn from. Naya-Film appears in none of them.
PAIRS = [(0, 0), (0, 1), (0, 3), (1, 0), (1, 2), (1, 3), (2, 1), (2, 2), (2, 3),
         (3, 4), (3, 5), (4, 4), (4, 6), (5, 5), (5, 6)]
u_idx = np.array([u for u, _ in PAIRS])
i_idx = np.array([i for _, i in PAIRS])
N = len(PAIRS)

D, LR, EPOCHS, TEMP = 4, 0.5, 400, 0.2
rng = np.random.default_rng(0)
U = rng.normal(0, 0.1, (len(USERS), D))      # user tower: one vector per person
W = rng.normal(0, 0.1, (len(FEATURES), D))   # item tower: tags -> vector


def normalise(M):
    return M / np.maximum(np.linalg.norm(M, axis=1, keepdims=True), 1e-9)


def through_normalise(M, dM):
    """Push a gradient back through the row-normalisation step."""
    n = np.maximum(np.linalg.norm(M, axis=1, keepdims=True), 1e-9)
    Mh = M / n
    return (dM - (dM * Mh).sum(axis=1, keepdims=True) * Mh) / n


print("epoch    loss")
for epoch in range(EPOCHS + 1):
    Vall = F @ W                             # the whole item tower, one matrix multiply
    Ub, Vb = normalise(U[u_idx]), normalise(Vall[i_idx])
    logits = Ub @ Vb.T / TEMP                # each pair scored against every other pair
    logits -= logits.max(axis=1, keepdims=True)
    P = np.exp(logits)
    P /= P.sum(axis=1, keepdims=True)
    if epoch % 100 == 0:
        loss = -np.mean(np.log(P[np.arange(N), np.arange(N)] + 1e-12))
        print(f"{epoch:5d}  {loss:.4f}")
    if epoch == EPOCHS:
        break

    G = (P - np.eye(N)) / (N * TEMP)         # gradient of the loss wrt the logits
    dU_raw = through_normalise(U[u_idx], G @ Vb)
    dV_raw = through_normalise(Vall[i_idx], G.T @ Ub)
    dU = np.zeros_like(U)
    np.add.at(dU, u_idx, dU_raw)             # several pairs share a user row
    U -= LR * dU
    W -= LR * (F[i_idx].T @ dV_raw)

Un, Vn = normalise(U), normalise(F @ W)
SEEN = {u: {i for uu, i in PAIRS if uu == u} for u in range(len(USERS))}

print("\ntop 2 items for each person, out of everything including the unseen film:")
for u in range(len(USERS)):
    scores = Vn @ Un[u]
    order = sorted(range(len(TITLES)), key=lambda i: (-scores[i], TITLES[i]))
    picks = [i for i in order if i not in SEEN[u]][:2]
    print(f"  {USERS[u]:<7} " + "   ".join(f"{TITLES[i]} {scores[i]:+.3f}" for i in picks))

print("\nNaya-Film was never in the training pairs. Its score anyway:")
j = TITLES.index("Naya-Film")
for u in range(len(USERS)):
    print(f"  {USERS[u]:<7} {float(Vn[j] @ Un[u]):+.3f}")

epoch    loss
    0  3.3444
  100  1.8388
  200  1.8211
  300  1.8169
  400  1.8147

top 2 items for each person, out of everything including the unseen film:
  Aarav   Naya-Film +0.785   3 Idiots +0.110
  Bhavna  Dangal +0.757   Naya-Film +0.757
  Chetan  Lagaan +0.547   Naya-Film +0.547
  Divya   Bhoot -0.154   3 Idiots -0.302
  Esha    Stree -0.112   Dangal -0.417
  Farhan  Tumbbad +0.135   3 Idiots -0.063

Naya-Film was never in the training pairs. Its score anyway:
  Aarav   +0.785
  Bhavna  +0.757
  Chetan  +0.547
  Divya   -0.889
  Esha    -0.417
  Farhan  -0.359

The last block is the whole point of the design#

Naya-Film never appeared in a single training pair. The model has no idea it exists as an entity.

It still scores +0.785 for Aarav and -0.889 for Divya. The item tower read its tags, and the tags were enough. Compare that with matrix factorisation, where an item with no interactions has a factor vector of exactly zero and scores zero for everybody.

That difference is why this architecture won.

Three things in the output that you should not gloss over#

The loss stops at 1.81, not near zero. That is not a bug and it is not a badly tuned learning rate. The batch contains fifteen pairs, and Sholay is a positive for three different users. When Aarav's row treats Bhavna's positive as a negative, the model is being told something false.

These are false negatives, and they are unavoidable with in-batch negatives. The floor on the loss is the price. Popular items suffer most, because they appear in the most rows.

Bhavna scores Dangal and Naya-Film at exactly +0.757. Look at their tag rows: identical. A tower that reads only features cannot distinguish two items with the same features. Production towers therefore concatenate a learned per-item vector onto the feature vector, which restores the distinction — and reintroduces a cold-start gap for that component.

Divya's best candidate scores -0.154, a negative number. Retrieval scores are not probabilities. Their sign and scale carry no meaning on their own; only the ordering does. Anyone who puts a threshold like score > 0.5 on a retrieval model's output has misread what the number is.

Line by line, for the parts that are not obvious#

logits = Ub @ Vb.T / TEMP — the in-batch negatives trick. An (N, D) @ (D, N) multiply scores every user in the batch against every item in the batch. Row n has one correct answer, at column n, and N-1 free negatives you did not have to sample. Bigger batches give harder training, which is why retrieval models are trained with batch sizes in the thousands.

/ TEMP — temperature. Dividing the logits sharpens the softmax. With cosine similarities capped at 1.0, the logits are tiny and the gradient is nearly flat without it. Set TEMP = 1.0 and training barely moves. Values between 0.05 and 0.2 are typical.

logits -= logits.max(axis=1, keepdims=True) — the standard numerically stable softmax. Skip it and large logits overflow np.exp. It changes nothing mathematically, because softmax is shift-invariant.

(P - np.eye(N)) / (N * TEMP) — the gradient of softmax cross-entropy is the predicted distribution minus the one-hot target. The identity matrix is the target here, since row n's answer is column n.

np.add.at(dU, u_idx, dU_raw) — accumulation, not assignment. Aarav appears in three pairs, so three gradients must be summed into his row. Writing dU[u_idx] = dU_raw keeps only the last one, and the bug is silent.

F[i_idx].T @ dV_raw — the chain rule through the item tower. Since V = F @ W, the gradient with respect to W is F.T @ dV. Every item sharing a tag pushes the same column of W, which is exactly the sharing that lets an unseen item inherit a sensible vector.

The same model in PyTorch#

# pip install torch  (CPU wheel, roughly 200 MB)
import torch, torch.nn as nn, torch.nn.functional as Fn

TEMP = 0.2
user_tower = nn.Embedding(6, 4)
item_tower = nn.Linear(6, 4, bias=False)
opt = torch.optim.Adam(list(user_tower.parameters()) + list(item_tower.parameters()), lr=0.05)

feats = torch.tensor([[1,0,0,1,0,0],[1,0,0,1,0,0],[1,0,1,0,0,0],[1,0,0,0,1,0],
                      [0,1,0,0,0,1],[0,1,1,0,0,0],[0,1,0,0,1,0]], dtype=torch.float)
pairs = torch.tensor([[0,0],[0,1],[0,3],[1,0],[1,2],[1,3],[2,1],[2,2],[2,3],
                      [3,4],[3,5],[4,4],[4,6],[5,5],[5,6]])

for step in range(400):
    u = Fn.normalize(user_tower(pairs[:, 0]), dim=1)
    v = Fn.normalize(item_tower(feats[pairs[:, 1]]), dim=1)
    logits = u @ v.T / TEMP
    loss = Fn.cross_entropy(logits, torch.arange(len(pairs)))   # the diagonal is the target
    opt.zero_grad(); loss.backward(); opt.step()

print(round(loss.item(), 4))

No output block for this one, deliberately. The initialisation and the optimiser differ from the NumPy version, so the printed number depends on your PyTorch version. What should match is the shape of the result: a loss that falls quickly and then flattens well above zero, for the false-negative reason above.

Common mistakes#

Forgetting to normalise. Without unit-length vectors, popular items grow long vectors and win every dot product regardless of direction. Normalising plus a temperature is the standard pairing, and neither works properly alone.

Ignoring the popularity bias in in-batch negatives. Negatives are sampled from the traffic distribution, so popular items appear as negatives far more often and get systematically pushed down. The standard correction subtracts log(sampling_probability) from each logit during training — see the researcher block.

Using retrieval scores as final scores. This model is a filter. Its job is recall of a few hundred candidates. Ordering the page is a separate, richer model's job.

Rebuilding the entire item index on every deploy. Item vectors change whenever the tower changes, so a new model version invalidates the whole index. Plan for that: version the index alongside the model, and never serve a user tower against an index built by a different one. Mismatched versions produce results that look plausible and are nonsense.

Judging it with nDCG. Retrieval is measured by recall@k for k in the hundreds or thousands. Whether the right item is at position 3 or 300 within the candidate set does not matter, because the ranker re-orders them anyway.

Try it yourself#

Set TEMP = 1.0 and re-run. The loss flattens at 2.1442 instead of 1.8147, and it stops improving after about a hundred epochs. Six users and seven items are easy enough that the recommendations survive anyway — the loss gap is the honest signal, and temperature is consistently one of the highest-impact hyperparameters in a retrieval model.

Then give Naya-Film the tags [0, 1, 0, 0, 0, 1] — horror and fantasy — and re-run. Every score flips sign: Aarav goes from +0.785 to -0.477, Divya from -0.889 to +0.699. Not one training pair changed.

Watching a model's opinion of an item move purely through its description is the clearest demonstration of what a feature-based tower buys you.

What to learn next#

• Vector databases — the nearest-neighbour index this design assumes.
• Sequence-aware recommendation — richer user towers that read order.
• Ranking metrics — why recall is the retrieval metric.