Title: LSH Softmax: Sub-Linear Learning and Inference of the Softmax Layer in Deep Architectures
Authors: Daniel Levy
Danlu Chan
Stefano Ermon
Venue: ICLR 2018
Keywords: LSH
softmax
deep
learning
sub
linear
efficient
GPU

Abstract: Log-linear models models are widely used in machine learning, and in particular are ubiquitous in deep learning architectures in the form of the softmax. While exact inference and learning of these requires linear time, it can be done approximately in sub-linear time with strong concentrations guarantees. In this work, we present LSH Softmax, a method to perform sub-linear learning and inference of the softmax layer in the deep learning setting. Our method relies on the popular Locality-Sensitive Hashing to build a well-concentrated gradient estimator, using nearest neighbors and uniform samples. We also present an inference scheme in sub-linear time for LSH Softmax using the Gumbel distribution. On language modeling, we show that Recurrent Neural Networks trained with LSH Softmax perform on-par with computing the exact softmax while requiring sub-linear computations.

TL;DR: we present LSH Softmax, a softmax approximation layer for sub-linear learning and inference with strong theoretical guarantees; we showcase both its applicability and efficiency by evaluating on a real-world task: language modeling.

Review discussion:
The paper proposes to use LSH to approximate softmax, which greatly speeds up classification with large output space. The paper is overall well-written.
However, similar ideas have been proposed before, such as "Deep networks with large output spaces" by Vijayanarasimhan et. al. (ICLR 2015). And this manuscript does not provide any comparison to any of those similar methods.
A few questions about the implementation,
(1) As stated in the manuscript, the proposed method contains three steps, hashing, lookup and distance. GPU is not good at lookup, so the manuscript proposes to do lookup on CPU. Does that mean the data should go back and forth between CPU and GPU? Would this significantly increase the overhead?
(2) At page 6, the LSH structure returns m list of C candidates. Is it a typo? C is the total number of classes. And how do you guarantee that each LSH query returns the same amount of candidates?
Experiment-wise, the manuscript leaves something to be desired.
(1) More baselines be evaluated and compared. In this manuscript, only IS and NS are compared. And pure negative sampling is actually rarely used in language modeling.
In addition to Vijayanarasimhan's LSH method, there are also a few other methods out there, such as hierarchical softmax, NCE, D-sothat ftmax ("Strategies for Training Large Vocabulary Neural Language Models" by Chen et. al. ACL 2016), adaptive softmax ("Efficient softmax approximation for GPUs" by Grave et. al).
(2) The results of the proposed method is not impressive. D-softmax and adaptive softmax can achieve 147 ppl on text 8 with 512 hidden units as described in other paper, while the proposed method can only achieve 224 ppl with 650 hidden units. Even the exact softmax have large difference in ppl. It looks like the authors do not tune the hyper-parameters well. With this suboptimal setting, it is hard to judge the significance of this manuscript.
(3) Why one billion word dataset is used in eval but not used for training? It is one of the best datasets to test the scalability of language models.
(4) We can see, as reported in the manuscript, that NS has bigger speedup than the proposed method.
So it would be nice to show ppl vs time curve for all methods.
Eventually, what we want is the best model given a fixed amount of training time. With the same amount of epochs, NS loses the advantage of being faster.
We first and foremost want to thank you for your time and valuable comments.
Comparison with Vijayanarasimhan et. al.:
It is true that the method from this work is similar in spirit to ours.
However, we wish to emphasize two key points.
First of all, their method is encompassed in ours by simply setting l=0.
Secondly, as shown in Mussmann et al. (UAI 2017), using the top-k largest values leads to highly biased gradients and significantly worse performance (Figure 4 and 5 of Mussman et al.).
Implementation
(1) It is important to note that the weight vectors are never copied over to CPU.
The data (i.e. the weight vectors for the classes) *never* needs to be copied over to CPU. Our method only requires copying to CPU the *hashed* batch of hidden states. This consists of a bit matrix of shape (batch_size x (k * L)). This is a small matrix and thus the copying overhead is minimal.
When copying back to GPU, one must simply copy the *indices* of the weight vectors for the gather operation, which is a small matrix (batch size x number of candidates).
(2) It is a typo. We fixed that in the text, thank you for pointing it out. We guarantee a fixed number of candidates by padding with uniform samples.
Experiment-wise
(1) We added several baselines in the text, namely another (unbiased) version of Importance Sampling and NCE. We decided against comparing against Hierarchical Softmax methods (such as D-Softmax and adaptive softmax) as these requires domain-knowledge and hand-engineering.
Furthermore, in contrast, they additionally enjoy no theoretical guarantees.
(2) On Text8, our models were trained for 3 epochs, whereas the cited methods were trained for 5 or 10 epochs. Our hyperparameters were chosen from the literature for good performance with exact softmax and not tuned additionally for the approximate softmaxes.
(3) We did not evaluate the One Billion Word dataset due to computational constraints but provided a computational comparison to show how our method could perform on even larger datasets.
Thank you once again for your time and comments, we hope this addresses your concerns and that you will reconsider your rating in light of this.
In this paper, the authors propose a new approximation of the softmax, based on approximate nearest neighbors search and sampling.
More precisely, they propose to approximate to partition function (which is the bottleneck to compute the softmax and its gradient), by using:
- the top-k classes (retrieved using LSH) ;
- uniform samples (to account for the tail of the distribution).
They describe how this technique can be used for learning, by performing sparse updates for the gradient (corresponding to the elements used to compute the partition function), and re-hashing the updated element of the softmax layers.
In section 5, they show how this method can be implemented on GPU, using standard operations available in neural networks framework such as TensorFlow or PyTorch.
Finally, they compare their approach to importance sampling and negative sampling, using language modeling as a benchmark.
They use 3 standards datasets to perform the evaluations: penn treebank, text8 and wikitext-2.
Pros:
- well written and easy to read paper
- interesting theoretical guarantees of the approximation
Cons:
- a bit incremental
- weak empirical evaluations
- no support for the claim of efficient GPU implementation
== Incremental ==
While the theoretical justification of the methods are interesting, these are not a contribution of the paper (but of previous work by Mussmann et al.).
In fact, the main contribution of this paper is to show how to apply the technique of Mussmann et al. in the setup of neural network.
The main difference with Mussmann et al. is the necessity of re-hashing the updated elements of the softmax at each step.
Other previous works have also proposed to use LSH to speed up computations in neural network, but are not discussed in the paper (see list of references).
== Weak evaluations ==
I believe that the empirical evaluation of section 6 are a bit weak.
First, there is a large gap between the perplexity obtained using the proposed method and the exact softmax (e.g. 97 v.s. 83 on ptb, 115 v.s. 95 on wikitext-2).
Thus, I do not believe that the experiments support the claim that the proposed method "perform on-par with computing the exact softmax".
Moreover, these numbers are pretty far from what other papers have reported on these datasets with similar models (I am wondering if the gap would be even larger with SOTA models).
Second, the authors do not report any runtime numbers for their method and the baselines on GPUs.
I believe that it would be more fair to plot the learning curves (Fig. 1) using the runtime instead of the number of epochs.
== Efficient implementation ==
In section 5, the authors claims that their approach can be efficiently implemented on GPUs.
However, several of the operations used by their approach are inefficient, especially when using mini-batches.
The authors state that only step 2 is inefficient, but I also believe that step 3 is (compared to sampling approaches).
Indeed, for their method, each example of a mini-batch uses a different set of elements to approximate the partition function (while for other sampling methods, the same set is used for the whole batch).
Thus a matrix-matrix multiplication is replaced by n matrix-vector multiplication (n is the batch size).
While these can be performed in parallel, it is much less efficient than a matrix-matrix multiplication.
Finally, the only runtime numbers provided by the authors comparing their approach to sampling is for a CPU implementation with a batch of size 1.
This setting is super favorable to their approach, but a bit unrealistic for most practical settings.
== Missing references ==
Scalable and Sustainable Deep Learning via Randomized Hashing
Ryan Spring, Anshumali Shrivastava
A New Unbiased and Efficient Class of LSH-Based Samplers and Estimators for Partition Function Computation in Log-Linear Models
Ryan Spring, Anshumali Shrivastava
Deep networks with large output spaces
Sudheendra Vijayanarasimhan, Jonathon Shlens, Rajat Monga & Jay Yagnik
We first and foremost want to thank you for your time and valuable comments.
Thank you for the additional references, we added those in the text along with a discussion.
== ``Incremental" and related work==
In addition to updating the MIPS structure with the updated weight vectors, we go above and beyond the experimental setup of Mussmann et al. 2017.
Indeed, while their experiments support their theoretical results, they are far from being close to a real-world setting and usable on a large-scale task. Building on this, we extend their theoretical results and introduce LSH Softmax, which is usable in a real-world setting and on a widespread task: language modeling.
Regarding the additional references: we added those in the text but we wanted to emphasize the following points:
- Regarding (Vijayanarasimhan et al.), their method is encompassed in ours by simply setting l=0.
Furthermore, as shown in Mussmann et al. (UAI 2017), using the top-k largest values leads to highly biased gradients and significantly worse performance (Figure 4 and 5 of Mussman et al.).
- Regarding Spring et al. (KDD), there are several significant differences.
First of all, their paper provides no theoretical guarantees which is a major difference with our work.
Secondly, their paper focuses on reducing memory footprint which is not the aim of our work.
- Regarding Spring et al. (arXiv), their estimator is indeed unbiased and efficient but, in contrast, provides no concentration guarantees. As with most importance sampling technique, their variance can get arbitrarily bad.
Finally, the results reported on PTB are worse than those of LSH Softmax and the ones for Text8 are comparable whilst being trained for 10 epochs (theirs) compared to 3 epoch (ours).
== Weak evaluations ==
The gap between exact and LSH is always within 20% whilst enjoying speed-ups up to 4.1x.
Regarding the exact softmax implementation on PTB, we used the hyperparameters provided by the standard PyTorch implementation.
While more complex models (HyperNetworks, PointerSentinel, Variational Dropout etc...) can provide better perplexity, our baseline (79.8 on the test set) is not weak by any mean (See [1] for a thorough evaluation of various models on PTB).
== Efficient Implementation ==
It is true that we do not provide a GPU comparisons as our implementation is not yet competitive with TensorFlow IS and NCE implementations.
However, since all of our operations are parallelizable, we posit that given professional engineering attention (which is the case for the TensorFlow IS and NCE) it should be competitive, especially given the theoretical runtime.
The CPU evaluation is meant to provide us with a reasonable FLOPS estimate; on that basis, we significantly outperform competing methods.
We hope that this addresses your comments and that you will reconsider your rating in light of these.
[1] Regularizing and Optimizing LSTM Language Models. Merity S. et al. 2017
Authors present LSH Softmax - a fast, approximate nearest neighbor search based, approach for computing softmax that utilizes the Gumbel distribution and it relies on an LSH implementation of the maximum inner product search.
In general the work presented in this paper is very interested and the proposed method is very appealing especially on large datasets. For the most part it draws from a previous work which is my main concern. It is very much inline with the previous work by Mussmman et al. and authors don’t really do a good job in emphasizing the relationship with this work which uses two datasets for their empirical analysis. This in turn gives the overall impression that their work is a simple addition to it.
With this in mind, my other concern is that their empirical analysis are only focused on a single task from the NLP domain (language modeling).
It would be good to see how well does the model generalizes across tasks in other domains outside of NLP.
How do the different softmax approaches perform across different model configurations? It appears that the analysis were performed using a single architecture.
What about a performance comparison on an extrinsic task?
Authors should discuss the performance of LSH Softmax on the PTB train set. It appears that it outperforms the exact (i.e. “full”) Softmax or perhaps it’s an overlook on my end.
Overall it feels that the paper was written really close to the conference deadline. Given the fact that the work is mostly based on the previous work by Mussmman et al. what would make the paper stronger and definitely ready to be presented at this conference is more in-depth performance analysis that would answer some of the above questions.
LSH is typically an abbreviation for “Locality Sensitive” rather than “Locally-Sensitive” Hashing.
At least this is the case with the original LSH approach.
For better clarity try rephrasing or splitting the first sentence in the second paragraph of the introduction.
I think the authors spent too much time in background section of the paper where they give an overview of concepts that should be well known to the reader (NNs and Softmax).
Theorem 3: Second sentence should be rephrased - “...and  $\mathcal{T}$, and $l$ uniform samples from…”
Theorem 3:  $\epsilon$ and $\delta$ should be formally introduced.
Section 5: pretty much covers well known concepts related to GPU implementations. Authors should spent more time focusing on the empirical analysis of their approach.
Section 6.1: “...to benchmark language modeling models...” should be rephrased.
How were the number of epochs chosen across the 3 collections?
Section 6.1.1: “...note that the final perplexities are evaluated with using the full softmax…” - This sentence is very confusing and it should be rephrased.
We first and foremost want to thank you for your time and valuable comments.
We updated the draft to address your comments and provide more specific answers below.
== Relationship with Mussmann et al. 2017 ==
While Mussmann et al. 2017 provides the theoretical grounding for our work, it is important to note that their experimental setup is very constrained.
While their experiments support their theoretical results, they are far from being close to a real-world setting and usable on a large-scale task. Building on this, we extend their theoretical results and introduce LSH Softmax, which is usable in a real-world setting and on a widespread task: language modeling.
== Tasks from different domain than NLP ==
We want to emphasize that our method is not a all domain-specific and conserves theoretical guarantees across domains. We evaluate our method on NLP task for two reasons: 1) they are particularly well-suited for evaluating our method (naturally large output spaces) 2) we did not dispose of the computational resources to tackle tasks from other domains such as vision (e.g. Flickr100M) which requires hundreds of GPUs for weeks. We briefly touched on that point in the introduction of Section 6.
== Architecture and hyperparameters cross-validation ==
First of all, it is important to note that our theoretical guarantees hold regardless of architecture, hyperparameters etc...
Secondly, we wanted to show that our technique performed well without further parameter tuning; to that end, we tuned all of our models for the EXACT softmax. We then evaluated the approximate softmaxes by simply swapping them in without further tuning. In our opinion, ease of tuning makes these methods used in practice.
== PTB Train set ==
The hyperparameters (and thus regularization strength) were heavily cross-validated for performance on PTB with the EXACT softmax. It thus makes sense that the generalization gap be as small as possible in that case; it is not clear how LSH Softmax interacts with those multiple regularization schemes and thus we did not pay particular attention to that lower training perplexity.
Thank you once again for your valuable comments and feedback, we hope to have addressed your concerns, and we hope you will reconsider your rating in light of this.

Meta-review:
role: PC
data: comment: The authors propose an efficient LSH-based method for computing unbiased gradients for softmax layers, building on (Mussmann et al. 2017). Given the somewhat incremental nature of the method, a thorough experimental evaluation is essential to demonstrating its value. The reviewers however found the experimental section weak and expressed concerns about the choice of baselines and their surprisingly poor performance.

Decision: Reject