Sampled Softmax

Training a retrieval model as classification over the catalogue means a softmax with one output per item. At ten million items, computing that denominator once per training example is the entire cost of training. Sampled softmax is the standard way of not paying it.

The mechanism
A full softmax normalises the user-item score over every item. Sampled softmax replaces that denominator with the positive item plus a small set of sampled negatives, and corrects each sampled logit by subtracting the log of its sampling probability. That correction is what keeps the gradient an approximately unbiased estimate of the full-softmax gradient; drop it and frequent items are systematically over-penalised. In a two-tower model the negatives are commonly the other items already in the minibatch (in-batch negatives), which costs one extra matrix multiply plus a logQ correction, because batch membership is itself popularity-weighted.

Choices that decide whether it works

• Negative distribution. Uniform negatives are easy and teach little; popularity-weighted negatives sharpen head items; hard negatives mined from an index sharpen the decision boundary at the cost of gradient noise.
• Count. More negatives tighten the approximation and raise recall, which is why large-batch training matters far more for retrieval than for ranking.
• False negatives. A sampled negative the user would have loved is a label error, and drawing one gets likelier as sampling tilts toward popularity.

Why it matters
This is what makes a retrieval tower trainable at catalogue scale, and the negatives shape the embedding space that FAISS or ScaNN later searches. Get it wrong and the index is fast, well-engineered, and full of the wrong neighbours.