In this paper, we use dynamical system to analyze the nonlinear weight dynamics of two-layered bias-free networks in the form of g(x;w) = ∑K j=1 σ(w T j x), where σ(·) is ReLU nonlinearity. We assume that the input x follow Gaussian distribution. The network is trained using gradient descent to mimic the output of a teacher network of the same size with fixed parameters w∗ using l2 loss. We first show that when K = 1, the nonlinear dynamics can be written in close form, and converges to w∗ with at least (1 − )/2 probability, if random weight initializations of proper standard derivation (∼ 1/ √ d) is used, verifying empirical practice [Glorot & Bengio (2010); He et al. (2015); LeCun et al. (2012)]. For networks with many ReLU nodes (K ≥ 2), we apply our close form dynamics and prove that when the teacher parameters {w∗ j}j=1 forms orthonormal bases, (1) a symmetric weight initialization yields a convergence to a saddle point and (2) a certain symmetry-breaking weight initialization yields global convergence to w∗ without local minima. To our knowledge, this is the first proof that shows global convergence in nonlinear neural network without unrealistic assumptions on the independence of ReLU activations. In addition, we also give a concise gradient update formulation for a multilayer ReLU network when it follows a teacher of the same size with l2 loss. Simulations verify our theoretical analysis.
∑K j=1 σ(w ᵀ j x),
where σ(·) is ReLU nonlinearity. We assume that the input x follow Gaussian distribution. The network is trained using gradient descent to mimic the output of a teacher network of the same size with fixed parameters w∗ using l2 loss. We first show that when K = 1, the nonlinear dynamics can be written in close form, and converges to w∗ with at least (1 − )/2 probability, if random weight initializations of proper standard derivation (∼ 1/ √ d) is used, verifying empirical practice [Glorot & Bengio (2010); He et al. (2015); LeCun et al. (2012)]. For networks with many ReLU nodes (K ≥ 2), we apply our close form dynamics and prove that when the teacher parameters {w∗j}Kj=1 forms orthonormal bases, (1) a symmetric weight initialization yields a convergence to a saddle point and (2) a certain symmetry-breaking weight initialization yields global convergence to w∗ without local minima. To our knowledge, this is the first proof that shows global convergence in nonlinear neural network without unrealistic assumptions on the independence of ReLU activations. In addition, we also give a concise gradient update formulation for a multilayer ReLU network when it follows a teacher of the same size with l2 loss. Simulations verify our theoretical analysis.
1 INTRODUCTION
Deep learning has made substantial progress in many applications, including Computer Vision [He et al. (2016); Simonyan & Zisserman (2015); Szegedy et al. (2015); Krizhevsky et al. (2012)], Natural Language Processing [Sutskever et al. (2014)] and Speech Recognition [Hinton et al. (2012)]. However, till now, how and why it works remains elusive due to a lack of theoretical understanding. First, how simple approaches like gradient descent can solve a very complicated non-convex optimization effectively. Second, how the deep models, especially deep convolutional models, achieve generalization power despite massive parameters.
In this paper, we focus on the first problem and use dynamical system to analyze the nonlinear gradient descent dynamics of certain two-layered nonlinear network in the following form:
g(x;w) = K∑ j=1 σ(wᵀj x) (1)
where σ(x) = max(x, 0) is the ReLU nonlinearity. We consider the following setting: a student network learns the parameters that minimize the l2 distance between its prediction and the supervision provided by the teacher network of the same size with a fixed set of parameters w∗. We assume all inputs x to follow Gaussian distribution and thus the network is bias-free. Eqn. 1 is highly nonconvex and could contain exponential number of symmetrically equivalent solutions.
To analyze this, we first derive novel and concise gradient update rules for multilayer ReLU networks (See Lemma 2.1) in the teacher-student setting under l2 loss. Then for K = 1, we prove that the nonlinear gradient dynamics of Eqn. 1 has a close form and converges to w∗ with at least (1 −
)/2 probability, if initialized randomly with standard derivation on the order of 1/ √ d, verifying commonly used initialization techniques [Glorot & Bengio (2010); He et al. (2015); LeCun et al. (2012)],. When K ≥ 2, we prove that when the teacher parameters {wj}Kj=1 form orthonormal bases, (1) a symmetric initialization of a student network gets stuck at a saddle point and (2) under a certain symmetric breaking weight initialization, the dynamics converges to w∗, without getting stuck into any local minima. Note that in both cases, the initialization can be arbitrarily close to the origin for a fixed ‖w∗‖, showing that such a convergence behavior is beyond the local convex structure at w∗. To our knowledge, this is the first proof of its kind.
Previous works also use dynamical system to analyze deep neural networks. [Saxe et al. (2013)] analyzes the dynamics of multilayer linear network, and [Kawaguchi (2016)] shows every local minima is global for multilinear network. Very little theoretical work has been done to analyze the dynamics of nonlinear networks, especially deep ones. [Mei et al. (2016)] shows the global convergence whenK = 1 with activation function σ(x) when its derivatives σ′, σ′′, σ′′′ are bounded and σ′ > 0. Similar to our approach, [Saad & Solla (1996)] also uses the student-teacher setting and analyzes the dynamics of student network when the teacher’s parameters w∗ forms a orthonomal bases; however, it uses σ(x) = erf(x) as the nonlinearity and only analyzes the local behaviors of the two critical points (the saddle point in symmetric initializations, and w∗). In contrast, we prove the global convergence behavior in certain symmetry-breaking cases.
Many previous works analyze nonlinear network based on the assumption of independent activations: the activations of ReLU (or other nonlinear) nodes are independent of the input and/or mutually independent. For example, [Choromanska et al. (2015a;b)] relate the nonlinear ReLU network with spin-glass models when several assumptions hold, including the assumption of independent activations (A1p and A5u). [Kawaguchi (2016)] proves that every local minimum in nonlinear network is global based on similar assumptions. [Soudry & Carmon (2016)] shows the global optimality of the local minimum in a two-layered ReLU network, by assuming small sample size and applying independent multiplicative Bernoulli noise on the activations. In practice, the activations are highly dependent due to their common input. Ignoring such dependency also misses important behaviors, and may lead to misleading conclusions. In this paper, no assumption of independent activation is made. For sigmoid activation, [Fukumizu & Amari (2000)] gives quite complicated conditions for a local minimum to be global when adding a new node to a two-layered network. [Janzamin et al. (2015)] gives guarantees on recovering the parameters of a 2-layered neural network learnt with tensor decomposition. In comparison, we analyze ReLU networks trained with gradient descent, which is a more popular setting in practice.
The paper is organized as follows. Sec. 2 introduces the basic formulation and some interesting novel properties of ReLU in multilayered ReLU networks. Sec. 3 and Sec. 4 then analyze the twolayered model Eqn. 1 for K = 1 and K ≥ 2, respectively. Sec. 5 shows that simulation results are consistent with theoretical analysis. Finally Sec. 7 gives detailed proofs for all theorems.
2 PRELIMINARY
2.1 NOTATION
Denote X as a N -by-d input data matrix and w∗ is the parameter of the teacher network with desired N -by-1 output u = g(X;w∗). Now suppose we have an estimator w and the estimated output v = g(X;w). We want to know with l2 loss E(w) = 12‖u − v‖ 2 = 12‖u − g(X;w)‖ 2, whether gradient descent will converge to the desired solution w∗.
The gradient descent update is w(t+1) = w(t) + η∆w(t), where ∆w(t) ≡ −∇E(w(t)). If we let η → 0, then the update rule becomes a first-order differential equation dw/dt = −∇E(w), or more concisely, ẇ = −∇E(w). In this case, Ė = ∇E(w)ᵀẇ = −‖∇E(w)‖2 ≤ 0, i.e., the function value E is nonincreasing over time. The key is to check whether there exist other critical points w 6= w∗ so that ∇E(w) = 0. In our analysis, we assume entries of inputX follow Gaussian distribution. In this situation, the gradient is a random variable and ∆w = −E [∇E(w)]. The expected E [E(w)] is also nonincreasing no matter whether we follow the expected gradient or the gradient itself, because
E [ Ė ] = −E [∇E(w)ᵀ∇E(w)] ≤ −E [∇E(w)]ᵀ E [∇E(w)] ≤ 0 (2)
Therefore, we analyze the behavior of expected gradient E [∇E(w)] rather than∇E(w).
2.2 PROPERTIES OF RELU
In this paper, we discover a few useful properties of ReLU that make our analysis much simpler. Denote D = D(w) = diag(Xw > 0) as a N -by-N diagonal matrix. The l-th diagnonal element of D is a binary variable showing whether the neuron is on for sample l. Using this notation, we could write σ(Xw) = DXw. Note that D only depends on the direction of w but not its magnitude.
Note that for ReLU,D is also “tranparent” on derivatives. For example, the Jacobian Jw[σ(Xw)] = σ′(Xw)X = DX at differentiable regions. This gives a very concise rule for gradient descent in ReLU network: suppose we have negative gradient inflow vector g (of dimension N -by-1) on the current ReLU node with weights w, then we can simply write the update ∆w as:
∆w = Jw[σ(Xw)] ᵀg = XᵀDg (3)
This can be easily applied to multilayer ReLU network. Denote j ∈ [c] if node j is in layer c, dc as the width of layer c, and uj and vj as the output of teacher network and student network, respectively. A simple deduction yields the following lemma:
Lemma 2.1 For neural network with ReLU nonlinearity and using l2 loss to match with a teacher network of the same size, the negative gradient inflow gj for node j at layer c has the following form:
gj = Lj ∑ j′ (L∗j′uj′ − Lj′vj′) (4)
where Lj and L∗j are N -by-N diagonal matrices. For any k ∈ [c+ 1], Lk = ∑ j∈[c] wjkDjLj and similarly for L∗k. For the first layer, L = L ∗ = I .
The intuition here is to start from g = u − v (true for l2 loss) at the top layer and use induction. With this formulation, we could write the finite dynamics for wc (all parameters in layer c). Denote the N -by-dc+1dc matrix Rc = [LjDj ]j∈[c]Xc and R∗c = [L ∗ jD ∗ j ]j∈[c]X ∗ c . Using gradient descent rules:
∆wj = X ᵀ cDjgj = X ᵀ cDjLj ∑ j′ L∗j′D ∗ j′X ∗ cw ∗ j′ − ∑ j′ Lj′Dj′Xcwj′  (5) = XᵀcDjLj (R ∗ cw ∗ c −Rcwc) (6)
Therefore we have: ∆wc = R ᵀ c (R ∗ cw ∗ c −Rcwc) (7)
3 SINGLE RELU CASE
Let’s start with the simplest case where there is only one ReLU node, K = 1. At iteration t, following Eqn. 3, the gradient update rule is:
∆w(t) = XᵀD(t)g(t) = XᵀD(t)(D∗Xw∗ −D(t)Xw(t)) (8)
Note here how the notation ofD(t) comes into play (andD(t)D(t) = D(t)). Indeed, when the neuron is cut off at sample l, then (D(t))ll is zero and will block the corresponding gradient component.
Linear case. In this situationD(t) = D∗ = I (no gating in either forward or backward propagation) and:
w(t+1) = w(t) + η
N XᵀX(w∗ −w(t)) (9)
where η/N is the learning rate. When it is sufficiently small so that the spectral radius ρ(I − η NX
ᵀX) < 1, w(t+1) will converge to w∗ when t→ +∞. Note that this convergence is guaranteed for any initial condition w(1), if XᵀX is full rank with suitable η. This is consistent with its convex nature. If entries of X follow i.i.d Gaussian distribution, then E [ 1 NX ᵀX ]
= I and the condition satisfies.
Nonlinear (ReLU) case. In this case, ∆w = XᵀD(D∗Xw∗ −DXw) in which D is a function of w. Intuitively, this term goes to zero when w → w∗, and should be approximated to be N2 (w
∗ − w) in the i.i.d Gaussian case, since roughly half of the samples are blocked. However, once we make such approximation, we lost the nonlinear behavior of the network and would draw the wrong conclusion of global convergence.
Then how should we analyze it? Notice that in ∆w, both of the two terms have the form F (e,w) = XᵀD(e)D(w)Xw. Using this form, E [∆w] = E [F (w/‖w‖,w∗)] − E [F (w/‖w‖,w)]. Here e is a unit vector called the “projected” weight. In the following, we will show that E [F (e,w)] has the following close form under i.i.d Gaussian assumption on X:
Lemma 3.1 Denote F (e,w) = XᵀD(e)D(w)Xw where e is a unit vector, X = [x1,x2, · · · ,xN ]ᵀ is N -by-d sample matrix and D(w) = diag(Xw > 0) is a binary diagonal matrix. If xi ∼ N(0, I) and are i.i.d (and thus bias-free), then:
E [F (e,w)] = N
2π [(π − θ)w + ‖w‖ sin θe] (10)
where θ = ∠(e,w) ∈ [0, π] is the angle between e and w.
Note that the expectation analysis smooths out the non-differentiable property of ReLU, leaving only one singularity at e = 0. The intuition is that expectation analysis involves an integration over the data distribution. With simple algebraic manipulation, E [∆w] takes the following closed form:
E [∆w] = N 2 (w∗ −w) + N 2π (α sin θw − θw∗) (11)
where α = ‖w∗‖/‖w‖ and θ ∈ [0, π] is the angle between w and w∗. The first term is expected while the last two terms show the nonlinear behavior. Using Lyapunov’s method, we show that the dynamics (if treated continuously) converges to w∗ when w(1) ∈ Ω = {w : ‖w −w∗‖ < ‖w∗‖}:
Lemma 3.2 When w(1) ∈ Ω = {w : ‖w −w∗‖ < ‖w∗‖}, following the dynamics of Eqn. 11, the Lyapunov function V (w) = 12‖w − w
∗‖2 has V̇ < 0 and the system is asymptotically stable and thus w(t) → w∗ when t→ +∞.
See Appendix for the proof. The intuition is to represent V as a 2-by-2 bilinear form of vector [‖w‖, ‖w∗‖], and the bilinear coefficient matrix is positive definite. One question arises: will the same approach show the dynamics converges when the initial conditions lie outside the region Ω, in particular for any region that includes the origin? The answer is probably no. Note that w = 0 is a singularity in which ∆w is not continuous (if approaching from different directions towards w = 0, ∆w is different). It is due to the fact that ReLU function is not differentiable at the origin. We could remove this singularity by “smoothing out” ReLU around the origin. This will yield ∆w→ 0 when w → 0. In this case, V̇ (0) = 0 so Lyapunov method could only tell that the dynamics is stable but not convergent. Note that for ReLU activation, σ′(x) = 0 for certain negative x even after a local smoothing, so the global convergence claim in [Mei et al. (2016)] for l2 loss does not apply.
Random Initialization. Then we study how to sample w(1) so that w(1) ∈ Ω. We would like to sample within Ω, but we don’t know where is w∗. Sampling around origin with big radius r ≥ 2‖w∗‖ is inefficient in particular in high-dimensional space. This is because when the sample is uniform, the probability of hitting the ball is proportional to (r/‖w∗‖)d ≤ 2−d, which is exponentially small.
A better idea is to sample around the origin with very small radius (but not at w = 0), so that the convergent hypersphere behaves like a hyperplane near the origin, and thus almost half of the samples is useful (Fig. 2(a)), as shown in the following theorem:
Theorem 3.3 The dynamics in Eqn. 11 converges to w∗ with probability at least (1 − )/2, if the initial value w(1) is sampled uniformly from Br = {w : ‖w‖ ≤ r} with r ≤ √ 2π d+1‖w ∗‖.
The intution here is to lower-bound the probability of the shaded area (Fig. 2(b)). From the proof, the conclusion could be made stronger to show r ∼ 1/ √ d, consistent with common initialization techniques [Glorot & Bengio (2010); He et al. (2015); LeCun et al. (2012)]. Fig. 2(c) shows an example in the 2D case, in which there is a singularity at the origin, and sampling towards w∗ yields the convergence. This is consistent with the analysis above.
4 MULTIPLE RELUS CASE
Now we are ready to analyze the network g(x) = ∑K j=1 σ(w ᵀ j x) for K ≥ 2 (Fig. 1(c)). Theoretical analysis of such networks is also the main topic in many previous works [Saad & Solla (1996); Soudry & Carmon (2016); Fukumizu & Amari (2000)]. In this case, Lj = L∗j = I for 1 ≤ j ≤ K. Then we have the following nonlinear dynamics from Eqn. 7:
∆wj = K∑ j′=1 f(wj ,wj′ ,w ∗ j′) (12)
where f = F (wj/‖wj‖,w∗j′)− F (wj/‖wj‖,wj′). Therefore, using Eqn. 10, its expectation is:
2π N E [ f(wj ,wj′ ,w ∗ j′) ] = (π − θ∗j ′ j )w ∗ j′ − (π − θ j′ j )wj′ + (‖w∗j′‖ ‖wj‖ sin θ∗j ′ j − ‖wj′‖ ‖wj‖ sin θj ′ j ) wj (13) where θ∗j ′ j ≡ ∠(wj ,w∗j′) and θ j′
j ≡ ∠(wj ,wj′). Eqn. 12 (and its expected version) gives very complicated nonlinear dynamics and could be hard to solve in general. Unlike K = 1, a similar approach with Lyaponov function does not yield a decisive conclusion. However, if we consider the symmetric case: wj = Pjw and w∗j = Pjw ∗ where Pj is a cyclic permutation matrix that maps index j′ + 1 to (j′ + j mod K) + 1 (and P1 is the identity matrix), then RHS of the expected version of Eqn. 12 can be simplified as follows:
E [∆wj ] = ∑ j′ E [ f(wj ,wj′ ,w ∗ j′) ] = ∑ j′ E [f(Pjw, Pj′w, Pj′w∗)]
= ∑ j′′ E [f(Pjw, PjPj′′w, PjPj′′w∗)] ({Pj}Kj=1 is a group)
= Pj ∑ j′′ E [f(w, Pj′′w, Pj′′w∗)] (‖Pw1‖ = ‖w1‖, ∠(Pw1, Pw2) = ∠(w1,w2))
= PjE [∆w1] (14)
which means that if all wj and w∗j are symmetric under the action of cyclic group, so does their expected gradient. Therefore, the trajectory {w(t)} keeps such cyclic structure. Instead of solving a system of K equations, we only need to solve one:
E [∆w] = K∑ j=1 E [f(w, Pjw, Pjw∗)] (15)
Surprisingly, there is another layer of symmetry in Eqn. 15 when {w∗j} forms an orthonomal basis (w∗j′ ᵀw∗j = δjj′ ). In this case, if we start with w (1) = xw∗ + y ∑ j 6=1 Pjw
∗ then we could show that the trajectory keeps this structure and Eqn. 15 can be further reduced into the following 2D nonlinear dynamics:
2π N E [ ∆x ∆y ] = − { [(π − φ)(x− 1 + (K − 1)y)] [ 1 1 ] + [ θ φ∗ − φ ] + φ [ x− 1 y ]} + [(K − 1)(α sinφ∗ − sinφ) + α sin θ] [ x y ] (16)
Here the symmetrical factor (α ≡ ‖w∗j′‖/‖wj‖, θ ≡ θ ∗j j , φ ≡ θ
j′
j , φ ∗ ≡ θ∗j
′
j ) are defined as follows:
α = (x2 + (K−1)y2)−1/2, cos θ = αx, cosφ∗ = αy, cosφ = α2(2xy+ (K−2)y2) (17) For this 2D dynamics, we thus have the following theorem:
Theorem 4.1 For any K ≥ 2, the 2D dynamics (Eqn. 16) shows the following behaviors:
(1) Symmetric case. If the initial condition x(1) = y(1) ∈ (0, 1], then the dynamics reduces to 1D and converges to a saddle point x = y = 1πK ( √ K − 1− arccos(1/ √ K) + π).
(2) Symmetry-Breaking. If (x(1), y(1)) ∈ Ω = {x ∈ (0, 1], y ∈ [0, 1], x > y}, then dynamics always converges to (x, y) = (1, 0).
From (x(t), y(t)) we could recover w(t)j = x (t)w∗j + y (t) ∑ j′ 6=j w ∗ j′ . Obviously, a convergence of Eqn. 16 to (1, 0) means Eqn. 12 converges to {w∗j}, i.e, the teacher parameters are recovered:
Corollary 4.2 For a bias-free two-layered ReLU network g(x;w) = ∑ j σ(w ᵀ j x) that takes Gaussian i.i.d inputs (Fig. 1), if the teacher’s parameters {w∗j} form orthogonal bases, then when the student parameters is initialized in the form of w(1)j = x (1)w∗j + y (1) ∑ j′ 6=j w ∗ j′ where (x(1), y(1)) ∈ Ω = {x ∈ (0, 1], y ∈ [0, 1], x > y}, then the dynamics (Eqn. 12) converges to {w∗j} without being trapped into local minima.
When symmetry is broken, since the closure of Ω includes the origin, there exists a path starting at arbitrarily small neighborhood of origin to w∗, regardless of how large ‖w∗‖ is. In contrast to traditional convex analysis that only gives the local parameter-dependent convergence basin around w∗j , here we obtain a convergence basin that is parameter-independent. In comparison, [Saad & Solla (1996)] uses a different activation function (σ(x) = erf(x)) and only analyzes local behaviors near the two fixed points (the symmetric saddle point and the teacher’s weights w∗), leaving symmetry breaking an empirical procedure. Here we show that it is possible to give global convergence analysis on certain symmetry breaking cases for two-layered ReLU network.
By symmetry, Corollary 4.1 immediately suggests that when w(1) = y(1) ∑K j=1 w ∗ j + (x
(1) − y(1))w∗j′ , then the dynamics will converge to Pj′w
∗. Since x > y but can be arbitrarily close, a slighest preturbation on the symmetric solution x = y leads to a different fixed point, which is a permutation of w∗. This is very similar to Spontaneously Symmetric-Breaking (SSB) procedure in physics, in which a high energy state with full symmetry goes to a low energy state and only retains part of the symmetry. In this case, the energy is the objective function E, the high energy state is the initialization that is almost symmetrical but with small fluctuation, and the low energy state is the fixed point the dynamics converges into.
From the simulation shown in Fig. 4, we could see that gradient descent takes a detour to reach the desired solution w∗, even when the initialization is aligned with w∗. This is because in the first stage, all ReLU nodes receive the residue and try to explain the data in the same way (both x and y increases); when the “obvious” component has been explained away, then the residue changes its direction and pushes some ReLU nodes to explain other components as well (x increases but y decreases).
Empirically this path also converges to w∗ under noise. We leave it a conjecture that the system converges in the presence of reasonably large noise. If this conjecture is true, then with high probability a random initialization stays in the convergence basin and converges to a permutation of w∗. The reason is that a random initialization almost never gives ties. Without a tie, there exists one leading component which will dominate the convergence.
Conjecture 4.3 When the initialization w(1) = x(1)w∗j + y(1) ∑ j′ 6=j w ∗ j′ + , where is Gaussian noise and (x(1), y(1)) ∈ Ω, then the dynamics Eqn. 12 also converges to w∗ without trapped into local minima.
5 SIMULATION
5.1 CLOSE FORM SOLUTION FOR ONE RELU NODE
We verify our close form expression of E [F (e,w)] = E [XᵀD(e)D(w)Xw] (Eqn. 10) with simulation. We randomly pick e and w so that their angle ∠(e,w) is uniformly distributed in [0, π]. We prepare the input data X with standard Gaussian distribution and compare the close form solution E [F (e,w)] with F (e,w), the actual data term in gradient descent without expectation. We use relative RMS error: err = ‖E [F (e,w)] − F (e,w)‖/‖F (e,w)‖. As shown in Fig. 3(a), The error distribution on angles shows the properties of the close-form solution. For small θ, D(w) and
very large noise is present. Both teacher and student networks use g(x) =
∑K
j=1 σ(w ᵀ j x). Each
experiment has 8 runs. Bottom row: Convergence when we use g2(x) = ∑K j=1 ajσ(w ᵀ j x). Here the top weights aj is fixed at different numbers (rather than 1). Large positive aj correponds to fast convergence. When aj has positive/negative components, the network does not converge to w∗.
D(e) overlaps sufficiently, giving a reliable estimation for the gradient. When θ → π, D(w) and D(e) tend not to overlap, leaving very few data involved in the gradient computation. As a result, the variance grows. Note that all our analysis operate on θ ∈ [0, π/2] and is not affected by this behavior. In the following, angles are sampled from [0, π/2].
Fig. 3(a) shows that the close form expression becomes more accurate with more samples. We also examine other zero-mean distributions of X , e.g., uniform distribution in [−1/2, 1/2]. As shown in Fig. 3(d), the close form expression still works for large d, showing that it could be quite general. Note that the error is computed up to a scaling constant, due to the difference in normalization constants among different distributions. We leave it to the future work to prove its usability for broader distributions.
5.2 CONVERGENCE FOR MULTIPLE RELU NODES
Fig. 4(a) and (b) shows the 2D vector field given by the 2D dynamics (Eqn. 16) and Fig. 4(c) shows the 2D trajectory towards convergence to the teacher’s parameters w∗. Interestingly, even when we initialize the weights as (10−3, 0), aligning with w∗, the gradient descent takes detours to reach the destination. One explanation is, at the beginning all nodes move similar direction trying to explain the data, once the data have been explained partly, specialization follows (y decreases).
Fig. 5 shows empirical convergence for K ≥ 2, when the initialization deviates from symmetric initialization in Thm. 4.1. Unless the deviation is large, gradient descent converges to w∗. We also check the convergence of a more general network g2(x) = ∑K j=1 ajσ(w ᵀ j x). When aj > 0 convergence follows; however, when some aj is negative, the network does not converge to w∗, even that the student network already knows the ground truth value of {aj}Kj=1.
6 CONCLUSION AND FUTURE WORK
In this paper, we analyze the nonlinear dynamical behavior of certain two-layered bias-free ReLU networks in the form of g(x;w) = ∑K j=1 σ(w ᵀ j x), where σ = max(x, 0) is the ReLU node. We assume that the input x follows Gaussian distribution and the output is generated by a teacher network with parameters w∗. In K = 1 we show a close-form nonlinear dynamics can be obtained and its convergence to w∗ can be proven, if we sample the initialization properly. Such initialization is consistent with common practice [Glorot & Bengio (2010); He et al. (2015)] and is independent of the value of w∗. ForK ≥ 2, when the teacher parameters {w∗j} form a orthonormal bases, we prove that the trajectory from symmetric initialization is trapped into a saddle point, while certain symmetric breaking initialization converges to w∗ without trapped into any local minima. Future work includes analysis of general cases (or symmetric case plus noise) for K ≥ 2, and a generalization to multilayer ReLU (or other nonlinear) networks.
7 APPENDIX
Here we list all detailed proof for all the theorems.
7.1 PROPERTIES OF RELU NETWORKS
Lemma 7.1 For neural network with ReLU nonlinearity and using l2 loss to match with a teacher network of the same size, the negative gradient inflow gj for node j at layer c has the following form:
gj = Lj ∑ j′ (L∗j′uj′ − Lj′vj′) (18)
where Lj and L∗j are N -by-N diagonal matrices. For any k ∈ [c+ 1], Lk = ∑ j∈[c] wjkDjLj and similarly for L∗k.
Proof We prove by induction on layer. For the first layer, there is only one node with g = u − v, therefore Lj = Lj′ = I . Suppose the condition holds for all node j ∈ [c]. Then for node k ∈ [c+1], we have:
gk = ∑ j wjkDjgj = ∑ j wjkDjLj ∑ j′ L∗j′uj′ − ∑ j′ Lj′vj′  =
∑ j wjkDjLj ∑ j′ L∗j′ ∑ k′ D∗j′w ∗ jk′uk′ − ∑ j′ Lj′ ∑ k′ Dj′wjk′vk′  =
∑ j wjkDjLj ∑ j′ L∗j′D ∗ j′ ∑ k′ w∗jk′uk′ − ∑ j wjkDjLj ∑ j′ Lj′Dj′ ∑ k′ wjk′vk′
= ∑ k′ ∑ j wjkDjLj ∑ j′ L∗j′D ∗ j′w ∗ jk′ uk′ −∑ k′ ∑ j wjkDjLj ∑ j′ Lj′Dj′wjk′ vk′ Setting Lk = ∑ j wjkDjLj and L ∗ k = ∑ j w ∗ jkD ∗ jL ∗ j (both are diagonal matrices), we thus have:
gk = ∑ k′ LkL ∗ k′uk′ − LkLk′vk′ = Lk ∑ k′ L∗k′uk′ − Lk′vk′ (19)
7.2 ONE RELU CASE
Lemma 7.2 Suppose F (e,w) = XᵀD(e)D(w)Xw where e is a unit vector and X = [x1,x2, · · · ,xN ]ᵀ is N -by-d sample matrix. If xi ∼ N(0, I) and are i.i.d, then:
E [F (e,w)] = N
2π ((π − θ)w + ‖w‖ sin θe) (20)
where θ ∈ [0, π] is the angle between e and w.
Proof Note that F can be written in the following form: F (e,w) = ∑
i:xᵀi e≥0,x ᵀ i w≥0
xix ᵀ iw (21)
where xi are samples so that X = [x1,x2, · · · ,xn]ᵀ. We set up the axes related to e and w as in Fig. 6, while the rest of the axis are prependicular to the plane. In this coordinate system, any vector x = [r sinφ, r cosφ, x3, . . . , xd]. We have an orthonomal set of bases: e, e⊥ = −e−w/‖w‖ cos θsin θ (and any set of bases that span the rest of the space). Under the basis, the representation for e and w is [1,0d−1] and [‖w‖ cos θ,−‖w‖ sin θ,0d−2]. Note that here θ ∈ (−π, π]. The angle θ is positive when e “chases after” w, and is otherwise negative.
Now we consider the quality R(φ0) = E [ 1 N ∑ i:φi∈[0,φ0] xix ᵀ i ] . If we take the expectation and use polar coordinate only in the first two dimensions, we have:
R(φ0) = E  1 N ∑ i:φi∈[0,φ0] xix ᵀ i  = E [xixᵀi |φi ∈ [0, φ0]]P [φi ∈ [0, φ0]] =
∫ +∞ 0 ∫∫ +∞ −∞ ∫ φ0 0 r sinφr cosφ. . . xd  [r sinφ r cosφ . . . xd] p(r)p(θ) d∏ k=3 p(xk)rdrdφdx3 . . . dxd
where p(r) = e−r 2/2 and p(θ) = 1/2π. Note that R(φ0) is a d-by-d matrix. The first 2-by-2 block can be computed in close form (note that ∫ +∞ 0
r2p(r)rdr = 2). Any off-diagonal element except for the first 2-by-2 block is zero due to symmetric property of i.i.d Gaussian variables. Any diagonal element outside the first 2-by-2 block will be P [φi ∈ [0, φ0]] = φ0/2π. Finally, we have:
R(φ0) = E  1 N ∑ i:φi∈[0,φ0] xix ᵀ i  = 1 4π [ 2φ0 − sin 2φ0 1− cos 2φ0 0 1− cos 2φ0 2φ0 + sin 2φ0 0 0 0 2φ0Id−2 ] (22)
= φ0 2π Id + 1 4π
[ − sin 2φ0 1− cos 2φ0 0 1− cos 2φ0 sin 2φ0 0
0 0 0
] (23)
With this equation, we could then compute E [F (e,w)]. When θ ≥ 0, the condition {i : xᵀi e ≥ 0,xᵀiw ≥ 0} is equivalent to {i : φi ∈ [θ, π]} (Fig. 6(a)). Using w = [‖w‖ cos θ,−‖w‖ sin θ,0d−2] and we have:
E [F (e,w)] = N (R(π)−R(θ))w (24)
= N
4π
( 2(π − θ)w − ‖w‖ [ − sin 2θ 1− cos 2θ 0 1− cos 2θ sin 2θ 0
0 0 0
][ cos θ − sin θ
0
]) (25)
= N
2π
( (π − θ)w + ‖w‖ [ sin θ 0 ]) (26)
= N
2π ((π − θ)w + ‖w‖ sin θe) (27)
For θ < 0, the condition {i : xᵀi e ≥ 0,x ᵀ iw ≥ 0} is equivalent to {i : φi ∈ [0, π + θ]} (Fig. 6(b)), and similarly we get
E [F (e,w)] = N (R(π + θ)−R(0))w = N 2π ((π + θ)w − ‖w‖ sin θe) (28)
Notice that by abuse of notation, the θ appears in Eqn. 20 is the absolute value and Eqn. 20 follows.
Lemma 7.3 In the region ‖w(1) −w∗‖ < ‖w∗‖, following the dynamics (Eqn. 11), the Lyapunov function V (w) = 12‖w−w
∗‖2 has V̇ < 0 and the system is asymptotically stable and thus w(t) → w∗ when t→ +∞.
Proof Denote that Ω = {w : ‖w(1) −w∗‖ < ‖w∗‖}. Note that
V̇ = (w −w∗)ᵀ∆w = −yᵀMy (29) where y = [‖w∗‖, ‖w‖]ᵀ and M is the following 2-by-2 matrix:
M = 1
2
[ sin 2θ + 2π − 2θ −(2π − θ) cos θ − sin θ
−(2π − θ) cos θ − sin θ 2π
] (30)
In the following we will show that M is positive definite when θ ∈ (0, π/2]. It suffices to show that M11 > 0, M22 > 0 and det(M) > 0. The first two are trivial, while the last one is:
4det(M) = 2π(sin 2θ + 2π − 2θ)− [(2π − θ) cos θ + sin θ]2 (31) = 2π(sin 2θ + 2π − 2θ)− [ (2π − θ)2 cos2 θ + (2π − θ) sin 2θ + sin2 θ ] (32)
= (4π2 − 1) sin2 θ − 4πθ + 4πθ cos2 θ − θ2 cos2 θ + θ sin 2θ (33) = (4π2 − 4πθ − 1) sin2 θ + θ cos θ(2 sin θ − θ cos θ) (34)
Note that 4π2 − 4πθ − 1 = 4π(π − θ) − 1 > 0 for θ ∈ [0, π/2], and g(θ) = sin θ − θ cos θ ≥ 0 for θ ∈ [0, π/2] since g(0) = 0 and g′(θ) ≥ 0 in this region. Therefore, when θ ∈ (0, π/2], M is positive definite.
When θ = 0, M(θ) = π[1,−1;−1, 1] and is semi-positive definite, with the null eigenvector being√ 2 2 [1, 1], i.e., ‖w‖ = ‖w
∗‖. However, along θ = 0, the only w that satisfies ‖w‖ = ‖w∗‖ is w = w∗. Therefore, V̇ = −yᵀMy < 0 in Ω. Note that although this region could be expanded to the entire open half-spaceH = {w : wᵀw∗ > 0}, it is not straightforward to prove the convergence inH, since the trajectory might go outsideH. On the other hand, Ω is the level set V < 12‖w
∗‖2 so the trajectory starting within Ω remains inside.
Theorem 7.4 The dynamics in Eqn. 11 converges to w∗ with probability at least (1 − )/2, if the initial value w(1) is sampled uniformly from Br = {w : ‖w‖ ≤ r} with:
r ≤ √ 2π
d+ 1 ‖w∗‖ (35)
Proof Given a ball of radius r, we first compute the “gap” δ of sphere cap (Fig. 2(b)). First cos θ = r
2‖w∗‖ , so δ = r cos θ = r2
2‖w∗‖ . Then a sufficient condition for the probability argument to hold, is to ensure that the volume Vshaded of the shaded area is greater than 1− 2 Vd(r), where Vd(r) is the volume of d-dimensional ball of radius r. Since Vshaded ≥ 12Vd(r)− δVd−1, it suffices to have:
1 2 Vd(r)− δVd−1 ≥ 1− 2 Vd(r) (36)
which gives
δ ≤ 2 Vd Vd−1
(37)
Using δ = r 2
2‖w∗‖ and Vd(r) = Vd(1)r d, we thus have:
r ≤ Vd(1) Vd−1(1) ‖w∗‖ (38)
where Vd(1) is the volume of the unit ball. Since the volume of d-dimensional unit ball is
Vd(1) = πd/2
Γ(d/2 + 1) (39)
where Γ(x) = ∫∞ 0 tx−1e−tdt. So we have
Vd(1) Vd−1(1) = √ π Γ(d/2 + 1/2) Γ(d/2 + 1) (40)
From Gautschi’s Inequality
x1−s < Γ(x+ 1)
Γ(x+ s) < (x+ s)1−s x > 0, 0 < s < 1 (41)
with s = 1/2 and x = d/2 we have:( d+ 1
2
)−1/2 < Γ(d/2 + 1/2)
Γ(d/2 + 1) <
( d
2
)−1/2 (42)
Therefore, it suffices to have
r ≤ √ 2π
d+ 1 ‖w∗‖ (43)
Note that this upper bound is tight when δ → 0 and d→ +∞, since all inequality involved asymptotically becomes equal.
7.3 TWO LAYER CASE
Lemma 7.5 For φ∗, θ and φ defined in Eqn. 17:
α ≡ (x2 + (K − 1)y2)−1/2 (44) cos θ ≡ αx (45) cosφ∗ ≡ αy (46) cosφ ≡ α2(2xy + (K − 2)y2) (47)
we have the following relations in the triangular region Ω 0 = {(x, y) : x ≥ 0, y ≥ 0, x ≥ y + 0} (Fig. 6(c)):
(1) φ, φ∗ ∈ [0, π/2] and θ ∈ [0, θ0) where θ0 = arccos 1√K .
(2) cosφ = 1− α2(x− y)2 and sinφ = α(x− y) √ 2− α2(x− y)2.
(3) φ∗ ≥ φ (equality holds only when y = 0) and φ∗ > θ.
Proof Propositions (1) and (2) are computed by direct calculations. In particular, note that since cos θ = αx = 1/ √ 1 + (K − 1)(y/x)2 and x > y ≥ 0, we have cos θ ∈ (1/ √ K, 1] and θ ∈ [0, θ0). For Preposition (3), φ∗ = arccosαy > θ = arccosαx because x > y. Finally, for x > y > 0, we have
cosφ cosφ∗ = α2(2xy + (K − 2)y2) αy = α(2x+ (K − 2)y) > α(x+ (K − 1)y) > 1 (48)
The final inequality is because K ≥ 2, x, y > 0 and thus (x + (K − 1)y)2 > x2 + (K − 1)2y2 > x2 + (K − 1)y2 = α−2. Therefore φ∗ > φ. If y = 0 then φ∗ = φ.
Theorem 7.6 For the dynamics defined in Eqn. 16, there exists 0 > 0 so that the trianglar region Ω 0 = {(x, y) : x ≥ 0, y ≥ 0, x ≥ y + 0} (Fig. 6(c)) is a convergent region. That is, the flow goes inwards for all three edges and any trajectory starting in Ω 0 stays.
Proof We discuss the three boundaries as follows:
Case 1: y = 0, 0 ≤ x ≤ 1, horizontal line. In this case, θ = 0, φ = π/2 and φ∗ = π/2. The y component of the dynamics in this line is:
f1 ≡ 2π N ∆y = −π 2 (x− 1) ≥ 0 (49)
So ∆y points to the interior of Ω.
Case 2: x = 1, 0 ≤ y ≤ 1, vertical line. In this case, α ≤ 1 and the x component of the dynamics is:
f2 ≡ 2π
N ∆x = −(π − φ)(K − 1)y − θ + (K − 1)(α sinφ∗ − sinφ) + α sin θ (50)
= −(K − 1) [(π − φ)y − (α sinφ∗ − sinφ)] + (α sin θ − θ) (51) Note that since α ≤ 1, α sin θ ≤ sin θ ≤ θ, so the second term is non-positive. For the first term, we only need to check whether (π − φ)y − (α sinφ∗ − sinφ) is nonnegative. Note that
(π − φ)y − (α sinφ∗ − sinφ) (52) = (π − φ)y + α(x− y) √ 2− α2(x− y)2 − α √ 1− α2y2 (53)
= y [ π − φ− α √ 2− α2(x− y)2 ] + α [ x √ 2− α2(x− y)2 − √ 1− α2y2 ]
(54)
In Ω we have (x − y)2 ≤ 1, combined with α ≤ 1, we have 1 ≤ √ 2− α2(x− y)2 ≤ √
2 and√ 1− α2y2 ≤ 1. Since x = 1, the second term is nonnegative. For the first term, since α ≤ 1,
π − φ− α √
2− α2(x− y)2 ≥ π − π 2 − √ 2 > 0 (55)
So (π − φ)y − (α sinφ∗ − sinφ) ≥ 0 and ∆x ≤ 0, pointing inwards.
Case 3: x = y + , 0 ≤ y ≤ 1, diagonal line. We compute the inner product between (∆x,∆y) and (1,−1), the inward normal of Ω at the line. Using φ ≤ π2 sinφ for φ ∈ [0, π/2] and φ
∗ − θ = arccosαy − arccosαx ≥ 0 when x ≥ y, we have:
f3(y, ) ≡ 2π
N
[ ∆x ∆y ]ᵀ [ 1 −1 ] = φ∗ − θ − φ+ [(K − 1)(α sinφ∗ − sinφ) + α sin θ] (56)
≥ (K − 1) [ α sinφ∗ − ( 1 +
π
2(K − 1)
) sinφ ] = α(K − 1) [√ 1− α2y2 − ( 1 + π
2(K − 1)
)√ 2− α2 2 ] Note that for y > 0:
αy = 1√
(x/y)2 + (K − 1) = 1√ (1 + /y)2 + (K − 1) ≤ 1√ K
(57)
For y = 0, αy = 0 < √ 1/K. So we have √ 1− α2y2 ≥ √ 1− 1/K. And √ 2− α2 2 ≤ √
2. Therefore f3 ≥ α(K−1)(C1− C2) withC1 ≡ √ 1− 1/K > 0 andC2 ≡ √ 2(1+π/2(K−1)) > 0. With = 0 > 0 sufficiently small, f3 > 0.
Lemma 7.7 (Reparametrization) Denote = x − y > 0. The terms αx, αy and α involved in the trigometric functions in Eqn. 16 has the following parameterization:
α [ y x ] = 1
K
[ β − β2
β + (K − 1)β2 Kβ2
] (58)
where β2 = √
(K − β2)/(K − 1). The reverse transformation is given by β =√ K − (K − 1)α2 2. Here β ∈ [1, √ K) and β2 ∈ (0, 1]. In particular, the critical point (x, y) = (1, 0) corresponds to (β, ) = (1, 1). As a result, all trigometric functions in Eqn. 16 only depend on the single variable β. In particular, the following relationship is useful:
β = cos θ + √ K − 1 sin θ (59)
Proof This transformation can be checked by simple algebraic manipulation. For example:
1
αK (β − β2) =
1
K
(√ K
α2 − (K − 1) 2 −
) = 1
K
(√ (Ky + )2 − ) = y (60)
To prove Eqn. 59, first we notice that K cos θ = Kαx = β + (K − 1)β2. Therefore, we have (K cos θ − β)2 − (K − 1)2β22 = 0, which gives β2 − 2β cos θ + 1 −K sin2 θ = 0. Solving this quadratic equation and notice that β ≥ 1, θ ∈ [0, π/2] and we get:
β = cos θ + √ cos2 θ +K sin2 θ − 1 = cos θ + √ K − 1 sin θ (61)
Lemma 7.8 After reparametrization (Eqn. 58), f3(β, ) ≥ 0 for ∈ (0, β2/β]. Furthermore, the equality is true only if (β, ) = (1, 1) or (y, ) = (0, 1).
Proof Applying the parametrization (Eqn. 58) to Eqn. 56 and notice that α = β2 = β2(β), we could write f3 = h1(β)− (φ+ (K − 1) sinφ) (62) When β is fixed, f3 now is a monotonously decreasing function with respect to > 0. Therefore, f3(β, ) ≥ f3(β, ′) for 0 < ≤ ′ ≡ β2/β. If we could prove f3(β, ′) ≥ 0 and only attain zero at known critical point (β, ) = (1, 1), the proof is complete.
Denote f3(β, ′) = f31 + f32 where
f31(β, ′) = φ∗ − θ − ′φ+ ′α sin θ (63) f32(β, ′) = (K − 1)(α sinφ∗ − sinφ) ′ (64)
For f32 it suffices to prove that ′(α sinφ∗ − sinφ) = β2 sinφ∗ − β2β sinφ ≥ 0, which is equivalent to sinφ∗ − sinφ/β ≥ 0. But this is trivially true since φ∗ ≥ φ and β ≥ 1. Therefore, f32 ≥ 0. Note that the equality only holds when φ∗ = φ and β = 1, which corresponds to the horizontal line x ∈ (0, 1], y = 0. For f31, since φ∗ ≥ φ, φ∗ > θ and ′ ∈ (0, 1], we have the following:
f31 = ′(φ∗ − φ) + (1− ′)(φ∗ − θ)− ′θ + β2 sin θ ≥ − ′θ + β2 sin θ ≥ β2 ( sin θ − θ
β
) (65)
And it reduces to showing whether β sin θ − θ is nonnegative. Using Eqn. 59, we have:
f33(θ) = β sin θ − θ = 1
2 sin 2θ +
√ K − 1 sin2 θ − θ (66)
Note that f ′33 = cos 2θ + √ K − 1 sin 2θ − 1 = √ K cos(2θ − θ0) − 1, where θ0 = arccos 1√K . By Prepositions 1 in Lemma 7.5, θ ∈ [0, θ0). Therefore, f ′33 ≥ 0 and since f33(0) = 0, f33 ≥ 0. Again the equity holds when θ = 0, φ∗ = φ and ′ = 1, which is the critical point (β, ) = (1, 1) or (y, ) = (0, 1).
Theorem 7.9 For the dynamics defined in Eqn. 16, the only critical point (∆x = 0 and ∆y = 0) within Ω is (y, ) = (0, 1).
Proof We prove by contradiction. Suppose (β, ) is a critical point other than w∗. A necessary condition for this to hold is f3 = 0 (Eqn. 56). By Lemma 7.8, > ′ = β2/β > 0 and
− 1 +Ky = 1 α (β2 − α+ β − β2) = β − α α = β − β2/ α > β − β2/ ′ α = 0 (67)
So − 1 +Ky is strictly greater than zero. On the other hand, the condition f3 = 0 implies that
((K − 1)(α sinφ∗ − sinφ) + α sin θ) = −1 (φ∗ − θ) + φ (68)
Using φ ∈ [0, π/2], φ∗ ≥ φ and φ∗ > θ, we have: 2π
N ∆y = −(π − φ)( − 1 +Ky)− (φ∗ − φ)− φy + ((K − 1)(α sinφ∗ − sinφ) + α sin θ) y
= −(π − φ)( − 1 +Ky)− (φ∗ − φ)− 1 (φ∗ − θ)y < 0 (69)
So the current point (β, ) cannot be a critical point.
Theorem 7.10 Any trajectory in Ω 0 converges to (y, ) = (1, 0), following the dynamics defined in Eqn. 16.
Proof We have Lyaponov function V = E [E] so that V̇ = −E [∆wᵀ∆w] ≤ −E [∆w]ᵀ E [∆w] ≤ 0. By Thm. 7.9, other than the optimal solution w∗, there is no other symmetric critical point, ∆w 6= 0 and thus V̇ < 0. On the other hand, by Thm. 7.6, the triangular region Ω 0 is convergent, in which the 2D dynamics isC∞ differentiable. Therefore, any 2D solution curve ξ(t) will stay within. By PoincareBendixson theorem, when there is a unique critical point, the curve either converges to a limit circle or the critical point. However, limit cycle is not possible since V is strictly monotonous decreasing along the curve. Therefore, ξ(t) will converge to the unique critical point, which is (y, ) = (1, 0) and so does the symmetric system (Eqn. 12).
Theorem 7.11 When x = y ∈ (0, 1], the 2D dynamics (Eqn. 16) reduces to the following 1D case: 2π
N ∆x = −πK(x− x∗) (70)
where x∗ = 1πK ( √ K − 1− arccos(1/ √ K) + π). Furthermore, x∗ is a convergent critical point.
Proof The 1D system can be computed with simple algebraic manipulations (note that when x = y, φ = 0 and θ = φ∗ = arccos(1/ √ K)). Note that the 1D system is linear and its close form solution is x(t) = x0 + Ce−K/2Nt and thus convergent.

