PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Walid Bendada, Guillaume Salha-Galvan

arXiv:2610.10483v1cs.LGcs.IRstat.ML

TL;DR

Exact softmax sampling is costly at scale, while 2LS offers a sublinear alternative but systematically misweights clusters. The paper proposes S-2LS and SD-2LS to correct size and dispersion biases, and reports improved fidelity with different latency trade-offs.

  • Problem

    Exact softmax sampling costs O(dN) per query, making it impractical for large item sets.

    Source: §1 · p. 1, §2.1 · p. 2

  • Method

    S-2LS corrects cluster-size imbalance, while SD-2LS also accounts for intra-cluster similarity dispersion.

    Source: §1 · p. 2, §4.3 · p. 6

  • Results

    Across natural corpora and temperatures, SD-2LS achieves the lowest KL and S-2LS is second-best; gains over 2LS range from 3–5× at τ = 0.05 to one to four orders of magnitude at τ ∈{0.1, 0.2}.

    Source: §5.2 · p. 9

  • Takeaways & Limitations

    S-2LS consistently improves fidelity over 2LS at no additional cost, while SD-2LS reduces approximation error with moderate overhead.

    Source: §6 · p. 9

    The SD-2LS analysis uses truncated moment-generating-function expansions that may be less accurate for embedding distributions with substantial higher-order moments.

    Source: §4.3 · p. 6

Abstract

from arXiv · show

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.

1 Introduction

Exact softmax sampling becomes impractical at large scale, motivating two-level softmax (2LS). The paper identifies systematic biases in 2LS and introduces corrected variants that improve fidelity across five large-scale datasets.

1 Introduction

  • Exact softmax sampling scales linearly with the item count, making it impractical for large settings such as recommender systems with millions of candidates.2LS offers a sublinear-time alternative by sampling a cluster and then an item within it.Source: §1 · p. 1
  • 2LS misweights clusters by ignoring both size imbalance and intra-cluster similarity dispersion, causing systematic sampling bias.It can assign different probabilities to items with equal query similarity when they belong to different clusters.Source: §1 · p. 2
  • S-2LS and SD-2LS correct these biases with negligible to non-existent additional computational overhead.The paper evaluates their sampling properties on five large-scale datasets.Source: §1 · p. 2
  • The authors report that S-2LS and SD-2LS consistently outperform 2LS in approximating softmax sampling.They recommend using the corrected variants in future work.Source: §1 · p. 2

2 Preliminaries

The preliminaries define exact softmax sampling and its computational cost, describe the two-stage 2LS alternative, and position both among other approximation and retrieval methods.

2.1 Softmax Sampling

  • Exact softmax weights items by query similarity, with temperature controlling distribution sharpness, but requires O(dN) time per query.The cost becomes prohibitive for large item sets.Source: §2.1 · p. 2

2.2 Two-Level Softmax (2LS) Sampling

  • 2LS samples a cluster using a centroid softmax, then an item using a softmax restricted to that cluster.Its cost is O(dK + d|Ck|), sublinear in N when K ≪ N and clusters are sufficiently balanced; the partition is typically computed offline and reused.Source: §2.2 · pp. 2–3

2.3 Related Work

  • Related methods trade off sampling scope, similarity use, and computation: truncation restricts candidates, ε-greedy ignores similarity, and hierarchical softmax uses a tree factorization.2LS shares a two-stage structure with hierarchical softmax but uses a cluster-based approximation and query-dependent sampling.Source: §2.3 · pp. 3–4

3 Bias Characterization of Two-Level Softmax (2LS) Sampling

This section characterizes 2LS distortion and traces it to cluster weighting. Its analysis identifies size imbalance and intra-cluster similarity dispersion as complementary sources of bias.

3.1 Sampling Ratio

  • The 2LS-to-softmax sampling ratio measures relative distortion: values above 1 indicate oversampling, while values below 1 indicate undersampling.The ratio depends on an item's cluster, not directly on its query similarity; items in the same cluster share the same distortion factor.Source: §3.1 · p. 4

3.2 Systematic Bias in 2LS Sampling

  • Under a mixture-distribution setting for item embeddings, the analysis studies how the sampling ratio behaves as the item count grows.The cluster-level expression depends on cluster probability, centroid similarity, and the expected exponential similarity within that component.Source: §3.2 · p. 4
  • Equal-similarity items can receive different probabilities under 2LS when they belong to different clusters.Tranheden et al. [36] had empirically noted an efficiency loss under cluster-size imbalance; this analysis also identifies dispersion as a bias source.Source: §3.2 · pp. 4–5
  • 2LS oversamples smaller clusters and undersamples larger ones, while also undersampling clusters with greater intra-cluster query-similarity variability.These are the two complementary cluster-misweighting axes identified by the analysis.Source: §3.2 · p. 5

4 Size- and Dispersion-Corrected Two-Level Softmax Sampling

