The fast softmax sampler that weighs a cluster like one item

Suppose a recommender has millions of candidate songs. Exact softmax scores every song, so a common shortcut samples a cluster first and a song second. Bendada and Salha-Galvan show that the shortcut's cluster softmax quietly treats a centroid as though it represented one item. A cluster of ten equally good songs can receive the same mass as a cluster of one.

Softmax assigns each item weight exp(score / temperature). The true mass of a cluster is therefore a sum of exponentials. Standard two-level softmax, abbreviated 2LS, replaces that sum with the exponential of the cluster's mean score. The missing multiplier is the number of items, and the missing shape term comes from variation around the mean.

S-2LS adds cluster size. SD-2LS also adds half the variance of scaled item scores along the current query. The variance term is the Gaussian moment-generating-function correction: if projected scores are normal, their exponential mean is exp(mean + variance / 2).

def cluster_probs(clusters, temperature, method):
    logits = []
    for scores in clusters:
        scaled = [score / temperature for score in scores]
        mean = sum(scaled) / len(scaled)
        variance = sum((x - mean) ** 2 for x in scaled) / len(scaled)
        logit = mean
        if method != "2ls": logit += log(len(scaled))
        if method == "sd-2ls": logit += variance / 2
        logits.append(logit)
    return softmax(logits)

The TypeScript core used by every chart implements this rule directly. The following deterministic toy holds size and mean fixed, then spreads the second cluster's scores symmetrically. It is illustrative, not paper data.

Mean and variance do not determine an arbitrary distribution's exponential mass. Here, two four-item clusters both have mean zero and variance one. One has scores [-1, -1, 1, 1]; the other has one large negative score and three smaller positive scores. SD-2LS sees identical summaries and assigns 50% to each, while exact softmax does not.

The paper evaluates 1.2 million GloVe words, 7.7 million YAMBDA tracks, 19.6 million VK-LSVD videos, and two million-item synthetic corpora. For each dataset and temperature it compares probabilities for 1,000 query embeddings. On the three natural corpora, SD-2LS has the lowest KL at every tested temperature. S-2LS matches ordinary 2LS latency within measurement noise; SD-2LS is slower because it evaluates cluster covariance quadratic forms.

At temperature 0.1, the reported SD-2LS KL is roughly 59 times smaller than ordinary 2LS on GloVe, 149 times smaller on VK-LSVD, and 60 times smaller on YAMBDA. But the synthetic balanced result at temperature 0.05 is a useful brake on universal claims: top-k is best there, and SD-2LS is slightly worse than ordinary 2LS. Structure determines which approximation wins.

First, stratify residual KL by cluster skewness and kurtosis rather than reporting only an average. Second, compare diagonal, low-rank, and full-rank covariance corrections under a fixed latency budget. Third, test whether errors multiply down deeper hierarchical trees, as the authors conjecture. Finally, measure end-to-end recommendation quality: probability fidelity is a mechanism metric, not proof of better user outcomes.