In this paper, the author analyzes the convergence dynamics of a single layer non-linear network under Gaussian iid input assumptions. The first half of the paper, dealing with a single hidden node, was somewhat clear, although I have some specific questions below. The second half, dealing with multiple hidden nodes, was very difficult for me to understand, and the final "punchline" is quite unclear. I think the author should focus on intuition and hide detailed derivations and symbols in an appendix.
In terms of significance, it is very hard for me to be sure how generalizable these results are: the Gaussian assumption is a very strong one, and so is the assumption of iid inputs. Real-world feature inputs are highly correlated and are probably not Gaussian. Such assumptions are not made (as far as I can tell) in recent papers analyzing the convergence of deep networks e.g. Kawaguchi, NIPS 2016. Although the author says the no assumption is made on the independence of activations, this assumption is shifted to the input instead. I think this means that the activations are combinations of iid random variables, and are probably Gaussian like, right? So I'm not sure where this leaves us.
Specific comments:
1. Please use D_w instead of D to show that D is a function of w, and not a constant. This gets particularly confusing when switching to D(w) and D(e) in Section 3. In general, notation in the paper is hard to follow and should be clearly introduced.
2. Section 3, statement that says "when the neuron is cut off at sample l, then (D^(t))_u" what is the relationship between l and u? Also, this is another example of notational inconsistency that causes problems to the reader.
3. Section 3.1, what is F(e, w) and why is D(e) introduced? This was unclear to me.
4. Theorem 3.3 suggests that (if \epsilon is > 0), then to have the maximal probability of convergence, \epsilon should be very close to 0, which means that the ball B_r has radius r -> 0? This seems contradictory from Figure 2.
5. Section 4 was really unclear and I still do not understand what the symmetry group really represents. Is there an intuitive explanation why this is important?
6. Figure 5: what is a_j ?
I encourage the author to rewrite this paper for clarity. In it's present form, it would be very difficult to understand the takeaways from the paper.


