LSH S OFTMAX : S UB -L INEAR L EARNING AND I NFERENCE OF THE S OFTMAX
L AYER
IN
D
EEP A RCHITECTURES
Anonymous authors
Paper under double-blind review
A BSTRACT
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
Locally-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.
I NTRODUCTION
Deep neural networks have achieved impressive successes in tasks spanning vision (He et al., 2016;
Krizhevsky et al., 2012), language (Bahdanau et al., 2014), speech (Graves et al., 2013; Oord et al., 2016) and videos (Abu-El-Haija et al., 2016).
While these models can vastly differ in architecture, activation functions, and presence of recurrence, they (almost) all share a common trait: the softmax layer.
The softmax layer, or log-linear model, is a widely used model in machine learning and statistics that transforms a feature vector into a distribution over the output space, modeling log-probabilities as a linear function of the feature vector.
For example, in object classiﬁcation, the softmax layer at the end of a deep convolutional network transforms a feature vector into a probability distribution over classes for the image; in language modeling using recurrent neural networks, it maps the hidden state to a distribution over next words.
While parameterizing for logits offers modeling ﬂexibility, inference and learning have linear runtime in the number of classes as it requires computing the un-normalized probability for every class to compute the partition function and retrieve an actual probability distribution.
Problems with large output spaces arise naturally in many areas like natural language processing (NLP), where the output space is a language’s vocabulary and can be on the order of hundreds of thousands of elements Jozefowicz et al.
; Jean et al.
This can also occur in computer vision (Joulin et al., 2016) when attempting tag prediction on massive, weakly-labeled datasets such as Flickr100M (Thomee et al., 2015).
Many solutions have been proposed to address this bottleneck, all revolving around two themes: approximation of the softmax probabilities or computation of exact probabilities for an approximate model.
Canonical examples of the former are importance sampling (IS) or noise contrastive estimation (NCE; Gutmann & Hyv ¨ arinen (2012)).
Instead of computing probabilities over the whole output space, these methods compute the softmax over a smaller, sampled vocabulary and re-weight the probabilities, providing an unbiased estimator.
An illustration of the latter is Hierarchical Softmax (Morin & Bengio, 2005), where the output classes are ﬁrst clustered such that you only need to compute the softmax over a smaller output space.
While the former is an unbiased estimate, it comes with no concentration guarantees, and it is often more art than science to craft proposal distributions which will provide low-variance estimators.
The latter, while efﬁcient, requires carefully hand-crafted clustering of the output space, at the risk of making mistakes from which there is no recovery.
More recently, estimators based on nearest neighbor search have been proposed for inference and learning in log-linear models (Mussmann & Ermon, 2016; Mussmann et al., 2017).
These estimators hinge on Maximum Inner Product Search using Locally-Sensitive to retrieve the largest logits of the distribution and account for the tail with uniformly sampled classes.
They boast strong theoretical guarantees and well-established concentration bounds.
However, they were constrained to toy settings and not directly applicable to real-world, large-scale, machine learning.
In this work, we build upon these estimators to make them amenable to deep learning practitioners, without losing any theoretical guarantees.
We ﬁrst show how they can be extended to be usable within training of deep learning models, then present our efﬁcient implementation, adapted to deep learning hardware and frameworks.
Finally, we show the applicability and efﬁciency of our method by evaluating on a real-world task: language modeling.
We show signiﬁcant perplexity gains against competing methods with signiﬁcant speed-ups.
Our contributions are as follows:
• We present a new deep learning layer, LSH Softmax, an efﬁcient replacement for the softmax layer based on Locally-Sensitive Hashing and the Gumbel distribution, for any deep learning architecture, with strong theoretical guarantees for sub-linear learning and inference.
• We provide details for efﬁcient implementation on deep learning hardware (GPUs) and modern deep learning frameworks (Abadi et al., 2016; Maclaurin et al.)).
• Empirically, we show, on several datasets, that training and sampling from LSH Softmax performs similarly to an exact softmax while requiring signiﬁcantly less FLOPS.
B ACKGROUND
In this section, we ﬁrst provide a quick overview of Neural Networks and the most popular classiﬁcation layer, the softmax layer.
We then present the Gumbel distribution (Gumbel & Lieblein, 1954) and introduce Locally-Sensitive Hashing (Indyk & Motwani, 1998), both of which our estimator is built upon for inference and learning.
Notationally, X is the input space, e.g. X (cid:44) R d and Y is a discrete output space:
Y (cid:44) { 1 , . .
, C } .
2. N EURAL N ETWORKS
Feedforward Networks Neural networks models are built hierarchically by applying linear and non-linear transformations in alternating fashion.
Formally, given input x ∈ X , an m -layer neural network with σ ( · ) non-linearity transforms x into h deﬁned as:
[Equation]
{ W i } i ≤ m and { b i } i ≤ m are learned weights of the network.
σ ( · ) denotes an element-wise nonlinearity such as ReLU ( max( · , 0) ) or sigmoid ( (1 + exp( −· )) − 1 ).
Recurrent Networks Recurrent Neural Networks (RNN) are an extension of the previous setting to arbitrarily long sequences by keeping an internal state h t .
Formally, given an input sequence ( x 1 , . . .
, x T ) , it can be written as a dynamical system of the form:
[Equation]
where U and V are learnable weight matrices.
In practice, this parametrization is not wellconditioned for optimization as it can be subject to vanishing or exploding gradients and in practice the Longer Short Term Memory (LSTM; Hochreiter & Schmidhuber (1997)) is preferred.
In both cases, these outputs are then given as input to a softmax layer which produces a distribution over the output space Y .
In the rest of this work, we denote by φ the parameters of the neural network.
2. S OFTMAX
The softmax layer is the common name given to a log-linear model for multi-classiﬁcation at the end of a neural network.
Let us consider the multi-classiﬁcation setting with inputs in X and outputs in Y .
Given a feature vector ψ ( x ) and C weight vectors { θ c } c ≤ C , the softmax layer parameterizes the following distribution:
[Equation]
In the particular case of neural networks, p ( y | x ; θ, φ ) ∝ exp( h T θ i ) .
{ h T θ i } i ≤ C are called the logits.
It is important to note that computing the distribution over the output space, for inference or learning, requires O ( C ) operations.
For the rest of this work, θ denotes the parameters of the softmax whereas φ denotes the parameters of the neural network (producing the feature).
2. G UMBEL DISTRIBUTION
First introduced by Gumbel & Lieblein (1954), the Gumbel distribution is deﬁned by the following cumulative distribution function: p ( G < s ) =
exp( − exp( − s )) .
More practically, one can sample from the Gumbel distribution by ﬁrst sampling U ∼ U [0 , 1] and returning G = − log( − log( U )) .
This distribution is particularly useful as it casts sampling as optimization.
Theorem 1 (Maddison et al.
Let { y i } i ≤ C be un-normalized probabilities (or logits) over
Y and let { G i } i ≤ C be i.i.d Gumbel variables.
Then:
[Equation]
2. MIPS AND L OCALLY -S ENSITIVE H ASHING
Nearest neighbor search is a task that arises in many ﬁelds, such as information retrieval.
Given a ﬁxed set of vectors S and a distance, this task consists of, given any incoming query q , returning the vectors closest to the query according to the speciﬁed distance.
In this work, we will be interested in the Maximum Inner Product Search (MIPS) task.
Let S =
{ s 1 , . . .
, s N } be a subset of R d .
Given a query q ∈ R d , MIPS aims at retrieving arg max s ∈S q T s .
This requires Θ( N ) operations as one has to compute the dot-product of q with all elements of S .
In the case where we assume that, for a given set S , it is needed to retrieve the nearest neighbor for a large numbers of queries, we can achieve amortized sub-linear time .
This problem is commonly addressed with space partitioning techniques, such as Locally-Sensitive
Hashing (LSH; Indyk & Motwani (1998)).
LSH leverages hashing to reduce the number of candidate vectors to evaluate, based on the idea that similar vectors will hash in the same bucket.
We have the following result:
Theorem 2 (Indyk & Motwani (1998)) .
Given a set S of size N , a similarity measure d ( · , · ) and a family of hash functions H s.t. for S > T and p > q :
• ∀ x, y ∈ S , d ( x, y ) ≥ S ⇒ p
[ h ( x ) =
h ( y )] ≥ p .
• ∀ x, y ∈ S , d ( x, y ) ≤ T ⇒ p
[ h ( x ) =
h ( y )] ≤ q .
we can construct a data structure s.t. given an incoming query q , a nearest neighbor can be retrieved, with high probability, in sub-linear time
O ( N ρ log N ) with ρ (cid:44) log p log q < 1 .
Recent work builds on top of LSH to either reduce the number of tables (Lv et al., 2007), or utilize more expressive hash functions (Andoni et al., 2015).
A common family of hash is the hyperplane hash, i.e. for v ∼ N (0 , I ) , h v ( x ) = sign (cid:0) v T x (cid:1) , also called Signed Random Projections (Charikar, 2002).
For the rest of this work, we denote b
the number of hashing bits (equivalently, the number of random vectors) per table, and L the number of tables.
L EARNING
In this section, we show how we can apply Theorem 3.5 of (Mussmann et al., 2017) to enable sublinear learning of softmax parameters in the context of deep models, i.e. where both weights and inputs can change.
This is crucial for real-world use.
Deep learning models for both classiﬁcation (Krizhevsky et al., 2012) and generation (Mikolov, 2012) are often trained with a maximum-likelihood objective.
Formally, given a training pair ( x, y ) ∈
X × Y , one aims at maximizing log p ( y | x ; θ, φ ) , where θ ∈ Θ and φ ∈ Φ are respectively the parameters of the softmax and of the neural network.
To optimize this model, the usual method is to use back-propagation (Rumelhart et al., 1988) to differentiate and then perform stochastic gradient descent (SGD; LeCun et al.
on θ and φ .
Let’s denote by f ( x ; φ ) (cid:44)
h the feature vector given as input to the softmax.
Given our notation, the objective is written as (cid:96) ( x, y, θ, φ ) (cid:44) log p ( y | x ; θ, φ ) =
h T θ y − log (cid:80) i ≤ C exp( h T θ i ) .
For backpropagation, we need to compute the gradient of (cid:96) w.r.t to both θ and h – the gradient w.r.t.
h is then passed down to compute the gradient w.r.t.
φ .
[Equation]
with Z θ ( h ) (cid:44) (cid:80) i exp( h T θ i ) .
Computing these gradients clearly requires O ( |Y| ) operations.
In practice, this constitutes a major bottleneck for large output spaces.
Mussmann et al.
(2017) shows how to compute expectation in in sub-linear time, with a well-concentrated estimator using an LSH structure.
Intuitively, we can build a good estimate of the partition function by retrieving the largest logits (using LSH) and accounting for the tail with uniform samples.
Applying this result, we can compute the expectations necessary to compute the softmax gradients in sub-linear time.
This is described in Theorem 3.
Theorem 3 (LSH Softmax for Learning) .
Let h =
f
( x ; φ ) be input to a softmax layer with parameters { θ c } c ≤ C and deﬁne (cid:96) ( x, y, θ, φ ) as previously.
Given S , the k -nearest neighbors of h in { θ c } c ≤ C and T , and l uniform samples from { 1 , . . .
, C } − S , let us deﬁne:
[Equation]
These estimators are well concentrated: for (cid:15), δ > 0 , if k = l = O (cid:16) n 23 1 (cid:15) (cid:113) 1 δ (cid:17) , then with probability greater than 1 − δ :
[Equation]
While Theorem 3 provides computation of the gradients in sub-linear time, it is only usable in a setting where the weights ( { θ i } i ≤ C ) are not updated.
Indeed, querying nearest neighbors in sublinear time assumes that an appropriate data structure (here LSH) was built in advance.
However, when training deep models, we are required to update the weights at every training step.
This necessitates online updating of the LSH structure.
To maintain the sub-linear runtime, we perform these updates in a sparse manner.
We describe in Algorithm 1 how this estimator can be used in a training loop, with weight updating and sparse LSH updates.
Proposition 4.
The softmax computations described in Algorithm 1 run in sub-linear time.
Inputs: Dataset D = { ( x ( i ) , y ( i ) } i ≤ N ⊂ X × Y , k , l , n iters number of training iterations.
Initialize θ and φ
Initialize the MIPS structure with { i
V
for j ≤ n iters do Sample an example
∆ θ ←
Compute h ← f ( x ; φ )
Find S , k -nearest-neighbors of h using the MIPS.
Deﬁne T as l indexes uniformly sampled from Y − S .
[Equation]
function estimate
Output ˆ (cid:96) =
h T θ y − log ˆ Z
[Equation]
( h
[Equation]
ˆ
[Equation]
Pass down ˆ g h for back-propagation.
Re-hash the updated vectors (at most ( k + l ) ) into the right buckets.
end for
Proof.
The softmax computations can be split into three parts: retrieving nearest neighbors, computing the forward/backward passes, and rehashing updated vectors.
With a sub-linear MIPS such as LSH, the ﬁrst part is guaranteed to be sub-linear.
For the second part, computing the partition function and the entire gradient estimator requires computing a ﬁnite number of sums over
O ( k + l ) =
O ( n 23 ) terms, which is sub-linear.
The third part consists of re-hashing updated vectors.
Re-hashing a vector is a constant operation (consisting of b × L dot-products) and thus, given that only a sub-linear number of vectors are updated, re-hashing is sub-linear.
I NFERENCE
In the last section, we presented a method to speed-up training time based on an LSH data structure.
In addition to these training time gains, LSH Softmax can be utilized for computational gains at inference time as well.
While MAP inference can be easily derived from the MIPS structure, sampling from the conditional distribution is often required (e.g. to generate diverse sentences in language modeling or machine translation).
These gains can be crucial for large-scale deployment.
This is a direct application of (Mussmann et al., 2017) that once again leverages a MIPS structure and the
Gumbel distribution.
By lazily evaluating Gumbel noise, once can devise an inference scheme which allows to sample from log-linear models in sub-linear time.
Theorem 5 (LSH Softmax for Inference) .
We reuse the same notations as the once in Theorem 3.
We deﬁne t (cid:44) − log( − log(1 − l/C )) .
Let { G i } i
≤ k be k samples from the Gumbel distribution.
We then proceed to sample m ∼ Binomial( C, l/C ) , and sample T , m points from Y − S with associated Gumbels { G (cid:48) i } i ≤ m s.t.
each G (cid:48)
i are larger than t .
Let us deﬁne:
[Equation]
We then have the two following results:
For k = l ≥ (cid:113) log 1 δ , ˆ y is a sample from p ( y | x ; θ, φ ) with probability greater than 1 − δ .
This inference scheme runs in sub-linear time.
Proof.
(Mussmann et al., 2017)
We denote by p Gumbel ( ·| h ; θ ) the implicit distribution over Y provided by this inference scheme.
While we can sample from p Gumbel , we note that the likelihood is intractable.
We also emphasize that this scheme can be utilized for any softmax model, regardless of the training method.
E FFICIENT I MPLEMENTATION
Recent successes of deep neural networks hinge on their efﬁcient implementation on specialized hardware: Graphics Processor Units (GPU), which enables training of large models in reasonable time.
Often, methods with theoretically faster runtime are dismissed by practitioners because of their incompatibility with the hardware, rendering them hard to implement efﬁciently and ultimately not widely used.
In this section, we ﬁrst detail how our method is indeed amenable to GPU implementation and can amount to wall-clock gains in practice, and explain why LSH Softmax is easy to implement in the context of modern deep learning frameworks who often provide a gradient computation API.
GPU Implementation Standard LSH implementations consist of three steps:
Hashing :
Given a query q ∈ R d , hash q into L tables i
e. computing b × L dot-product with (random) hyperplanes.
Look-up : Given L signatures in { 0 , 1 } b , retrieve candidates in each of the L tables.
Let us denote C q the number of candidates retrieved.
Distances :
Given those candidates { x 1 , . . .
, x C q } ⊂ R d , compute the distances { q T x i } i ≤ C q
and only return the closest one.
It is also important to note that deep learning models are often trained using minibatch optimization; let us describe how each of these steps can be computed efﬁciently and in the minibatch setting.
The ﬁrst step is amenable to the GPU setting; a batch of queries { q i } i ≤ m ⊂ R d can be represented by Q ∈ R m × d .
Given that the hyperplanes are similarly presented in matrix form i.e.
H ∈ R d × ( b × L ) , the hashing step is equivalent to sign ( Q · H ) ∈ { 0 , 1 } m × ( b × L ) .
This is the type of operations that GPUs excel at: matrix-multiply followed by element-wise function.
The second step, while not as compatible with GPU, is still massively parallelizable using multithreading on CPU.
Given the computed signatures, one can run parallelism at the query level (i.e. each thread retrieves candidates for a given query), rendering that step efﬁcient.
It also allows for more memory-efﬁcient look-up such as (Lv et al., 2007).
The last operation is, once again, very amenable to GPU.
It simply consists of a gather (i.e. building a matrix with the appropriate indexes from the candidates) into a 3 -d tensor.
Indeed, after the previous step, the LSH structure returns m lists of C candidates, and the gather step returns the appropriate vectors from the vocabulary into a 3 -d tensor of shape
R m × C × d .
As the batched queries can be also seen as a 3 -d tensor R m × d × 1 , computing the exact distances then reduces to a batch matrix-multiply which is a very efﬁcient operation on GPU.
Software Implementation Another crucial point for practitioners is the ability to rely on frameworks automatically providing gradients, such as (Abadi et al., 2016; Maclaurin et al.), to implement deep learning models; this abstracts away the need to write down the exact gradients which can be both cumbersome and error-prone.
An additional advantage of our estimator is that it can be effortlessly implemented in these frameworks.
Indeed, given logits computed over the nearest-neighbors and the additional uniformly sampled indexes, one can compute the estimate of the partition function and thus an estimate of the loss.
Computing the gradient estimators now reduces to differentiating this loss, which can be very simply done using the framework’s differentiation API.
E XPERIMENTS
After having presented our new layer LSH Softmax, we now proceed to show its applicability and efﬁciency in a real-world setting for deep learning practitioners, speciﬁcally towards language modeling.
We ﬁrst show that our method signiﬁcantly outperforms approximate softmax baselines while performing within 20% of the performance of the exact softmax.
We then provide a computational comparison.
While we evaluate our method on NLP tasks, we want to emphasize that it is directly applicable to other domains, such as vision.
However, public vision benchmark datasets with large output spaces require signiﬁcantly more computational resources (e.g. 98 GPU nodes for 8 days for
Flickr100M (Thomee et al., 2015)) which is outside the scope of this paper.
6. L ANGUAGE M ODELING
Language modeling is the task of, given a sequence of words ( w 1 , . . .
, w T ) in a vocabulary V , estimating p ( w 1 , . . .
, w T ) =
(cid:81) t ≤ T p ( w t | w <t ) .
Substantial work has been done to model these distributions using non-parametric n -gram counts with additional smoothing techniques, but can fail to model long histories because of an exponential number of sequences.
Recently, parametric models using RNNs have shown impressive success on these tasks (Mikolov, 2012).
In this setting, large output spaces arise naturally, as the vocabulary size can range from 10 4 to 10 6 .
We ﬁrst describe our experimental protocol, and then report perplexity (ppl) of LSH Softmax against a set of baselines on this task for several datasets.
Datasets We evaluate our method on three standard datasets for Language Modeling with varying number of characters and vocabulary size:
• Penn TreeBank
(PTB): We follow the pre-processing described by (Mikolov, 2012), which results in 929 k training tokens, 73 k validation and 82 k test tokens with a 10 k vocabulary size.
• Text8 is a dataset consisting of the ﬁrst 100 millions characters of Wikipedia, and has a vocabulary size of 44 k .
This dataset has been used recently to benchmark language modeling models (Xie et al., 2017) as well as to benchmark softmax performance (Grave et al., 2016).
We use the 90 M ﬁrst words for training and split the remaining between the validation and test set.
• Wikitext-2.
First introduced in Merity et al.
(2016), this is a selected corpus of Wikipedia articles.
It has a vocabulary size of 33 k and contains 217 k tokens.
As previously, we split between a training, validation and testing set.
Baselines We evaluate the performance of models trained with (1) exact softmax i.e. computed over the entire output space, (2) Importance Sampled softmax (IS), as presented in (Jozefowicz et al., 2016), which consists of sub-sampling the vocabulary according to a proposal distribution based on unigram counts, and (3) Negative Sampling (NS), proposed in Mikolov et al.
(2013), equivalent to (IS) with a uniform distribution.
These baselines are what practitioners canonically use to circumvent the bottleneck of large output spaces.
Implementation Details Our architecture is a 2 -layer RNN with LSTM cells and 650 hidden units.
Weights are initialized uniformly within [ − 0 . 1 , 0 . 1] .
Our models are trained using SGD using gradient clipping, with an initial learning rate of 20 .
This learning rate is annealed when the validation perplexity plateaus.
Our models are trained for 40 epochs for PTB, 3 epochs for Text8,epochs for Wikitext-2.
With the notations of Theorem 3, for LSH Softmax, we choose k = 10 (cid:112) |V| and l = (cid:112) |V| .
For the IS and NS baselines, we choose to sample k + l classes from the output space for a fair comparison.
We choose the number of bits per signature b (cid:44) log 2 |V| and choose L , number of tables, to have sufﬁcient recall for the MIPS task.
6.1. L EARNING
We report perplexity for a ﬁxed architecture but comparing different softmax evaluations; we present both learning curves and perplexity on each set.
It is important to note that the ﬁnal perplexities are evaluated with using the full softmax.
Perplexities are reported in Table 1 and learning curves in
Figure 1.
We see that LSH Softmax consistently outperforms the approximate baselines by a fair margin while performing a similar number of operations, showcasing the strength of this estimator.
We also observe from the training curves that approximate methods’ performances tend to plateau, as IS and NS cannot target the proper classes to push down.
In constrast, LSH Softmax does not.
6. C OMPUTATIONAL C OMPARISON
Having established that the proposed estimator performs very well on real-world tasks, we now proceed to evaluate the computation gains.
It is important to note that for models with large output spaces, the softmax computation can amount to about 80% of the total computation (Joulin et al.,
[Table]
[Figure]
2016; Ji et al., 2015); we thus choose to only evaluate computational gains in the softmax layer.
We evaluate our method in CPU, with a batch size of 1 , to have an accurate estimation of the ratio of FLOPS.
We report both speed-up and validation perplexity (ppl) relative difference with the exact softmax for LSH Softmax and NS.
Note that NS requires the same number of operations as importance sampling (IS) but outperforms it in all tasks.
Additionally, we show the speed-ups one can achieve on the One Billion Word dataset (Chelba et al., 2013), whose ppl was not evaluated due to computational constraints.
We report the results in Table 2.
We observe that, while faster,
NS performs signiﬁcantly worse than LSH Softmax.
Furthermore, its performance deteriorates signiﬁcantly when increasing the size of the output space, contrary to LSH Softmax which always performs in the same relative range.
[Table]
C ONCLUSION
In this work, we presented LSH Softmax, a softmax approximation layer for large output spaces with sub-linear learning and inference cost (in the number of states) and strong theoretical guarantees.
We showcased both its applicability and efﬁciency by evaluating LSH on a common NLP task, language modeling.
On several datasets for this task, we report perplexity closest to exact training among all baselines, as well as signiﬁcant speed-ups.
Our hope is that, for any architecture, this layer could be chosen in lieu of softmax, when the output space is sufﬁciently large to warrant the approximation.
To that end, we plan to release source-code with the camera-ready version.