Adaptive Sampling for Efficient Softmax Approximation

This paper introduces an efficient algorithm called AdaptiveSoftmax, designed to compute the top k outputs of the softmax function more effectively than the traditional full softmax computation. The key innovation lies in its adaptive approach, which leverages multi-armed bandit techniques to prioritize computations for the largest input elements, significantly reducing sample complexity. The paper provides PAC (Probably Approximately Correct) guarantees for AdaptiveSoftmax, demonstrating its theoretical soundness. Empirical results on both real and synthetic datasets, including the EuroSAT dataset, validate the algorithm's efficiency, showing substantial reductions in computational overhead, often achieving tenfold improvements or more. Additionally, the proposed method for estimating the softmax partition function offers potential applications beyond the current scope, such as in kernel density estimation.

This paper focuses on the efficient approximation of the softmax function. The authors propose an algorithm named AdaptiveSoftmax, which aims to reduce the computational cost of the softmax function in high-dimensional environments. Inspired by the multi-armed bandit problem, the algorithm adaptively allocates computational resources to important output values, efficiently calculating the top k values of the softmax function. It also provides PAC (Probably Approximately Correct) guarantees, demonstrating its effectiveness both theoretically and experimentally.

The softmax function is a widely used tool, e.g., as an activation in the final layer of a classifier. Hence, cutting down on its computational costs can have a significant impact across the AIML field. This paper aims at this by introducing an adaptive variant of computing the softmax for the top $k$ values by estimating the normalization and the index of the likely highest value in the input by employing a multi-armed bandit setting. Furthermore, theoretical guarantees on the accuracy of the output of the adaptive algorithm are provided and can be controlled by an additional parameter $\delta$.

This paper introduces an algorithm named AdaptiveSoftmax, designed to efficiently compute the top-k softmax values rather than the full softmax computation. This algorithm is particularly useful in high-dimensional settings where the computation of the full softmax function can be prohibitively expensive. The authors provide theoretical guarantees for the efficiency and accuracy of AdaptiveSoftmax and demonstrate its effectiveness through empirical results on real and synthetic datasets.

We would like to thank all the reviewers for their careful reading of our manuscript. We were pleased to see that all reviewers appreciated the novel PAC guarantees provided by this work for efficient, instance adaptive softmax computation. We have addressed all the comments and suggestions made by the reviewers, which has helped improve the quality and clarity of the paper. We discuss some common points of concern below, and look forward to the upcoming discussion period.

1. Assumption of sub-Gaussianity: This is the only assumption that we make in this paper. It is one of the weakest assumptions possible (does not assume that the arms are Bernoulli or Gaussian), and is a common assumption in the multi-armed bandit and adaptive computation literature [4]. Unfortunately, without such an assumption, no nontrivial results are possible; consider the case where we do not have preprocessing access to $A$, the vector $x$ is all ones, and $A$ is all $1$s except for one randomly selected entry which has value $2$. In this case, any algorithm for PAC computation of softmax$(Ax)$ with $\delta < 1-1/n$ (even just identification of the largest entry of $Ax$) requires $\Omega(nd)$ samples. More practically though, these vectors are the result of a machine learning pipeline, and not of adversarial construction. As shown by our simulations, this worst case scenario never occurs in practice, and arm pulls are generally well approximated by a Gaussian (see Figure 1 in the attached pdf, now added to the Appendix). Additionally, note that for any fixed problem instance (fixed matrix $A$ and $x$), all arm pulls are bounded, and are thus sub-Gaussian. We have added this more detailed discussion to the Appendix, and have added a reference to it in the main text.

2. Comparison: While many algorithms have been devised for accelerating softmax computation, no existing methods, to our knowledge, provide $(\epsilon,\delta)$-PAC guarantees (as we formulate in equation (4)) save for exact computation. Current approximations use some combination of Locality Sensitive Hashing (LSH) to cluster the vocabulary, truncation of the tail to approximate the normalization factor, and sketching to approximate the matrix multiplication. However, these methods are not truly instance adaptive, and fall prey to many common flaws in machine learning pipelines. For example, reviewer wft8 references a method which learns a quantization of the matrix $A$ during training [2]. The method proposed in this paper has no theoretical guarantees, requires prior knowledge of the correct output label, and only adapts to the training data, not the actual data presented at inference. Since it focuses on learning a good quantization over the training data, it can suffer from potential distributional shifts of the training and test datasets (concerns raised by reviewer kEa9, that our method avoids). See $**Comparison**$ response to Reviewer kEa9 for more details.

