@inproceedings{
baharav2024adaptive,
title={Adaptive Sampling for Efficient Softmax Approximation},
author={Tavor Baharav and Ryan Kang and Colin Sullivan and Mo Tiwari and Eric Sager Luxenberg and David Tse and Mert Pilanci},
booktitle={The Thirty-eighth Annual Conference on Neural Information Processing Systems},
year={2024},
url={https://openreview.net/forum?id=XsNA2b8GPz}
}
The softmax function is ubiquitous in machine learning and optimization applications. Computing the full softmax evaluation of a matrix-vector product can be computationally expensive in high-dimensional settings. In many applications, however, it is sufficient to calculate only the top few outputs of the softmax function. In this work, we present an algorithm, dubbed AdaptiveSoftmax, that adaptively computes the top k softmax values more efficiently than the full softmax computation, with probabilistic guarantees. We demonstrate the sample efficiency improvements afforded by AdaptiveSoftmax on real and synthetic data to corroborate our theoretical results. AdaptiveSoftmax yields >10x gain over full softmax computation on most datasets, yielding up to 30x improvement for Mistral7B evaluated on the Wikitext dataset. The adaptive method we propose for estimating the partition function (the softmax denominator) is of independent interest and can be used in other applications such as kernel density estimation.
Adaptive Sampling for Efficient Softmax Approximation
Tavor Baharav Ryan Kang Colin Sullivan Mo Tiwari Eric Sager Luxenberg David Tse Mert Pilanci