In-Batch Sampled Softmax + logQ Correction

Background

Training a two-tower retrieval model with the full softmax over $10^9$ items is infeasible. In-batch sampled softmax treats the other items in a training batch as negatives. The logQ correction de-biases for the fact that those negatives are not uniform — popular items appear more often than uniform, so their logits are reduced by $\log Q(j)$ to keep the gradient unbiased.

Problem statement

Implement sampled_softmax_loss(user_embs, item_embs, sampling_log_probs). Compute:

$$\ell_i \;=\; -\log\!\frac{\exp(u_i \cdot v_i - \log Q(i))}{\sum_j \exp(u_i \cdot v_j - \log Q(j))}$$

Return the mean across the batch.

Input

• user_embs — np.ndarray shape $(N, d)$.
• item_embs — np.ndarray shape $(N, d)$.
• sampling_log_probs — np.ndarray length $N$, the empirical $\log Q(j)$ for each item in the batch.

Output

Returns a Python float >= 0.

Examples

Example 1 — identity embeddings → small loss

Input:  user_embs == item_embs == 10 * I_4; sampling_log_probs = zeros
Output: loss small (diagonal dominates the softmax)

Example 2 — logQ correction changes the loss

Input:  any embeddings, two different sampling distributions
Output: with non-zero logQ, the loss differs from the no-correction version

Example 3 — shape mismatch

Input:  user_embs (3, 4), item_embs (5, 4)
Output: ValueError

Constraints

• user_embs, item_embs, and sampling_log_probs must agree on $N$.
• Use a numerically stable logsumexp (subtract row max before exp).
• Returns a plain Python float.

Notes

• Why subtract logQ. Without the correction, popular items appear too often as negatives and their logits get suppressed at train time — the model under-weights them at retrieval.
• In-batch vs sampled negatives. In-batch is free (same batch, no extra forward passes). Sampled negatives (drawn from a popularity distribution) can give larger effective $N$ at the cost of more work per step.