We propose an approximate strategy to efficiently train neural network based language models over very large vocabularies. Our approach, called adaptive softmax, circumvents the linear dependency on the vocabulary size by exploiting the unbalanced word distribution to form clusters that explicitly minimize the expectation of computational complexity. Our approach further reduces the computational cost by exploiting the specificities of modern architectures and matrix-matrix vector operations, making it particularly suited for graphical processing units. Our experiments carried out on standard benchmarks, such as EuroParl and One Billion Word, show that our approach brings a large gain in efficiency over standard approximations while achieving an accuracy close to that of the full softmax. The code of our method is available at https://github.com/facebookresearch/adaptive-softmax.
We propose an approximate strategy to efficiently train neural network based language models over very large vocabularies. Our approach, called adaptive softmax, circumvents the linear dependency on the vocabulary size by exploiting the unbalanced word distribution to form clusters that explicitly minimize the expectation of computational complexity. Our approach further reduces the computational cost by exploiting the specificities of modern architectures and matrix-matrix vector operations, making it particularly suited for graphical processing units. Our experiments carried out on standard benchmarks, such as EuroParl and One Billion Word, show that our approach brings a large gain in efficiency over standard approximations while achieving an accuracy close to that of the full softmax. The code of our method is available at https://github.com/facebookresearch/adaptive-softmax.
1 INTRODUCTION
This paper considers strategies to learn parametric models for language modeling with very large vocabularies. This problem is key to natural language processing, with applications in machine translation (Schwenk et al., 2012; Sutskever et al., 2014; Vaswani et al., 2013) or automatic speech recognition (Graves et al., 2013; Hinton et al., 2012). In particular, Neural Network Language Models (NNLMs) have received a renewed interest in recent years, by achieving state of the art performance on standard benchmarks (Jozefowicz et al., 2016; Mikolov et al., 2010). These approaches are more costly but generalize better than traditional non-parametric models (Bahl et al., 1983; Kneser & Ney, 1995).
Statistical language models assign a probability to words given their history (Bahl et al., 1983). They are evaluated by objective criteria such as perplexity (ppl), which directly measures the ability of the system to determine proper probabilities for all the words. This potentially makes parametric models prohibitively slow to train on corpara with very large vocabulary. For instance, the vocabulary of the One Billion Word benchmark (Chelba et al., 2013) contains around 800K words. In standard NNLMs, such as feedforward networks (Bengio et al., 2003a) or recurrent networks (Mikolov et al., 2010), computing this probability over the whole vocabulary is the bottleneck. Many solutions have been proposed to reduce the complexity of this expensive step (Bengio et al., 2003b; Goodman, 2001a; Gutmann & Hyvärinen, 2010). We distinguish (i) the methods that consider the original distribution and aim at providing approximations of the probabilities, or of a subset of them (Bengio et al., 2003b; Ji et al., 2015), from (ii) the approaches that compute exact probabilities for an approximate model yielding a lower computational cost, such as the popular hierarchical softmax (Goodman, 2001a; Mnih & Hinton, 2009; Morin & Bengio, 2005).
Our approach, called adaptive softmax, belongs to the second category. More specifically, it is inspired by the hierarchical softmax and its subsequent variants. In contrast to previous works and motivated by the trend that GPUs are comparatively more and more performant than CPUs, our design is oriented towards efficient processing on GPUs. In this context, our paper makes the following points:
• We define a strategy to produce an approximate hierarchical model. It departs from previous ones in that it explicitly takes into account the complexity of matrix-matrix multiplications on modern architectures, which is not trivially linear in the dimensions of the matrices.
• We conduct an empirical complexity analysis of this model on recent GPUs. This leads us to define a realistic complexity model that is incorporated in the proposed optimization;
• Our approach provides a significant acceleration factor compared to the regular softmax, i.e., 2× to 10× speed-ups. Equivalently we improve the accuracy under computational constraints. Importantly, on the largest corpus, this higher efficiency empirically comes at no cost in accuracy for a given amount of training data, in contrast to concurrent approaches improving the efficiency.
This paper is organized as follows. Section 2 briefly reviews the related work and Section 3 provides some background on the language modeling task that we consider. Section 4 describes our proposal, which is subsequently evaluated in Section 5 on typical benchmarks of the language modeling literature, including Text8, Europarl and One Billion Word datasets.
2 RELATED WORK
Many methods have been proposed to approximate the softmax efficiently (Bengio et al., 2003b; Goodman, 2001a; Gutmann & Hyvärinen, 2010; Morin & Bengio, 2005). We briefly describe the most popular ones below and point the reader to Chen et al. (2015) for a comparative study. For the sake of completeness, we refer the reader to other strategies that can speed-up the training of language models in complementary manners (Mikolov et al., 2011b).
Loss function approximation. The Hierarchical Softmax (HSM) is an approximation of the softmax function introduced by Goodman (2001a). This approach is generally used with a two-level tree (Goodman, 2001a; Mikolov et al., 2011c) but has also been extended to deeper hierarchies (Morin & Bengio, 2005; Mnih & Hinton, 2009). In general, the hierarchy structure is built on word similarities (Brown et al., 1992; Le et al., 2011; Mikolov et al., 2013) or frequency binning (Mikolov et al., 2011c). In particular, Mikolov et al. (2013) proposes an optimal hierarchy by constructing a Huffman coding based on frequency. However this coding scheme does not take into account the theoretical complexity reduction offered by matrix-matrix multiplication and distributed computation, in particular with modern GPUs.
Similar to our work, Zweig & Makarychev (2013) constructs their hierachy in order to explicitly reduce the computational complexity. They also solve the assignment problem with dynamic programming. However, they only consider hierachies where words are kept in the leaves of the tree, leading to a significant drop of performance (reported to be around 5− 10%), forcing them to also optimize for word similarity. In our case, allowing classes to be stored in the internal node of the tree leads to almost no drop of performance. Also, they assume a linear cost for the vector-matrix operation which significantly limits the use of their approach on distributed system such as GPU.
The idea of keeping a short-list of the most frequent words has been explored before (Le et al., 2011; Schwenk, 2007). In particular, Le et al. (2011) combines a short-list with a hierachical softmax based on word representation. In contrast, the word hierarchy that we introduce in Section 4 explicitly aims at reducing the complexity.
Our work also shares similarities with the d-softmax introduced by Chen et al. (2015). They assign capacity to words according to their frequency to speed up the training. Less frequent words have smaller classifiers than frequent ones. Unlike our method, their formulation requires accessing the whole vocabulary to evaluate the probability of a word.
Sampling based approximation. Sampling based approaches have been successfully applied to approximate the softmax function over large dictionaries in different domains, such as language modeling (Jozefowicz et al., 2016), machine translation (Jean et al., 2015) and computer vision (Joulin et al., 2015). In particular, importance sampling (Bengio & Senécal, 2008; Bengio et al., 2003b) selects a subset of negative targets to approximate the softmax normalization. Different schemes have been proposed for sampling, such as the unigram and bigram distribution (Bengio et al., 2003b) or more recently, a power-raised distribution of the unigram (Ji et al., 2015; Mikolov et al., 2013). While this approach often leads to significant speed-up at train time, it still requires to evaluate the full softmax at test time.
Self-normalized approaches. Self-normalized approaches aim at learning naturally normalized classifier, to avoid computing the softmax normalization. Popular methods are Noise Contrastive Estimation (Gutmann & Hyvärinen, 2010; Mnih & Teh, 2012; Vaswani et al., 2013) or a penalization on the normalization function (Andreas & Klein, 2014; Devlin et al., 2014). Noise Contrastive Estimation (Gutmann & Hyvärinen, 2010) replaces the softmax by a binary classifier distinguishing the original distribution form a noisy one. While the original formulation still requires to compute the softmax normalization, Mnih & Teh (2012) shows that good performance can be achieved even without it.
Finally, Vincent et al. (2015) have also proposed an efficient way to train model with high dimensional output space. Their approach is exact and leads to a promising speed-up but it cannot be directly applied to the softmax function, limiting its potential application to language modeling.
3 PRELIMINARIES ON LANGUAGE MODELING
The goal of language modeling is to learn a probability distribution over a sequence of words from a given dictionary V . The joint distribution is defined as a product of conditional distribution of tokens given their past (Bahl et al., 1983). More precisely, the probability of a sequence of T words w1, . . . , wT ∈ VT is given as
P (w1, . . . , wT ) = T∏ t=1 P (wt | wt−1, . . . , w1). (1)
This problem is traditionally addressed with non-parameteric models based on counting statistics (Goodman, 2001b). In particular, smoothed N-gram models (Bahl et al., 1983; Katz, 1987; Kneser & Ney, 1995) achieve good performance in practice (Mikolov et al., 2011a), especially when they are associated with cache models (Kuhn & De Mori, 1990). More recently, parametric models based on neural networks have gained popularity for language modeling (Bengio et al., 2003a; Jozefowicz et al., 2016; Mikolov et al., 2010). They are mostly either feedforward networks (Bengio et al., 2003a) or recurrent networks (Mikolov et al., 2010).
Feedforward network. In a standard feedforward network for language modeling, we fix a window of length N and predict the next words according to the words appearing in this window. In the simplest case, this probability is represented by a 2-layer neural network acting on an input xt ∈ VN , defined as the concatenation of the one-hot representation of the N previous words, wt−N+1, . . . , wt. The state ht of the hidden layer and subsequently the vector of scores yt associated with the next token wt+1 are computed as
ht = σ(APxt), (2) yt = f(Bht), (3)
where σ is a non linearity, e.g., the pointwise sigmoid function σ(z) = 1/(1+exp(−z)), and f is the softmax function discussed in the next section. This model is parameterized by the weight matrices P , A and B and is routinely learned with an optimization scheme such as stochastic gradient descent or Adagrad (Duchi et al., 2011).
Recurrent network. A Recurrent network (Elman, 1990) extends a feedforward network in that the current state of the hidden layer also depends on its previous state. The hidden state ht is updated according to the equation ht = σ(Awt + Rht−1), where R is a weight matrix and xt is the one-hot representation of the current word wt. Computing the exact gradient for this model is challenging but it is possible to compute an efficient and stable approximation of it, using a truncated back-propagation through time (Werbos, 1990; Williams & Peng, 1990) and norm clipping (Mikolov et al., 2010).
Since the model introduced by Elman (1990), many extensions have been proposed, such as Longer Short Term Memory (LSTM) (Hochreiter & Schmidhuber, 1997), Gated recurrent units (Chung et al., 2014) or structurally constrained network (Mikolov et al., 2014). These models have been successfully used in the context of language modeling (Jozefowicz et al., 2016; Mikolov et al., 2010; Mikolov & Zweig, 2012). In this work, we focus on the standard word level LSTM architecture since it has obtained state of the art performance on the challenging One Billion Word Benchmark (Jozefowicz et al., 2016).
Class-based hierarchical softmax. In neural language modeling, predicting the probability of the next word requires to compute scores for every word in the vocabulary and to normalize them to form a probability distribution. This is typically achieved by applying a softmax function to the unnormalized score zw associate with each word w, where the softmax function is defined as
f(zw) = exp(zw)∑
w′∈V exp(zw′) . (4)
For a vocabulary comprising k = |V| words, this function requires O(k) operations once the scores are computed. In the case of neural networks, the overall complexity is O(dk), where d is the size of the last hidden layer. When the vocabulary is large, this step is computationally expensive and often dominates the computation of the whole model (Jozefowicz et al., 2016; Mikolov et al., 2014), as discussed in introduction and related work. A simple approach (Goodman, 2001a) to reduce this computational cost is to assign each word w of the vocabulary to a unique class C(w) and to factorize the probability distribution over words as
p(wt | ht) = p1(C(wt) | ht)× p2(wt | C(wt), ht),
where p1 and p2 are obtained using the softmax function (Eq. 4). If each class contains √ k words,
the computational cost is reduced from O(dk) to O(d √ k).
4 OUR APPROACH: THE ADAPTIVE SOFTMAX
In this section, we propose the adaptive softmax, a simple speedup technique for the computation of probability distributions over words. The adaptive softmax is inspired by the class-based hierarchical softmax, where the word classes are built to minimize the computational complexity. Our method is designed to be efficient for GPUs, which are commonly used to train neural networks. For the sake of clarity, we first present the intuition behind our method in the simple case where we simply split our dictionary in two distinct clusters, before analyzing a more general case.
4.1 GPU COMPUTATIONAL MODEL
The bottleneck of the model described in the previous section is the matrix multiplication between the matrix representing the hidden states (of size B × d, where B denotes the batch size), and the matrix of word representations, of size d× k. For a fixed size d of the hidden layer, we denote by g(k,B, d) the complexity of this multiplication, and simplify the notation wherever some parameters are fixed. Figure 1 reports empirical timings as a function of k for typical parameters of B and d for two GPU models, namely K40 and Maxwell. We observe that the complexity g(k) is constant for low values of k, until a certain inflection point k0 ≈ 50, and then becomes affine for values k > k0. This suggests a cost function of the form
g(k) = max(c+ λk0, c+ λk) = cm +max [ 0, λ(k − k0) ] . (5)
Empirically, cm = 0.40ms on a K40 and 0.22 ms on a Maxwell. We observe the same behavior when measuring the timings as a function of the other parameters, i.e., it is inefficient to matrix-multiply when one of the dimensions is small. This observation suggests that hierarchical organizations of words with a low number of children per node, such as binary Huffman codes, are highly suboptimal.
4.2 INTUITION: THE TWO-CLUSTERS CASE
In natural languages, the distribution of the words notoriously follows a Zipf law (Zipf, 1949). Most of the probability mass is covered by a small fraction of the dictionary, e.g., 87% of the document is covered by only 20% of the vocabulary in the Penn TreeBank. Similar to the frequency binning hierarchical softmax (Mikolov et al., 2011c), this information can be exploited to reduce the computation cost.
A simple strategy to reduce the overall complexity is to partition the dictionary V into two clusters as Vh and Vt, where Vh denotes the head of the distribution consisting of the most frequent words, and where Vt is the tail associated with a large number of rare words. The classifier frequently accesses the head, which motivates the fact that it should be computed efficiently. In contrast, the tail occurs
������
�� �
��� �
���� �
��
� �� �� �� �� �� �� �� �� ��� ��� ��� ��� ��� ���
���� ����
���� ������
���� ���
�������������������
��� �������
Figure 1: GPU timings for multiplying two matrices in the dominant step of the RNN model. We consider matrices of size 2560 × 2048 and 2048 × k representing hidden states and word representations. We report the timings as a function of k (number of word representations), and we compute the averages (circles) over 1000 measures, and the minima and maxima for the K40. The standard deviation does not exceed 5% of each timing.
less frequently and the corresponding computation can be slower. This suggests to define clusters with unbalanced cardinalities |Vh| |Vt| and probabilities P (Vh) P (Vt), where P (A) = ∑ w∈A pi is the probability of a word to occur in the set Vi. For instance, one may define the head would only contain 20% of the vocabulary (covering for 87% on PennTree Bank). These two clusters can be organized in two different ways: either they are both leaves of a 2-level tree (Mikolov et al., 2011c), or the head cluster is kept as a short-list in the root node (Le et al., 2011). We now analyze what is the best structure and how to split the vocabulary by determining the corresponding complexities, assuming that the head consists of the most frequent words. The next subsection shows the optimality of this choice.
Given a vocabulary of k words, we are looking for the number kh = |Vh| of words from the head of the distribution to be assigned to the first cluster. These words will cover for ph of the distribution. The tail cluster will then contain the rest of the vocabulary, made of kt = k − kh words and covering for pt = 1− ph of the overall distribution. We denote by g(k, d) the computational complexity of computing the softmax function over k words with d dimensional input features. Figure 1 shows an example of this function for a fixed d. The complexity of putting the head of the distribution in the root of the tree is g(kh + 1, d) + ptg(kt, d), while the complexity associated with putting both cluster in leaves is g(2, d) + phg(kh, d) + ptg(kt, d). Depending on the distribution of a corpus, it is then simple to choose the best assignment of words into the two clusters. For example, on PennTree Bank, with a hidden layer of size d = 128, the optimal configuration is to keep a short-list of 1400 classes in the root node, leading to an average cost of 0.33 ms per batch of size 512, while it takes 0.36 ms when both clusters are in the leaves. In comparison, the full softmax takes 0.80 ms for the same configuration, leading to a ×2.4 speed-up.
Adapting the classifier capacity for each cluster. Each cluster is accessed independently of each other, they thus do not need to have the same capacity. Frequent words need high capacity to be predicted correctly. In contrast, rare words cannot be learned very well, since we only see them a few times. It would then be wasteful to associate them with high capacity. Like in Chen et al. (2015), we exploit this observation to further reduce the computational cost of our classifier. Unlike Chen et al. (2015), we share the state of hidden layer across clusters and simply reduce the input size of the classifiers by applying a projection matrix. Typically, the projection matrix for the tail cluster reduces the size from d to dt = d/4, reducing the complexity from g(kt, d) to g(dt, d) + g(kt, dt).
Compromising between efficiency and accuracy. We observe empirically that putting all the clusters in the leaves of the tree leads to a significant drop of performance (around 5− 10% performance drop, Mikolov et al., 2011c; Zweig & Makarychev, 2013). The reason is that the probability of every word w belonging to a cluster c is multiplied by the probability of its class, i.e., it is equal to P (c | h)P (w | c, h), while attaching a frequent word directly to the root associates it directly to the probability P (w | h) making its inference sharper. For this reason, unless there is a significant difference in computational complexity, we favor using a short-list, over the standard 2-level hierarchical softmax.
Vh
V1
V2
V3
4.3 GENERAL CASE
Let us now consider the more general case where the dictionary is partitioned as V = Vh ∪ V1 . . .VJ , Vi ∩ Vj = ∅ if i 6= j. We consider the hierarchical model depicted in Figure 2, where the subdictionary Vh is accessed at the first level, and the others in the second level. We now investigate the computational cost C of the forward (equivalently, backward) pass of this approximate softmax layer. For the time being, we fix the batch size B and the dimensionality d of the hidden layer, in order to analyze the complexity as a function of the sub-dictionary sizes and probabilities. We denote by pi = ∑ w∈Vi p(w) the probability P (w ∈ Vi) and ki = |Vi| the cardinality of each cluster.
The expected cost C is decomposed as C = Ch + ∑ i Ci, where
Ch = ph g(J + kh) and ∀i, Ci = pi [ g(J + kh) + g(ki) ] , (6)
leading to C = g(J + kh) + ∑ i pi g(ki). (7)
We add the constraint k ≥ k0 to ensure that there is no penalty induced by the constant part of the cost model of Equation 5, the previous equation simplifies as
C = c+ λ(J + kh) + ∑ i pi(c+ λki) (8)
= c(2− ph) + λ [ J + kh + ∑ i pi ki ] . (9)
Let us discuss this equation, by first considering that the cardinalities of the sub-vocabularies are fixed. The right-most term is the only one that depends on the word probabilities. For two distinct clusters Vi and Vj , we can re-write pjkj as (pi+j − pi)kj , where pi+j = pi + pj , so that
piki + pjkj = pi(ki − kj) + pi+jkj . (10)
Without loss of generality, we assume that ki > kj . The quantities pi+j , ki and kj being fixed, the second term of the right-hand side of this equation is constant, and the best strategy is trivially to minimize the probability of the largest cluster Vi. In other terms, an optimal solution for Equation 9 requires that the most frequent words are assigned to the smallest cluster. This remark is true for any tuple (i, j), and we easily see that this point also holds for the head cluster. As a consequence, for a fixed number of clusters of given sizes, the best strategy is to assign the words by decreasing probabilities to clusters of increasing size. Note, this analysis remains valid as long as the g is monotonically increasing in k.
Determining ki with J fixed: dynamic programming. We now assume that the number of clusters is fixed. Following our analysis above, the optimization solely depends on the cardinalities ki for all clusters, which perfectly determines how to split the list of words ordered by frequency. We solve this problem by dynamic programming.
Finding the number of clusters. The only remaining free variable in our optimization is J , since the other parameters are then determined by the aforementioned optimizations. For this step, the cost of Equation 9 over-estimates the number of clusters because we have neglected the effect of the non-linearity of the batch size: in the second layer, the batches are typically smaller than the inflection
full softmax 144 83 min sampling 166 41 min HSM (freq) 166 34 min D-softmax 195 53 min D-softmax [*] 147 54 min
Ours 147 30 min
point k0. In practice, we optimize over small values of J = 1, 2, 3, 4 and empirically determine the best compromise speed/perplexity on training data. Note, having a lower number of clusters with numerous frequent words on the first level has another flavor: we empirically observe that it offers a better perplexity than word hierarchy with a large number of clusters. It is comparable to that of the exact softmax for large corpora, as shown later by our experiments.
5 EXPERIMENTS
This section provides a set of experiments aiming at analyzing the trade-off between actual complexity and effectiveness of several strategies, in particular the approach presented in the previous section. First we describe our evaluation protocol, then we evaluate some of the properties of our model and finally we compare it on standard benchmark against standard baselines.
Datasets. We evaluate our method on standard datasets, and use the perplexity (ppl) as an evaluation metric, as the function of the training time or of the number of training data (epochs). The datasets have varying vocabulary sizes, in different languages, which allows us to better understand the strengths and weaknesses of the different approaches.
• Text81 is a standard compression dataset containing a pre-processed version of the first 100 million characters from Wikipedia in English. It has been recently used for language modeling (Mikolov et al., 2014) and has a vocabulary of 44k words.
• Europarl2 is a machine translation corpus, containing 20 languages (Koehn, 2005). For most languages, there are 10M–60M tokens and the vocabulary is in between 44k and 250k words.
• One Billion Word 3 is a massive corpus introduced by Chelba et al. (2013). It contains 0.8B tokens and a vocabulary comprising almost 800k words.
Implementation details. We use an LSTM with one layer in all our experiments. On Text8 and Europarl, the models have d = 512 hidden units and are regularized with weight decay (λ = 10−6).
1http://mattmahoney.net/dc/textdata 2http://www.statmt.org/europarl/ 3https://code.google.com/archive/p/1-billion-word-language-modeling-benchmark/
On the One Billion Word benchmark, we use d = 2048 hidden units and no regularization. The dimension of the input word embeddings is set to 256, so that large models fit in GPU memory. For the backpropagation through time, we unroll the models for 20 steps. We use Adagrad (Duchi et al., 2011), with a step size of 0.1 and 5 epochs, and we clip the norm of the gradients to 1. The batch size B is set to 128, except on the Finnish portion of Europarl where B=64 due to memory constraints. All the experiments were run on the same GPU with the Maxwell architecture.
Baselines. Our method is compared to: (1) the full softmax, (2) the hierarchical softmax (HSM) with frequency binning (Mikolov et al., 2011b), (3) importance sampling (Bengio et al., 2003b; Bengio & Senécal, 2008) and (4) the differentiated softmax (Chen et al., 2015). For HSM, we tried different strategies for the binning. We observe that using the square root function on the count before computing the word bins is the most efficient. For the negative sampling method, we used a number of samples equal to 20% of the size of the vocabulary (Chen et al., 2015). For the differentiated softmax (D-softmax), we used the same partitions for the vocabulary as for our approach. We tried two version of the differentiated softmax. The first is the one described by Chen et al. (2015), where each word cluster uses a disjoint subset of the hidden representation. We also present an improved version, referred to as D-softmax [*], which uses our choice to have the whole hidden representation mapped to the different word clusters using projection matrices of different sizes.
Comparison with the state of the art. Table 1 reports the results that we achieve on Text8. On this small vocabulary, approximate methods are comparatively less interesting. Our approach is the only one to approach the result of the full soft-max (below by 3 points of perplexity), while being the fastest. Our improved variant D-softmax [*] of the work by Chen et al. (2015) obtains similar results but is slower by a factor ×1.8. On Europarl, we first present the convergence properties of our approach compared to other approximate strategies in Figure 3 show the perplexity (ppl) as a function of training time. Our approach significantly outperforms all competitors by a large margin. For reference, we also show the performance (D-softmax [*]) obtained by improving the D-softmax, to make it more comparable to our method. Our method is 2× to 3× faster than this improved competitor, which demonstrates how critical is our optimization strategy. Similar conclusions are drawn from Table 3 for other languages from the Europal corpus.
Table 2 gives the test perplexity on One Billion Word benchmark: Our method achieves a perplexity of 43.9 after five epochs, taking less than three days to train on a single GPU. In comparison, only Jozefowicz et al. (2016) achieves a lower perplexity, but with a model 8× bigger than ours and trained over 32 GPUs during 3 weeks. We also note that for models of similar size, we achieve similar perplexity than the method introduced by Jozefowicz et al. (2016). As far as we know, ours the first method to achieve a perplexity lower than 50 on a single GPU.
6 CONCLUSION
In this paper, we have proposed a simple yet efficient approximation of the softmax classifier. To our knowledge, it is the first speed optimizing approximation that obtains performance on par with the exact model. This is achieved by explicitly taking into account the computational complexity of parallel systems and combining it with a few important observations, namely keeping a short-list of frequent words in the root node (Schwenk, 2007) and reducing the capacity of rare words (Chen et al., 2015). In all our experiments on GPU, our method consistently maintains a low perplexity while enjoying a speed-up going from 2× to 10× compared to the exact model. This type of speed-up allows to deal with extremely large corpora in reasonable time and without the need of a large number of GPUs. We believe our approach to be general enough to be applied to other parallel computing architectures and other losses, as well as to other domains where the distributions of the class are unbalanced.
ACKNOWLEDGMENTS
The authors would like to thank Jeff Johnson for his help with GPU benchmarking and Tomas Mikolov for insightful discussions.

