Two-level softmax sampling is shown to be systematically biased; S-2LS and SD-2LS corrections cut KL divergence by up to four orders of magnitude
Synopsis
The work proves that two-level softmax (2LS) sampling systematically misweights clusters by ignoring cluster size imbalance and intra-cluster similarity dispersion, breaking softmax's invariance for equally similar items; the authors propose S-2LS (size-weighted cluster scores) and SD-2LS (adding a second-order covariance correction), which on five large-scale datasets match exact softmax more closely at comparable or slightly higher cost than 2LS.
Figure 1 : Sampling ratio on VK-LSVD ( τ = 0.1 \tau=0.1 ) as a function of cluster size, averaged over queries. Standard 2LS undersamples large clusters and oversamples small ones relative to exact softmax ( R i ( N ) ( q ) ≠ 1 R_{i}^{(N)}(q)\neq 1 , p-value p < 0.01 p<0.01 ), while S-2LS and SD-2LS correct these biases.
arXivInterpretation
The paper gives the first formal characterization of 2LS sampling bias: the sampling ratio depends only on the cluster containing an item and not on its similarity, so all items in a cluster are over- or undersampled by the same factor and equally similar items in different clusters receive different probabilities. 2LS is widely deployed but its bias had not been formally characterized; the closest prior connection is an empirical observation of efficiency loss under cluster size imbalance by Tranheden et al. Derived analytically in Proposition 1 and empirically validated on five datasets including VK-LSVD via sampling-ratio plots against cluster size and intra-cluster variance.
The bias acts along two complementary axes: 2LS ignores cluster size imbalance, oversampling small clusters and undersampling large ones relative to their true softmax mass, and ignores intra-cluster similarity dispersion, undersampling clusters whose items vary more in similarity to the query. Decomposes the distortion into two separately correctable sources and gives a simplified ratio form under a Gaussian approximation. Propositions 2 and 3 establish the asymptotic behavior under an i.i.d. mixture setting via the strong law of large numbers; experiments isolate the dispersion bias once S-2LS removes the size bias.
S-2LS incorporates cluster size into the cluster-level score at the same sampling complexity as 2LS, while SD-2LS additionally approximates the moment generating function of the similarity via a second-order term using intra-cluster covariance matrices, adding an O(Kd^2) cost and asymptotically recovering exact softmax under Gaussian assumptions. Both corrections are drop-in replacements: S-2LS is free, SD-2LS trades moderate computation for stronger guarantees. Propositions 4, 5 and 6 provide the asymptotic analysis; measured latency shows S-2LS matching 2LS within noise and SD-2LS within roughly 1-2x on higher-dimensional data and up to about 3x on lower-dimensional data.
Across five large-scale datasets, SD-2LS attains the lowest KL on all natural corpora, with gains over 2LS ranging from one to two orders of magnitude at low temperature to one to four orders of magnitude at high temperature; S-2LS is consistently second best, already reducing KL by one to two orders of magnitude in the low-temperature regime most relevant to recommendation. Top-k truncation shows large errors even at low temperature and hierarchical softmax is markedly less accurate, particularly on YAMBDA, so the corrected methods give a strictly better fidelity-latency trade-off. Table 1 reports KL across five datasets and three temperatures; Table 2 reports per-query latency on single-threaded CPU; synthetic balanced and unbalanced corpora isolate the role of cluster size imbalance.
Perspective
The results target inference regimes with a fixed offline partition and repeated online sampling, such as sampling over millions of candidate items in recommendation or large-vocabulary generation in language models, where the dominant cost is inference for a fixed partition and the sublinear advantage holds. S-2LS suits highly imbalanced clusters with critical latency constraints; SD-2LS suits settings where intra-cluster dispersion is non-negligible and resources allow, with the paper noting the extra cost is amortized when K is much smaller than N. The authors recommend S-2LS as an unconditional replacement for standard 2LS and SD-2LS when intra-cluster variance is non-negligible.
SD-2LS relies on a truncated expansion of the moment generating function of the similarity, and the paper notes the second-order approximation may be less accurate for embedding distributions with high skewness or kurtosis, leaving higher-order corrections to future work. The bias analysis targets 2LS, and the authors note the mismatch recurs at each level of a tree-based hierarchical softmax and accumulates multiplicatively along the path, with per-node corrections left open. Latency is measured single-threaded on CPU with sequential queries, so overhead under other hardware and concurrency settings remains for readers to assess. The main-text figures focus on VK-LSVD due to space limits, with the full analysis for the other four datasets in Appendix C.