U NBIASED SCALABLE SOFTMAX OPTIMIZATION
Anonymous authors
Paper under double-blind review
A BSTRACT
Recent neural network and language models have begun to rely on softmax distributions with an extremely large number of categories.
In this context calculating the softmax normalizing constant is prohibitively expensive.
This has spurred a growing literature of efﬁciently computable but biased estimates of the softmax.
In this paper we present the ﬁrst two unbiased algorithms for maximizing the softmax likelihood whose work per iteration is independent of the number of classes and datapoints (and does not require extra work at the end of each epoch).
We compare our unbiased methods’ empirical performance to the state-of-the-art on seven real world datasets, where they comprehensively outperform all competitors.
I NTRODUCTION
Under the softmax model 1 the probability that a random variable y takes on the label (cid:96) ∈ { 1 , ..., K } , is given by
[Equation]
where x ∈ R D is the covariate, w k ∈ R D is the vector of parameters for the k -th class, and
W = [ w 1 , w 2 , ..., w K ] ∈ R D
× K is the parameter matrix.
Given a dataset of N label-covariate pairs
D = { ( y i , x i ) } Ni =1 , the ridge-regularized maximum log-likelihood problem is given by
[Equation]
where (cid:107) W (cid:107) 2 denotes the Frobenius norm.
This paper focusses on how to maximize (2) when N, K, D are all large.
Having large N, K, D is increasingly common in modern applications such as natural language processing and recommendation systems, where N, K, D can each be on the order of millions or billions (Partalas et al., 2015;
Chelba et al., 2013; Bhatia et al.).
A natural approach to maximizing L ( W ) with large N, K, D is to use Stochastic Gradient Descent (SGD), sampling a mini-batch of datapoints each iteration.
However if K, D are large then the
O ( KD ) cost of calculating the normalizing sum (cid:80)
Kk =1 e x (cid:62) i w k in the stochastic gradients can still be prohibitively expensive.
Several approximations that avoid calculating the normalizing sum have been proposed to address this difﬁculty.
These include tree-structured methods (Bengio et al., 2003;
Daume III et al., 2016; Grave et al., 2016), sampling methods (Bengio & Sen´ecal, 2008; Mnih &
Teh, 2012; Joshi et al., 2017) and self-normalization (Andreas & Klein, 2015).
Alternative models such as the spherical family of losses (de Br´ebisson & Vincent, 2015; Vincent et al., 2015) that do not require normalization have been proposed to sidestep the issue entirely (Martins & Astudillo, 2016).
Krishnapuram et al.
(2005) avoid calculating the sum using a maximization-majorization approach based on lower-bounding the eigenvalues of the Hessian matrix.
All 2 of these approximations are computationally tractable for large N, K, D , but are unsatisfactory in that they are biased and do not converge to the optimal W ∗
=
argmax L ( W ) .
Recently Raman et al.
(2016) managed to recast (2) as a double-sum over N and K .
This formulation is amenable to SGD that samples both a datapoint and class each iteration, reducing the per iteration cost to O
( D ) .
The problem is that vanilla SGD when applied to this formulation is unstable, in that the gradients suffer from high variance and are susceptible to computational overﬂow.
Raman et al.
(2016) deal with this instability by occasionally calculating the normalizing sum for all datapoints at a cost of O ( NKD ) .
Although this achieves stability, its high cost nulliﬁes the beneﬁt of the cheap
O ( D ) per iteration cost.
The goal of this paper is to develop robust SGD algorithms for optimizing double-sum formulations of the softmax likelihood.
We develop two such algorithms.
The ﬁrst is a new SGD method called
U-max, which is guaranteed to have bounded gradients and converge to the optimal solution of (2) for all sufﬁciently small learning rates.
The second is an implementation of Implicit SGD, a stochastic gradient method that is known to be more stable than vanilla SGD and yet has similar convergence properties (Toulis et al., 2016).
We show that the Implicit SGD updates for the doublesum formulation can be efﬁciently computed and has a bounded step size, guaranteeing its stability.
We compare the performance of U-max and Implicit SGD to the (biased) state-of-the-art methods for maximizing the softmax likelihood which cost O ( D ) per iteration.
Both U-max and Implicit SGD outperform all other methods.
Implicit SGD has the best performance with an average log-loss 4.29 times lower than the previous state-of-the-art.
In summary, our contributions in this paper are that we:
Provide a simple derivation of the softmax double-sum formulation and identify why vanilla SGD is unstable when applied to this formulation (Section 2).
Propose the U-max algorithm to stabilize the SGD updates and prove its convergence (Section 3.1).
Derive an efﬁcient Implicit SGD implementation, analyze its runtime and bound its step size (Section 3.2).
Conduct experiments showing that both U-max and Implicit SGD outperform the previous state-of-the-art, with Implicit SGD having the best performance (Section 4).
C ONVEX DOUBLE - SUM FORMULATION 2. D ERIVATION OF DOUBLE - SUM
In order to have an SGD method that samples both datapoints and classes each iteration, we need to represent (2) as a double-sum over datapoints and classes.
We begin by rewriting (2) in a more convenient form,
[Equation]
The key to converting (3) into its double-sum representation is to express the negative logarithm using its convex conjugate:
[Equation]
4) where u = − log( − v ) and the optimal value of u is u ∗ ( a )
= log( a ) .
Applying (4) to each of the logarithmic terms in (3) yields
[Equation]
where f ( u, W ) =
N (cid:88) i = (cid:88)
[Equation]
is our double-sum representation that we seek to minimize and the optimal solution for u i is u ∗ i ( W )
=
log(1 + (cid:80) k
(cid:54) =
y i e x (cid:62)
i
[Equation]
w k − w yi ) ) ≥ 0 .
Clearly f is a jointly convex function in u and W .
In Appendix A we prove that the optimal value of u and W is contained in a compact convex set and that f is strongly convex within this set.
Thus performing projected-SGD on f is guaranteed to converge to a unique optimum with a convergence rate of O (1 /T ) where T is the number of iterations (Lacoste-Julien et al., 2012).
2. I NSTABILITY OF VANILLA SGD
The challenge in optimizing f using SGD is that it can have problematically large magnitude gradients.
Observe that f = E ik [ f ik ] where i ∼ unif ( { 1 , ..., N } ) , k ∼ unif ( { 1 , ..., K } − { y i } ) and f ik ( u, W ) =
[Equation]
(6) where β j =
N n j +( N − n j )( K − 1) is the inverse of the probability of class
j being sampled either through i or k , and n j = |{ i : y i =
j, i = 1 , ..., N }| .
The corresponding stochastic gradient is:
[Equation]
If u i equals its optimal value u ∗ i ( W )
=
log(1+ (cid:80)
k
[Equation]
=
y i e
[Equation]
i ( w k
[Equation]
w yi ) )
then e x (cid:62) i ( w k
[Equation]
w yi ) − u i ≤ 1 and the magnitude of the N ( K − 1) terms in the stochastic gradient are bounded by N ( K − 1) (cid:107) x
i (cid:107) 2 .
However if u i (cid:28) x (cid:62) i ( w k − w y i ) , then e x (cid:62) i ( w k − w yi ) − u i (cid:29) 1 and the magnitude of the gradients can become extremely large.
Extremely large gradients lead to two major problems: (a) the gradients may computationally overﬂow ﬂoating-point precision and cause the algorithm to crash, (b) they result in the stochastic gradient having high variance, which leads to slow convergence 3 .
In Section 4 we show that these problems occur in practice and make vanilla SGD both an unreliable and inefﬁcient method 4 .
The sampled softmax optimizers in the literature (Bengio & Sen ´ ecal, 2008; Mnih & Teh, 2012;
Joshi et al., 2017) do not have the issue of large magnitude gradients.
Their gradients are bounded by N ( K − 1) (cid:107) x
i (cid:107) 2 due to their approximations to u ∗ i ( W ) always being greater than x (cid:62) i ( w k − w y i ) .
For example, in one-vs-each (Titsias, 2016), u ∗ i ( W ) is approximated by log(1 + e x
(cid:62) i ( w k − w yi ) )
> x (cid:62) i ( w k − w y i ) .
However, as they only approximate u ∗ i ( W ) they cannot converge to the optimal W ∗ .
The goal of this paper is to design reliable and efﬁcient SGD algorithms for optimizing the doublesum formulation f
( u, W ) in (5).
We propose two such methods: U-max (Section 3.1) and an implementation of Implicit SGD (Section 3.2).
But before we introduce these methods we should establish that f is a good choice for the double-sum formulation.
2. C HOICE OF DOUBLE - SUM FORMULATION
The double-sum in (5) is different to that of Raman et al.
Their formulation can be derived by applying the convex conjugate substitution to (2) instead of (3).
The resulting equations are
L
[Equation] [Equation]
and the optimal solution for ¯
u i is ¯ u ∗
i ( W ∗ )
= log( (cid:80)
Kk =1 e x (cid:62) i w ∗ k ) .
Although both double-sum formulations can be used as a basis for SGD, our formulation tends to have smaller magnitude stochastic gradients and hence faster convergence.
To see this, note that typically x (cid:62) i w y i =
argmax k
{ x (cid:62) i w k } and so the ¯ u i , x (cid:62) i
w y i and e x (cid:62) i w yi
− ¯ u i terms in (8) are of the greatest magnitude.
Although at optimality these terms should roughly cancel, this will not be the case during the early stages of optimization, leading to stochastic gradients of large magnitude.
In contrast the function f ik in (6) only has x (cid:62) i w y
i appearing as a negative exponent, and so if x (cid:62) i w y i is large then the magnitude of the stochastic gradients will be small.
In Section 4 we present numerical results conﬁrming that our double-sum formulation leads to faster convergence.
S TABLE SGD METHODS 3. U- MAX METHOD
As explained in Section 2.2, vanilla SGD has large gradients when u i (cid:28) x (cid:62) i ( w k − w y i ) .
This can only occur when u i is less than its optimum value for the current W , since u ∗ i ( W )
=
log(1 + (cid:80) j (cid:54) =
y i e x (cid:62) i ( w k − w yi ) )
≥ x (cid:62) i ( w k − w y i ) .
A simple remedy is to set u i
= log(
+ e x (cid:62) i ( w k − w yi ) )
whenever u i (cid:28) x (cid:62) i ( w k − w y i ) .
Since log(1 + e x (cid:62) i ( w k − w yi ) ) > x (cid:62) i ( w k − w y i ) this guarantees that u i > x (cid:62) i ( w k − w y i ) and so the gradients are bounded.
It also brings u i closer 5 to its optimal value for the current W and thereby decreases the the objective f ( u, W ) .
This is exactly the mechanism behind the U-max algorithm — see Algorithm 1 in Appendix C for its pseudocode.
U-max is the same as vanilla SGD except for two modiﬁcations: (a) u i is set equal to log(1 + e x (cid:62) i ( w k − w yi ) )
whenever u i ≤ log(1 + e x (cid:62) i ( w k − w yi ) )
− δ for some threshold δ > 0 , (b) u i is projected onto [0 , B u ] , and W onto { W : (cid:107) W (cid:107) 2 ≤ B W } , where B u and B W are set so that the optimal u ∗ i ∈ [0 , B u ] and the optimal W ∗ satisﬁes (cid:107) W ∗ (cid:107) 2 ≤ B W .
See Appendix A for more details on how to set B u and B W .
Theorem 1.
Let B f = max (cid:107) W (cid:107) 22 ≤ B 2 W , 0 ≤ u ≤ B u , max ik (cid:107)∇ f ik ( u, W ) (cid:107) 2 .
Suppose a learning rate η t ≤ δ 2 / (4 B 2 f ) , then U-max with threshold δ converges to the optimum of (2), and the rate is at least as fast as SGD with same learning rate, in expectation.
Proof.
The proof is provided in Appendix D.
U-max directly resolves the problem of extremely large gradients.
Modiﬁcation (a) ensures that δ ≥ x (cid:62) i ( w k − w y i )
− u i
(otherwise u i would be increased to log(1 + e x (cid:62) i ( w k − w yi ) ) ) and so the magnitude of the U-max gradients are bounded above by N ( K − 1)
e δ (cid:107) x
i (cid:107) 2 .
In U-max there is a trade-off between the gradient magnitude and learning rate that is controlled by δ .
For Theorem 1 to apply we require that the learning rate η t ≤ δ 2 / (4 B 2 f ) .
A small δ yields small magnitude gradients, which makes convergence fast, but necessitates a small η t , which makes convergence slow.
3. I MPLICIT SGD
Another method that solves the large gradient problem is Implicit SGD 6 (Bertsekas, 2011; Toulis et al., 2016).
Implicit SGD uses the update equation
[Equation]
[Equation] where θ ( t ) is the value of the t th iterate, f is the function we seek to minimize and ξ t is a random variable controlling the stochastic gradient such that ∇ f ( θ ) =
E ξ t [ ∇ f ( θ, ξ t )] .
The update (9) differs from vanilla SGD in that θ ( t +1) appears on both the left and right side of the equation,
[Equation]
[Equation]
whereas in vanilla SGD it appears only on the left side.
In our case θ = ( u, W ) and ξ t =
( i t , k t ) with ∇ f ( θ ( t +1) , ξ t ) = ∇ f i t ,k t ( u ( t +1) , W ( t +1) ) .
Although Implicit SGD has similar convergence rates to vanilla SGD, it has other properties that can make it preferable over vanilla SGD.
It is known to be more robust to the learning rate (Toulis et al., 2016), which important since a good value for the learning rate is never known a priori.
Another property, which is of particular interest to our problem, is that it has smaller step sizes.
Proposition 1.
Consider applying Implicit SGD to optimizing f ( θ ) =
E ξ [ f ( θ, ξ )]
where f ( θ, ξ ) is m -strongly convex for all ξ .
Then
[Equation]
and so the Implicit SGD step size is smaller than that of vanilla SGD.
Proof.
The proof is provided in Appendix E.
The bound in Proposition 1 can be tightened for our particular problem.
Unlike vanilla SGD whose step size magnitude is exponential in x (cid:62) i ( w k − w y i ) − u
i , as shown in (7), for Implicit SGD the step size is asymptotically linear in x (cid:62) i ( w k − w y i ) − u i .
This effectively guarantees that Implicit SGD cannot suffer from computational overﬂow.
Proposition 2. Consider the Implicit SGD algorithm where in each iteration only one datapoint i and one class k (cid:54) =
y i is sampled and there is no ridge regularization.
The magnitude of its step size in w is O
( x (cid:62) i ( w k − w y i )
− u i ) .
Proof.
The proof is provided in Appendix F.2.
The difﬁculty in applying Implicit SGD is that in each iteration one has to compute a solution to (9).
The tractability of this procedure is problem dependent.
We show that computing a solution to (9) is indeed tractable for the problem considered in this paper.
The details of these mechanisms are laid out in full in Appendix F.
Proposition 3.
Consider the Implicit SGD algorithm where in each iteration n datapoints and m classes are sampled.
Then the Implicit SGD update θ ( t +1) can be computed to within (cid:15) accuracy in runtime O ( n ( n + m )( D + n log( (cid:15) − 1 ))) .
Proof.
The proof is provided in Appendix F.3.
In Proposition 3 the log( (cid:15) − 1 ) factor comes from applying a ﬁrst order method to solve the strongly convex Implicit SGD update equation.
It may be the case that performing this optimization is more expensive than computing the x (cid:62) i w k inner products, and so each iteration of Implicit SGD may be signiﬁcantly slower than that of vanilla SGD or U-max.
However, in the special case of n = m =we can use the bisection method to give an explicit upper bound on the optimization cost.
Proposition 4.
Consider the Implicit SGD algorithm with learning rate η where in each iteration only one datapoint i and one class k (cid:54) =
y i is sampled and there is no ridge regularization.
Then the Implicit SGD iterate θ ( t +1) can be computed to within (cid:15) accuracy with only two D -dimensional vector inner products and at most log 2 ( (cid:15) − 1 )+log 2 ( | x (cid:62) i ( w k − w y i ) − u
i | +2 ηN (cid:107) x i (cid:107) 22 +log( K − 1)) bisection method function evaluations.
Proof.
The proof is provided in Appendix F.
For any reasonably large dimension D , the cost of the two D -dimensional vector inner products will outweigh the cost of the bisection, and Implicit SGD will have roughly the same speed per iteration as vanilla SGD or U-max.
In summary, Implicit SGD is robust to the learning rate, does not have overﬂow issues and its updates can be computed in roughly the same time as vanilla SGD.
[Table]
E XPERIMENTS
Two sets of experiments were conducted to assess the performance of the proposed methods.
The ﬁrst compares U-max and Implicit SGD to the state-of-the-art over seven real world datasets.
The second investigates the difference in performance between the two double-sum formulations discussed in Section 2.3.
We begin by specifying the experimental setup and then move onto the results.
4. E XPERIMENTAL SETUP
Data.
We used the MNIST, Bibtex, Delicious, Eurlex, AmazonCat-13K, Wiki
[Table]
and Wiki-small datasets 7 , the properties of which are summarized in Table 1.
Most of the datasets are multi-label and, as is standard practice (Titsias, 2016), we took the ﬁrst label as being the true label and discarded the remaining labels.
To make the computation more manageable, we truncated the number of features to be at most 10,000 and the training and test size to be at most 100,000.
If, as a result of the dimension truncation, a datapoint had no non-zero features then it was discarded.
The features of each dataset were normalized to have unit L 2 norm.
All of the datasets were pre-separated into training and test sets.
We only focus on the performance on the algorithms on the training set, as the goal in this paper is to investigate how best to optimize the softmax likelihood, which is given over the training set.
Algorithms.
We compared our algorithms to the state-of-the-art methods for optimizing the softmax which have runtime O
( D ) per iteration 8 .
The competitors include Noise Contrastive Estimation (NCE) (Mnih & Teh, 2012), Importance Sampling (IS) (Bengio & Sen´ecal, 2008) and One-Vs-Each (OVE) (Titsias, 2016).
Note that these methods are all biased and will not converge to the optimal softmax MLE, but something close to it.
For these algorithms we set n = 100 , m = 5 , which are standard settings 9 .
For Implicit SGD we chose to implement the version in Proposition 4 which has n = 1 , m = 1 .
Likewise for U-max we set n = 1 , m = 1 and the threshold parameter δ = 1 .
The ridge regularization parameter µ was set to zero for all algorithms.
Epochs and losses.
Each algorithm is run for 50 epochs on each dataset.
The learning rate is decreased by a factor of 0.9 each epoch.
Both the prediction error and log-loss (2) are recorded at the end of 10 evenly spaced epochs over the 50 epochs.
Learning rate.
The magnitude of the gradient differs in each algorithm, due to either under- or overestimating the log-sum derivative from (2).
To set a reasonable learning rate for each algorithm on
We also experimented setting n = 1 , m = 1 in these methods and there was virtually no difference except the runtime was slower.
For example, in Appendix G we plot the performance of NCE with n = 1 , m =and n = 100 , m = 5 applied to the Eurlex dataset for different learning rates and there is very little difference between the two.
[Table]
[Figure]
x-axis is the number of epochs and the y-axis is the log-loss from (2) calculated at the current value of W .
each dataset, we ran them on 10% of the training data with initial learning rates η = 10 0 , ± 1 , ± 2 , ± 3 .
The learning rate with the best performance after 50 epochs is then used when the algorithm is applied to the full dataset.
The tuned learning rates are presented in Table 2.
Note that vanilla SGD requires a very small learning rate, otherwise it suffered from overﬂow.
4. R ESULTS
Comparison to state-of-the-art.
Plots of the performance of the algorithms on each dataset are displayed in Figure 1 with the relative performance compared to Implicit SGD given in Table 3.
The
Implicit SGD method has the best performance on virtually all datasets.
Not only does it converge faster in the ﬁrst few epochs, it also converges to the optimal MLE (unlike the biased methods that prematurely plateau).
On average after 50 epochs, Implicit SGD’s log-loss is a factor of 4.29 lower than the previous state-of-the-art.
The U-max algorithm also outperforms the previous state-of-theart on most datasets.
U-max performs better than Implicit SGD on AmazonCat, although in general
Implicit SGD has superior performance.
Vanilla SGD’s performance is better than the previous state-of-the-art but worse than U-max and Implicit SGD.
The difference in performance between vanilla SGD and U-max can largely be explained by vanilla SGD requiring a smaller learning rate to avoid computational overﬂow.
[Table]
[Figure]
The sensitivity of each method to the initial learning rate can be seen in Appendix G, where the results of running each method on the Eurlex dataset with learning rates η = 10 0 , ± 1 , ± 2 , ± 3 is presented.
The results are consistent with those in Figure 1, with Implicit SGD having the best performance for most learning rate settings.
For learning rates η = 10 3 , 4 the U-max log-loss is extremely large.
This can be explained by Theorem 1, which does not guarantee convergence for U-max if the learning rate is too high.
Comparison of double-sum formulations.
Figure 2 illustrates the performance on the Eurlex dataset of U-max using the proposed double-sum in (6) compared to U-max using the double-sum of Raman et al.
(2016) in (8).
The proposed double-sum clearly outperforms for all 10 learning rates η = 10 0 , ± 1 , ± 2 , − 3 , − 4 , with its 50 th -epoch log-loss being 3 .
08 times lower on average.
This supports the argument from Section 2.3 that SGD methods applied to the proposed double-sum have smaller magnitude gradients and converge faster.
C ONCLUSION
In this paper we have presented the U-max and Implicit SGD algorithms for optimizing the softmax likelihood.
These are the ﬁrst algorithms that require only O ( D ) computation per iteration (without extra work at the end of each epoch) that converge to the optimal softmax MLE.
Implicit SGD can be efﬁciently implemented and clearly out-performs the previous state-of-the-art on seven real world datasets.
The result is a new method that enables optimizing the softmax for extremely large number of samples and classes.
So far Implicit SGD has only been applied to the simple softmax, but could also be applied to any neural network where the ﬁnal layer is the softmax.
Applying Implicit SGD to word2vec type models, which can be viewed as softmaxes where both x and w are parameters to be ﬁt, might be particularly fruitful.
A P ROOF OF VARIABLE BOUNDS AND STRONG CONVEXITY
We ﬁrst establish that the optimal values of u and W are bounded.
Next, we show that within these bounds the objective is strongly convex and its gradients are bounded.
Lemma 1 ((Raman et al., 2016)) .
The optimal value of W is bounded as (cid:107) W ∗ (cid:107) 22 ≤ B 2 W where
B 2 W = 2 µ N log( K ) .
Proof.
[Equation]
Rearranging gives the desired result.
Lemma 2.
The optimal value of u i is bounded as u ∗ i ≤ B u
where B u = log(
+ ( K − 1) e 2 B x B w ) and B x = max i {(cid:107)
x i (cid:107) 2 }
Proof.
[Equation]
Lemma 3.
If (cid:107) W (cid:107) 22 ≤ B 2 W and u i ≤ B u then f ( u, W ) is strongly convex with convexity constant greater than or equal to min { exp( − B u ) , µ } .
Proof.
Let us rewrite f as
[Equation]
where θ =
( u (cid:62) , w (cid:62) 1 , ..., w (
cid:62) k ) ∈
R N + KD
with a i and b ik being appropriately deﬁned.
The Hessian of f is
[Equation]
where e i is the
i th canonical basis vector, 0 N is an N -dimensional vector of zeros and 1 KD is a
KD -dimensional vector of ones.
It follows that
[Equation]
Lemma 4.
If (cid:107) W (cid:107) 22 ≤ B 2 W and u i ≤ B u then the 2-norm of both the gradient of f and each stochastic gradient f ik are bounded by
[Equation]
Proof.
By Jensen’s inequality
[Equation]
Using the results from Lemmas 1 and 2 and the deﬁnition of f ik from (6),
[Equation]
and for j indexing either the sampled class k (cid:54) =
y i or the true label
y i ,
[Equation]
Letting
[Equation]
we have
[Equation]
In conclusion:
[Equation]
B S TOCHASTIC C OMPOSITION O PTIMIZATION
We can write the equation for L ( W ) from (3) as (where we have set µ = 0 for notational simplicity),
L ( W ) = −
N (cid:88) i
= log(1 + (cid:88)
k (cid:54) =
y i e x (cid:62) i ( w k − w yi ))
=
E i [ h i ( E k [ g k ( W )])] where i ∼ unif ( { 1 , ..., N } ) , k ∼ unif ( { 1 , ..., K } ) , h i ( v ) ∈
R , g k ( W ) ∈ R N and h i ( v ) =
− N log(
+ e (cid:62)
[Equation]
if k
[Equation]
= y
[Equation]
otherwise .
Here e (cid:62) i v = v i ∈ R is a variable that is explicitly kept track of with v i
≈ E k [ g k ( W )]
i = (cid:80) k
(cid:54) =
y i e x (cid:62) i ( w k − w yi ) (with exact equality in the limit as t → ∞ ).
Clearly v i in stochastic composition optimization has a similar role as u i has in our formulation for f in (5).
If i, k are sampled with k (cid:54) =
y i in stochastic composition optimization then the updates are of the form (Wang et al., 2016)
[Equation]
where z k is a smoothed value of w k .
These updates have the same numerical instability issues as vanilla SGD on f in (5): it is possible
that e x (cid:62) i zk
+ v i (cid:29) 1 where ideally we should have 0 ≤ e x (cid:62)
i
zk
+ v i ≤ 1 .
C U- MAX PSEUDOCODE
Input : Data D = { ( y i , x i ) :
y i ∈ { 1 , . . .
, K } , x i ∈ R d } Ni
=1 , number of classes K , number of datapoints N , learning rate η t , class sampling probability β k
=
N n k +( N − n k )( K − 1) , threshold parameter δ > 0 , bound B W on W such that (cid:107) W (cid:107) 2 ≤ B W and bound B u on u such that u i ≤ B u for i = 1 , ..., N
Output: W 1 Initialize 2 for k = 1 to K do w k ←4 end 5 for i = 1 to N do u
i ← log( K ) 7 end 8 Run SGD 9 for t = 1 to T do
Sample indices i ∼ unif ( { 1 , ..., N } )
k ∼ unif ( { 1 , ..., K } − { y i } )
Increase u i if u i < log(1 + e x (cid:62) i ( w k − w yi ) )
− δ then u i ← log(1 + e x (cid:62) i ( w k − w yi ) )
SGD step w k ← w k − η t [ N ( K − 1) e x (cid:62) i ( w k − w yi )
− u
i x i + µβ k w k ] w y i ← w y i − η t [ − N
( K − 1) e x (cid:62) i ( w k − w yi )
− u i x i
+ µβ y
i
w y
i ] u i ← u
i − η t
[ N (1 − e − u
i − ( K − 1) e x (cid:62)
i ( w k − w yi ) − u i )]
Projection w k ← w k · min { 1 , B W / (cid:107) w k (cid:107) 2 } w y i ← w y i · min { 1 , B W / (cid:107)
w y i (cid:107) 2 } u
i ← max { 0 , min { B u , u i }} 24 end
D P ROOF OF CONVERGENCE OF U- MAX METHOD
In this section we will prove the claim made in Theorem 1, that U-max converges to the softmax optimum.
Before proving the theorem, we will need a lemma.
Lemma 5.
For any δ > 0 , if u i ≤ log(1+ e x (cid:62) i ( w k − w yi ) ) − δ then setting u i
= log(
+ e x (cid:62) i ( w k − w yi ) ) decreases f ( u, W ) by at least δ 2 / 2 .
Proof.
As in Lemma 3, let θ =
( u (cid:62) , w (cid:62) 1 , ..., w (
cid:62) k ) ∈ R N + KD .
Then setting u i
= log(
+ e x (cid:62) i ( w k − w yi ) ) is equivalent to setting θ = θ + ∆ e i where e i is the
i th canonical basis vector and ∆ =
log(1 + e x (cid:62) i ( w k − w yi ) )
− u i ≥ δ .
By a second order Taylor series expansion
[Equation]
(10) for some λ ∈ [0 , 1] .
Since the optimal value of u i for a given value of W is u ∗ i ( W )
=
log(1 + (cid:80) k
(cid:54) =
y i e x (cid:62) i ( w k − w yi ) )
≥ log(
+ e x (cid:62) i ( w k − w yi ) ) , we must have ∇ f ( θ +∆ e i ) (cid:62) e i ≤ 0 .
From Lemma we also know that
[Equation]
Putting in bounds for the gradient and Hessian terms in (10),
[Equation]
Now we are in a position to prove Theorem 1.
Proof of Theorem 1.
Let θ ( t ) =
( u ( t ) , W ( t ) )
∈
Θ denote the value of the t th iterate.
Here Θ = { θ : (cid:107) W (cid:107) 22 ≤ B 2 W , u i ≤ B u } is a convex set containing the optimal value of f ( θ ) .
Let π ( δ ) i ( θ ) denote the operation of setting u i
= log(
+ e x (cid:62) i ( w k − w yi ) )
if u i ≤ log(1 + e x (cid:62) i ( w k − w yi ) )
− δ .
If indices i, k are sampled for the stochastic gradient and u i ≤ log(1 + e x (cid:62) i ( w k − w yi ) ) − δ , then the value of f at the t + 1 st iterate is bounded as
[Equation]
since η t ≤ δ 2 / (4 B 2 f ) by assumption.
Alternatively if u i ≥ log(1 + e x (cid:62) i ( w k − w yi ) ) − δ then
[Equation]
Either way f
( θ ( t +1) )
≤ f ( θ ( t ) − η t ∇ f ik ( θ ( t ) )) .
Taking expectations with respect to i, k ,
[Equation]
Finally let P denote the projection of θ onto Θ .
Since Θ is a convex set containing the optimum we have f ( P ( θ )) ≤ f ( θ ) for any θ , and so
[Equation]
] , which shows that the rate of convergence in expectation of U-max is at least as fast as that of standard
SGD.
E P ROOF OF GENERAL I MPLICIT SGD GRADIENT BOUND
Proof of Theorem 2.
Let f ( θ, ξ ) be m -strongly convex for all ξ .
The vanilla SGD step size is η t (cid:107)∇ f ( θ ( t ) , ξ t ) (cid:107) 2 where η t is the learning rate for the t th iteration.
The Implicit SGD step size is η t (cid:107)∇ f ( θ ( t +1) , ξ t ) (cid:107) 2 where θ ( t +1) satisﬁes θ ( t +1)
= θ ( t ) − η t ∇ f
( θ ( t +1) , ξ t ) , or equivalently,
[Equation]
argmin
[Equation]
In Implicit SGD it must be the case that ∇ f ( θ ( t +1) , ξ t ) is parallel to θ ( t ) − θ ( t +1) , that is ∇ f ( θ ( t +1) , ξ t ) (cid:62) ( θ ( t ) − θ ( t +1) )
=
(cid:107)∇ f ( θ ( t +1) , ξ t ) (cid:107) 2 (cid:107) θ ( t ) − θ ( t +1) (cid:107) 2 .
Imagine, to the contrary, that ∇ f ( θ ( t +1) , ξ t ) =
α ( θ ( t ) − θ ( t +1) )
+ βg where β > 0 , (cid:107) g (cid:107) 2 > 0 and g (cid:62) ( θ ( t ) − θ ( t +1) )
= 0 .
Then the g − directional derivative of (11) at θ = θ ( t +1) is 2 η t β > 0 .
Since f is convex it must be continuous, and so taking a sufﬁciently small step in the direction − g is guaranteed to improve the objective in (11).
However θ ( t +1) is the optimal value of (11).
Thus the assumption that β > 0 must be false for all g perpendicular to θ ( t ) − θ ( t +1) .
It must also be the case that α > 0 , otherwise the objective could be improved by taking a small step in the direction θ ( t ) − θ ( t +1) .
This proves that ∇ f ( θ ( t +1) , ξ t ) is parallel to θ ( t ) − θ ( t +1) .
Our desired result follows:
[Equation]
where the ﬁrst inequality is by Cauchy-Schwarz, the second inequality by strong convexity and the ﬁrst equality is from ∇ f ( θ ( t +1) , ξ t ) being parallel to θ ( t ) − θ ( t +1) .
F U PDATE EQUATIONS FOR IMPLICIT SGD
In this section we will derive the updates for Implicit SGD.
We will ﬁrst consider the simplest case where only one datapoint ( x i , y i ) and a single class is sampled in each iteration with no regularizer.
Then we will derive the more complicated update for when there are multiple datapoints and sampled classes with a regularizer.
F. S INGLE DATAPOINT , SINGLE CLASS , NO REGULARIZER
Equation (6) for the stochastic gradient for a single datapoint and single class with µ = 0 is
[Equation]
The Implicit SGD update corresponds to ﬁnding the variables optimizing
[Equation]
where η is the learning rate and the tilde refers to the value of the old iterate (Toulis et al., 2016, Eq.
Since f ik is only a function of u i , w k , w y i the optimization reduces to
[Equation]
The optimal value of w k , w y i must deviate from the old value ˜ w k , ˜ w y i in the direction of x i .
Furthermore we can observe that the deviation of w k must be exactly opposite that of w y i , that is:
[Equation]
for some a ≥ 0 .
The optimization problem reduces to min
[Equation] [Equation]
We’ll approach this optimization problem by ﬁrst solving for a as a function of u i and then optimize over u i .
Once the optimal value of u i has been found, we can calculate the corresponding optimal value of a .
Finally, substituting a into (12) will give us our updated value of W .
Solving for a
We solve for a by setting its derivative equal to zero in (13)
[Equation]
The solution for a can be written in terms of the principle branch of the Lambert W function P ,
[Equation]
Substituting the solution to a ( u i ) into (13)
[Equation]
we now only need minimize over u i :
[Equation]
where we used the fact that e − P ( z )
=
P ( z ) /z .
The derivative with respect to u i in (16) is
[Equation]
where to calculate ∂ u i a ( u i ) we used the fact that ∂ z P ( z ) =
P ( z ) z (1+ P ( z )) and so
[Equation]
i
[Equation]
Bisection method for u i
We can solve for u i using the bisection method.
Below we show how to calculate the initial lower and upper bounds of the bisection interval and prove that the size of the interval is bounded (which ensures fast convergence).
Start by calculating the derivative in (17) at u i =
˜ u
i .
If the derivative is negative then the optimal u i is lower bounded by
˜ u
i .
An upper bound is provided by u
[Equation]
=
[Equation]
In the ﬁrst inequality we set a ( u i )
= 0 ,
since by the envelop theorem the gradient of u i is monotonically increasing in a .
In the second inequality we used the assumption that u i is lower bounded by
˜ u
i .
Thus if the derivative in (17) is negative at u i =
˜ u
i then ˜ u i ≤ u i ≤ log(
+( K − 1) e x (cid:62) i ( ˜ w k − ˜ w yi ) ) .
If ( K − 1) e x (cid:62) i ( ˜ w k − ˜ w yi ) ≤ 1 then the size of the interval must be less than log(2) , since ˜ u
i ≥ 0 .
Otherwise the gap must be at most log(2( K − 1)
e x
[Equation]
i ( ˜ w k − ˜ w yi ) )
− ˜ u i =
[Equation]
Either way, the gap is upper bounded by log(2( K − 1))
+ | x (cid:62)
i
[Equation]
−
[Equation]
u i | .
Now let us consider if the derivative in (17) is positive at u i =
˜ u
i .
Then u i is upper bounded by ˜ u
i .
Denoting a (cid:48) as the optimal value of a , we can lower bound u i using (13)
[Equation]
where the ﬁrst inequality comes dropping the ( u i − ˜ u i ) 2 term due to the assumption that u i < ˜ u i .
Recall (14),
[Equation]
The solution for a (cid:48) is strictly monotonically increasing as a function of the right side of the equation.
Thus replacing the right side with an upper bound on its value results in an upper bound on a (cid:48) .
Substituting the bound for u i ,
[Equation] [Equation]
Substituting this bound for a (cid:48) into (18) yields
[Equation]
Thus if the derivative in (17) is postive at u i =
˜
u
i then log( K − 1) + x (cid:62) i ( ˜ w k − ˜ w y i ) − 2 ηN (cid:107) x
i (cid:107) 22 ≤ u i ≤ ˜ u
i .
The gap between the upper and lower bound is ˜ u
i − x (cid:62) i ( ˜ w k − ˜ w y i )
+2 ηN (cid:107) x
i (cid:107) 22 − log( K − 1) .
In summary, for both cases of the sign of the derivative in (17) at u i =
˜
u
i we are able to calculate a lower and upper bound on the optimal value of u i such that the gap between the bounds is at most | ˜
u i − x (cid:62) i ( ˜ w k − ˜ w y i ) | + 2 ηN (cid:107) x
i (cid:107) 22 + log( K − 1) .
This allows us to perform the bisection method where for (cid:15) > 0 level accuracy we require only log 2 ( (cid:15) − 1 )+log 2 ( | ˜ u i − x (cid:62) i
( ˜ w k − ˜ w y i ) | +2 ηN (cid:107) x
i (cid:107) 2 2 + log( K − 1)) function evaluations.
F.
B OUND ON STEP SIZE
Here we will prove that the step size magnitude of Implicit SGD with a single datapoint and sampled class with respect to w is bounded as O ( x (cid:62) i ( ˜ w k − ˜ w y i ) − ˜ u i ) .
We will do so by considering the two cases u (cid:48)
i ≥ ˜
u
i and u (cid:48)
i < ˜ u
i separately, where u (cid:48)
i denotes the optimal value of u i in the Implicit SGD update and ˜ u
i is its value at the previous iterate.
Case: u (cid:48)
i ≥
[Equation]
Let a (cid:48) denote the optimal value of a in the Implicit SGD update.
From (15)
[Equation]
Now using the fact that P ( z ) =
O (log( z )) ,
[Equation]
Case: u
[Equation]
i < ˜ u
i
If u (cid:48) i < ˜ u
i then we can lower bound a (cid:48) from (19) as
[Equation]
Combining cases
Putting together the two cases,
[Equation]
The actual step size in w is ± a x
i 2 (cid:107) x i (cid:107) 22 .
Since a is O ( x (cid:62) i ( ˜ w k − ˜ w y i ) − ˜ u i ) , the step size magnitude is also O
( x (cid:62) i ( ˜ w k − ˜ w y i )
− ˜ u i ) .
F. M ULTIPLE DATAPOINTS , MULTIPLE CLASSES
The Implicit SGD update when there are multiple datapoints, multiple classes, with a regularizer is similar to the singe datapoint, singe class, no regularizer case described above.
However, there are a few signiﬁcant differences.
Firstly, we will require some pre-computation to ﬁnd a low-dimensional representation of the x values in each mini-batch.
Secondly, we will integrate out u i for each datapoint (not w k ).
And thirdly, since the dimensionality of the simpliﬁed optimization problem is large, we’ll require ﬁrst order or quasi-Newton methods to ﬁnd the optimal solution.
F.3. D EFINING THE MINI - BATCH
The ﬁrst step is to deﬁne our mini-batches of size n .
We will do this by partitioning the datapoint indices into sets S 1 , ..., S J with S j = { j (cid:96) : (cid:96) = 1 , ..., n } for j = 1 , ..., (cid:98) N/n (cid:99) , S J = { J (cid:96) : (cid:96) = 1 , ..., N mod n } , S i ∩ S j = ∅ and ∪ J j =1 S j = { 1 , ..., N } .
Next we deﬁne the set of classes C j which can be sampled for the j th mini-batch.
The set C j is deﬁned to be all sets of m distinct classes that are not equal to any of the labels y for points in the mini-batch, that is,
C j = { ( k 1 , ..., k m ) : k i ∈ { 1 , ..., K } , k i (cid:54) =
k (cid:96) ∀ (cid:96) ∈ { 1 , ..., m } − { i } , k
i
[Equation]
= y (cid:96) ∀ (cid:96) ∈
S j } .
Now we can write down our objective from (5) in terms of an expectation of functions corresponding to our mini-batches:
[Equation]
where j is sampled with probability p j
= | S j | /N and C is sampled uniformly from C j and f j,C ( u, W ) =
p − 1 j
[Equation]
The value of the regularizing constant β k is such that E
[ I [ k ∈ C ∪ S j ] β k ] = 1 , which requires that
[Equation]
F.3. S IMPLIFYING THE I MPLICIT SGD UPDATE EQUATION
The Implicit SGD update corresponds to solving
[Equation]
where η is the learning rate and the tilde refers to the value of the old iterate (Toulis et al., 2016,
Eq. 6).
Since f j,C is only a function of u S j = { u i : i ∈ S j } and W j,C = { w k : k ∈ S j ∪ C } the optimization reduces to
[Equation]
The next step is to analytically minimize the u S j terms.
The optimization problem in (22) decomposes into a sum of separate optimization problems in u i for i ∈
S
j ,
[Equation]
where
[Equation]
Setting the derivative of u i equal to zero yields the solution
[Equation]
where P is the principle branch of the Lambert W function.
Substituting this solution into our optimization problem and simplifying yields
[Equation]
where we have used the identity e − P ( z )
=
P ( z ) /z .
We can decompose (20) into two parts by splitting W j,C = W (cid:107) j,C + W ⊥ j,C , its components parallel and perpendicular to the span of { x i : i ∈ S j } respectively.
Since the leading term in (20) only depends on W (cid:107)
j
[Equation]
C , the two resulting sub-problems are
[Equation]
Let us focus on the perpendicular component ﬁrst.
Simple calculus yields the optimal value w ⊥ k
=
+ µβ k / 2 ˜
w ⊥ k for k ∈ S j ∪ C .
Moving onto the parallel component, let the span of { x i : i ∈ S j } have an orthonormal basis
V j
=
( v j 1 , ..., v jn
) ∈ R D × n with x i =
V j b i for some b i ∈ R n .
With this basis we can write w
[Equation]
k +
[Equation]
j a k for a k ∈ R n which reduces the parallel component optimization problem tomin
[Equation]
22) where A j,C = { a k : k ∈ S j ∪ C } ∈ R ( n + m ) × n and
[Equation]
The e b (cid:62) i ( a k − a yi ) factors come from
[Equation]
since V j is an orthonormal basis.
F.3. O PTIMIZING THE I MPLICIT SGD UPDATE EQUATION
To optimize (22) we need to be able to take the derivative:
[Equation]
where we used that ∂ z P ( z )
=
P
[Equation]
z )
z
[Equation]
1+ P
[Equation]
z )) and e − P ( z ) =
P ( z ) /z .
To complete the calculation of the derivate we need,
[Equation]
In order to calculate the full derivate with respect to A j,C we need to calculate b (cid:62) i a k for all i ∈ S j and k ∈ S j ∪ C .
This is a total of n ( n + m ) inner products of n -dimensional vectors, costing
O ( n 2 ( n + m )) .
To ﬁnd the optimum of (22) we can use any optimization procedure that only uses gradients.
Since (22) is strongly convex, standard ﬁrst order methods can solve to (cid:15) accuracy in
O (log( (cid:15) − 1 )) iterations (Boyd & Vandenberghe, 2004, Sec. 9.3).
Thus once we can calculate all of the terms in (22), we can solve it to (cid:15) accuracy in runtime O ( n 2 ( n + m ) log( (cid:15) − 1 )) .
Once we have solved for A j,C , we can reconstruct the optimal solution for the parallel component of w k as w (cid:107) k
= ˜ w (cid:107) k
+ V j a k .
Recall that the solution to the perpendicular component is w ⊥ k =
[Equation]
k / 2 ˜ w ⊥ k .
Thus our optimal solution is w k = ˜
w
[Equation]
If the features x i are sparse, then we’d prefer to do a sparse update to w , saving computation time.
We can achieve this by letting
[Equation]
where γ k is a scalar and r k a vector.
Updating w k = ˜ w (cid:107) k
+ V j
[Equation]
+
[Equation]
k is equivalent to
[Equation]
Since we only update r k along the span of { x i : i ∈ S j } , its update is sparse.
F.3. R UNTIME
There are two major tasks in calculating the terms in (22).
The ﬁrst is to calculate x (cid:62) i ˜ w k for i ∈ S
j and k ∈ S
j ∪ C .
There are a total of n ( n + m ) inner products of D -dimensional vectors, costing
O ( n ( n + m ) D ) .
The other task is to ﬁnd the orthonormal basis V j of { x i : i ∈ S j } , which can be achieved using the Gram-Schmidt process in O ( n 2 D ) .
We’ll assume that { V j : j = 1 , ..., J } is computed only once as a pre-processing step when deﬁning the mini-batches.
It is exactly because calculating { V j : j = 1 , ..., J } is expensive that we have ﬁxed mini-batches that do not change during the optimization routine.
Adding the cost of calculating the x (cid:62) i ˜ w k inner products to the costing of optimizing (22) leads to the claim that solve the Implicit SGD update formula to (cid:15) accuracy in runtime O ( n ( n + m ) D + n 2 ( n + m ) log( (cid:15) − 1 ))
=
O ( n ( n + m )( D + n log( (cid:15) − 1 ))) .
F.3. I NITIALIZING THE I MPLICIT SGD OPTIMIZER
As was the case in Section F.1, it is important to initialize the optimization procedure at a point where the gradient is relatively small and can be computed without numerical issues.
These numerical issues arise when an exponent x (cid:62) i ( ˜ w k − ˜ w y i ) − ˜ u
i + b (cid:62) i ( a k − a y i )
(cid:29) 0 .
To ensure that this does not occur for our initial point, we can solve the following linear problem,
[Equation]
Note that if k = y
i then the constraint 0 ≥ x (cid:62) i ( ˜ w k − ˜ w y i )
− ˜ u
i + b (cid:62) i ( a k − a y i ) =
− ˜ u
i is automatically fulﬁlled since ˜ u i ≥ 0 .
Also observed that setting a k = − V (cid:62) j ˜ w k satisﬁes all of the constraints, and so
[Equation]
We can use the solution to (23) to gives us an upper bound on (22).
Consider the optimal value
A ( R ) j,C of the linear program in (23) with the value of the minimum being R .
Since A ( R ) j,C satisﬁes the constrain in (23) we have z ijC ( A ( R ) j,C ) ≤ Kηp − 1 j exp( ηp
[Equation]
1 j ) .
Since P ( z ) is a monotonically increasing function that is non-negative for z ≥
we also have (1 + P ( z ijC ( A ( R ) j,C )))
2 ≥ (1 +
P ( Kηp − 1 j exp( ηp − 1 j ))) 2 .
Turning to the norms, we can use the fact that (cid:107) a (cid:107) 2 ≤ (cid:107) a (cid:107) 1 for any vector a to bound (cid:88) k
[Equation]
Putting the bounds together we have that the optimal value of (22) is upper bounded by its value at the solution to (23)
[Equation]
which in turn is upper bounded by
[Equation]
This bound is guarantees that our initial iterate will be numerically stable
[Equation]
G L EARNING RATE PREDICTION AND LOSS
Here we present the results of using different learning rates for each algorithm applied to the Eurlex dataset.
In addition to the Implicit SGD, NCE, IS, OVE and U-max algorithms, we also provide results for NCE with n = 1 , m = 1 , denoted as NCE (1,1) .
NCE and NCE (1,1) have near identical performance.
[Figure]