Adaptive Sampling for Efficient Softmax Approximation

Authors: tavor z baharav
ryan  kang
colin  sullivan
mo  tiwari
eric  luxenberg
david  tse
mert  pilanci
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 05:35:18

Abstract: 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.

Main paper: 
 introduction: 
 The softmax function appears in many different fields and applications. It is often used in multiclass classification problems, as the final operation in a neural network to obtain a probability distribution over classes, in reinforcement learning to obtain a probability distribution over possible actions, and in statistical mechanics to derive various thermodynamic quantities.In machine learning applications, the softmax function often appears as the final operation in classification models and in attention layers. Crucially, the softmax function takes a vector of weights †denotes equal contribution as input and returns a probability distribution defined by those weights. Formally, the softmax function for a given temperature parameter β ∈ R is defined as:EQUATIONwhere µ ∈ R n is the input vector of weights, also referred to as the logits. Usually, the logits are the result of a matrix-vector product (e.g., in a fully connected layer where the softmax is used as the nonlinear activation function). The output of the softmax function (Equation 1) is a probability distribution that is a "soft" version of the max operator that is differentiable with respect to the logits. The softmax function can thus be used in gradient-based algorithms as a proxy for the non-differentiable max function.Intuitively, the temperature parameter β controls the peakiness of the softmax output. A larger β corresponds to a peakier distribution and a "harder" max. The choice of β = 1 corresponds to the canonical softmax function σ, and the choice of β = ∞ corresponds to the "hard" argmax. The denominator of Equation (1) , j e βµj , is called the partition function and is denoted by Z β .The softmax function is critical in many popular, recent machine learning applications like large language models (LLMs). However, it can present a computational bottleneck in high-dimensional applications. During the training of neural networks, for example, each training example x requires the computation of the softmax function σ β (x), the partition function Z β (x), and their gradients. During inference of these models, the number of possible labels for next-token prediction corresponds to the vocabulary size, which can be in the hundreds of thousands for common languages such as English. As such, there has been significant recent interest in accelerating the computation of the softmax function and its derivatives [37, 14, 15] .Key Observations: In many applications, we are only interested in identifying the top few outputs of the softmax function; in these settings, it is unnecessary to compute the smaller entries. This suggests that some of the computation of the full softmax function may be unnecessary and motivates our study. First, we observe that when the input vector to the softmax function is the result of a matrix-vector product, we can approximate the intermediary computation instead of exactly computing it. This, in turn, allows us to approximate the output of the softmax function and converts the problem of computing the softmax function from a computational one to a statistical one. We also note that the softmax output is heavily influenced by the largest input elements which suggests that we can allocate computation adaptively to larger input elements to estimate them with greater certainty. This procedure is inspired by recent work in multi-armed bandits that converts computational problems to statistical ones [4] .Outline: We begin this study with a summary of related work in Section 2. In Section 3, we formalize the reduction of computing the softmax function to a statistical estimation problem. In Section 4, we propose the AdaptiveSoftmax algorithm based on this reduction. In the same section, we provide probably approximately correct (PAC) guarantees for AdaptiveSoftmax and prove that it is more efficient than brute force computation. Crucially, AdaptiveSoftmax allocates greater computational resources towards important output values. In Section 5, we demonstrate the empirical advantages of our algorithm in several real-world applications, including in a multiclass classification setting and in large language models. In Section 6, we conclude with a discussion of further applications, potential limitations, and directions for future work. 
related work: 
 Recent work has identified the computational complexity of the softmax function as a significant bottleneck in recent machine learning applications [37] . Well before the attention revolution, [36] proposed methods to accelerate softmax computation via a hierarchical model. In their work, a binary tree representing possible outputs as leaves is used and, at each step, the model must predict which path to traverse to a leaf. For a balanced tree with n leaves, the computational complexity of the softmax function is reduced to O(log n) from O(n), at the expense of O(n) internal classifiers and providing only an approximate output dictated by the quality of the clustering. Google's word2vec models used Huffman trees instead of vanilla binary trees in a similar tree-based approach [37] . Other approaches include target sampling, noise contrastive estimation and self normalization (summarized [14] ), but all of these methods reduce the complexity in terms of the vocabulary size n, rather than by the dimension d. Additionally, our proposed algorithm provides direct PAC guarantees on the original softmax output, instead of approximating the softmax in a sequence of steps without provable accuracy guarantees. Independently, some works have developed fast methods to approximate the softmax gradients during training by using importance sampling over the classes, improving scaling with respect to n [11, 10] leading to a sampled softmax. This is in contrast with AdaptiveSoftmax , which utilizes importance sampling in an orthogonal direction to subsample the features efficiently, enabling gains in d at both train and test time. These sampled softmax methods were later specialized to kernel-based sampling, resulting in provably bounded bias [12, 38] . However, these and other optimized methods [23] typically require prior knowledge about the desired output label frequencies, leaving them susceptible to phenomena like distribution shift between training and inference data, where the frequency distribution changes at the time the model is evaluated. Unlike these approaches, AdaptiveSoftmax does not require auxiliary knowledge, is adaptive on a per instance basis, and provides provable guarantees for the true softmax computation directly rather than a proxy.Our algorithm is inspired by adaptive sampling techniques from the multi-armed bandit literature.Randomized algorithms based on multi-armed bandit algorithms have seen a surge of recent work, due to their ability to provide instance-adaptive guarantees for a variety of problems. This idea was first formalized in the specific case of Monte Carlo Tree Search [26] and later studied in the context of hyper-parameter tuning [31] . Recent work has formalized this approach into the framework of Bandit-Based Monte Carlo Optimization [4] , where the computational task is reduced to one of statistical estimation that is solved efficiently with adaptivity. Applications of this framework include finding the medoid of a dataset [5, 7, 41] , k-nearest neighbor graph construction [30, 34] , Monte Carlo permutation-based multiple testing [46] , and an adaptive singular value decomposition (SVD) [24] . Most relevant is the recent work of [8] , where the authors provide a general framework for adaptive sampling to approximate function evaluation at unknown but estimable points. This work provides general guarantees, but requires a bound on the Lipschitz factor of the function's gradients as input and has potentially poor performance on specific function classes due to its generality.A sub-problem in our softmax approximation is identifying the index of the largest component; this is equivalent to the Maximum Inner Product Search (MIPS) problem on the preceding matrix-vector product. MIPS is a common problem that has inspired many directions of research [22, 32] . Many of these algorithms focus on specific use cases and place restrictive assumptions on the data (e.g., that all elements of the matrix-vector product are positive), require preprocessing, are not adaptive to the underlying data distribution, or lack PAC guarantees. One large family of MIPS algorithms are based on locality-sensitive hashing (LSH) [19, 2] . In addition to significant preprocessing overhead and practical implementation issues, a shortcoming of these LSH-based approaches is that the maximum dot product is often small compared to the vector norms in high dimensions, which necessitates many hashes and significant storage space (often orders of magnitude more than the data itself). Promising LSH-based algorithms have recently been applied to the problem of softmax computation [1, 15] . These methods focus on intensive preprocessing and work primarily by attaining gains in terms of n.In contrast, AdaptiveSoftmax subsamples matrix-vector products and obtains gains with respect to d. Furthermore, AdaptiveSoftmax provides an instance-adaptive algorithm with no required preprocessing and still has PAC guarantees. 
problem formulation: 
 In this work, we focus on the problem of identifying the k largest entries of the softmax of an input vector that is the result of a computationally expensive matrix-vector multiplication. Specifically, we analyze the setting where the input vector µ is the result of a matrix vector product Ax, as is common in the final linear layer of neural networks (such scenarios frequently arise in other machine learning problems as well [4] ). Our objective is to design an algorithm that can with probability at least 1 -δ estimate the top k values to multiplicative accuracy ε, where ε and δ are given input parameters.For clarity of exposition, we focus on the case of k = 1, i.e., identifying and estimating the largest component. All our theoretical results, however, easily extend to the setting k > 1 (discussed in Section 4.2).We use [n] to denote the set {1, 2, . . . , n} and ∥ • ∥ to denote the vector ℓ 2 norm, unless otherwise specified. We use ∥ • ∥ Ψ2 to denote the Orlicz norm (i.e., the sub-Gaussianity parameter or variance proxy) of a random variable; this is discussed in greater detail in Appendix A.2 [44] . For matrix A and vector x, we denote the resulting product as µ = Ax. Assuming for notational simplicity that the arms are in sorted order µ 1 > µ 2 ≥ . . . µ n , we define the gaps between the entries of µ as ∆ i = µ 1 -µ i . We use the convention from the best-arm identification literature that ∆ 1 = ∆ 2 and assume that ∆ 2 > 0 (this assumption is easily relaxed). Furthermore, we define α i = e βµi and γ i = e βµi/2 , which are proportional to the optimal first (respectively, second) order sampling frequencies; these are discussed further in Section B.2. Finally, we define p as the softmax output, and i * as its largest entry (assumed to be unique), i.e.,σ β (Ax) = p, i * = argmax i∈[n] p i .Our goal is to design an algorithm which efficiently outputs the best index i * and an estimate its value where, with probability at least 1 -δ, the best index is correct and the estimated value is within a factor of ϵ multiplicative accuracy. Mathematically, defining the algorithm's outputs as i * ∈ [n] and pi * ∈ [0, 1], we define the success events E id , E est where the algorithm identifies the largest entry, and where it estimates its value to within multiplicative accuracy ϵ. We define the algorithm as providing (ε, δ)-PAC guarantees if these events happen simultaneously with probability at least 1 -δ, with respect to the randomness of the algorithm.EQUATIONEQUATIONOur objective them becomes to design an algorithm that satisfies Equation ( 4) and minimizes the requisite sample complexity. 
