Two-Level Softmax Sampling Done Right: Correcting Bias from Size Imbalance and Dispersion
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.
Related
- Mapping urban heat from citizen-collected data: a scalable Gaussian process approach to correcting for preferential sampling
- Decarbonising price formation: unit-level evidence on battery storage and the imbalance price in the GB Balancing Mechanism
- Coherent and robust Bayesian inference under informative sampling via sample-level weight models
- Notes on Uniform Policy Score Bounds and Lipschitz Continuity for Truncated Gaussian, Product Cauchy, Boltzmann, and Softmax Policies
- Softmax Attention on Gaussian Mixtures: Linear When It Can, Selective When It Must
BibTeX
@misc{bendada2026twolevel,
title = {{Two-Level Softmax Sampling Done Right: Correcting Bias from Size Imbalance and Dispersion}},
author = {Walid Bendada and Guillaume Salha-Galvan},
howpublished = {arXiv},
year = {2026},
month = oct,
doi = {10.48550/arxiv.2610.10483},
eprint = {2610.10483},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2610.10483},
}
RIS
TY - UNPB
TI - Two-Level Softmax Sampling Done Right: Correcting Bias from Size Imbalance and Dispersion
AU - Walid Bendada
AU - Guillaume Salha-Galvan
PY - 2026
DA - 2026/10/07
PB - arXiv
DO - 10.48550/arxiv.2610.10483
UR - https://arxiv.org/abs/2610.10483
AB - 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.
ER -