The section introduces two corrected 2LS methods: S-2LS addresses cluster-size imbalance, while SD-2LS additionally accounts for intra-cluster dispersion.

4.1 Size-Corrected Two-Level Softmax (S-2LS) Sampling

  • S-2LS removes dependence on cluster size and empirically avoids undersampling items in larger clusters or oversampling items in smaller clusters.Its within-cluster distribution remains unchanged, and its sampling complexity is identical to 2LS.Source: §4.1 · p. 5, §4 · p. 5

4.2 Size- and Dispersion-Corrected Two-Level Softmax (SD-2LS) Sampling

  • SD-2LS adds a variance-based correction for intra-cluster similarity dispersion to the cluster-size correction in S-2LS.It approximates the moment generating function with a truncated expansion; the approximation is exact under Gaussian assumptions and is reported effective on real-world data beyond that assumption.Source: §4.2 · pp. 5–6
  • SD-2LS costs O(d2K + d|Ck|), with the additional cost described as negligible in practice when d is much smaller than N.The method evaluates quadratic forms q⊤Σkq for all K clusters; the within-cluster step is unchanged.Source: §4.2 · pp. 5–6

4.3 Discussion

  • S-2LS retains 2LS sampling complexity, whereas SD-2LS incurs covariance-evaluation overhead that may matter in high-dimensional settings.The paper favors S-2LS when low latency is critical and clusters are highly imbalanced, and SD-2LS when dispersion is substantial and resources suffice.Source: §4.3 · p. 6
  • SD-2LS’s truncated moment expansion may be less accurate for embedding distributions with substantial higher-order moments, such as skewness or kurtosis.The paper identifies efficient higher-order corrections as future work.Source: §4.3 · p. 6

5 Experimental Analysis

The experiments evaluate whether 2LS biases appear across large-scale datasets and whether S-2LS and SD-2LS improve fidelity without sacrificing the efficiency of approximate sampling. The reported results show that the corrections reduce bias and improve the fidelity–latency trade-off, with SD-2LS providing the strongest fidelity gains and a moderate latency cost.

5.1 Experimental Setup

  • The evaluation uses five large-scale datasets and compares sampling ratio, KL divergence from exact softmax, and per-query latency across 1,000 queries per dataset.The datasets include three real-world embedding corpora and two synthetic Gaussian-mixture corpora; methods are also evaluated at three temperatures.Source: §5.1 · p. 7

5.2 Results and Discussion

  • On VK-LSVD, 2LS undersamples large clusters and oversamples small ones, while S-2LS and SD-2LS correct the cluster-size bias.Similar patterns were observed on the other four datasets.Source: Figure 1 · p. 8, §5.2 · p. 8
  • High-variance clusters are undersampled and low-variance clusters oversampled; S-2LS leaves this bias, which SD-2LS corrects.S-2LS's correction of cluster-size bias isolates the remaining variance effect.Source: Figure 2 · p. 8, §5.2 · p. 8
  • For items with identical query similarity, standard 2LS assigns widely different probabilities, while S-2LS and SD-2LS markedly reduce that dispersion.The figure plots per-item sampling probability against query-item dot-product similarity for a fixed query.Source: Figure 3 · p. 8
  • SD-2LS achieves the lowest KL on all natural corpora across temperatures, while S-2LS is consistently second best; top-k and hierarchical softmax are less accurate.The table reports sampling fidelity relative to exact softmax, averaged over 1,000 queries with 95% confidence intervals.Source: Table 1 · p. 9, §5.2 · p. 9
  • S-2LS matches 2LS latency within measurement noise, while SD-2LS is 1.2–1.8× slower on higher-dimensional datasets and up to ∼3× slower on lower-dimensional ones.SD-2LS's latency overhead accompanies KL reductions of several orders of magnitude.Source: §5.2 · p. 9

6 Conclusion

The paper concludes that 2LS has systematic sampling biases and that S-2LS and SD-2LS improve fidelity while preserving 2LS's computational efficiency. It recommends S-2LS as a no-additional-cost improvement, and SD-2LS when intra-cluster variance matters and its moderate overhead is acceptable.

6 Conclusion

  • Theoretical and experimental results show that S-2LS improves on 2LS at no cost, while SD-2LS further reduces approximation error with moderate overhead.The paper recommends S-2LS as a replacement for standard 2LS and SD-2LS when intra-cluster variance is non-negligible and the cost is acceptable.Source: §6 · p. 9
  • Because SD-2LS relies on a second-order approximation, higher-order moments remain a proposed direction for improving accuracy on embedding distributions with high skewness or kurtosis.The appendix provides detailed proofs for the theoretical results and additional experimental results and technical details.Source: §6 · p. 9, Appendix · p. 13

Appendix

  • The appendix organizes its supplementary material into proofs for Sections 3 and 4, plus results and technical details complementing Section 5.Source: Appendix · p. 13