adaptivesoftmax algorithm: 
 We now introduce the AdaptiveSoftmaxAlgorithm, which approximates the output of the softmax in Algorithm 1. First, AdaptiveSoftmax approximates the softmax normalization constant Z β to a multiplicative accuracy of ϵ/4 via NormalizationEstimation (Algorithm 2). Next, AdaptiveSoftmax identifies the best arm (or top k arms, depending on the setting) using a standard multi-armed bandit algorithm, BestArmId (Algorithm 3). In our setting, "arms" correspond to different rows of A and pulling arm i corresponds to computing a coordinate-wise scalar product A i,j x j for some coordinate j (we provide a more formal overview of the best-arm identification problem and the associated algorithm in Appendix A). Finally, AdaptiveSoftmax estimates the value of the identified best arm (or top k arms) to a multiplicative accuracy of ϵ/4 by sampling each arm a sufficient number of times via EstimateArm (Algorithm 10).We prove (ε, δ)-PAC guarantees for AdaptiveSoftmax by union-bounding over the error probabilities in each step of Algorithm 1. Our results will show that, with probability at least 1 -δ, AdaptiveSoftmax is able to identify the largest output of the softmax function and estimate its value to multiplicative accuracy ϵ.Algorithm 1 Adaptive Softmax x have bounded entries. For algorithmic simplicity we utilize the following assumption. Assumption 1. We assume that we are given a variance proxy bound σ 2 for the sub-Gaussian parameters of the constructed estimators:σ 2 ≥ ∥A iJ x J ∥ Ψ2 ∀ i, for J ∼ Unif([n]).We provide theoretical guarantees for AdaptiveSoftmax under Assumption 1. Recall that we defined our optimal first and second order sampling frequencies α i = e βµi and γ i = e βµi/2 (see Appendix B.2). We first show in Proposition 1 that our softmax normalization estimation algorithm (Algorithm 2) obtains the desired guarantees. Proposition 1. For input ε ∈ (0, 1/2), δ ∈ (0, 1), and σ satisfying Assumption 1, Algorithm 2 will, with probability at least 1 -δ, estimate Z β = j e βµj to a multiplicative accuracy of ϵ. On this success event, Algorithm 2 requires at mostCβ 2 σ 2   n log n δ + log n δ   j γ j   2   ϵ j α j   -1 + log 1 δ ε -2   samples for some absolute constant C, where non-asymptotic bounds with numerical constants are provided in Appendix B.With the sample complexity of Algorithm 2 bounded, the complexity of best arm identification and the cost of estimating the best arm to a target accuracy are readily available from the multi-armed bandit literature. This enables us to state an overall result for AdaptiveSoftmax (Algorithm 1) in the following Theorem. Theorem 1. For input ε ∈ (0, 1/2), δ ∈ (0, 1), and σ satisfying Assumption 1, Algorithm 1 identifies the largest component in σ β (Ax) and estimates its value to a multiplicative accuracy of ϵ with probability at least 1 -δ, as in (4). On this success event, the algorithm uses T samples whereT ≤ Cσ 2   β 2 n log n δ + n i=1 log n log d δ ∆ 2 i + β 2 log n δ j γ j 2 ϵ j α j + β 2 log(1/δ) ε 2    ,for some absolute constant C. Tighter bounds with non-asymptotic numerical constants are provided in Appendix B.The proofs of these two results are detailed in Appendix B; we provide some intuition and brief sketches of the proofs here.For Proposition 1, we first show that we can estimate the quantities {α i }, {γ i }, to constant multiplicative error with high probability. This allows us to construct a sampling scheme based off of the asymptotically optimal sampling frequencies, and guarantee that each arm is sampled at least half of what this asymptotically optimal frequency requires. Then, sampling each arm i enough so that the first order Taylor expansion of e β μi is sufficiently accurate, we can sample further according to these determined frequencies to guarantee PAC estimation. This is an improved and specialized modification of [8] that exploits the structure of the softmax function to remove the assumption of Lipschitz gradients and yield improved sample complexity (this is discussed further in Appendix B.3).Next, we utilize a classical best-arm identification algorithm to identify the best arm with high probability, leveraging standard results in Bandit-Based Monte Carlo Optimization [6] . Finally, we sample the identified best arm enough times to estimate its value to multiplicative accuracy ϵ/4 with high probability. By union bounding over these error probabilities, we achieve the desired PAC guarantees.We now simplify and further interpret the sample complexity results in Theorem 1. First, note that the ε -2 dependence exhibited by NormalizationEstimation (Algorithm 2) is optimal: it is inherent even in estimating the mean of the best arm to accuracy ε. The cost stemming from the second order error, which scales as ε -1 , is bounded between β 2 σ 2 log(n/δ)ε -1 and nβ 2 σ 2 log(n/δ)ε -1 , where in the case where one arm is much better than the rest this will match the first term. Concretely, we analyze the setting where the minimum gap is ∆, i.e. µ 1 -µ i = ∆ i ≥ ∆ for all i. Corollary 1. Under the conditions of Theorem 1, when the minimum gap is at least ∆, Algorithm 1 identifies and provides (ε, δ)-PAC estimation (Equation (4)) of the largest softmax entry, usingT ≤ C β 2 σ 2 log n δ n + ε -1 n 2 n + e β∆ + β 2 σ 2 ε -2 log(1/δ) + nσ 2 log n log d δ ∆ -2samples for some universal constant C. In the case where the gap is large (∆ ≥ 2 β log n), β is not too small, and d < e e n (see Equation (47) for precise statement), this can be simplified toT ≤ Cβ 2 σ 2 log n δ n + ε -1 + ε -2 log(1/δ) .where all sample complexities are for the 1 -δ success event.The proof of this upper bound is in Appendix B.1. Note that this directly implies that when the gap is large (i.e. there is a clear largest output element) and ε is constant, the sample complexity is nearly linear in n and is upper bounded by Cβ 2 σ 2 n log(n/δ).There are many techniques that we can use to extend and improve AdaptiveSoftmax in practice. We discuss changes from the written algorithm in detail in Appendix C.The variance-proxy bound σ 2 of the arms plays a large factor in the AdaptiveSoftmax algorithm's sample complexity. The underlying sub-Gaussianity parameter of these estimators can be improved using techniques from randomized numerical linear algebra, such as the randomized Hadamard transform [42] . If a small number of entries dramatically increase the variance of the estimator, then the randomized Hadamard transform will make the coordinates more uniform. We provide theoretical guarantees for this approach in Appendix A.3.1.Top-k Identification: Extending our algorithmic results from best arm identification (top 1) to identifying multiple components (top k) follows directly from existing multi-armed bandit algorithms. Numerous algorithms have been developed for this setting [21] , and variants for computational settings have been developed and studied in [6] . For simplicity and clarity, we focused on the top 1 identification in this paper, but the top k extension readily follows. Furthermore, in numerical experiments we observe estimating the normalization constant Z β dominates the sample complexity, and the increase in cost from identifying the top k arms and estimating their values to multiplicative accuracy ε/4 is minimal.Relaxing Assumption of Known sub-Gaussian Parameter σ 2 : Assumptions regarding known arm concentration parameters are common in multi-armed bandit works and simplify theoretical exposition. These results can naturally be extended in several directions. One simple extension is to the setting where we have a separate sub-Gaussianity parameter σ 2 i for each arm, i.e., heterogeneous variances. A more practical extension is to the setting where we do not have a bound on the sub-Gaussianity parameter for each arm but know that the arm returns are bounded. In this setting, a common multi-armed bandit approach is to utilize the empirical variance [35] . These approaches are discussed further in [8] .Improved Estimators μ: Naïvely, the AdaptiveSoftmax algorithm samples coordinates uniformly at random with replacement from the set of coordinates {1, . . . , d} to estimate each j A ij x j . This procedure can be improved in several ways. For example, we may utilize importance sampling and sample each coordinate with probability z j ∝ |x j |. Furthermore, we can sample coordinates without replacement; this is known to yield tighter confidence intervals than sampling with replacement [9] . We can combine these techniques and compute the effective variance as in [18] . Sampling without replacement can be achieved in a computationally efficient manner via Gumbel sampling [27] . We discuss these details further in Appendix A.3; these details may be of independent interest. 
