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

2026-10-07

AI Summary

This study addresses the systematic bias inherent in conventional two-level sampling methods for large-scale Softmax sampling, which arises from neglecting cluster size imbalance and dispersion heterogeneity. To mitigate this, we propose two correction algorithms, S-2LS and SD-2LS, that rigorously quantify and rectify these biases through probabilistic analysis, achieving unbiased sampling while preserving sublinear time complexity. This work provides the first theoretical elimination of size and dispersion biases in standard two-level sampling, yielding provably superior approximations with negligible computational overhead. Extensive experiments across five large-scale datasets demonstrate that the proposed methods substantially enhance the accuracy of Softmax distribution approximation.

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.

Problem

• Softmax Sampling
• Two-Level Sampling
• Sampling Bias
• Size Imbalance
• Dispersion

Innovation

• Two-level Softmax Sampling
• Bias Correction
• Size Imbalance
• Dispersion
• Sublinear-time Sampling