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

arXiv: https://arxiv.org/abs/2610.10483 · HTML: https://mlsanity.com/paper/2610.10483/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 7 of its 38 references that are in the map.
- First posted on arXiv 7 Oct 2026.

## Builds on (most cited first, 7 of 7)

- [CSV-Decode: Certifiable Sub-Vocabulary Decoding for Efficient Large Language Model Inference](https://mlsanity.com/paper/2511.21702/) (2025)
- [DynaSpec: Context-aware Dynamic Speculative Sampling for Large-Vocabulary Language Models](https://mlsanity.com/paper/2510.13847/) (2025)
- [Doubly Sparse: Sparse Mixture of Sparse Experts for Efficient Softmax Inference](https://mlsanity.com/paper/1901.10668/) (2019)
- [Adaptive Sampled Softmax with Inverted Multi-Index: Methods, Theory and Applications](https://mlsanity.com/paper/2501.08563/) (2025)
- [VQ-Logits: Compressing the Output Bottleneck of Large Language Models via Vector Quantized Logits](https://mlsanity.com/paper/2505.10202/) (2025)
- [FlashHead: Efficient Drop-In Replacement for the Classification Head in Language Model Inference](https://mlsanity.com/paper/2603.14591/) (2026)
- [Efficient Item ID Generation for Large-Scale LLM-based Recommendation](https://mlsanity.com/paper/2509.03746/) (2025)

## Abstract (arXiv)

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.