3. Converting FLOP gains to wall clock gains: The focus of this paper is to develop the first provably adaptive softmax algorithm with PAC guarantees, highlighting its dramatically improved sample complexity across a wide variety of models and tasks. The eventual goal of this method is to enable wall-clock improvements in hardware implementations. We provided a brief discussion of this in the Limitations and future work section (Lines 333-352), but have now added additional discussion (below) there and to the supplement. These next steps of converting our provably and empirically performant method into a hardware optimized wall-clock efficient algorithm is an exciting direction of future work, which we detail below. In most modern-day transformer architectures, memory I/O serves as the primary bottleneck [1]. AdaptiveSoftmax already presents an opportunity to significantly scale down the number of entries of the matrix $A$ that must be loaded at inference time, and, in the future - if memory remains the bottleneck - improve model bandwidth by a similar factor. This objective appears in reach, since we have designed the components of AdaptiveSoftmax to be amenable to tiling and parallelization. Most notably, our implementation of AdaptiveSoftmax uses the same column to generate an importance-weighted sample for each active arm. The reasons for this implementation decision are two-fold. First, it takes advantage of the locality of entries in the same column to load samples faster, and, second, it removes intra-column correlation, which can yield theoretically improved performance [Baharav and Tse 2019]. Adjacent column samples can also be combined by simply summing their respective importance weights - admitting a simple tiling of our matrix $A$ that could easily be sized particularly to fit individual tiles into SRAM on a GPU along with a copy of the vector $x$ and the current mean/variance estimates for each arm. Then, we can dynamically load these tiles into SRAM based on the arm estimates as we do currently. The successive elimination bandits algorithm utilized by AdaptiveSoftmax is also, by choice, quite easily parallelizable. We may also store two copies of our matrix $A$ — one with wider tiles and one with taller tiles — to take advantage of our tiling at all stages of the AdaptiveSoftmax algorithm: both when a larger number of samples is necessary for fewer arms, in later stages of adaptive estimation, and when a smaller number of samples is necessary for many arms, in earlier stages of the adaptive estimation. This said, we observe in our experiments that the bulk of compute is invested in our early samples of many arms. Just using basic parallelization to speed up this step could result in the desired speed improvements.

[1] Ivanov et al. "Data movement is all you need: A case study on optimizing transformers" 2021

[2] Joulin et al. "Efficient softmax approximation for GPUs." 2017.

Computing full softmax can be expensive. This paper proposes a new algorithm to efficiently estimate the top-k softmax values. The algorithm is inspired by the multi-armed bandit problem: it adaptively allocates resources to important output values. The work provides PAC guarantees and empirical results.

The final recommendations from the reviewers are unanimously acceptance (7, 5, 5, 7), albeit with varying degrees.

The AC agrees with the reviewers that the problem studied is important, and the proposed method is novel and theoretically sound.

One major concern from the reviewers with less favorable rating is regarding the sub-Gaussianity assumptions. The AC agrees with the authors that this assumption is not strict.

On the other hand, the AC finds that the concerns regarding more discussions and comparisons with other method to be valid (qEfP, wft8).

The topic studied is highly related to sampled softmax. From this point of view, the literature review is not complete . The original/early papers discussing sampled softmax are not cited:

Bengio and Senecal, Quick Training of Probabilistic Neural Nets by Importance Sampling 2003

Bengio and Senecal, Adaptive Importance Sampling to Accelerate Training of a Neural Probabilistic Language Model 2008

There has been many ways of improving the sampling distribution (in the sense of importance sampling framework) since Bengio and Senegal 2008: different types of hashing methods, kernel-based sampling etc. for example the following work and the works it reviewed.

Blanc and Rendle Adaptive Sampled Softmax with Kernel Based Sampling 2017

The claim that “no other methods provide theoretical guarantees“ seems to be overstated. Sampled softmax comes with theoretical guarantees from importance sampling directly. See

Rawat et al Sampled Softmax with Random Fourier Features 2019

We therefore suggest that the authors work on a better and more comprehensive literature review and slightly tone down the claims in the final version.

Overall, the AC finds that this is an interesting work to be shared with the community with the above revisions.