Two-Level Softmax Sampling Done Right: Correcting Bias from Size Imbalance and Dispersion

AI summary

Overview

Research area: Machine learning — scalable sampling algorithms, with applications in recommendation systems, language modeling, and information retrieval.

Technical level: Advanced. The paper builds on asymptotic analysis, moment generating functions, and Gaussian approximations, though its practical conclusions are easy to state.

Scope: The paper characterizes the sampling bias of two-level softmax (2LS) sampling, proposes two bias-corrected variants (S-2LS and SD-2LS), and validates them on five large-scale datasets.

What This Paper Is About

Softmax sampling — picking items in proportion to the exponential of their similarity to a query — costs time linear in the number of items, which is impractical when a system has millions of candidates. Two-level softmax (2LS) sampling avoids this by first sampling a cluster and then an item inside it, giving sublinear time. This paper shows that 2LS silently distorts the sampling distribution: it oversamples small clusters, undersamples large ones, and undersamples clusters whose items vary a lot in similarity to the query. The authors derive corrections for both distortions and test whether the corrected samplers match true softmax more faithfully.

Key Contributions

1. Formal characterization of 2LS bias. The authors define the sampling ratio R_i^(N)(q) = p_2LS(i|q) / p(i|q) and prove (Proposition 1) that this ratio depends only on the cluster containing item i, not on the item's own similarity q^T x_i. This means 2LS breaks softmax's invariance: two items with identical similarity to the query can be sampled with different probabilities if they sit in different clusters.

2. Decomposition of the bias into two axes. Proposition 2 gives the asymptotic ratio R_k^(∞)(q) → c(q) · exp(q^T μ_k/τ) / (π_k · E_{x~D_k}[exp(q^T x/τ)]), and Proposition 3 shows that under a Gaussian approximation R_k^(∞)(q) ∝ 1 / (π_k exp(½ σ_k²(q))). This identifies two separate sources of misweighting: cluster size imbalance (the π_k term) and intra-cluster similarity dispersion (the σ_k²(q) term).

3. Two corrected samplers. Size-Corrected 2LS (S-2LS) multiplies the cluster score by |C_k| with no change in sampling complexity; Size- and Dispersion-Corrected 2LS (SD-2LS) additionally adds ½ σ_k²(q) to the cluster score. Proposition 4 shows S-2LS removes the cluster-size dependence, and Proposition 6 shows that under Gaussian assumptions SD-2LS recovers the exact softmax asymptotically (R_{i,SD-2LS}^(N)(q) → 1).

4. Empirical validation on five large-scale datasets. Experiments on VK-LSVD, YAMBDA, GloVe-100, Synth-balanced, and Synth-unbalanced confirm the predicted biases and show improved fidelity to softmax, with code released at https://github.com/bendadaw/neurips-2ls.

Main Findings

• 2LS breaks softmax invariance in practice. On VK-LSVD at τ = 0.1, standard 2LS assigns widely different sampling probabilities to items with identical query-item dot-product similarity, while S-2LS and SD-2LS markedly reduce this dispersion (Figure 3).

• Small clusters are oversampled, large clusters undersampled. Figure 1 (VK-LSVD, τ = 0.1) shows the sampling ratio deviating from 1 as a function of cluster size for 2LS, with p < 0.01. S-2LS and SD-2LS concentrate the ratio tightly around 1.

• High-dispersion clusters are undersampled. Figure 2 isolates the variance effect: S-2LS corrects the cluster-size bias but not the variance bias, revealing that clusters with high intra-cluster similarity variance are undersampled and low-variance clusters oversampled relative to softmax (p < 0.01).

• SD-2LS achieves the lowest KL divergence on all natural corpora. At τ = 0.05 on GloVe-100, SD-2LS reaches KL = 0.0288 ± 0.0076 versus 0.0961 ± 0.0110 for 2LS. On VK-LSVD at the same temperature, SD-2LS reaches 0.0129 ± 0.0008 versus 0.0541 ± 0.0023 for 2LS. On YAMBDA, 0.0177 ± 0.0056 versus 0.0485 ± 0.0018.

• Gains grow with temperature. The improvement over 2LS ranges from 3–5× at τ = 0.05 to one to four orders of magnitude at τ ∈ {0.1, 0.2}. At τ = 0.2 on GloVe-100, SD-2LS reaches 1.2e-06 ± 9e-08 versus 0.0677 ± 0.0001 for 2LS.

• S-2LS is consistently second best. It reduces KL by one to two orders of magnitude in the low-temperature regime most relevant for recommendation, while matching 2LS latency within measurement noise.

• Competing baselines are less accurate. Top-k truncation (k = 1000, via Faiss IVFFlat) shows large errors on natural corpora — for instance 5.4103 ± 0.0217 on GloVe-100 at τ = 0.2. Hierarchical softmax remains significantly less accurate, particularly on YAMBDA (1.4484 ± 0.0563 at τ = 0.05), likely from accumulated routing errors along the tree.

• Synthetic data confirms the mechanism. On Synth-balanced, where cluster sizes are nearly uniform, the advantage of the corrections is limited and top-k can perform well at low temperature (KL 0.0003 ± 0.0001 at τ = 0.05). On Synth-unbalanced, which mirrors heavy-tailed natural corpora, SD-2LS dominates and the gap to 2LS reaches an order of magnitude.

• Overhead is small. S-2LS matches 2LS latency. SD-2LS, which uses full-rank covariance matrices Σ_k ∈ R^(d×d), incurs a 1.2–1.8× slowdown on higher-dimensional datasets and up to ~3× on lower-dimensional ones.