experiments: 
 In this section, we demonstrate the empirical advantages of AdaptiveSoftmax over the brute-force softmax computation in terms of sample complexity. All of our results are reproducible via a 1-line script, publicly available on GitHub at github.com/ThrunGroup/adaptiveSoftmax.Crucially, the AdaptiveSoftmax algorithm scales sublinearly in d. More precisely, Corollary 1 implies that, for fixed ε and δ, the sample complexity of the AdaptiveSoftmax algorithm scales as O(n log n). We now empirically validate this behavior.We first run the AdaptiveSoftmax algorithm on two synthetic datasets. In each dataset, we generate A and x with n = 100 and vary d.In the first synthetic dataset, we set x to be a d-dimensional vector of all 1s. We draw each element of A i.i.d.∼ N (0, 1) and add the vector of all 1s to the first row of A, thereby planting a signal. In expectation, the first row of A will have inner product d with x whereas all other rows will have inner product 0 with x. Furthermore, all arms have expected variance σ 2 i that scales with d. In the second synthetic dataset, we draw each element of A i.i.d.∼ N (0, 1) and set x to be |A 1,: |, the entrywise absolute value of the first row of A. Here, arms have expected variance σ 2 i = Θ(1).Figures 1(a ) and 1(b) demonstrates the scaling of the AdaptiveSoftmax Algorithm on each of the two datasets. On the first synthetic dataset, the AdaptiveSoftmax algorithm scales with d because the variance proxies σ 2 i do. On the second synthetic dataset, the AdaptiveSoftmax algorithm does not exhibit significant scaling with d. On both datasets, the AdaptiveSoftmax algorithm significantly outperforms the naïve brute-force computation of the softmax function.Multinomial logistic regression (MNL) is a form of multiclass classification in which the final operation performed by the classifier is of the form:EQUATIONi.e., the probabilities that datapoint x belongs to each class c is given by the softmax applied to the vector W h(x), where W is the matrix containing rows w 1 , . . . , w c and h(x) is a latent representation of x (i.e., the forward pass of some neural network on x).The multinomial logistic regression is naturally amenable to accelerated softmax computation in Equation ( 5). In many real-world settings, both the number of classes C and the dimension of the latent representation h(x) (and therefore the dimensionality of each w c ) can be very large, motivating the usage of AdaptiveSoftmax to identify and estimate the probability of the most likely class. However, the application of AdaptiveSoftmax extends far past vanilla MNL. For instance, the final layer of any neural network classifier utilizing softmax can also be viewed as an MNL problem. We now provide several such practical settings for which we demonstrate the benefits of applying the AdaptiveSoftmax algorithm.We now demonstrate the performance of the AdaptiveSoftmax algorithm on several real-world datasets. For each setting, we provide the sample complexity gain relative to the sample complexity of the brute-force, naïve softmax computation sample complexity nd. We also provide the success rate of our algorithm in each setting, i.e., the proportion of times the AdaptiveSoftmax algorithm correctly identifies the maximum likelihood output (i.e. î⋆ = i ⋆ ) and estimates its probability p i ⋆ within a multiplicative error of ε = 30%.We consider the application of AdaptiveSoftmax to CNN classifiers on two distinct image classification datasets:1. The MNIST dataset, containing black and white images of handwritten digits as input and ten output classes representing all ten possible digits. 2. The EuroSAT dataset, containing RGB satellite imagery as input and ten output classes, representing possible land types (e.g., river, residential, etc)On both of these datasets and for distinct architectures, we show that AdaptiveSoftmax provides a drastic improvement in sample efficiency.MNIST For the MNIST dataset, we train a shallow CNN from scratch with two convolutional blocks (Conv2d, ReLu, MaxPool, BatchNorm). This model achieves over 99% accuracy on the test set. The matrix A is obtained by extracting the weight matrix of the model's final linear layer. The vector x is extracted as the output of the final hidden layer (the layer before the final linear layer) constructed by passing the MNIST image through the trained model and flattening the result. The dimensionality of x is adjusted by changing the number of output channels of the convolution blocks.The sample complexity of our algorithm is measured by running the algorithm on 1000 different images in test set with same matrix A. The empirical error rate δ is calculated as the fraction of experiments where the adaptive algorithm fails to identify the same class, or fails to estimate the probability to accuracy ϵ, as assigned by exact computation.EuroSAT We also utilize a larger pre-trained CNN classifier fine-tuned on the EuroSAT dataset, to show that AdaptiveSoftmax works with larger more sophisticated CNNs. Specifically, we freeze all convolution blocks of VGG-19 (pretrained on ImageNet) and changed the final output dimension to 10 classes for EuroSAT without freezing the weights. The resulting model achieves 92% accuracy on the test set. As before, the matrix A can be extracted from the weights of the final linear layer and the vector x represents the final hidden layer activations. The empirical error rate δ is calculated in the same manner as for MNIST.Dataset (Model) δ = 10% δ = 5% δ = 1% EuroSAT (VGG-19)5.18x (80.62%) 5.16x (83.00%) 4.54x (98.37%) MNIST (Shallow CNN) 8.95x (92.25%) 8.81x (93.75%) 8.13x (99.38%) Table 1 : Performance improvement and success rate afforded by AdaptiveSoftmax for multinomial logistic regression on two different real-world datasets. We used a total of q = 800 test queries to measure success rate.We also apply the AdaptiveSoftmax algorithm to LLMs using HuggingFace's AutoModelForCausalLM module for the task-generation task [45] . The matrix A is the lm-head layer for each model, and the queries x are the final hidden states of the model that is extracted by running a forward pass of the model on the given dataset with a window moving at a certain stride. The context window and stride is modified to generate a desired number of queries. 2 : Performance improvement and success rate afforded by AdaptiveSoftmax for LLM inference (improvement for final softmax layer). Experiment details in Section 5.3.2. We used q = 1000 unseen test queries to measure δ-accuracy.Dataset (Model) δ = 10% δ = 5% δ = 1% Wikitext (GPT-2) 8.Our matrix A is the extracted lm-head from HuggingFace's AutoModelForCausalLM for the four models: GPT-2 (n = 50257, d = 768), Llama3-7B (n = 128256, d = 4096), Mistral7B (n = 32000, d = 4096), and Gemma7B (n = 256000, d = 3072). Our task is task-generation, and we generate our queries x by using two datasets (Wikitext and Penn Treebank) with a sliding window of certain stride. Stride and context window is set to get q = 1000 number of queries. Constants and confidence intervals given by theory are empirically quite loose, so we tuned algorithm parameters (constant coefficients for stage length and confidence interval width) on initial training data, described in Appendix C. An aggressive tuning strategy was undertaken in order to demonstrate the potential gains in sample complexity provided by AdaptiveSoftmax . Specifically, constant multiples were applied to variance estimate within Algorithm 3 and Algorithm 2. Due to the limited sample set, this approach occasionally overoptimized the constants on training data, yielding lower success rates than targeted. However, from the results, it is clear that this target parameter still provides users sufficient control over the tradeoff between true success rate and sample complexity. 