A.1 Proof of Proposition 1

The proofs characterize the 2LS sampling ratio within clusters, establish asymptotic limits, and examine Gaussian and size-weighted cluster-level scores.

A.1 Proof of Proposition 1

  • The 2LS-to-softmax sampling ratio is constant across items in a cluster because its expression no longer depends directly on item-query similarity.The proof obtains this after the similarity terms cancel.Source: §A.1 · p. 13

A.2 Proof of Proposition 2

  • For each cluster, its empirical proportion converges almost surely to the corresponding mixture weight.Source: §A.2 · p. 14

A.3 Proof of Proposition 3

  • The proof models each component’s query projection as Gaussian and uses its moment generating function; the resulting c(q) does not depend on the cluster index.Source: §A.3 · p. 14

B.1 Proof of Proposition 4

  • The Proposition 4 proof changes the cluster-level weighting to include cluster cardinality, then applies the asymptotic limits used in Proposition 2.Source: §B.1 · pp. 14–15

B.2 Proof of Proposition 5

These proof passages describe a cluster score with size and dispersion corrections, a Gaussian-approximation cancellation, and the memory constraint shaping the hierarchical-softmax comparison.

B.2 Proof of Proposition 5

  • The Proposition 5 proof’s cluster-level score includes both cluster cardinality and a dispersion correction.Source: §B.2 · p. 15

B.3 Proof of Proposition 6

  • Under a Gaussian approximation of query projections within components, the cluster-dependent factor cancels in the Proposition 6 proof.Source: §B.3 · pp. 15–16

C.1 Technical Details on Hierarchical Softmax Implementation

  • A one-item-per-leaf binary tree would add roughly N internal-node embeddings and double index memory, so it is impractical for the tested tens-of-millions-scale corpora.The comparison instead uses a recursive binary k-means tree, terminating at 1024 items per leaf; its routing requires O(log2(N/ℓ)) decisions.Source: §C.1 · p. 16

C.2 Per-Query Latency Across Sampling Methods

Table 2 reports per-query latency and speed-up comparisons across sampling methods, with results organized by dataset.

C.2 Per-Query Latency Across Sampling Methods

  • Per-query latency is reported in milliseconds as a mean with 95% confidence intervals.Source: Table 2 · p. 16
  • Measurements average over temperatures τ ∈ {0.05, 0.1, 0.2}, with K = 1024 clusters and 1000 queries.Source: Table 2 · p. 16
  • Speed-up is computed relative to Exact, and the table marks the best and second-best results per dataset.Source: Table 2 · p. 16

C.3 Cluster Size Distributions Across Datasets

At K = 1024, cluster-size distributions vary across the five datasets, with natural and synthetic corpora illustrating different size regimes.

C.3 Cluster Size Distributions Across Datasets

  • VK-LSVD, YAMBDA, and GloVe-100 have substantial cluster-size imbalance, spanning more than an order of magnitude.Source: Figure 4 · p. 17
  • Synth-balanced has tightly concentrated cluster sizes, while Synth-unbalanced reproduces the heavy-tailed distributions seen in natural corpora.Source: Figure 4 · p. 17
  • Despite equal points per component in its generating Gaussian mixture, Synth-balanced has some unequal observed cluster sizes because k-means does not exactly recover the components.Source: Figure 4 · p. 17

C.4 Complete Bias Analysis Across All Datasets

The appendix extends the main-paper bias analyses from VK-LSVD to YAMBDA, GloVe-100, Synth-Balanced, and Synth-Unbalanced. The results are consistent with the main paper.

C.4 Complete Bias Analysis Across All Datasets

  • Across the five datasets, the 2LS biases are observable and substantial in practice, and S-2LS and SD-2LS effectively correct them.The appendix reports that its results for the four additional datasets are consistent with the main-paper findings.Source: §C.4 · p. 17
  • Across all five datasets, standard 2LS undersamples large clusters and oversamples small ones, while S-2LS and SD-2LS correct these biases.The reported comparison is against exact softmax; the captions report p < 0.01 for the cluster-size sampling-ratio deviation.Source: Figure 5 · p. 18, Figure 8 · p. 19, Figure 11 · p. 20
  • For items with identical query-item similarity, standard 2LS assigns widely different sampling probabilities, while S-2LS and SD-2LS markedly reduce this dispersion across all five datasets.Source: Figure 7 · p. 18, Figure 10 · p. 19, Figure 13 · p. 20
  • Across VK-LSVD, YAMBDA, GloVe-100, and both synthetic datasets, S-2LS exposes high-variance undersampling and low-variance oversampling; SD-2LS corrects most synthetic cases.For Synth-Balanced and Synth-Unbalanced, SD-2LS corrects the variance-related biases for the vast majority of points.Source: Figure 6 · p. 18, Figure 9 · p. 19, Figure 12 · p. 20
Loading 2610.10483v1…