Title: UNBIASED SCALABLE SOFTMAX OPTIMIZATION

Abstract: 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 efficiently computable but biased estimates of the softmax. In this paper we present the first 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.

1. INTRODUCTION:

Under the softmax model1 the probability that a random variable y takes on the label ∈ {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 )} N i=1 , the ridge-regularized maximum log-likelihood problem is given by

L(W ) = N i=1

x i w yi -log(

EQUATION

where W 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 K k=1 e x 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 difficulty. These include tree-structured methods (Bengio et al., 2003; Daume III et al., 2016; Grave et al., 2016) , sampling methods (Bengio & Senécal, 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ébisson & 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. All2 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 overflow. Raman et al. (2016) deal with this instability by occasionally calculating the normalizing sum for all datapoints at a cost of O(N KD). Although this achieves stability, its high cost nullifies the benefit 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 first 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 sufficiently 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 efficiently 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:

1. Provide a simple derivation of the softmax double-sum formulation and identify why vanilla SGD is unstable when applied to this formulation (Section 2). 2. Propose the U-max algorithm to stabilize the SGD updates and prove its convergence (Section 3.1). 3. Derive an efficient Implicit SGD implementation, analyze its runtime and bound its step size (Section 3.2). 4. 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).

2. CONVEX DOUBLE-SUM FORMULATION:



2.1. DERIVATION 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

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

L(W ) = N i=1 max ui∈R {-u i -e -ui (1 + k =yi e x i (w k -wy i ) ) + 1} - µ 2 W 2 2 = -min u≥0 {f (u, W )} + N, where f (u, W ) = N i=1 k =yi u i + e -ui K -1 + e x i (w k -wy i )-ui + µ 2 W 2 2 (5)

is our double-sum representation that we seek to minimize and the optimal solution for u i is

u * i (W ) = log(1 + k =yi e x i (w k -wy i ) ) ≥ 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.2. INSTABILITY OF VANILLA SGD:

The challenge in optimizing f using SGD is that it can have problematically large magnitude gradients. Observe that

EQUATION

where

β j = N nj +(N -nj )(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 + k =yi e x i (w k -wy i ) ) then e x i (w k -wy i )-ui ≤ 1 and the magnitude of the N (K -1) terms in the stochastic gradient are bounded by N (K -1) x i 2 . However if u i

x i (w k -w yi ), then e x i (w k -wy i )-ui 1 and the magnitude of the gradients can become extremely large.

Extremely large gradients lead to two major problems: (a) the gradients may computationally overflow floating-point precision and cause the algorithm to crash, (b) they result in the stochastic gradient having high variance, which leads to slow convergence3 . In Section 4 we show that these problems occur in practice and make vanilla SGD both an unreliable and inefficient method4 .

The sampled softmax optimizers in the literature (Bengio & Senécal, 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) x i 2 due to their approximations to u * i (W ) always being greater than x i (w k -w yi ). For example, in one-vs-each (Titsias, 2016), u * i (W ) is approximated by log(1 + e x i (w k -wy i ) ) > x i (w k -w yi ). 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 efficient 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.3. CHOICE OF DOUBLE-SUM FORMULATION:

The double-sum in ( 5) is different to that of Raman et al. (2016) . Their formulation can be derived by applying the convex conjugate substitution to (2) instead of (3). The resulting equations are

L(W ) = min ū 1 N N i=1 1 K-1 k =yi fik (ū, W ) + N where fik (ū, W ) = N ūi -x i w yi + e x i wy i -ūi + (K -1)e x i w k -ūi + µ 2 (β yi w yi 2 2 + β k w k 2 2 ) (8)

and the optimal solution for ūi is ū *

i (W * ) = log( K k=1 e x 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 i w yi = argmax k {x i w k } and so the ūi , x i w yi and e x i wy i -ū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 i w yi appearing as a negative exponent, and so if x i w yi is large then the magnitude of the stochastic gradients will be small. In Section 4 we present numerical results confirming that our double-sum formulation leads to faster convergence.

3. STABLE SGD METHODS:



3.1. U-MAX METHOD:

As explained in Section 2.2, vanilla SGD has large gradients when u i x i (w k -w yi ). This can only occur when u i is less than its optimum value for the current W , since u * i (W ) = log(1 + j =yi e x i (w k -wy i ) ) ≥ x i (w k -w yi ). A simple remedy is to set u i = log(1 + e x i (w k -wy i ) ) whenever u i

x i (w k -w yi ). Since log(1 + e x i (w k -wy i ) ) > x i (w k -w yi ) this guarantees that u i > x i (w k -w yi ) and so the gradients are bounded. It also brings u i closer5 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 modifications: (a) u i is set equal to log(1 + e x i (w k -wy i ) ) whenever u i ≤ log(1 + e x i (w k -wy i ) ) -δ for some threshold δ > 0, (b) u i is projected onto [0, B u ], and W onto {W : W 2 ≤ B W }, where B u and B W are set so that the optimal u * i ∈ [0, B u ] and the optimal W * satisfies W * 2 ≤ B W . See Appendix A for more details on how to set B u and B W .

EQUATION

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. Modification (a) ensures that δ ≥ x i (w k -w yi ) -u i (otherwise u i would be increased to log(1 + e x i (w k -wy i ) )) and so the magnitude of the U-max gradients are bounded above by N (K -1)e δ x i 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 /(4B 2 f ). A small δ yields small magnitude gradients, which makes convergence fast, but necessitates a small η t , which makes convergence slow.

3.2. IMPLICIT SGD:

Another method that solves the large gradient problem is Implicit SGD6 (Bertsekas, 2011; Toulis et al., 2016) . Implicit SGD uses the update 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, 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 it,kt (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 ∇f (θ (t+1) , ξ t ) 2 ≤ ∇f (θ (t) , ξ t ) 2 -m θ (t+1) -θ (t) 2

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 i (w k -w yi ) -u i , as shown in ( 7), for Implicit SGD the step size is asymptotically linear in x i (w k -w yi ) -u i . This effectively guarantees that Implicit SGD cannot suffer from computational overflow.

Proposition 2. Consider the Implicit SGD algorithm where in each iteration only one datapoint i and one class k = y i is sampled and there is no ridge regularization. The magnitude of its step size in

w is O(x i (w k -w yi ) -u i ).

Proof. The proof is provided in Appendix F.2.

The difficulty 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 accuracy in runtime O(n(n + m)(D + n log( -1 ))).

Proof. The proof is provided in Appendix F.3.

In Proposition 3 the log( -1 ) factor comes from applying a first 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 i w k inner products, and so each iteration of Implicit SGD may be significantly slower than that of vanilla SGD or U-max. However, in the special case of n = m = 1 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 = y i is sampled and there is no ridge regularization. Then the Implicit SGD iterate θ (t+1) can be computed to within accuracy with only two D-dimensional vector inner products and at most log 838 (3, 993)  5, 000  15, 709 (2, 919)  10, 000 (203, 882)  100, 000 (1, 186, 239)  Wiki10  4, 021 (30, 938)  10, 000 (101, 938 

2 ( -1 )+log 2 (|x i (w k -w yi )-u i |+2ηN x i

4. EXPERIMENTS:

Two sets of experiments were conducted to assess the performance of the proposed methods. The first 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.1. EXPERIMENTAL SETUP:

Data. We used the MNIST, Bibtex, Delicious, Eurlex, AmazonCat-13K, Wiki10, 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 first 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écal, 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 7 All of the datasets were downloaded from http://manikvarma.org/downloads/XC/ XMLRepository.html, except Wiki-small which was obtained from http://lshtc.iit. demokritos.gr/.

8 Raman et al. (2016) have runtime O(N KD) per epoch, which is equivalent to O(KD) per iteration. This is a factor of K slower than the methods we compare against. 9 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 = 1 and n = 100, m = 5 applied to the Eurlex dataset for different learning rates and there is very little difference between the two.

Table 2 : Tuned initial learning rates for each algorithm on each dataset. The learning rate in 10 0,±1,±2,±3 with the lowest log-loss after 50 epochs using only 10% of the data is displayed. Vanilla SGD applied to AmazonCat, Wiki10 and Wiki-small suffered from overflow with a learning rate of 10 -3 , but was stable with smaller learning rates (the largest learning rate for which it was stable is displayed).

DATASET OVE NCE IS:

Vanilla U-max Implicit MNIST 10 1 10 1 10 1 10 -2 10 1 10 -1 Bibtex 10 2 10 2 10 2 10 -2 10 -1 10 1 Delicious 10 1 10 3 10 3 10 -3 10 -2 10 -2 Eurlex 10 -1 10 2 10 2 10 -3 10 -1 10 1 AmazonCat 10 1 10 3 10 3 10 -5 10 -2 10 -3 Wiki10 10 -2 10 3 10 2 10 -4 10 -2 10 0 Wiki-small 10 3 10 3 10 3 10 -4 10 -3 10 -3 0 10 20 30 40 50 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 overflow.

4.2. RESULTS:

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 first 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 overflow. 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 all10 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.

5. CONCLUSION:

In this paper we have presented the U-max and Implicit SGD algorithms for optimizing the softmax likelihood. These are the first 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 efficiently 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 final 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 fit, might be particularly fruitful.

A PROOF OF VARIABLE BOUNDS AND STRONG CONVEXITY:

We first 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

W * 2 2 ≤ B 2 W where B 2 W = 2 µ N log(K).

Proof.

-

N log(K) = L(0) ≤ L(W * ) ≤ - µ 2 W * 2 2

Rearranging gives the desired result.

Lemma 2. The optimal value of u i is bounded as

u * i ≤ B u where B u = log(1 + (K -1)e 2BxBw ) and B x = max i { x i 2 } Proof. u * i = log(1 + k =yi e x i (w k -wy i ) ) ≤ log(1 + k =yi e xi 2( w k 2+ wy i 2 ) ) ≤ log(1 + k =yi e 2BxBw ) = log(1 + (K -1)e 2BxBw ) Lemma 3. If W 2 2 ≤ 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

f (u, W ) = N i=1 u i + e -ui + k =yi e x i (w k -wy i )-ui + µ 2 W 2 2 = N i=1 a i θ + e -ui + k =yi e b ik θ + µ 2 W 2 2 .

where θ = (u , w 1 , ..., w k ) ∈ R N +KD with a i and b ik being appropriately defined. The Hessian of f is

∇ 2 f (θ) = N i=1 e -ui e i e i + k =yi e b ik θ b ik b ik + µ • diag{0 N , 1 KD }

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

∇ 2 f (θ) I • min{ min 0≤u≤Bu {e -ui }, µ} = I • min{exp(-B u ), µ} 0. Lemma 4. If W 2 2 ≤ 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

B f = N max{1, e Bu -1} + 2(N e Bu B x + µ max k {β k }B W ).

Proof. By Jensen's inequality max

W 2 2 ≤B 2 W ,0≤u≤Bu ∇f (u, W ) 2 = max W 2 2 ≤B 2 W ,0≤u≤Bu ∇E ik f ik (u, W ) 2 ≤ max W 2 2 ≤B 2 W ,0≤u≤Bu E ik ∇f ik (u, W ) 2 ≤ max W 2 2 ≤B 2 W ,0≤u≤Bu max ik ∇f ik (u, W ) 2 .

Using the results from Lemmas 1 and 2 and the definition of f ik from (6),

∇ ui f ik (u, W ) 2 = N 1 -e -ui -(K -1)e x i (w k -wy i )-ui ) 2 = N |1 -e -ui (1 + (K -1)e x i (w k -wy i ) )| ≤ N max{1, (1 + (K -1)e xi 2( w k 2+ wy i 2 ) ) -1}

≤ N max{1, e Bu -1} and for j indexing either the sampled class k = y i or the true label y i ,

∇ wj f ik (u, W ) 2 = ± N (K -1)e x i (w k -wy i )-ui x i + µβ j w j 2 ≤ N (K -1)e xi 2( w k 2+ wy i 2) x i 2 + µβ j w j 2 ≤ N e Bu B x + µ max k {β k }B W . Letting B f = N max{1, e Bu -1} + 2(N e Bu B x + µ max k {β k }B W ) we have ∇f ik (u, W ) 2 ≤ ∇ ui f ik (u, W ) 2 + ∇ w k f ik (u, W ) 2 + ∇ wy i f ik (u, W ) 2 = B f . In conclusion: max W 2 2 ≤B 2 W ,0≤u≤Bu ∇f (u, W ) 2 ≤ max W 2 2 ≤B 2 W ,ui≤Bu, max ik ∇f ik (u, W ) 2 ≤ B f .

B STOCHASTIC COMPOSITION OPTIMIZATION:

We can write the equation for L(W ) from ( 3) as (where we have set µ = 0 for notational simplicity),

L(W ) = - N i=1 log(1 + k =yi e x i (w k -wy i )) = 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(1 + e i v) [g k (W )] i = Ke x i (w k -wy i ) if k = y i 0 otherwise .

Here

e i v = v i ∈ R is a variable that is explicitly kept track of with v i ≈ E k [g k (W )] i =

k =yi e x i (w k -wy i ) (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 = y i in stochastic composition optimization then the updates are of the form (Wang et al., 2016) 

w yi = w yi + η t N K e x i (z k -zy i ) 1 + v i x i w k = w k -η t N K e x i (z k -zy i ) 1 + v i x i ,

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 i z k 1+vi 1 where ideally we should have 0 ≤ e x i z k 1+vi ≤ 1.

C U-MAX PSEUDOCODE:

Algorithm 1: U-max

Input : Data D = {(y i , x i ) : y i ∈ {1, . . . , K}, x i ∈ R d } N i=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 W 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 3 w k ← 0 4 end 5 for i = 1 to N do 6 u i ← log(K) 7 end 8 Run SGD 9 for t = 1 to T do 10 Sample indices 11 i ∼ unif ({1, ..., N }) 12 k ∼ unif ({1, ..., K} -{y i }) 13 Increase u i 14 if u i < log(1 + e x i (w k -wy i ) ) -δ then 15 u i ← log(1 + e x i (w k -wy i ) ) 16 SGD step 17 w k ← w k -η t [N (K -1)e x i (w k -wy i )-ui x i + µβ k w k ] 18 w yi ← w yi -η t [-N (K -1)e x i (w k -wy i )-ui x i + µβ yi w yi ] u i ← u i -η t [N (1 -e -ui -(K -1)e x i (w k -wy i )-ui )] 20 Projection 21 w k ← w k • min{1, B W / w k 2 } 22 w yi ← w yi • min{1, B W / w yi 2 } 23 u i ← max{0, min{B u , u i }} 24 end

D PROOF 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 i (w k -wy i ) )-δ then setting u i = log(1+e x i (w k -wy i ) ) decreases f (u, W ) by at least δ 2 /2.

Proof. As in Lemma 3, let θ = (u , w 1 , ..., w k ) ∈ R N +KD . Then setting u i = log(1 + e x i (w k -wy i ) ) is equivalent to setting θ = θ + ∆e i where e i is the i th canonical basis vector and ∆ = log(1 + e x i (w k -wy i ) ) -u i ≥ δ. By a second order Taylor series expansion

EQUATION

for some λ ∈ [0, 1]. Since the optimal value of u i for a given value of W is u * i (W ) = log(1 + k =yi e x i (w k -wy i ) ) ≥ log(1+e x i (w k -wy i ) ), we must have ∇f (θ+∆e i ) e i ≤ 0. From Lemma 3 we also know that

e i ∇ 2 f (θ + λ∆e i )e i = exp(-(u i + λ∆)) + k =yi e x i (w k -wy i )-(ui+λ∆) = exp(-λ∆)e -ui (1 + k =yi e x i (w k -wy i ) ) = exp(-λ∆) exp(-(log(1 + e x i (w k -wy i ) ) -∆))(1 + k =yi e x i (w k -wy i ) ) ≥ exp(∆ -λ∆) ≥ exp(∆ -∆) = 1.

Putting in bounds for the gradient and Hessian terms in (10),

f (θ) -f (θ + ∆e i ) ≥ ∆ 2 2 ≥ δ 2 2 .

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 Θ = {θ :

W 2 2 ≤ 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(1 + e x i (w k -wy i ) ) if u i ≤ log(1 + e x i (w k -wy i ) ) -δ.

If indices i, k are sampled for the stochastic gradient and u i ≤ log(1 + e x i (w k -wy i ) ) -δ, then the value of f at the t + 1 st iterate is bounded as

f (θ (t+1) ) = f (π i (θ (t) ) -η t ∇f ik (π i (θ (t) ))) ≤ f (π i (θ (t) )) + max θ∈Θ η t ∇f ik (π i (θ)) 2 max θ∈Θ ∇f (θ) 2 ≤ f (π i (θ (t) )) + η t B 2 f ≤ f (θ (t) ) -δ 2 /2 + η t B 2 f ≤ f (θ (t) -η t ∇f ik (θ (t) )) -δ 2 /2 + 2η t B 2 f ≤ f (θ (t) -η t ∇f ik (θ (t) )), since η t ≤ δ 2 /(4B 2 f ) by assumption. Alternatively if u i ≥ log(1 + e x i (w k -wy i ) ) -δ then f (θ (t+1) ) = f (π i (θ (t) ) -η t ∇f ik (π i (θ (t) ))) = f (θ (t) -η t ∇f ik (θ (t) )).

Either way f (θ (t+1) ) ≤ f (θ (t) -η t ∇f ik (θ (t) )). Taking expectations with respect to i, k,

E ik [f (θ (t+1) )] ≤ E ik [f (θ (t) -η t ∇f ik (θ (t) ))].

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

E ik [f (P (θ (t+1) ))] ≤ E ik [f (θ (t) -η t ∇f ik (θ (t) ))],

which shows that the rate of convergence in expectation of U-max is at least as fast as that of standard SGD.

E PROOF OF GENERAL IMPLICIT SGD GRADIENT BOUND:

Proof of Theorem 2. Let f (θ, ξ) be m-strongly convex for all ξ. The vanilla SGD step size is η t ∇f (θ (t) , ξ t ) 2 where η t is the learning rate for the t th iteration. The Implicit SGD step size is η t ∇f (θ (t+1) , ξ t ) 2 where θ (t+1) satisfies θ (t+1) = θ (t) -η t ∇f (θ (t+1) , ξ t ). Rearranging, ∇f (θ (t+1) , ξ t ) = (θ (t) -θ (t+1) )/η t and so it must be the case that ∇f (θ (t+1) , ξ t ) (θ (t) -θ (t+1) ) = ∇f (θ (t+1) , ξ t ) 2 θ (t) -θ (t+1) 2 . Our desired result follows:

∇f (θ (t) , ξ t ) 2 ≥ ∇f (θ (t) ) (θ (t) -θ (t+1) ) θ (t) -θ (t+1) 2 ≥ ∇f (θ (t+1) ) (θ (t) -θ (t+1) ) + m θ (t) -θ (t+1) 2 2 θ (t) -θ (t+1) 2 = ∇f (θ (t+1) ) 2 θ (t) -θ (t+1) 2 + m θ (t) -θ (t+1) 2 2 θ (t) -θ (t+1) 2 = ∇f (θ (t+1) ) 2 + m θ (t) -θ (t+1) 2

where the first inequality is by Cauchy-Schwarz and the second inequality by strong convexity.

F UPDATE EQUATIONS FOR IMPLICIT SGD:

In this section we will derive the updates for Implicit SGD. We will first 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.1 SINGLE DATAPOINT, SINGLE CLASS, NO REGULARIZER Equation ( 6) for the stochastic gradient for a single datapoint and single class with µ = 0 is f ik (u, W ) = N (u i + e -ui + (K -1)e x i (w k -wy i )-ui ).

The Implicit SGD update corresponds to finding the variables optimizing

min u,W 2ηf ik (u, W ) + u -ũ 2 2 + W -W 2 2 ,

where η is the learning rate and the tilde refers to the value of the old iterate (Toulis et al., 2016, Eq. 6) . Since f ik is only a function of u i , w k , w yi the optimization reduces to

min ui,w k ,wy i 2ηf ik (u i , w k , w yi ) + (u i -ũi ) 2 + w yi -wyi 2 2 + w k -wk 2 2 = min ui,w k ,wy i 2ηN (u i + e -ui + (K -1)e x i (w k -wy i )-ui ) + (u i -ũi ) 2 + w yi -wyi 2 2 + w k -wk 2 2 .

The optimal value of w k , w yi must deviate from the old value wk , wyi in the direction of x i . Furthermore we can observe that the deviation of w k must be exactly opposite that of w yi , that is:

w yi = wyi + a x i 2 x i 2 2 w k = wk -a x i 2 x i 2 2 (11)

for some a ≥ 0. The optimization problem reduces to min ui,a≥0 2ηN (u i + e -ui + (K -1)e x i ( wk

EQUATION

We'll approach this optimization problem by first 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 (11) will give us our updated value of W .

Solving for a:

We solve for a by setting its derivative equal to zero in ( 12)

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 ( 12), we now only need minimize over u i :

min ui 2ηN u i + 2ηN e -ui + 2ηN (K -1)e x i ( wk -wy i ) e -a(ui)-ui + (u i -ũi ) 2 + a(u i ) 2 1 2 x i 2 2 = min ui 2ηN u i + 2ηN e -ui + a(u i ) x i -2 2 + (u i -ũi ) 2 + a(u i ) 2 1 2 x i 2 2 (15)

where we used the fact that e -P (z) = P (z)/z. The derivative with respect to u i in ( 15) is

∂ ui 2ηN u i + 2ηN e -ui + a(u i ) x i -2 2 + (u i -ũi ) 2 + a(u i ) 2 1 2 x i 2 2 = 2ηN -2ηN e -ui + ∂ ui a(u i ) x i -2 2 + 2(u i -ũi ) + 2a(u i )∂ ui a(u i ) 1 2 x i 2 2 = 2ηN -2ηN e -ui - a(u i ) 1 + a(u i ) x i -2 2 + 2(u i -ũi ) - a(u i ) 2 (1 + a(u i )) x i 2 2 (16)

where to calculate ∂ ui a(u i ) we used the fact that ∂ z P (z) = P (z) z(1+P (z)) and so

∂ ui a(u i ) = - a(u i ) e x i ( wk -wy i )-ui+log(2ηN (K-1) xi 2 2 ) (1 + a(u i )) e x i ( wk -wy i )-ui+log(2ηN (K-1) xi 2 2 ) = - a(u i ) 1 + a(u i ) .

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 ( 16) at u i = ũi . If the derivative is negative then the optimal u i is lower bounded by ũi . An upper bound is provided by

u i = argmin ui 2ηN (u i + e -ui + (K -1)e x i ( wk -wy i ) e -a(ui)-ui ) + (u i -ũi ) 2 + a(u i ) 2 2 x i 2 2 ≤ argmin ui 2ηN (u i + e -ui + (K -1)e x i ( wk -wy i ) e -ui ) + (u i -ũi ) 2

≤ argmin ui 2ηN (u i + e -ui + (K -1)e x i ( wk -wy i ) e -ui ) = log(1 + (K -1)e x i ( wk -wy i ) ).

In the first 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 ũi . Thus if the derivative in ( 16) is negative at u i = ũi then ũi ≤ u i ≤ log(1 + (K -1)e x i ( wk -wy i ) ). If (K -1)e x i ( wk -wy i ) ≤ 1 then the size of the interval must be less than log(2), since ũi ≥ 0. Otherwise the gap must be at most log(2(K -1)e x i ( wk -wy i ) )-ũi = log(2(K -1))+x i ( wk -wyi )-ũi . Either way, the gap is upper bounded by log(2(K -1)) + |x i ( wk -wyi ) -ũi |.

Now let us consider if the derivative in ( 16) is positive at u i = ũi . Then u i is upper bounded by ũi . Denoting a as the optimal value of a, we can lower bound u i using ( 12)

EQUATION

where the first inequality comes dropping the (u i -ũi )2 term due to the assumption that u i < ũi . Recall ( 13), a e a = 2ηN (K -1) x i 2 2 e x i ( wk -wy i )-ui . The solution for a 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 . Substituting the bound for u i ,

a ≤ min{a : ae a = 2ηN (K -1) x i 2 2 e x i ( wk -wy i )-(log(K-1)+x i ( wk -wy i )-a) } = min{a : a = 2ηN x i 2 2 } = 2ηN x i 2 2 .

(18) Substituting this bound for a into (17) yields

u i ≥ log(K -1) + x i ( wk -wyi ) -2ηN x i 2 2 .

Thus if the derivative in ( 16) is postive at u i = ũi then log(K -1) + x i ( wk -wyi ) -2ηN x i 2 2 ≤ u i ≤ ũi . The gap between the upper and lower bound is ũi -x i ( wk -wyi )+2ηN x i 2 2 -log(K-1). In summary, for both cases of the sign of the derivative in ( 16) at u i = ũ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 |ũ i -x i ( wk -wyi )| + 2ηN x i 2 2 + log(K -1). This allows us to perform the bisection method where for > 0 level accuracy we require only log 2 ( -1 )+log 2 (|ũ i -x i ( wk -wyi )|+2ηN x i 2 2 + log(K -1)) function evaluations.

F.2 BOUND 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 i ( wk -wyi ) -ũi ). We will do so by considering the two cases u i ≥ ũi and u i < ũi separately, where u i denotes the optimal value of u i in the Implicit SGD update and ũi is its value at the previous iterate.

Case: u i ≥ ũi Let a denote the optimal value of a in the Implicit SGD update. From ( 14) a = a(u i ) = P (e x i ( wk -wy i )-u i +log(2ηN (K-1) xi 2 2 ) ) = P (e x i ( wk -wy i )-ũi+log(2ηN (K-1) xi 2 2 ) ). Now using the fact that P (z) = O(log(z)), a = O(x i ( wk -wyi ) -ũi + log(2ηN (K -1) x i Case: u i < ũi If u i < ũi then we can lower bound a from (18) as

a ≤ 2ηN x i 2 2 .

Combining cases:

Putting together the two cases,

a = O(max{x i ( wk -wyi ) -ũi , 2ηN x i 2 2 }) = O(x i ( wk -wyi ) -ũi ).

The actual step size in w is ±a xi

2 xi 2 2

. Since a is O(x i ( wk -wyi ) -ũi ), the step size magnitude is also O(x i ( wk -wyi ) -ũi ).

F.3 MULTIPLE 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 significant differences. Firstly, we will require some pre-computation to find 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 simplified optimization problem is large, we'll require first order or quasi-Newton methods to find the optimal solution.

F.3.1 DEFINING THE MINI-BATCH:

The first step is to define 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 : = 1, ..., n} for j = 1, ..., N/n , S J = {J : = 1, ..., N mod n}, S i ∩ S j = ∅ and ∪ J j=1 S j = {1, ..., N }.

Next we define the set of classes C j which can be sampled for the j th mini-batch. The set C j is defined 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 = k ∀ ∈ {1, ..., m} -{i}, k i = y ∀ ∈ S j }.

Now we can write down our objective from (5) in terms of an expectation of functions corresponding to our mini-batches:

f (u, W ) = E[f j,C (u, W )]

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 i∈Sj   u i + e -ui + k∈Sj -{i} e x i (w k -wy i )-ui + K -n m k∈C e x i (w k -wy i )-ui   + µ 2 k∈C∪Sj β k w k 2 2 .

The value of the regularizing constant

β k is such that E[I[k ∈ C ∪ S j ]β k ] = 1, which requires that β -1 k = 1 - 1 J J j=1 I[k = S j ](1 - m K -|S j |

).

F.3.2 SIMPLIFYING THE IMPLICIT SGD UPDATE EQUATION:

The Implicit SGD update corresponds to solving

min u,W 2ηf j,C (u, W ) + u -ũ 2 2 + W -W 2 2 ,

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 Sj = {u i : i ∈ S j } and W j,C = {w k : k ∈ S j ∪ C} the optimization reduces to

min u S j ,W j,C 2ηf j,C (u Sj , W j,C ) + u Sj -ũSj 2 2 + W j,C -Wj,C 2 2 .

The next step is to analytically minimize the u Sj terms. The optimization problem in (21) decomposes into a sum of separate optimization problems in u i for i ∈ S j ,

min ui 2ηp -1 j (u i + e -ui d i ) + (u i -ũi ) 2

where

d i (W j,C ) = 1 + k∈Sj -{i} e x i (w k -wy i ) + K -n m k∈C e x i (w k -wy i ) .

Setting the derivative of u i equal to zero yields the solution

u i (W j,C ) = ũi -ηp -1 j + P (ηp -1 j d i (W j,C ) exp(ηp -1 j -ũi ))

where P is the principle branch of the Lambert W function. Substituting this solution into our optimization problem and simplifying yields min

EQUATION

where we have used the identity e -P (z) = P (z)/z. We can decompose (19) into two parts by splitting W j,C = W 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 (19) only depends on W j,C , the two resulting sub-problems are min

EQUATION

Let us focus on the perpendicular component first. Simple calculus yields the optimal value w

⊥ k = 1 1+µβ 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 basis11 

V j = (v j1 , ..., v jn ) ∈ R D×n with x i = V j b i for some b i ∈ R n .

With this basis we can write

EQUATION

where

A j,C = {a k : k ∈ S j ∪ C} ∈ R (n+m)×n and z ijC (A j,C ) = ηp -1 j exp(ηp -1 j ) exp(-ũ i ) + k∈Sj -{i} e x i ( wk -wy i )-ũi e b i (a k -ay i ) + K -n m k∈C e x i ( wk -wy i )-ũi e b i (a k -ay i )

.

The e b i (a k -ay i ) factors come from

x i w k = x i ( w k + a k V j ) = x i wk + (V j b i ) V j a k = x i wk + b i V j V j a k = x i wk + b i a k ,

since V j is an orthonormal basis.

F.3.3 OPTIMIZING THE IMPLICIT SGD UPDATE EQUATION:

To optimize ( 21) we need to be able to take the derivative:

∇ a   i∈Sj (1 + P (z ijC (A j,C ))) 2 + k∈Sj ∪C (1 + µβ k 2 ) a k 2 2 + µβ k w k V j a k   = i∈Sj 2(1 + P (z ijC (A j,C )))∂ z ijC (A j,C ) P (z ijC (A j,C ))∇ a z ijC (A j,C ) + (2 + µβ )a + µβ w V j = i∈Sj 2(1 + P (z ijC (A j,C ))) P (z ijC (A j,C )) z ijC (A j,C )(1 + P (z ijC (A j,C ))) ∇ a z ijC (A j,C ) + (2 + µβ )a + µβ w V j = i∈Sj 2 P (z ijC (A j,C )) z ijC (A j,C ) ∇ a z ijC (A j,C ) + (2 + µβ )a + µβ w V j = i∈Sj 2e -P (z ijC (A j,C )) ∇ a z ijC (A j,C ) + (2 + µβ )a + µβ w V j

where we used that ∂ z P (z) = P (z) z(1+P (z)) and e -P (z) = P (z)/z. To complete the calculation of the derivate we need,

∇ a z ijC (A j,C ) = ∇ a ηp -1 j exp(ηp -1 j ) exp(-ũ i ) + k∈Sj -{i} e x i ( w -wy i )-ũi e b i (a -ay i ) + K -n m k∈C e x i ( w -wy i )-ũi e b i (a -ay i ) = ηp -1 j exp(ηp -1 j )b i • I[ ∈ S j -{i}]e x i ( w -wy i )-ũi e b i (a -ay i ) + I[ ∈ C] K -n m e x i ( w -wy i )-ũi e b i (a -ay i ) -I[ = y i ] k∈Sj -{i} e x i ( w -wy i )-ũi e b i (a -ay i ) + K -n m k∈C e x i ( w -wy i )-ũi e b i (a -ay i ) .

In order to calculate the full derivate with respect to A j,C we need to calculate b 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 find the optimum of ( 21) we can use any optimization procedure that only uses gradients. Since ( 21) is strongly convex, standard first order methods can solve to accuracy in O(log( -1 )) iterations (Boyd & Vandenberghe, 2004, Sec. 9.3) . Thus once we can calculate all of the terms in (21), we can solve it to accuracy in runtime O(n 2 (n + m) log( -1 )).

Once we have solved for A j,C , we can reconstruct the optimal solution for the parallel component of w k as w k = w k + V j a k . Recall that the solution to the perpendicular component is w

⊥ k = 1 1+µβ k /2 w⊥ k . Thus our optimal solution is w k = w k + V j a k + 1 1+µβ k /2 w⊥ k .

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

w k = γ k • r k

where γ k is a scalar and r k a vector. Updating

w k = w k + V j a k + 1 1+µβ k /2 w⊥ k is equivalent to γ k = γk • 1 1 + µβ k /2 r k = rk + µβ k /2 • r k + γ-1 k (1 + µβ k /2) • V j a k .

Since we only update r k along the span of {x i : i ∈ S j }, its update is sparse.

F.3.4 RUNTIME:

There are two major tasks in calculating the terms in ( 21). The first is to calculate x i wk 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 find 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 defining the mini-batches. It is exactly because calculating {V j : j = 1, ..., J} is expensive that we have fixed mini-batches that do not change during the optimization routine.

Adding the cost of calculating the x i wk inner products to the costing of optimizing (21) leads to the claim that solve the Implicit SGD update formula to accuracy in runtime O(n(n + m)D + n 2 (n + m) log( -1 )) = O(n(n + m)(D + n log( -1 ))).

F.3.5 INITIALIZING THE IMPLICIT 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 i ( wk -wyi ) -ũi + b i (a k -a yi ) 0. To ensure that this does not occur for our initial point, we can solve the following linear problem,13 R = min

EQUATION

Note that if k = y i then the constraint 0 ≥ x i ( wk -wyi )-ũi +b i (a k -a yi ) = -ũ i is automatically fulfilled since ũi ≥ 0. Also observed that setting a k = -V j wk satisfies all of the constraints, and so

R ≤ k∈C∪Sj V j wk 1 ≤ (n + m) max k∈C∪Sj V j wk 1 .

We can use the solution to (22) to gives us an upper bound on (21). Consider the optimal value A j exp(ηp -1 j ). Since P (z) is a monotonically increasing function that is non-negative for z ≥ 0 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 a 2 ≤ a 1 for any vector a to bound

k∈Sj ∪C (1 + µβ k 2 ) a k 2 2 + µβ k w k V j a k ≤ k∈Sj ∪C (1 + µβ k 2 ) a k 2 1 + µβ k w k V j 1 a k 1 ≤ 1 + µ • max k∈Sj ∪C {β k }/2 k∈Sj ∪C a k 2 1 + µ max k∈Sj ∪C {β k w k V j 1 } k∈Sj ∪C a k 1 ≤ 1 + µ • max k∈Sj ∪C {β k }/2 R 2 + µ max k∈Sj ∪C {β k } max k∈Sj ∪C { w k V j 1 }R ≤ 1 + µ • max k∈Sj ∪C {β k }/2 (n + m) max k∈C∪Sj V j wk 1 2 + µ max k∈Sj ∪C {β k } max k∈Sj ∪C { w k V j 1 } (n + m) max k∈C∪Sj V j wk 1 ≤ (1 + µ • max k∈Sj ∪C {β k })(n + m) 2 max k∈C∪Sj V j wk 2 1 ≤ (1 + µ • max k∈Sj ∪C {β k })(n + m) 2 max k∈C∪Sj wk 2 1 .

Putting the bounds together we have that the optimal value of ( 21) is upper bounded by its value at the solution to ( 22), which in turn is upper bounded by

n(1 + P (Kηp -1 j exp(ηp -1 j ))) 2 + (1 + µ • max k∈Sj ∪C {β k })(n + m) 2 max k∈C∪Sj wk 2 1 .

This bound is guarantees that our initial iterate will be numerically stable.

G LEARNING 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. 

:

Also known as the multinomial logit model.

The method ofKrishnapuram et al. (2005) does converge to the optimal MLE, but has O(N D) runtime per iteration which is not feasible for large N, D.

The convergence rate of SGD is inversely proportional to the second moment of its gradients(Lacoste- Julien et al., 2012).

4 The same problems arise if we approach optimizing (3) via stochastic composition optimization(Wang et al., 2016). As is shown in Appendix B, stochastic composition optimization yields near-identical expressions for the stochastic gradients in (7) and has the same stability issues.

Since ui < x i (w k -wy i ) < log(1 + e x i (w k -wy i ) ) < log(1 + j =y i e x i (w k -wy i ) ) = u * i (W ).

Also known to as an "incremental proximal algorithm"(Bertsekas, 2011).

+log(K -1)) bisection method function evaluations.Proof. The proof is provided in Appendix F.1For 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 overflow issues and its updates can be computed in roughly the same time as vanilla SGD.

The learning rates η = 10 3,4 are not displayed in the Figure2for visualization purposes. It had similar behavior as η = 10 2 .

)) = O(x i ( wk -wyi ) -ũi )

We have assumed here that dim(span({xi : i ∈ Sj})) = n, which will be most often the case. If the dimension of the span is lower than n then let Vj be of dimension D × dim(span({xi : i ∈ Sj})).

Note that we have used wk instead of w k in writing the parallel component optimization problem. This does not make a difference as wk always appears as an inner product with a vector in the span of {xi : i ∈ Sj}.

Instead bounding the constraints on the right with 0, we could also have used any small positive number, like 5.

Figures and Tables:

Figure1: The x-axis is the number of epochs and the y-axis is the log-loss from (2) calculated at the current value of W .

Figure2: Log-loss of U-max on Eurlex for different learning rates with our proposed double-sum formulation and that ofRaman et al. (2016).

of the linear program in (22) with the value of the minimum being R. Since A (R) j,C satisfies the constrain in (22) we have z ijC (A (R) j,C ) ≤ Kηp -1

Figure 3: Log-loss on Eurlex different learning rates.

Datasets with a summary of their properties. Where the number of classes, dimension or number of examples has been altered, the original value is displayed in brackets.

Table content: <table><tr><td>DATASET</td><td>CLASSES</td><td>DIMENSION</td><td>EXAMPLES</td></tr><tr><td>MNIST</td><td>10</td><td>780</td><td>60,000</td></tr><tr><td>Bibtex</td><td>147 (159)</td><td>1,836</td><td>4,880</td></tr><tr><td>Delicious</td><td>350 (983)</td><td>500</td><td>12,920</td></tr><tr><td>Eurlex</td><td/><td/><td/></tr></table>

Relative log-loss. The values for each dataset are normalized by dividing by the corresponding Implicit SGD log-loss. The lowest log-loss for each dataset is in bold.

Table content: <table><tr><td colspan="4">DATASET OVE NCE</td><td>IS</td><td colspan="3">VANILLA-SGD U-MAX. IMPLICIT SGD</td></tr><tr><td>MNIST</td><td>5.73</td><td colspan="2">6.05</td><td>5.74</td><td>1.43</td><td>1.56</td><td>1.00</td></tr><tr><td>Bibtex</td><td colspan="4">29.03 28.52 32.71</td><td>15.18</td><td>9.77</td><td>1.00</td></tr><tr><td>Delicious</td><td>1.90</td><td colspan="2">1.91</td><td>1.89</td><td>1.25</td><td>1.11</td><td>1.00</td></tr><tr><td>Eurlex</td><td>6.47</td><td colspan="2">6.39</td><td>6.38</td><td>3.59</td><td>2.11</td><td>1.00</td></tr><tr><td colspan="2">AmazonCat 2.12</td><td colspan="2">2.15</td><td>2.12</td><td>1.47</td><td>0.98</td><td>1.00</td></tr><tr><td>Wiki10</td><td>7.04</td><td colspan="2">7.13</td><td>6.97</td><td>5.99</td><td>2.39</td><td>1.00</td></tr><tr><td>Wiki-small</td><td>1.02</td><td colspan="2">1.36</td><td>1.35</td><td>1.15</td><td>1.03</td><td>1.00</td></tr><tr><td>Average</td><td>7.62</td><td colspan="2">7.64</td><td>8.17</td><td>4.29</td><td>2.71</td><td>1.00</td></tr><tr><td/><td/><td>10 4</td><td colspan="3">Double-sum formulations</td><td/></tr><tr><td/><td>Log-likelihood</td><td>10 1 10 2 10 3</td><td/><td/><td>Learning rates</td><td/></tr><tr><td/><td/><td>0</td><td/><td>20 Epochs</td><td>40</td><td/></tr></table>

Peer review (iclr2018):
 The paper develops an interesting approach for solving multi-class classification with softmax loss. The key idea is to reformulate the problem as a convex minimization of a "double-sum" structure via a simple conjugation trick. SGD is applied to the reformulation: in each step samples a subset of the training samples and labels, which appear both in the double sum. The main contributions of this paper are: "U-max" idea (for numerical stability reasons) and an ""proposing an "implicit SGD" idea. Unlike the first review, I see what the term "exact" in the title is supposed to mean. I believe this was explained in the paper. I agree with the second reviewer that the approach is interesting. However, I also agree with the criticism (double sum formulations exist in the literature; comments about experiments); and will not repeat it here. I will stress though that the statement about Newton in the paper is not justified. Newton method does not converge globally with linear rate. Cubic regularisation is needed for global convergence. Local rate is quadratic. I believe the paper could warrant acceptance if all criticism raised by reviewer 2 is addressed. I apologise for short and late review: I got access to the paper only after the original review deadline.