discussion, limitations, and future work: 
 In this work, we proposed a theoretically novel and practically efficient algorithm for approximating the softmax function. We provided theoretical guarantees on the accuracy of our approximation and demonstrated that, with fewer samples than exact computation, we can approximate the softmax function to the desired accuracy. We further demonstrated the viability of our proposed algorithm in two real-world settings, multinomial logistic regression and LLM inference.A potential limitation of our proposed algorithm is that it is most beneficial when the inner dimension of the matrix vector product is high dimensional; its benefits over exact computation are more modest when the inner dimension is small. In particular, the exact computation of the matrix-vector product preceding a softmax operation is usually performed efficiently using BLAS (Basic Linear Algebra Subroutines, which are highly optimized). Adaptivity at its core is inherently sequential, whereas BLAS operations take advantage of batch computation. In this work we proposed minimally adaptive algorithms, with only a logarithmic number of rounds of adaptivity, but there are important directions of future work to realize these theoretical gains in practice.Limitations: Theoretical sample complexity bounds are useful for understanding the fundamental properties of an algorithm, but in practice, wall-clock time is often the metric of interest. Many of the steps in our algorithm can be batched and made BLAS efficient, yielding comparable wall clock times to brute force computation. However, in general adaptivity is the opposite of batching, as can be seen when we modify our algorithm to adapt to arm specific variances. In this case, we must sample each arm individually, as the number of samples required for each arm is different. This is a trade-off between adaptivity and wall-clock time, and in practice, the choice of which to prioritize depends on the specific application (energy efficiency, computational resources, etc.). There are also possible theoretical analyses, where we can e.g. create batches of arms with similar empirical variance and sample all arms within a batch together, leading to a trade-off between adaptivity and batched computational efficiency. Additionally, in large language models, the final softmax layer is often not a computationally significant step, so while such a method may greatly accelerate multinomial logistic regression, more work may be required to have this accelerate LLMs.Given the ubiquity of the softmax function in today's machine learning workflows, we hope that our algorithm will help pave the way for an optimized adaptive softmax that can accelerate a wide class of machine learning models. An interesting direction of future work is trying to combine this multi-armed bandit approach with LSH [15] to obtain (for the attention case) subquadratic complexity in n, and sublinear complexity in d. The adaptive method we propose for estimating the normalizing constant of the softmax function is of independent interest, and holds potential for applications in kernel density estimation and other machine learning tasks.To make this work accessible to a broad audience, we provide a self contained introduction to the multi-armed bandit setting.We consider a stochastic multi-armed bandit problem [13, 40, 29] with n arms (distributions), where each arm i has an unknown mean reward µ i . At each time step t, the algorithm selects an arm I t ∈ [n] and receives a reward X It,t drawn from the distribution of arm I t . Early work in the multi-armed bandit literature focused on the regret minimization setting, where the goal is to maximize the cumulative reward (sum of arm pulls observed so far), motivated by applications such as online advertising and gambling [28] . Recent work has seen increased interest in the best-arm identification setting, where the goal is to identify the arm with the highest mean reward with high probability, motivated by applications such as clinical trials. Significant research has been devoted to obtaining optimal logarithmic factors, but for the sake of clarity we highlight here a simpler and empirically well performing algorithm, multi-round ϵ-arm from [25] .Algorithm 3 BestArmId (modification of best-arm identification from [25] )Input: n arms, error probability δ, variance proxy σ 2 Output: Best arm i ⋆ with probability at least1 -δ S 0 ← [n] r ← 0 t 0 ← 0 while |S r | > 1 do r ← r + 1 ϵ r ← 2 -r t r ← ⌈8σ 2 ϵ -2 r log(4nr 2 /δ)⌉ for all arms i ∈ S r-1 doPull arm i t r -t r-1 times and observe rewards X i,tr-1+1 , . . . , X i,tr μi,r ← 1 tr tr s=1 X i,s ▷ Update mean estimates C i,r ← 2σ 2 log(4nr 2 /δ)/t r ▷ Compute confidence interval width end for Set S r ← {i ∈ S r-1 : μi,r + C i,r ≥ max j∈Sr-1 μj,r -C j,r } ▷ Filter "bad" arms end while return i * , the only element in S r ▷ We assume i * is unique (this assumption is easily relaxed)The algorithm proceeds in rounds, maintaining a set of arms S r that are still in contention for being the best arm. At each round r the algorithm pulls each surviving arm such that we can construct a high probability confidence interval of width ϵ r /2 around the mean of each arm. Then, any arms whose empirical mean plus confidence interval is less than the maximum empirical mean minus confidence interval are eliminated. If this is the case, then that arms mean is with high probability less than the maximum mean, and so it is eliminated from contention. This preserves the best arm with high probability, and the algorithm terminates when only one arm remains (this best arm).Following the exposition of [44] , for a strictly increasing convex function ψ : R + → R + where ψ(0) = 0, the ψ-Orlicz norm of a random variable X is defined as:∥X∥ ψ ≜ inf t > 0|E ψ t -1 |X|| ≤ 1 .∥X∥ ψ is infinite if the expectation E ψ t -1 |X|| does not exist for any finite t. The sub-Gaussian parameter of a random variable is defined as the ψ 2 -Orlicz norm, where ψ 2 (u) = e u 2 -1.Standard results (Hoeffding's lemma) provide that for a random variable X such that a ≤ X ≤ b almost surely, that ∥X∥ ψ2 ≤ (b-a) 2. Figure 2 : Distribution of arm pulls for the best, middle, and worst arms for a random query. The Gaussian fit plotted for each arm is computed using the empirical variance of the given arm, and can be seen to closely match the empirical distribution, indicating that the arm pull distributions are well approximated by a Gaussian. This merits the assumption of sub-Gaussianity. Arm pulls represent importance weighted samples based on the magnitude of the query vector, so for an arm i and sample j, the value is A i,j * sgn(x j ) * d. σ is computed across all samples for an arm, and windows are truncated at the lowest and highest ends of the 3σ ranges across arms for viewing clarity.Hoeffding's inequality provides a useful concentration bound, where for X with ∥X∥ ψ2 ≤ σ 2 ,P (|X -E[X]| ≥ t) ≤ exp - t 2 2σ 2 .The assumption of sub-Gaussianity 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 1s, and A is all 1s except for a randomly selected entry which has value 2. In this case, any algorithm for PAC computation of softmax(Ax) with δ = 1 -1/n (even just identification of the largest entry of Ax) requires Ω(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 2 ). Additionally, note that for any fixed problem instance, all arm pulls are bounded, and are thus sub-Gaussian.To provide theoretical guarantees in multi-armed bandit problems, stringent assumptions are often required, e.g. Assumption 1. This is so that we can provide high probability guarantees on the concentration of the estimator constructed as the empirical mean of the observed samples. Often, these assumptions are phrased as either that random variables corresponding to arm pulls are bounded in [0, 1] a.s., or that they are σ 2 sub-Gaussian, with a known bound on σ 2 . Such analyses have been generalized to bounded random variables with a known bound, where the algorithm is able to adapt to the unknown variance [33, 35] . In this work, as is often done to make multi-armed bandit algorithms more performant [4] , we instead directly use the empirical variance of an estimator as its sub-Gaussian parameter. In the Gaussian case, a random variable's sub-Gaussian parameter σ 2 matches its variance: as our arm pulls are constructed as the sum of many (d) terms, which can be thought of as weakly dependent, our arm pulls can be thought of as Gaussian random variables with variance σ 2 . Since in practice, we do not have a good bound on the magnitude of these arm pulls, we directly use Hoeffding's concentration inequalities [44] with the empirical variance σ2 , as opposed to an empirical Bernstein type concentration inequality [35] .We discuss applying a randomized Hadamard transform to reduce the sub-Gaussian parameter of our estimators. Define the rotation matrix R = 1 √ d DH, where D is a diagonal matrix with diagonal entries equiprobably ±1. H is a Hadamard matrix (d must be a power of 2, 0 padded if necessary). Then, we have that applying the transform R to A and R ⊤ to x yields arms with better sub-Gaussian parameters. Concretely, define Z = AR and y = R ⊤ x. Analyzing y first, we have that no entry of y is too large, as:EQUATION)The first inequality is a union bound over the n points and plugs in for y = R ⊤ x. The second equality uses the fact that the Hadamard matrix multiplied by the random diagonal ±1 matrix D makes R i i.i.d. ±1, and we use ξ j to denote these Rademacher i.i.d. ±1 random variables. Next we use Hoeffding's inequality. Finally, we simplify.Thus, with probability at least 1 -δ, ∥y∥ ∞ < ∥x∥ 2 2 log(2n/δ) d.A similar argument can be made for Z = AR, showing that each entry in the i-th row is upper bounded by ∥A i ∥ 2 2 log(2nd/δ) d , holding simultaneously for all i, j.Since Ax = Zy, we can use bandits to approximate Zy instead of Ax. With the above analysis, we have a bound on the maximum entry of Z iJ y J , giving us a bound on the sub-Gaussian parameter of each arm with probability ≥ 1 -δ.EQUATIONWhereas before:EQUATIONIn practice, we did not observe that this transformation yielded much improved variance, as opposed to simply importance sampling. Thus, we do not utilize this in our main algorithm.Instead of naively sampling each coordinate uniformly, we can construct improved estimators using importance sampling. Concretely, consider the unbiased estimator Z where P(Z = 1 pj A j x j ) = p j for some probability distribution {p j } where p j > 0 for all j and j p j = 1. Now, we are left with the design choice of how to construct {p j }; naively, this is 1/d. Unpacking the variance of our estimator, we see:EQUATION)Consider the case where A j i.i.d.∼ Q for some distribution Q with E[Q 2 ] = λ 2 (an empirically not unreasonable assumption), where {x j } are fixed constants. We assume x is known, but A isn't. Then, we simplify the expected variance with respect to this randomness in A asEQUATIONSimplifying this for Z naive where p j = 1 d , and for Z opt where p j ∝ |x j |, we haveEQUATIONEQUATIONIn the case where we are using the same matrix A over many vectors x (as is often case e.g. in LLM inference, or multinomial logistic regression), we can make our leverage scores a function of precomputed quantities based off of A, not just x. In this case, it makes sense to consider forms of leverage sampling; in this work we consider taking p j ∝ |x j || i |A ij |.Weighted sampling with replacement becomes wasteful in larger sample size regimes, for which the same high-weight elements are sampled repeatedly. It therefore becomes desirable to remove sampled elements from consideration after they are sampled, re-weighting the remaining elements accordingly. We could naively repeat this iterative process of sampling, removing, and re-weighting our elements until we ended up with a sample of the desired size, say k. However, this process is sequential and quite slow. Fortunately, as noted in [43] , sampling with replacement according to a set of weights is equivalent to perturbing the logits λ 1...n of our desired sample weights with draws from the i.i.d standard Gumbel distribution and taking the elements with the top-k perturbed logits as our sample, as detailed in Algorithm 6 and 7. This process is easily batched and is much faster as a result.Further, taking the (k + 1)-th perturbed logit as an empirical threshold τ , the inclusion of an element j in our sample is solely dependent on whether or not its perturbed logit value exceeded this threshold. This derivation is detailed in [27] , and gives us the following expression (Gumbel CDF) for the inclusion probability of element j in a set S of size k drawn without replacement according to weights w:EQUATIONEQUATIONThese empirical estimates of the marginal probabilities of selection for each column allow us to generate a sequence of estimators for each arms' mean, with improved variance discussed in the next section.The Gumbel sampling trick used in A.3.3 with the fixed empirical threshold also provides us a different lens on our sampling process. Namely, we can compute these closed form inclusion probabilities π, and by setting this empirical threshold, we may treat the inclusion of separate elements as independent. Given these observations (sampled S), an unbiased estimator for the variance of μi , constructed as the importance sampling weighted mean of the observations, can be computed from [18] as:EQUATIONEQUATIONEQUATION)EQUATIONFollowing the analysis of [18] and assuming that our threshold τ is fixed, we may conclude that the estimate of the variance of μi , the quantity Var( μi ; τ ), is an unbiased estimate of the true variance Var( μi ). Further, in all datasets we analyzed, the entries of A were generally symmetric, zero-mean, and not correlated with the corresponding entries of x, as seen in Figure 3 (a). Since τ and the values of π are selected solely based on the values of x, these two further assumptions make the second order term(A i,j A i,k x j x k ) π j,k -πj π k π j,k πj π k zero in expectation.Thus, for simplicity, we discard this latter summation and treat Var( μi ) as an unbiased estimate of the variance throughout our implementation, which can be computed in linear instead of quadratic time (updated in constant vs linear time). In practice, as observed in Figure 4 , this variance estimator Var( μi ) (solid green line) provides a far better estimate of the Gumbel sampler's true variance than other methods. (A i,j A i,k x j x k ) π j,k -πj π k π j,k πj π k .We begin by proving a standard best arm identification result, for a slightly modified version of the round-based algorithm from [16] . Lemma 1 (Best-arm identification). With probability at least 1 -δ, Algorithm BestArmId identifies the top softmax value correctly with probability using a number of observations at mostn i=1 min 32σ 2 ln 4n δ log 2 2 (4/∆ i ) ∆ 2 i , d .Proof. Following the proof from [16] , since an arm's mean is exactly computed after d pulls, we have that the best arm will be identified with probability at least 1 -δ requiring for arm i a number of pulls at mostEQUATIONSumming over arms yields the desired result, noting that the best arm is pulled at most as many times as it takes to eliminate every other arm (i.e. the second best arm, ∆ 1 = ∆ 2 ).Note that if ∆ i ≤ C √ d , then the second term (d) will be selected, for a sufficiently small absolute constant C. Thus,n i ≤ Cσ 2 log n δ log(d) ∆ 2 i (25)Hence, the total sample complexity on this success event is at mostT = n i=1 n i ≤ n i=1 min 32σ 2 ln 4n δ log 2 2 (4/∆ i ) ∆ 2 i , d ≤ C n i=1 min σ 2 log n δ log(d) ∆ 2 i , dWe additionally require a lemma for estimating the mean of the best-arm in a PAC sense. Lemma 2 (Exponential best arm estimation). Sampling an arm usingT = 32σ 2 β 2 log(2/δ) ϵ 2samples guarantees that e β μk estimates e β maxi µi to multiplicative accuracy ϵ, with probability at least 1 -δ.Proof. We estimate the mean of arm k after T draws using the plug-in estimator μk . For simplicity, assume β = 1, where in the end we scale the sample complexity by β 2 . Sub-Gaussian concentration provides that with probability at least 1 -δ,|μ k -µ k | ≤ 2σ 2 log(2/δ)/T = ϵ/4which in turn implieslog(1 -ϵ) ≤ -ϵ/4 ≤ μk -µ k ≤ ϵ/4 ≤ log(1 + ϵ)for 0 < ϵ < 1. Then, exponentiating both sides yields1 -ϵ ≤ e μk -µ k ≤ 1 + ϵ ⇐⇒ (1 -ϵ)e µ k ≤ e μk ≤ (1 + ϵ)e µ k .Scaling the number of samples by β 2 yields the desired result.With these two steps in place, we are now ready to tackle the novel technical challenge of this work; estimating the normalization constant of the softmax. We denote the softmax normalization constant asf (µ) = σ β (µ) = i e βµi .Proposition 2 (Softmax normalization estimation: restatement of Proposition 1). Under Assumption 1, Algorithm 2 will, with probability at least 1 -δ, estimate f β (µ) = j e βµj to a multiplicative accuracy of ϵ, using a number of samples at mostT = 2nT 0 + T 1 = 34β 2 σ 2 log(6n/δ)n + 91σ 2 β 2 log(6n/δ) ( n i=1 γ i ) 2 ϵ i α i + 16β 2 σ 2 log(12/δ) ϵ 2Proof. To prove Proposition 1, we want to show that with high probability, we can upper and lower bound our plugin estimator as(1 -ϵ)f (µ) ≤ f (μ) ≤ (1 + ε)f (µ),giving us our desired multiplicative error bound. We construct several success events that collectively guarantee our bound holds, and that occur with high probability.• E 1 is the event where our estimated optimal sampling frequencies are not too far from the unknown optimal frequencies, i.e.αi ≥ α i /2, βC i < 1, i = 1, . . . , n.On this event, we will sample arms sufficiently in the second round.• E 2 is the event where all estimators μi are within their 2-sided confidence intervals in stage 2, i.e.|μ i -µ i | ≤ 2σ 2 log(12n/δ)/n i , i = 1, . . . , n.On this event, we can bound the error in the exponentiated estimator.• E 3 is the event where the first and second order errors are small, i.e.- ϵ 2 f (µ) ≤ i e βµi β(µ i -μi ) ≤ ϵ 2 f (µ) and i e βµi β 2 (µ i -μi ) 2 ≤ ϵ 2 f (µ).These two terms arise from bounding f (μ).We now show that if E 1 and E 2 and E 3 all occur, then our desired bound holds.Lemma 3 shows that when E 2 holds,EQUATIONEQUATION.Note that if E 3 holds as well, the expression is upper bounded by (1 + ϵ)f (µ), since each of ( 1) and ( 2) is bounded above by (ϵ/2)f (µ).All that remains is to show the lower bound also holds. We use the global inequality 1 + x ≤ e x to lower boundEQUATION.Again, under E 2 , ( 1) is lower bounded by -ϵ/2f (µ), which implies our desired lower bound(1 -ϵ)f (µ) ≤ f (μ).Thus, all three events holding guarantees our desired bound holds. Now that we know our bound holds on the joint success event, all that remains is to show it holds with sufficiently high probability. To do so, we invoke our lemmas which characterize the probability of each event.The probability of E 1 , E 2 , and E 3 all holding isP(E 1 E 2 E 3 ) = P(E 1 )P(E 2 |E 1 )P (E 3 |E 2 E 1 ).Lemma 3 says that if each arm is sampled T 0 times, P (E 1 ) ≥ 1 -δ/3. Lemma 4 says P (E 2 |E 1 ) ≥ 1 -δ/3. Lemma 5 and Lemma 6 together show P (E 3 |E 2 E 1 ) ≥ 1 -δ/3. Thus, taken together we haveP(E 1 E 2 E 3 ) = P(E 1 )P(E 2 |E 1 )P (E 3 |E 2 E 1 ) ≥ (1 -δ/3) 3 ≥ 1 -δ,as desired, and so we are done.The lemmas and their proofs follow. Lemma 3. For n i = T 0 = 17β 2 σ 2 log( 6n δ ) for all i, we have that the eventE 1 , αi ≥ α i /2, βC i < 1, i = 1, . . . , noccurs with probability at least 1 -δ/3.Proof. By a standard Chernoff bound, with σ 2 a bound on the sub-Gaussian parameter of all arms, we have that with probability at least 1 -δ/3 that for all i = 1, . . . , n,αi = e β(μi-Ci) j e β(μj -Cj ) ≥ e β(µi-2Ci) j e βµj ≥ 1 2 α i ,where α i = e βµ i j e βµ j , andC i = 2σ 2 log(6/nδ)is the Chernoff confidence interval width constructed such that C i < log(2)/2β and so βC i < 1. To simplify constants, we use that 8/ ln 2 (2) < 17.Lemma 4. Sampling as n i ≥ T 0 guarantees that conditioned on E 1 , E 2 occurs with probability at least 1 -δ/3 and that onE 2 , f (μ) ≤ f (µ) + i β(µ i + μi ) + i β 2 (µ i -μi ) 2 .Proof. Suppose we sample each arm n i times. Note that by a Chernoff bound on each arm,|μ i -µ i | ≤ 2σ 2 log(6n/δ)/n iholds on each arm independently with probability at least 1 -δ/3n, so all arms are within the two-sided bound with probability at least 1 -δ/3.We upper bound the plugin estimator f (μ).EQUATIONEQUATIONEQUATIONwhere in (28) we use the upper bound e x ≤ 1 + x + x 2 for x ≤ 1.79 on the event E 2 , since on E 2 , (µ i -μi ) ≤ 1/β. This is because n i ≥ T 0 , and T 0 samples already guarantees this.Lemma 5 (First-order error concentration). On the event E 2 ∩ E 1 , the first order errorEQUATIONProof. First, defining E f as the failure event |G| ≥ ϵ 2 i e βµi , note thatP (E f ) = P (E f |E 2 )P (E 2 ) + P (E f |E 2 )P (E 2 ) implies P (E f |E 2 ) = P (E f ) -P (E f |E 2 )P (E 2 ) P (E 2 ) ≤ 2P (E f ),where we use that since δ < 1, P (E 2 ) ≥ 1/2. Thus, it suffices to show thatP (E f ) ≤ δ/6G is a sum of independent sub-Gaussian random variables, each scaled by a constant. Thus, we have the two-sided tail bound that with probability at least 1 -δ/6i e βµi β(µ i -μi ) ≤ 2B 2 log(12/δ)with the sum having sub-Gaussian parameterB 2 = i e 2βµi β 2 σ 2 /n iPlugging in our value of B 2 we find i e βµi β(µ i -μi )≤ 2 log(12/δ) i e 2βµi β 2 σ 2 /n i ≤ βσ 2 log(12/δ)   i e 2βµi 1 e βµ i T 2 j e βµ i   1/2 ≤ βσ 2 log(12/δ)     2 j e βµi   i e βµi 1 T   1/2 = 2βσ log(12/δ)f (µ)/T 1 2, where in (30) we use that on E 1 , αi ≥ 1 2 α i , and soEQUATIONis sufficient to yield the desired multiplicative error of ε/2 with probability at least 1 -δ/3.Lemma 6. If arm i is pulled at least 4 √ 2σ 2 β 2 log(6n/δ)γi n j=1 γj ϵ j αjtimes, then on the success events E 1 , E 2 where the confidence intervals hold,EQUATIONand Algorithm 2 will require a number of arm pulls at mostEQUATIONProof. Bounding the second order error, we utilize the fact that we have sampled proportional to γi . On the event E 2 (where n i ≥ γi √ 2 T ) the second order error can be bounded asEQUATIONEQUATIONEQUATIONEQUATIONEQUATIONWe want this second order error to be at most ϵ 2 i e βµi , and so require T to satisfy the inequality belowEQUATIONAlgorithmically this is not a valid T to use, since it depends on the unknown µ i . However, since αi , γi are close to their true values on the good event E 1 , we can use these estimates. Thus, we take as our budget T for this second order error:EQUATIONwhich on the event E 1 is larger than (38) . This is a random quantity, so to analyze the requisite sample complexity, we use the fact than on E 1 , we can bound the quantity in (39) asEQUATIONNote 64 √ 2 < 91 so we use this simple constant in the statement of the lemma.This can be directly compared to the case where we only sample according to α i , not γ i , which would yield a sample complexity ofEQUATIONsamples. Note that since ( n i=1 γ i ) 2 ≤ n i α i by Cauchy-Schwarz, this is always an improvement (up to absolute constants) up to a factor of n. Now we can provide a proof of the main theorem.Proof of Theorem 1. We utilize the adaptive approximation subroutine with error probability δ ′ = δ/3 and error ε ′ = ε/4 from Lemma 6. With 1 -δ ′ probability, it requires a number of samples at mostT ≤ 34β 2 σ 2 log(18n/δ)n + 363σ 2 β 2 log(18n/δ) ( n i=1 γ i ) 2 ϵ i α i + 256β 2 σ 2 log(36/δ) ε 2 .Best arm identification is called with error probability δ ′ = δ/3. From Lemma 1 this requires sample complexityEQUATIONWe then estimate the best arms mean using Lemma 2 to accuracy ε ′ = ε/4 with error probability δ ′ = δ/3. This requires a number of samples at mostEQUATION)By a union bound, all these algorithms succeed with probability at least 1 -δ. Analyzing the multiplicative error, we have that[1/(1 + ε ′ ), 1/(1 -ε ′ )] ∈ [1 -2ε, 1 + 2ε] for 0 < ε ≤ 1/2, and that (1 ± ε 1 )(1 ± ε 2 ) ∈ (1 ± (ε 1 + ε 2 + ε 1 ε 2 ). Thus, on these success events, the numerator is approximated to accuracy ε/4, and the denominator to accuracy ε/4. The denominator error converts to a multiplicative error of ε/2 in the numerator, which combines to yield an error of 3ε/4+ε 2 /8 < ε, as ε < .5. This allows us to simplify the total sample complexity asT ≤ 512σ 2 β 2 log(6/δ) ϵ 2 + 256β 2 σ 2 log(36/δ) ε 2 ≤ 768σ 2 β 2 log(36/δ) ε 2 Thus T ≤34β 2 σ 2 log(18n/δ)n + n i=1 min 32σ 2 β 2 ln 12n δ log 2 2 (4/∆ i ) ∆ 2 i , d + 363σ 2 β 2 log(18n/δ) ( n i=1 γ i ) 2 ϵ i α i + 768σ 2 β 2 log(36/δ)ϵ -2 ≤C        β 2 σ 2 n log n δ + n i=1 min   log n log d δ ∆ 2 i , d   + log n δ n j=1 γ j 2 ϵ j α j + log(1/δ) ε 2       Algorithmically, in the second stage, arm i needs to be pulled a number of timesn i =17β 2 σ 2 log(6n/δ) + 16 √ 2σ 2 β 2 log(6n/δ) ( n i=1 γi ) 2 ϵ i αi γi j γj + 16β 2 σ 2 log(12/δ)ϵ -2 αi j αj ≤17β 2 σ 2 log(6n/δ) + 64 √ 2σ 2 β 2 log(6n/δ) j γ j γ i ϵ j α j +16β 2 σ 2 log(12/δ)α i ϵ 2 j α j where we upper bound on the success event E 2 .Combining this with the initial T 0 pulls, the pulls from best arm identification, and the pulls from estimating the value of the best arm, we have that the overall sample complexity for arm i is upper bounded as:ñi = 34β 2 σ 2 log(6n/δ) + 64 √ 2σ 2 β 2 log(6n/δ) j γ j γ i ϵ j α j + 16β 2 σ 2 log(12/δ)α i ϵ 2 j α j + 32σ 2 ln 12n δ log 2 2 (4/∆ i ) ∆ 2 i + 512σ 2 β 2 log(6/δ) ϵ 2 1{i = 1}And so, the total algorithmic sample complexity on the success event isEQUATIONWe work to provide a simplified (looser) bound on the sample complexity when the minimum gap is ∆. The worst case sample complexity in this case is when the best arm has mean µ 1 , and all the rest have mean µ 1 -∆.This allows us to simplify the overall sample complexity asT ≤ C σ 2 β 2 n log n δ + n i=1 log n log d δ ∆ 2 i + β 2 log n δ n j=1 γ j 2 ϵ j α j + β 2 log(1/δ) ε 2 (45) ≤ C β 2 σ 2 log n δ n + ε -1 1 + n 2 e -β∆ 1 + ne -β∆ + β 2 σ 2 ε -2 log(1/δ) + nσ 2 log n log d δ ∆ -2where we use the fact that for n > 2:n j=1 γ j 2 j α j = 1 + (n -1)e -β∆/2 2 1 + (n -1)e -β∆ ≤ 2 1 + (n -1) 2 e -β∆ 1 + (n -1)e -β∆ ≤ C 1 + n 2 e -β∆ 1 + ne -β∆ .Evaluating our sample complexity in (45) , when ∆ > 2 β log n, this term is bounded by a constant, and the sample complexity can be more simply bounded as.EQUATIONIn the last line we require the condition thatEQUATIONi.e. d is not doubly exponential in n and β is not too small. Assumption 2 (large gap, moderate β, moderate d). We assume that (47) holds, and that ∆ > 2 β log n.Under the conditions of Assumption 2, the sample complexity in Theorem 1 can be simplified as in Corollary 1.Following the approach of [8] , we can show the asymptotic optimality of our sampling frequencies (sampling proportional to α i for minimizing the first order error, and γ i for minimizing our bound on the second order error).Considering a plug-in estimator μ; analyzing the first order taylor expansion of its error, we have thatf (μ) -f (µ) = ∇f (µ) ⊤ (μ -µ) + O ∥μ -µ∥ 2 2 .Thus in the high accuracy regime (ε → 0), we can consider only the first order term. Assuming Gaussian noise in our arm pulls (an identical result holds for sub-Gaussian noise), the first order error can be bounded as (assuming we use T pulls, and sample arm i, p i T times for a probability distribution p:∇f (µ) ⊤ (μ -µ) ∼ N 0, n i=1 β 2 e 2βµiT p iOptimizing over probability distributions p, using Sion's minimax theorem and strong duality as in [8] gives us thatα ⋆ = argmin p n i=1 β 2 e 2βµi T p i = argmin p max λ n i=1 e 2βµi T p i + λ i p i -1utilizing strong duality, we have thatEQUATION)Note that in the limit as ε → 0, multiplicative and additive error objectives are equivalent.With the first order error term in hand, the second order error term is left to be analyzed:EQUATIONWe minimize a bound on this given by our confidence intervals |µ i -μi | ≤ c √ ni , where c is some constant. We identify what sampling distribution p minimizes this second order bound, defining n i = p i T .By a similar argument as for the first order analysis:γ ⋆ = argmin p n i=1 β 2 e βµi 1 T p i = argmin p max λ n i=1 e βµi 1 p i + λ i p i -1utilizing strong duality, we have thatEQUATIONGains by sampling according to γ i Sampling according to γ gives a second order error bounded byEQUATIONas opposed to sampling according to α which gives a second order error bounded byn i e βµi = n∥γ∥ 2 2 (52)By standard norm inequalities, sampling according to γ is always at least as good, and up to a factor of n improvement in the case where one entry in the softmax is much larger than the rest; exactly the case of interest.In [8] , the problem of estimating a real valued function f to additive accuracy ε with probability at least 1 -δ is studied, under the assumption that the function has L-Lipschitz gradients. For the case of softmax estimation, the gradients are not Lipschitz due to the unbounded nature of the exponential. However, if we evaluate the norm of the gradient at a point µ, we obtain∥∇f (µ)∥ 2 = lim c→0 max ∥u∥≤c c -2 ∥∇f (µ + u) -∇f (µ)∥ 2 2 = lim c→0 max ∥u∥≤c c -2 i βe β(µi+ui) -βe βµi 2 = lim c→0 max ∥u∥≤c c -2 β 2 i e 2βµi e βui -1 2 = β 4 max i e 2βµi .Theorem 1 of [8] states that the number of samples required to achieve ε additive error with probability at least 1 -δ isEQUATIONwhere the noise variance σ 2 is assumed to be 1. Since the error in our setting is multiplicative, we are interested in ε ′ = ϵf (µ) = ε i e βµi . Additionally, ∥∇f (µ∥ 2 1 = β 2 i e βµi 2 . Thus, the number of samples required isT = O ∥∇f (µ∥ 2 1 log(1/δ) ε 2 + n 2 L log(n/δ) ε = O β 2 i e βµi 2 log(1/δ) ε 2 ( i e βµi ) 2 + n 2 β 4 max i e 2βµi log(n/δ) ε i e βµi = O β 2 log(1/δ) ε 2 + n 2 β 4 max i e 2βµi log(n/δ) ε i e βµiThis is to be compared with the sample complexity of the proposed algorithm in this paper, for the specific setting of softmax normalization estimation, which isT ≤ Cβ 2   n log n δ +   log n δ   n j=1 γ j   2      ϵ j α j   -1 + log 1 δ ε 2    , taking σ 2 = 1 to compare results.The constant term independent of ε is to linearize the exponential. The ε -2 term matches between the two settings, as asymptotically the optimal strategy is indeed to sample according to the first derivative. The term scaling with ε -1 improves dramatically on that of prior work. Note that in the case of f (µ) = ∥µ∥ 2 2 , the second order term (scaling with ε -1 ) can be improved to O(n 3/2 Lε -1 ), as the mean of the second order error can be removed. Thus, we can see the massive improvement afforded by our more refined algorithm, tailored for the specific structure of the softmax function.Adapting bandit algorithms to settings with heterogeneous variances has been done in both the standard regret [3] and best arm identification [33] settings.For best arm identification, sacrificing log factors for the sake of clarity, empirical-Bernstein-based confidence intervals [35] can be constructed where we iteratively pull each arm once, try and eliminate, and progress. Union bounding over the nd possible pulls naively upper bounds this, yielding a complexity ofEQUATIONThis assumes that all arms are bounded in [0, 1] with variance σ 2 i . We additionally require a lemma for estimating the mean of the best-arm in a PAC sense. Lemma 7 (Exponential best arm estimation). Sampling arm iT = 32σ 2 i β 2 log(2/δ) ϵ 2samples guarantees that e β μi estimates e βµi to multiplicative accuracy ϵ, with probability at least 1 -δ.This trivially follows from the proof of Lemma 2.For the softmax normalization estimation, we know from [8] that the optimal first order sampling frequencies are to sample arm i a number of times proportional toEQUATIONHowever, sampling like this yields an additive first order that scales as i σ i e βµi , which cannot be easily related to i e βµi , as we would need to get multiplicative error bounds. Thus, we instead use suboptimal target first order sampling frequencies, scaling with σ 2 i , to avoid this analysis issueEQUATIONScaling the number of pulls for each arm by σ 2 i yields: Proposition 3 (Softmax normalization estimation: heterogeneous variances variant of Proposition 1). Under Assumption 1, Algorithm 2 will, with probability at least 1 -δ, estimate f β (µ) = j e βµj to a multiplicative accuracy of ϵ, using a number of samples for arm i at mostn i = 34β 2 σ 2 i log(6n/δ)n + 91σ 2 i β 2 log(6n/δ) ( n i=1 γ i ) 2 ϵ i α i + 16β 2 σ 2 i log(12/δ) ϵ 2The proof of this proposition follows similarly to Proposition 2, as sampling proportional to σ 2 i cancels the differing variances, essentially reducing the problem to the homogeneous setting (suboptimally).In this appendix, we present the implementation details of our algorithm. In Algorithm 4, we provide pseudocode with greater detail about our implementation of the Adaptive Softmax algorithm. We note that Algorithm 4 contains some implementation differences from the original Algorithm 1 presented in Section 4. None of these changes materially affect the output of the algorithm; nonetheless, we provide a discussion of them here to enable reproducibility of our experimental results. Our results are also reproducible via a 1-line script in our code submission. In the following subsections, we describe each of these implementation details.As an important note, we consider all variables global unless stated otherwise. This is to say, calling "pull arms" updates the state of arm mean estimates μ, their variance estimates σ, the number of pulls per arm {n i }, inclusion probabilities π, etc. We define the estimators based on Gumbel sampling according to importance weights as the set of arms A. This idea of treating an arm simply as a sequence of estimators with confidence intervals was pioneered for the computational setting in [4] , and saw further usage in [24] .In our theoretical analysis in Appendix B, each phase of Algorithm 1 is handled independently. This allows us to union bound the error probabilities of each phase of the algorithm. In our implementation, Update variance estimate σ2 i ← 4: PullToVariance({i}, targetVar) 5: return μi however, we re-use arm pulls across different parts of Algorithm 1. Intuitively, once an arm return has been observed (corresponding to a scalar multiplication of an element of A with an element of x), it can be used to warm-start estimates of μ and σ2 in later stages of the algorithm. In practice, we observe that the re-use of arm pulls does not affect the correctness of our algorithm and yields significant sample complexity improvements compared to cold-starting each stage of the algorithm independently.EQUATIONIn Line 9 of Algorithm 4, we compute the mean of our estimated best arm i * exactly, and set the "estimate" μi = µ i . This allows us to reduce the approximation fidelity required by NormalizationEstimation (Algorithm 2) to ϵ 2 instead of ϵ 4 and saves a constant factor in sample complexity. Furthermore, this computation of µ i is efficiently computed as a vector-vector dot product.In practice, we found that Algorithm 1 usually required over d samples for the best arm. As such, performing the computation µ i = ⟨A i• , x⟩ (reusing coordinate-wise samples from previous stages of the algorithm) did not significantly increase sample complexity.In NormalizationEstimation (Algorithm 2), the initial number of arm pulls T 0 depends on σ 2 . However, as discussed in Appendix A.3, this variance proxy is often unknown a priori. As such, we set T 0 = d 10 . We observe that this choice of T 0 works well in experiments, across multiple datasets.We note that despite the changes made above, there still exists some looseness in practice. To remedy this, we scale our variance estimates by a constant factor to reduce the amount of pulls needed to reach target variances, and improving the gains of our algorithm as a result. We generate these constant factors for each dataset/model by tuning on a separate "training" group of queries. Tuning is performed, generally, via bisection to discover the minimal factor which still satisfies our provided failure probability parameter δ. This bisection is performed in geometric space and terminates when the log 10 difference between the low and high end of our interval is within 10 -2 . The range of factors we consider is [10 -6 , 1]. We first tune the constant factor independently for the variance estimates used in bandits such that the algorithm successfully identifies the best arm with a rate of at least 1 -δ on our training set. Next, we tune the constant factor for the variance estimate used in log norm estimation such that the entire SFTM algorithm succeeds with a rate of at least 1 -δ on our training set.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. 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 [20] . AdaptiveSoftmax already presents an opportunity to significantly scale down the number of entries of the matrix 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 [7] . Adjacent column samples can also be combined by simply summing their respective importance weights -admitting a simple tiling of our matrix that could easily be sized particularly to fit individual tiles into SRAM on a GPU along with a copy of the vector 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 -one with wider tiles and one with taller tiles -to take advantage of this structure 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.The user-desired parameters can take a wide range. We kept ϵ = 30% constant across all simulations because we observed varying ϵ did not result in significant changes to the performance of adaSoftmax. Further, we suspect that for most users, δ will fall in the range we considered: 90 -99%. However, to assuage any concerns and verify our assertion that adaSoftmax is not sensitive to choice of ϵ or δ, we include here adaSoftmax with a much wider range of parameters on the MNIST dataset in Table 3 . 3 : Success rate (%) and FLOP gains (x) for adaSoftmax with varied δ and ε on the MNIST dataset, showing the improved performance across a wide range of parameters, and that raising ε past 0.1 causes minimal difference in performance.ε = 0.001 ε = 0.01 ε = 0.1 ε = 1.0 δ = 0.Temperature is treated as a fixed constant (fixed parameter for the problem at hand, not tunable by the algorithm). This is because tuning the temperature fundamentally changes the problem. With higher temperatures, the only arms that matter are the best and second best arms, and so adaptivity is extremely helpful. At low temperatures, the output will be essentially the uniform distribution, and the computation is trivial and adaptivity unhelpful. With respect to other parameters, the error probability and FLOP gains of adaSoftmax are insensitive to changes in ε and vary most with the choice of δ. We demonstrate this trend on the MNIST dataset in Table 3 .The analysis posed in this paper focuses on identifying the largest entry in the softmax output and estimating its associated probability. As discussed above, this naturally extends to identifying the k largest elements in the output vector by replacing the bandit best arm identification algorithm with any top-k identification algorithm [39] . However, in LLM inference, the goal is often to draw a sample from the softmax output distribution via nucleus sampling [17] . Nucleus sampling avoids specifying k directly; instead, it provides cumulative probability p and requires the identification of the top k such that k is the smallest value such that the sum of the probabilities of the top k elements is greater than p. The next token is then sampled from the renormalization probability distribution on these k elements. Our adaptive sampling algorithm naturally applies to the nucleus sampling setting. AdaptiveSoftmax can maintain a predicted set of arms S such that the sum of the arm probabilities p is greater than p based on pessimistic arm mean estimates. Then, we iteratively sample: a) arms in S, sampling both the arm with the lowest mean minus LCB (in an attempt to verify the boundary), as well as the arm with the widest confidence interval (in order to better estimate p), and b) sampling the top arm in [n] \ S, to see if it belongs in S. For simplicity and concreteness, in this work we focus on identifying and estimating the probability of the top-1 element, but this is an exciting direction of future work. 
 
 Appendix:  


Decision: 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.

Reviews: review: paper_summary: paper_summary: 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.
strength_weakness: strength_weakness: a)The algorithm relies on assumptions about the data distribution and variance proxies, which may not hold in all practical scenarios.
b)The adaptive nature of the algorithm introduces implementation complexity, particularly in balancing computational resources and ensuring efficient sampling.
c)While the algorithm is theoretically sound, its practical benefits are most pronounced in scenarios with very high-dimensional data. In lower-dimensional settings, the gains may be less significant.
d)It is better to give detailed explanation of each step in Algorithms. What is EstimateArm in Algorithm 1 step 8?
e)How sensitive is the algorithm's performance to the choice of parameters such as the temperature parameter?
f)Besides comparing with the full softmax computation, is there a detailed comparison with other approximate softmax algorithms?

