Cosine search, contrastive loss and negative sampling

Similarity by angle

Two embeddings are compared by the cosine of the angle between them, u·v / (‖u‖‖v‖), so that a vector's length, which often tracks frequency rather than meaning, does not decide the ranking. Normalise every row once, and cosine similarity becomes a dot product: searching a matrix of n normalised vectors for a query is one matrix-vector product and a partial sort, np.argpartition, rather than a full sort.

A zero vector has no direction. Decide in advance what its similarity is (zero is the usual answer) instead of dividing by zero.

Training by contrast: InfoNCE

Given n pairs that belong together, a query and its matching key, contrastive learning asks each query to pick out its own key among all n. The logits are the cosine similarities divided by a temperature τ, and the loss is the cross-entropy of the diagonal:

logits = (Qn @ Kn.T) / tau                    # (n, n), rows normalised first
loss = mean_i( logsumexp(logits[i]) - logits[i, i] )

The temperature matters. Small values sharpen the distribution and make the loss focus on the hardest negatives; they also make the logits large, which is why the log-sum-exp must subtract the row maximum. CLIP uses this loss in both directions, queries to keys and keys to queries, and averages the two.

Making it cheap: negative sampling

Word2vec's skip-gram model would need a softmax over the whole vocabulary for every training pair. Negative sampling replaces it with k binary questions: is this the true context word (push σ(u·v) towards 1), and are these k random words not context words (push σ(u_k·v) towards 0)? The loss is

L = -log σ(u·v) - Σ_k log σ(-u_k·v)

and its gradients have a simple form, because the derivative of −log σ(x) is σ(x) − 1. Write log σ(x) as −softplus(−x) and compute softplus stably; a large negative dot product otherwise makes log σ overflow to −∞.

What interviewers listen for

• Normalise once, then use dot products.
• A stated rule for zero vectors and for ties in the ranking.
• The row maximum subtracted inside every log-sum-exp.
• Gradients derived from σ′ = σ(1 − σ), and checked against finite differences.