he authors provide an interesting, computational-complexity-driven approach for efficient softmax computation for language modeling based on GPUs. An adaptive softmax approach is proposed based on a hierarchical model. Dynamic programming is applied to optimize the structure of the hierarchical approach chosen here w.r.t. computational complexity based on GPUs.
However, it remains unclear, how robust the specific configuration obtained from dynamic programming is w.r.t. performance/perplexity. Corresponding comparative results with perplexity-based clustering would be desirable. Especially, in Sec. 5, Paragraph Baselines, and Table 1, respectively, it would be interesting to see a result on HSM(PPL) (cf. Zweig et al. 2013).
AFAIK, the first successful application of an LSTM-based language model for large vocabulary was published by Sundermeyer et al. 2012 (see below), which is missing in the sumary of prior work on the bottom of p. 3.
Mainly, the paper is well written and accessible, though notation in some cases should be improved, see detailed comments below.
Prior work on LSTM language modeling:
 - Sundermeyer et al.: LSTM Neural Networks for Language Modeling, Interspeech, pp. 194-197, 2012.
Notation:
 - use of g(k) vs. g(k,B,d): g(k) should be clearly defined (constant B and d?)
 - notation should not be reused (B is matrix in Eq. (3), and batch size in Sec. 4.1).
 - notation p_{i+j} (Eq. (10) and before) is kind of misleading, as p_{i+j} is not the same as p_{(i+j)}
Minor comments:
 - p. 1, item list at bottom, first item: take -> takes
 - p. 5, second paragraph: will then contained -> will then contain
 - p. 5, third paragaph: to associated -> to associate
 - Sec. 4.3, first paragraph: At the time being -> For the time being
 - below Eq. (9): most-right -> right-most
 - below Eq. (10): the second term of this equation -> the second term of the right-hand side of this equation
 - p. 6, second to last line: smaller that the -> smaller than the
 - p. 7, Sec. 5, itemize, first item: 100 millions -> 100 million
 - p. 8, last sentence: we are the -> ours is the