Typo:
Line 5: “we present present an …” should be “we present an …
Line 135: “Our objective them becomes …” should be “Our objective becomes …”
questions: questions: a)See weakness
limitations: limitations: 8.The limitations is discussed in the paper.
paper_summary: paper_summary: 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$.
strength_weakness: strength_weakness: - Overall, a more extensive evaluation with improved clarity is highly desirable. It would be important to see the effect of changing the $\epsilon$ parameter as well when conducting the evaluation, which is, to my understanding, set to a constant value of $30%$ throughout the experiments. Furthermore, reporting the effects of these changes could be made more comprehensive if presented in a plot with a wider range of choices for $\delta$, rather than just the three.
- As has been discussed in the limitations section, there can be a considerable trade-off between more involved, adaptive methods and easy-to-batch and parallelize brute-force computations. To my understanding, the experiments do not consider this trade-off, e.g., by reporting the wall-clock time next to the gains in sample efficiency.
questions: questions: I have no questions for the authors.
limitations: limitations: The authors have addressed the limitations in a dedicated section.
paper_summary: paper_summary: 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.
strength_weakness: strength_weakness: In practical implementation, parameter tuning may be necessary. An ablation study on the sensitivity and justification of default hyperparameters would enhance the paper. Additionally, discussing scenarios where the proposed method may not be suitable would strengthen the paper. For instance, while the performance degradation due to approximation is shown to be limited on average, what about the worst-case scenarios, such as cases with significant data distributional shift?
questions: questions: How was the performance of AdaptiveSoftmax evaluated? For example, are the speedup metrics in Tables 1 and 2 theoretical values or execution time comparisons? If it is the latter, what benchmarks were used? PyTorch’s torch.nn.CrossEntropyLoss utilizes log_softmax and nll_loss, which are highly optimized for GPU performance. Has AdaptiveSoftmax been compared with these in practical settings on GPUs?
limitations: limitations: The authors appropriately mention the limitations of the proposed algorithm, noting that it is most beneficial in high-dimensional settings. However, in practical settings involving LLMs, methods highly optimized for GPUs are usually used. Therefore, the question remains as to whether a fair comparison has been made regarding the current effectiveness of the proposed method in this setting.
paper_summary: paper_summary: 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.
strength_weakness: strength_weakness: 1. The method relies on the assumption of a variance proxy bound for the sub-Gaussian parameters of the constructed estimators. While the paper discusses loosening this assumption, its practical implications and the extent to which it holds in various scenarios are not thoroughly explored, potentially limiting the algorithm's applicability in more varied or less controlled environments.