Methodology in Plain English

The authors start from the observation that 2LS splits the sampling probability into two factors: one over cluster centroids and one over items inside the chosen cluster. Because the cluster-level softmax only sees centroids, it cannot know how many items a cluster holds or how spread out they are — it compares clusters as if each were a single point.

To measure the damage precisely, the authors define a ratio comparing what 2LS does to what exact softmax would do for the same item. They prove this ratio is constant across all items in a cluster, then study what happens to that ratio as the number of items grows large, modeling item embeddings as draws from a mixture of K component distributions. This yields an asymptotic formula in which cluster size (π_k) and dispersion (the variance σ_k²(q), via a moment generating function) appear as separate, interpretable terms. Assuming roughly Gaussian similarities turns the formula into a simple ratio ∝ 1/(π_k exp(½ σ_k²(q))).

Armed with that formula, they design two fixes at the cluster-selection step only. S-2LS multiplies each cluster's centroid score by the cluster's size. SD-2LS multiplies by the size and also adds ½ σ_k²(q), the variance of the query-projected similarities within that cluster. Because the variance can be written as q^T Σ_k q using full-rank intra-cluster covariance matrices, both corrections can be computed without ever touching individual items.

Empirically, they cluster five datasets with k-means into K = 1024 clusters, share that partition across 2LS, S-2LS, and SD-2LS, and sample 1000 query embeddings per dataset. They compare against exact softmax, top-k softmax (k = 1000), and hierarchical softmax built as a recursive binary k-means tree, at temperatures τ ∈ {0.05, 0.1, 0.2}. Fidelity is measured by KL divergence in the approximate-to-exact direction (which stays finite when an approximate method assigns zero probability to some items, as top-k does), alongside per-query latency in milliseconds on a single-threaded CPU with sequential queries.

Why This Matters

Sampling is a bottleneck operation in any system whose candidate set is too large for exact softmax, and the paper shows that a widely deployed approximation has a systematic, previously uncharacterized flaw. The corrections are drop-in replacements: they reuse the same partition and centroids, and the only change is the cluster sampling rule.

Real-world applications:

• Recommender systems. Both VK-LSVD (19.6M videos, d = 64) and YAMBDA (7.7M music tracks, d = 128) come from industrial recommenders, where sampling items by user similarity decides what a user actually sees. The sublinear regime matters because a fixed partition is reused across many queries online.
• Large language model text generation. 2LS-style samplers are used to sample next tokens from vocabularies too large for exact softmax; biases in cluster weighting translate directly into skewed token distributions.
• Information retrieval and approximate nearest neighbor search. IVF indexing, implemented in libraries such as Faiss, mirrors the 2LS two-stage structure, so the bias analysis is relevant to retrieval pipelines that rely on cluster-restricted search.
• Machine translation and text or image classification. These are listed among the adopted applications of 2LS and closely related variants.

Industry relevance: S-2LS offers a strict Pareto improvement over 2LS — the same latency and complexity with lower KL divergence — which makes adoption essentially free. SD-2LS adds a d²K term instead of dK, negligible when d ≪ N, giving a stronger correction for teams with compute headroom. The paper's conclusion recommends replacing standard 2LS with these variants in future work.

Future Directions

• Higher-order corrections for non-Gaussian embeddings. The SD-2LS analysis relies on a truncated expansion of the moment generating function. It is exact under Gaussian assumptions and empirically effective on real data, but may be less accurate for embedding distributions with significant skewness or kurtosis. Adding higher-order terms while preserving computational efficiency is left open.
• Extending the corrections to hierarchical softmax. The authors note that the mismatch identified in 2LS recurs at each level of a root-to-leaf hierarchy and accumulates multiplicatively along the path, suggesting node-level corrections as a follow-up.
• Choosing between S-2LS and SD-2LS. The paper describes the regimes — heavy cluster imbalance and latency-critical versus substantial dispersion with available resources — but does not provide a decision rule or an estimate of when the d² overhead is justified.
• Characterizing accuracy beyond the Gaussian regime. The relationship between the truncated MGF approximation error and measurable properties of real corpora is not quantified, leaving the practitioner with empirical results rather than a diagnostic.

Target Audience

This paper suits machine learning engineers and researchers who deploy or study scalable sampling for large item sets or vocabularies — particularly those working on recommender systems, large language model inference, or approximate nearest neighbor search. It is also relevant to researchers interested in hierarchical softmax and other softmax approximations, since the analysis is framed against those alternatives. Readers need comfort with softmax, k-means clustering, covariance matrices, and asymptotic arguments to follow the proofs, though the practical recommendation is straightforward: use S-2LS instead of 2LS at no cost, and SD-2LS when the extra computation is affordable.

Authors’ abstract

Sampling from a softmax distribution is a fundamental operation in machine learning, but its linear complexity in the number of items makes exact sampling impractical at scale. Two-level softmax (2LS) sampling is a popular alternative enabling sublinear-time sampling. Assuming items are partitioned into clusters, 2LS first samples a cluster and then an item within it. In this paper, we show that, despite its advantages, 2LS introduces systematic and undesirable sampling biases, which arise from misweighting clusters by ignoring both cluster size imbalance and intra-cluster similarity dispersion. We propose two sampling methods, Size-Corrected 2LS (S-2LS) and Size- and Dispersion-Corrected 2LS (SD-2LS), which correct these biases and provide provably better softmax approximations with negligible to non-existent computational overhead. In-depth experiments on five large-scale datasets validate the improved sampling properties of our methods. We recommend their consistent use in place of standard 2LS in future work.