2. The paper does not compare quantitatively to other adaptive Softmax methods, like A

[A] Joulin, Armand, Moustapha Cissé, David Grangier, and Hervé Jégou. "Efficient softmax approximation for GPUs." In International conference on machine learning, pp. 1302-1310. PMLR, 2017.
questions: questions: 1. Please discuss if the assumptions are a constraint for the environment, does it have an impact on the effectiveness of the algorithm?
2. It would be nice to see some quantitative comparison of the proposed method with other existing algorithms, eg., speed improvements.
limitations: limitations: The authors have provided an extensive explanation of the limitations in the final section of the paper.

Comments: title: Official Comment by Reviewer kEa9
comment: Thank you again for your detailed comments and valuable reference, which have helped improve the quality of the paper. We would like to follow up, and see if you have any remaining questions or concerns that we can address during this discussion period.
Thank you again for the detailed review and helpful comments. In our rebuttal we worked to answer and incorporate, in a point by point manner, all the concerns and ideas that you raised. Please let us know if you have any additional questions or concerns regarding this manuscript, we would be happy to discuss them in the remaining 1.5 days of the discussion period.
I thank the authors for the clarifications, I have no further questions. I would like to keep my positive score.
Thank you for the clarifying responses to the reviews. 

With the added insights from the rebuttal, I am happy to improve my rating.
Thank you for the swift response. We have incorporated yours and the other reviewers’ feedback and believe it has improved the quality of our paper. Further, it appears that the rebuttal has properly addressed the points raised in your review. If there are no further concerns, we would appreciate if your score were updated to reflect this. Please let us know if there is any additional information we can provide to help clarify.
Thank you for the insightful discussion!