L-SR1: A SECOND ORDER OPTIMIZATION METHOD
We describe L-SR1, a new second order method to train deep neural networks. Second order methods hold great promise for distributed training of deep networks. Unfortunately, they have not proven practical. Two significant barriers to their success are inappropriate handling of saddle points, and poor conditioning of the Hessian. L-SR1 is a practical second order method that addresses these concerns. We provide experimental results showing that L-SR1 performs at least as well as Nesterov’s Accelerated Gradient Descent, on the MNIST and CIFAR10 datasets. For the CIFAR10 dataset, we see competitive performance on shallow networks like LeNet5, as well as on deeper networks like residual networks. Furthermore, we perform an experimental analysis of L-SR1 with respect to its hyperparameters to gain greater intuition. Finally, we outline the potential usefulness of L-SR1 in distributed training of deep neural networks.
1 MOTIVATION
Second order methods hold great potential for distributing the training of deep neural networks. Due to their use of curvature information, they can often find good minima in far fewer steps than first order methods such as stochastic gradient descent (SGD). Moreover, stochastic second order methods can benefit from larger mini-batches (Le et al., 2011). This is because they estimate second derivatives via differences between estimated gradients. The gradient estimates need to have less variance, so that when we take their differences, the result has low variance. As a result they provide a different trade-off between number of steps and mini-batch size than do SGD-like methods. This trade-off is interesting, because while steps must be evaluated sequentially, a mini-batch may be evaluated in parallel. Thus, second order methods present an opportunity to extract more parallelism in neural network training. In particular, when mini-batches are sufficiently large, their evaluation may be distributed. Furthermore, there are relatively fewer hyperparameters to tune in second order methods, compared to variants of stochastic gradient descent.
L-BFGS (Nocedal, 1980; Liu & Nocedal, 1989) is perhaps the most commonly used second order method in machine learning. BFGS is a quasi-Newton method that maintains an approximation to the inverse Hessian of the function being optimized. L-BFGS is a limited memory version of BFGS that stores the most recent updates to the inverse Hessian approximation and can therefore be used practically for large scale problems. L-BFGS is typically combined with a line search technique to choose an appropriate step size at each iteration. L-BFGS has been used to good effect in convex optimization problems in machine learning, but has not found effective use in large scale non-convex problems such as deep learning.
Three critical weaknesses have been identified. First, we know that training deep neural networks involves minimizing non-convex error functions over continuous, high dimensional spaces. It has been argued that the proliferation of saddle points in these problems presents a deep and profound difficulty for quasi-Newton optimization methods (Dauphin et al., 2014). Furthermore, it has been argued that curvature matrices generated in second order methods are often ill-conditioned, and these need to be carefully repaired. A variety of approaches to this have been suggested, including the use of an empirical Fisher diagonal matrix (Martens, 2016). Finally, popular quasi-Newton
approaches, such as L-BFGS (in their default form), require line search to make parameter updates, which requires many more gradient and/or function evaluations.
We propose L-SR1, a second order method that addresses each of these concerns. SR1 (Symmetric Rank One) is a quasi-Newton method that uses a rank one update for updating the Hessian approximation of the function being optimized (Nocedal & Wright, 2006). Unlike BFGS, the SR1 update does not guarantee positive definiteness of the updated matrix. This was considered a major problem in the early days of nonlinear optimization when only line search iterations were used, and possibly led to the obscurity of SR1 outside the optimization community. However, with the development of trust-region methods, the SR1 updating formula is potentially very useful, and its ability to generate indefinite Hessian approximations can actually prove to be advantageous.
We believe that it is possible to overcome saddle points using rank-one update based second order methods. The more common rank-two methods, e.g. L-BFGS, maintain a positive definite approximation to the inverse of the Hessian, by design (Nocedal & Wright, 2006). At saddle-points, the true Hessian cannot be well approximated by a positive definite matrix, causing commonly used second order methods to go uphill (Dauphin et al., 2014). On the other hand, rank-one approaches such as SR1 don’t maintain this invariant, so they can go downhill at saddle points. Numerical experiments (Conn et al., 1991) suggest that the approximate Hessian matrices generated by the SR1 method show faster progress towards the true Hessian than those generated by BFGS. This suggests that a limited memory SR1 method (L-SR1, if you like) could potentially outperform L-BFGS in the task of high dimensional optimization in neural network training. The building blocks needed to construct an L-SR1 method have been suggested in the literature (Byrd et al., 1994; Khalfan et al., 1993). To the best of our knowledge, however, there is no complete L-SR1 method previously described in the literature 1. This prompted us to develop and test the approach, specifically in the large scale non-convex problems that arise in deep learning.
Two other insights make L-SR1 practical by removing the requirement for a line search and addressing the conditioning problem. First, we replace the line search using a trust region approach. While L-BFGS using line search is well studied, recently, an L-BFGS method that uses a trustregion framework has also been proposed (Burke et al., 2008). Second, we combine L-SR1 with batch normalization. Batch normalization is a technique of normalizing inputs to layers of a neural network, used to address a phenomenon known as internal covariate shift during training (Ioffe & Szegedy, 2015). Our hypothesis is that batch normalization may cause parameters of a neural network to be suitably scaled so that the Hessian becomes better conditioned. We tested this hypothesis empirically and outline the results below.
2 RELATED WORK
We now briefly summarize some other second order approaches that have been suggested in the literature, in order to place our approach in context. Pearlmutter (1994) derived a technique that directly calculated the product of the Hessian with an arbitrary vector, and applied this technique to a few variants of backpropagation, thereby showing a way to use the full Hessian without needing to compute and store it. Martens (2010) used a generalization of this technique, introduced by Schraudolph (2002), to develop a second order optimization method based on the “Hessian-free” approach, using it to train deep auto-encoders (Martens, 2010), as well as recurrent neural networks (Martens & Sutskever, 2011). The “Hessian-free” approach is essentially a line search Newton-CG (Conjugate Gradient) method, also known as the truncated Newton method (Nocedal & Wright, 2006), in which the search direction is computed by applying CG to the Newton method, and terminating it once it has made sufficient progress. This approach differs from ours in its use of line search instead of a trust region method. Moreover, it computes Hessian-vector products using finite differencing, as opposed to the limited-memory symmetric rank one update with trust region method, used in our approach. The cost of skipping the Hessian calculation in a truncated Newton method is one additional gradient evaluation per CG iteration (Nocedal & Wright, 2006). As mentioned previously, Dauphin et al. (2014) argue, that in high dimensional problems of practical interest, the proliferation of saddle points poses greater difficulty than local minima. In a bid to escape these saddle points, they propose second order optimization method called the saddle-free Newton method. Key to this
1 The reference Brust et al. (2016) describes an approach to solve the trust region sub-problem encountered in an L-SR1 method, but does not describe the L-SR1 method itself.
approach is the definition of a class of generalized trust region methods. This class extends classical trust region methods in a couple of ways. A first order Taylor expansion of the function is minimized, instead of the second order Taylor expansion. Moreover, the constraint on the step norm is replaced by generalized constraint on the distance between consecutive iterates. Our approach, by contrast, uses a a classical trust-region method. Rather than compute the Hessian exactly, Dauphin et al. (2014) use an approach similar Krylov subspace descent (Vinyals & Povey, 2012). The function is optimized in a lower-dimensional Krylov subspace, which is determined through Lanczos iteration of the Hessian (Vinyals & Povey, 2012). The Lanczos method may be considered a generalization of the CG method that can be applied to indefinite systems, and may be used to aid the CG method by gathering negative curvature information (Nocedal & Wright, 2006). The Lanczos method also involves finding an approximate solution to a trust-region subproblem in the range of a Krylov basis that it generates. This trust region problem differs from the one we solve, in that the Krylov basis generated has a special structure due to its mapping to a tridiagonal matrix (Nocedal & Wright, 2006).
It is worth noting that several approaches have been proposed to overcome the weaknesses of LBFGS. First, it has been proposed to initialize L-BFGS with a number of SGD steps. However, this diminishes the potential for parallelism (Dean et al., 2012; Le et al., 2011). Second, it has been proposed to use “forgetting”, where every few (say, for example, 5) steps, the history for L-BFGS is discarded. However, this greatly reduces the ability to use second order curvature information. There has also been a recent spurt of work on stochastic quasi-Newton methods for optimization. Byrd et al. (2016) propose a stochastic quasi-Newton method which uses the classical L-BFGS formula, but collects curvature information pointwise, at regular intervals, through sub-sampled Hessian vector products, rather than at every iteration. Mokhtari & Ribeiro (2014) propose RES, a regularized stochastic version of BFGS to solve convex optimization problems with stochastic objectives, and prove its convergence for bounded Hessian eigenvalues. Mokhtari & Ribeiro (2015) propose an online L-BFGS method for solving optimization problems with strongly convex stochastic objectives, and establish global almost sure convergence of their approach for bounded Hessian eigenvalues of sample functions. In the case of nonconvex stochastic optimization, Wang et al. (2014) propose, based on a general framework, two concrete stochastic quasi-Newton update strategies, namely stochastic damped-BFGS update and stochastic cyclic Barzilai-Borwein-like update, to adaptively generate positive definite Hessian approximations. They also analyze the almost sure convergence of these updates to stationary points. Keskar & Berahas (2015) propose ADAQN, a stochastic quasiNewton algorithm for training RNNs. This approach retains a low per-iteration cost while allowing for non-diagonal scaling through a stochastic L-BFGS updating scheme. The method also uses a novel L-BFGS scaling initialization scheme and is judicious in storing and retaining L-BFGS curvature pairs. Finally, Curtis (2016) proposes a variable-metric algorithm for stochastic nonconvex optimization which exploits fundamental self-correcting properties of BFGS-type updating, and uses it to solve a few machine learning problems. As one may notice, all of these approaches adapt the BFGS-style rank two updates in different ways to solve convex and non-convex problems. In contrast, our approach uses SR1-type updates, which we think can help better navigate the pathological saddle points present in the non-convex loss functions found in deep learning, by not constraining the Hessian approximation to be positive definite, as in the case of BFGS-style updates. Comparison of our approach with one of these recent stochastic second order methods is an interesting next step. In the Appendix, we provide a brief primer on line search and trust region methods, as well as on quasi-Newton methods and their limited memory variants.
3 ALGORITHM
Our algorithm is synthesized as follows. We take the basic SR1 algorithm described in Nocedal & Wright (2006) (Algorithm 6.2), and represent the relevant input matrices using the limited-memory representations described in Byrd et al. (1994). The particular limited-memory representations used in the algorithm vary, depending on whether we use trust region or line search methods as subroutines to make parameter updates, as does some of the internal logic. For instance, if k updates are made to the symmetric matrix B0 using the vector pairs {si, yi}k−1i=0 and the SR1 formula, the resulting matrix Bk can be expressed as (Nocedal & Wright, 2006)
Bk = B0 + (Yk −B0Sk)(Dk + Lk + LTk − STk B0Sk)−1(Yk −B0Sk)T
where Sk, Yk, Dk, and Lk are defined as follows:
Sk = [so, · · · , sk−1], andYk = [y0, · · · , yk−1] ,
(Lk)i,j = { sTi−1yj−1 if i > j 0 otherwise
Dk = diag[sT0 y0, · · · , sTk−1yk−1] The self-duality of the SR1 method (Nocedal & Wright, 2006) allows the inverse formula Hk to be obtained simply by replacing B, s, and y by H , y, and s, respectively, using standard matrix identities. Limited-memory SR1 methods can be derived exactly like in the case of the BFGS method. Additional details are present in the pseudocode provided in the Appendix. The algorithm we develop is general enough to work with any line search or trust region method. While we tested the algorithm with line search approaches described in Dennis Jr. & Schnabel (1983), and with the trust region approach described in Brust et al. (2016), in this paper, we focus our experimental investigations on using the trust region approach, and the advantage that provides over using other first and second order optimization methods.
We also make a note here about the space and time complexity of our algorithm. We respectively denote by m and n, the memory size, and parameter dimensions. We assume m << n. As discussed in Section 7.2 of Nocedal & Wright (2006), the limited-memory updating procedure of Bk requires approximately 2mn+O(m3) operations, and matrix vector products of the form Bkv can be performed at a cost of (4m+ 1)n+O(m2) multiplications. Moreover, the Cholesky and eigenvalue decompositions we perform within our trust-region method form×mmatrices requireO(m3) operations. It follows quite easily2 from this that the space complexity of our algorithm is O(mn), and the per iteration time complexity of our algorithm is O(mn).
4 EXPERIMENTS
In the following, we summarize the results of training standard neural networks on the MNIST and CIFAR10 datasets using our approach, and benchmarking the performance with respect to other first and second order methods. First, we compared our L-SR1 (with trust region) approach, with Nesterov’s Accelerated Gradient Descent (NAG), L-BFGS with forgetting every 5 steps, default SGD, AdaDelta, and SGD with momentum, by training small standard networks on the MNIST and CIFAR10 datasets. On these problems, we also studied the effect of varying the minibatch size, for L-SR1, Adam (Kingma & Ba, 2014), and NAG. Next, we compared our L-SR1 with trust region approach with default hyperparameters, with a benchmark SGD with momentum, and Adam, by training a 20-layer deep residual network on the CIFAR10 dataset. Following that, we varied each hyperparameter of the L-SR1 with trust region approach to observe its effect on training the residual network on CIFAR10.
4.1 LENET-LIKE NETWORKS
For each approach, and for each dataset, we considered the case where our networks had batch normalization layers within them, and the case where they did not. The parameters of the networks were randomly initialized. All experiments were repeated 10 times to generate error bars.
4.1.1 MNIST
We considered the LeNet5 architecture in this case, which comprised 2 convolutional layers, followed by a fully connected layer and an outer output layer. Each convolutional layer was followed by a max-pooling layer. In the case where we used batch-normalization, each convolutional and fully connected layer was followed by a spatial batch normalization layer. We used a mini-batch size of 20 for the first order methods like NAG, SGD, AdaDelta and SGD with momentum, and a mini-batch size of 400 for the second order methods like L-SR1 and L-BFGS. The memory size was set to 5 for both L-SR1 and L-BFGS. The networks were trained for 20 epochs. Further details on the network architecture and other parameter settings are provided in the Appendix.
2Deep neural networks typically have paramater dimensions in the tens of millions, while the memory size typically does not exceed 10. So n is indeed several orders of magnitude larger than m.
4.1.2 CIFAR10
We considered a slight modification to the ‘LeNet5’ architecture described above. We used a minibatch size of 96 for NAG, SGD, AdaDelta and SGD with momentum. The other mini-batch sizes and memory sizes for L-SR1 and L-BFGS were as above. As above, the networks were trained for 20 epochs. Further details on the network architecture and other parameter settings are provided in the Appendix.
4.1.3 VARIATION OF MINIBATCH SIZE
We also compared the variation of test loss between L-SR1, Adam and NAG, as we varied the mini-batch size from 500 to 1000 to 10000, in the presence of batch normalization. The network architectures were as above. For minibatch sizes 500 and 1000, we trained the networks for 50 epochs, while for the minibatch size of 10000, the networks were trained for 200 epochs.
4.1.4 DISCUSSION
Our first set of experiments (Figures 1, 2) suggest that L-SR1 performs as well as, or slightly better than all the first order methods on both the MNIST and CIFAR10 datasets, with or without batch normalization. L-SR1 is substantially better than L-BFGS in all settings, with or without forgetting. Forgetting appears to be necessary in order to get L-BFGS to work. Without forgetting, the approach appears to be stuck where it is initialized. For this reason, the plots for L-BFGS without forgetting have not been included. Batch normalization appears to improve the performance of all approaches, particularly the early performance of second order approaches like L-SR1 and L-BFGS.
The experiments with variation of minibatch sizes (Figures 3, 4), seem to provide compelling evidence of the potential for distributed training of deep networks, as may be seen from Table 1. First, we note that first order methods like NAG are not as sensitive to size of the minibatch, as commonly understood. For example, a 20 fold increase in minibatch size did not decrease the speed of convergence by the same or higher order of magnitude. Furthermore, approaches like L-SR1 and Adam appear to be much less sensitive to increasing minibatch size than NAG. This strengthens the case for their application to distributed training of deep neural networks. Finally, while Adam makes much faster initial progress than the other approaches, its final test loss by the end of training is worse than in the case of L-SR1.
One of the limitations of SR1 updating is that the denominator in the update can vanish. The literature however suggests that this happens rarely enough that the updates can be skipped when this phenomenon occurs, without affecting performance. In this regard, we had some interesting observations from our experiments. While in most cases, updates were either never skipped, or skipped less than 2.5% of the time, the cases of MNIST training with batch normalization, yielded abnor-
mally high levels of skipped updates, ranging all the way from 7% to higher than 60% (for minibatch size 10000). While this did not seem to affect performance adversely, it certainly warrants future investigation. Moreover, a better understanding of the interplay between batch normalization and optimization could help inform potential improvements in optimization approaches.
4.2 RESIDUAL NETWORKS
We next considered a deeper residual network architecture described in section 4.2 of He et al. (2015b), with n = 3. This led to a 20-layer residual network including 9 shortcut connections. As in He et al. (2015b), we used batch normalization (Ioffe & Szegedy, 2015) and the same initialization method (He et al., 2015a).
4.2.1 COMPARISON WITH SGD WITH MOMENTUM, AND ADAM
We trained the residual network using the benchmark SGD with momentum, and other parameter settings as described in He et al. (2015b). We also trained the network using L-SR1 with default settings. These included, a memory size of 5, a trust-region radius decrease factor of 0.5, and a trust-region radius increase factor of 2.0. Finally, we also compared with Adam, with default settings (Kingma & Ba, 2014). We used the same mini-batch size of 128 for all algorithms. Based on the learning rate schedule used, the learning rate was equal to 0.1 through the first 80 epochs, 0.01 up to 120 epochs, and 0.001 thereafter, for SGD with momentum. Figure 5 shows variation of test loss, over epochs, and by time. It needs to be noted that default L-SR1, with no parameter tuning at all, has a superior final test loss to Adam, and is competitive with SGD with momentum, which used custom parameters that were tuned carefully. L-SR1 does make slower progress over time, which can be further optimized. Finally, we note that the test loss for L-SR1 bounces around a lot more than the test loss for the other algorithms. This bears further exploration.
4.2.2 VARIATION OF L-SR1 HYPERPARAMETERS
We varied the hyperparameters of L-SR1 in turn, keeping the remaining fixed. In each case, we trained the network for 200 epochs. We first considered varying the increase and decrease factors together. We considered a trust-region radius decrease factor of 0.2, 0.5 and 0.8, and a trust-region radius increase factor 1.2 and 2.0. The respective default values of these factors are 0.5 and 2.0 respectively. This led to six different combinations of decrease and increase factors. We kept the memory size and mini-batch size fixed at 5 and 128 respectively. Next, we considered memory sizes of 2 and 10 (in addition to 5, which we tried earlier), keeping the mini-batch size, decrease factor, and increase factor fixed at 128, 0.5, and 2.0 respectively. Finally, we considered mini-batch sizes of 512, 2048 and 8192 (in addition to 128, which we tried earlier), keeping the memory size, decrease factor, and increase factor fixed at 5, 0.5, and 2.0 respectively. Figure 6 shows the results.
The following may be noted, based on the experiments with L-SR1 for training a residual network on CIFAR10. While there is potential value in increasing and decreasing the trust region radius at different rates, our experiments suggest that it may not be necessary to tune these hyperparameters. There is no noticeable performance gain from using a higher memory size in L-SR1. Furthermore, using a smaller memory size performs at least as well as in the default case. This is good news, due to the consequent savings in storage and computational resources. L-SR1 is relatively insensitive to a 4-fold increase in mini-batch size from 128 to 512, and a further 4-fold increase to 2048. The minibatch sensitivity of L-SR1 seems to be higher in the case of the residual network, compared
with the Le-Net like networks seen earlier. Finally, we found the proportion of skipped updates in the case of residual networks to be less than 0.5% in all cases.
5 CONCLUSIONS
In this paper, we have described L-SR1, a new second order method to train deep neural networks. Our experiments suggest that this approach is at the very least, competitive, with other first order methods, and substantially better than L-BFGS, a well-known second order method. Our experiments also appear to validate our intuition about the ability of L-SR1 to overcome key challenges associated with second order methods, such as inappropriate handling of saddle points, and poor conditioning of the Hessian. Our experimentation with the hyperparameters of L-SR1 suggested that it is relatively robust with respect to them, and requires minimal tuning. Furthermore, we have evidence to suggest that L-SR1 is much more insensitive to larger minibatch sizes than a first order method like NAG. This suggests that L-SR1 holds promise for distributed training of deep networks, and we see our work as an important step toward that goal.
APPENDIX
BACKGROUND
In the following, we provide a brief primer on line search and trust region methods, as well as on quasi-Newton methods and their limited memory variants. Further details may be found in Nocedal & Wright (2006).
LINE SEARCH AND TRUST REGION METHODS
In any optimization algorithm, there are two main ways of moving from the current point xk to a new iterate xk+1. One of them is line search. In it, the algorithm picks a descent direction pk and searches along this direction from the current iterate xk for a new iterate with a lower function value. The distance α to move along pk can be found by solving the following one-dimensional minimization problem:
min α>0 f(xk + αpk)
Instead of an exact minimization which may be expensive, the line search algorithm generates a limited number of trial step lengths until it finds one that generates a sufficient decrease in function
value. At the new point, the process of computing the descent direction and step length is repeated. The other way is to use a trust region method. In a trust region method, the information about f is used to construct a model function mk, which is supposed to approximate f near the current point xk. Since the model mk may not approximate f well when x is far from xk, the search for a minimizer of mk is restricted to some trust region within a radius ∆k around xk. To wit, the candidate step p approximately solves the following sub-problem:
min p:‖p‖≤∆k mk(xk + p),
If the candidate solution does not produce a sufficient decrease in f , the trust region is considered too large for the model function to approximate f well. So we shrink the trust region and re-solve. Essentially, the line search and trust region approaches differ in the order in which they choose the direction and magnitude of the move to the next iterate. In line search, the descent direction pk is fixed first, and then the step length αk to be taken along that direction is computed. In trust region, a maximum distance equal to the trust-region radius ∆k is first set, and then a direction is determined within this radius, that achieves the best improvement in the objective value. If such a direction does not yield sufficient improvement, the model function is determined to be a poor approximation to the function, and the trust-region radius ∆k is reduced until the approximation is deemed good enough. Conversely, as long as the model function appears to approximate the objective function well, the trust region radius is increased until the approximation is not good enough.
LIMITED MEMORY QUASI-NEWTON METHODS
Quasi-Newton methods are a useful alternative to Newton’s method in that they do not require computation of the exact Hessian, and yet still attain good convergence. In place of the true Hessian ∇2fk, they use an approximation Bk, which is updated after each step based on information gained during the step. At each step, the new Hessian approximation Bk+1 is required to satisfy the following condition, known as the secant equation:
Bk+1sk = yk
where sk = xk+1 − xk, yk = ∇fk+1 −∇fk
Typically, Bk+1, is also required to be symmetric (like the exact Hessian), and the difference between successive approximations Bk and Bk+1 is constrained to have low rank. One of the most popular formulae for updating the Hessian approximation Bk is the BFGS formula, named after its inventors, Broyden, Fletcher, Goldfarb, and Shanno, which is defined by
Bk+1 = Bk − Bksks
T kBk
sTkBksk + yky
T k yTk sk
A less well known formula, particularly in the machine learning community, is the symmetric-rankone (SR1) formula, defined by
Bk+1 = Bk + (yk −Bksk)(yk −Bksk)T
(yk −Bksk)T sk The former update is a rank-two update, while the latter is a rank-one update. Both updates satisfy the secant equation and maintain symmetry. The BFGS update always generates positive definite approximations whenever the initial approximation B0 is positive definite and sTk yk > 0. Often, in practical implementations of quasi-Newton methods, the inverse Hessian approximation Hk is used instead of the Bk,and the corresponding update formulae can be generated using the ShermanMorrison-Woodbury matrix identity (Hager, 1989).
Limited-memory quasi-Newton methods are useful for solving large problems where computation of Hessian matrices is costly or when these matrices are dense. Instead of storing fully dense n× n approximations, these methods save only a few vectors of length n that capture the approximations. Despite these modest storage requirements, they often converge well. The most popular limited memory quasi-Newton method is L-BFGS, which uses curvature information from only the most recent iterations to construct the inverse Hessian approximation. Curvature information from earlier
iterations, which is less likely to be useful to modeling the actual behavior of the Hessian at the current iteration, is discarded in order to save memory.
Limited-memory quasi-Newton approximations can be used with line search or trust region methods. As described in Byrd et al. (1994), we can derive efficient limited memory implementations of several quasi-Newton update formulae, and their inverses.
NETWORK ARCHITECTURES AND HYPERPARAMETER SETTINGS
MNIST
The layers of the LeNet5 architecture used, are described below. All the batch normalization layers were removed, in the ‘without batch normalization’ case.
• Convolutional Layer - filter size 5 × 5, 20 feature maps, stride 1, padding 0, and a ReLU activation function with bias 0 and Gaussian noise with mean 0 and standard deviation 0.1
• Spatial Batch Normalization Layer • Max Pooling Layer - filter size 2 • Convolutional Layer - filter size 5 × 5, 50 feature maps, stride 1, padding 0, and a ReLU
activation function with bias 0 and Gaussian noise with mean 0 and standard deviation 0.1 • Spatial Batch Normalization Layer • Max Pooling Layer - filter size 2 • Fully Connected Layer - 500 hidden units, and a tangent hyperbolic activation function • Spatial Batch Normalization Layer • Outer Output Layer - 10 outputs and output standard deviation of 0.1
Additionally, the network was trained with L2 regularization with parameter 0.0001. Training loss was measured as softmax cross entropy, while test loss was measured as multi-class error count. In the case of the first order methods, the learning rate was set to 0.003 where needed, and the momentum was set to 0.9, where needed. AdaDelta did not take any parameters.
CIFAR10
The layers of the architecture used, are described below. All the batch normalization layers were removed, in the ‘without batch normalization’ case.
• Convolutional Layer - filter size 5 × 5, 32 feature maps, stride 1, padding 2, and a ReLU activation function with bias 0 and Gaussian noise with mean 0 and standard deviation 0.01
• Spatial Batch Normalization Layer • Max Pooling Layer - filter size 2 • Activation Layer - ReLU activation function with bias 0 and Gaussian noise with mean 0
and standard deviation 0.1 • Convolutional Layer - filter size 5 × 5, 32 feature maps, stride 1, padding 2, and a ReLU
activation function with bias 0 and Gaussian noise with mean 0 and standard deviation 0.01 • Spatial Batch Normalization Layer • Max Pooling Layer - filter size 2 • Convolutional Layer - filter size 5 × 5, 64 feature maps, stride 1, padding 2, and a ReLU
activation function with bias 0 and Gaussian noise with mean 0 and standard deviation 0.01 • Spatial Batch Normalization Layer • Max Pooling Layer - filter size 2 • Fully Connected Layer - 64 hidden units, and a ReLU activation function with bias 0 and
Gaussian noise with mean 0 and standard deviation 0.1 • Spatial Batch Normalization Layer
• Outer Output Layer - 10 outputs and output standard deviation of 0.1
Additionally, the network was trained with L2 regularization with parameter 0.001. Training loss was measured as softmax cross entropy, while test loss was measured as multi-class error count. In the case of the first order methods, the learning rate was set to 0.01 where needed, and the momentum was set to 0.9, where needed. AdaDelta did not take any parameters.
PSEUDOCODE
Algorithm 1 provides the pseudocode for L-SR1 with trust region method, while Algorithm 2 provides the pseudocode for L-SR1 with line search.
Algorithm 1 L-SR1 with Trust Region Method Require: Sk = [s0, · · · , sk−1], Yk = [y0, · · · , yk−1], starting point x0 ∈ Rn, limited memory size
m << n, initial Hessian approximation B0 (a diagonal matrix, typically γIn, γ 6= 0), initial trust-region radius ∆0 = ‖∇f(x0)‖2, convergence tolerance t > 0, maximum iterations K, parameters η ∈ (0, 10−3), r ∈ (0, 1), and column dimension colDim; 1: k ← 0 2: while k < K and ‖∇f(xk)‖2 > t and ‖sk‖2 > t do 3: if k = 0 or Sk.colDim = 0 then 4: sk ← −∇f(xk) 5: Bk ← B0 6: else 7: sk ← TrustRegionMethod(Ψk,M−1k ,∇f(xk),∆k, γ, B0) (Solve the trust-region sub-
problem) 8: Bksk ← B0sk + ΨkMk(ΨTk sk) 9: end if
10: pred← − ( ∇f(xk)T sk + 12s T kBksk ) (predicted reduction) 11: ared← f(xk)− f(xk + sk) (actual reduction) 12: yk ← ∇f(xk + sk)−∇f(xk) 13: if ared/pred > η then 14: xk+1 ← xk + sk 15: else 16: xk+1 ← xk 17: end if 18: ∆k+1 ← ∆k 19: if ared/pred > u then (u = 0.75 by default) 20: if ‖sk‖ > ρ∆k then (ρ = 0.8 by default) 21: ∆k+1 ← 2∆k 22: end if 23: else if ared/pred < l then (l = 0.1 by default) 24: ∆k+1 ← 0.5∆k 25: end if 26: if |sTk (yk −Bksk)| ≥ r‖sk‖‖yk −Bksk‖ then 27: Sk+1 ← [Sk, sk], Yk+1 ← [Yk, yk] 28: if Sk+1.colDim > m then 29: Sk+1 ← Sk+1[, 2 : m+ 1], Yk+1 ← Yk+1[, 2 : m+ 1] 30: end if 31: while Sk+1.colDim > 0 do 32: Ψk+1 ← Yk+1 −B0Sk+1 33: M−1k+1 ← STk+1Yk+1 − STk+1B0Sk+1 34: if (ΨTk+1Ψk+1 0 and |M −1 k+1| 6= 0) then 35: break 36: else 37: Remove the first columns of Sk+1 and Yk+1 38: end if 39: end while 40: end if 41: k ← k + 1 42: end while
Algorithm 2 L-SR1 with Line Search Require: Sk = [s0, · · · , sk−1], Yk = [y0, · · · , yk−1], starting point x0 ∈ Rn, limited memory size
m << n, initial inverse Hessian approximationH0 (a diagonal matrix, typically In), initial step length λ0 = 1, convergence tolerance t > 0, maximum iterations K, r ∈ (0, 1), and column dimension colDim;
1: k ← 0 2: while k < K and (‖∇f(xk)‖2 > t or ‖sk‖2 > t) do 3: if k = 0 or Sk.colDim = 0 then 4: dk ← −∇f(xk) 5: Hk ← H0 6: else 7: dk ← −H0∇f(xk)−ΨkMk(ΨTk∇f(xk)) 8: end if 9: if k > 0 then
10: λ0 = min{1, 2λk−1} 11: end if 12: λk ←computeStepLength(f, xk, dk, λ0) (perform line search) 13: sk ← λkdk 14: xk+1 ← xk + sk 15: yk ← ∇f(xk+1)−∇f(xk) 16: if k > 0 and Sk.colDim > 0 then 17: Hkyk ← H0yk + ΨkMk(ΨTk yk) 18: end if 19: if |yTk (sk −Hkyk)| ≥ r‖yk‖‖sk −Hkyk‖ then 20: Sk+1 ← [Sk, sk], Yk+1 ← [Yk, yk] 21: if Sk+1.colDim > m then 22: Sk+1 ← Sk+1[, 2 : m+ 1], Yk+1 ← Yk+1[, 2 : m+ 1] 23: end if 24: while Sk+1.colDim > 0 do 25: Ψk+1 ← Sk+1 −H0Yk+1 26: M−1k+1 ← Y Tk+1Sk+1 − Y Tk+1H0Yk+1 27: if (|M−1k+1| 6= 0) then 28: break 29: else 30: Remove the first columns of Sk+1 and Yk+1 31: end if 32: end while 33: end if 34: k ← k + 1 35: end while

The paper proposes a new second-order method L-SR1 to train deep neural networks. It is claimed that the method addresses two important optimization problems in this setting: poor conditioning of the Hessian and proliferation of saddle points. The method can be viewed as a concatenation of SR1 algorithm of Nocedal & Wright (2006) and limited-memory representations Byrd et al. (1994). First of all, I am missing a more formal, theoretical argument in this work (in general providing more intuition would be helpful too), which instead is provided in the works of Dauphin (2014) or Martens. The experimental section in not very convincing considering that the performance in terms of the wall-clock time is not reported and the advantage over some competitor methods is not very strong even in terms of epochs. I understand that the authors are optimizing their implementation still, but the question is: considering the experiments are not convincing, why would anybody bother to implement L-SR1 to train their deep models? The work is not ready to be published.


Recurrent neural networks have been very successful at predicting sequences of words in tasks such as language modeling. However, all such models are based on the conventional classification framework, where the model is trained against one-hot targets, and each word is represented both as an input and as an output in isolation. This causes inefficiencies in learning both in terms of utilizing all of the information and in terms of the number of parameters needed to train. We introduce a novel theoretical framework that facilitates better learning in language modeling, and show that our framework leads to tying together the input embedding and the output projection matrices, greatly reducing the number of trainable variables. Our framework leads to state of the art performance on the Penn Treebank with a variety of network models.
1 INTRODUCTION
Neural network models have recently made tremendous progress in a variety of NLP applications such as speech recognition (Irie et al., 2016), sentiment analysis (Socher et al., 2013), text summarization (Rush et al., 2015; Nallapati et al., 2016), and machine translation (Firat et al., 2016).
Despite the overwhelming success achieved by recurrent neural networks in modeling long range dependencies between words, current recurrent neural network language models (RNNLM) are based on the conventional classification framework, which has two major drawbacks: First, there is no assumed metric on the output classes, whereas there is evidence suggesting that learning is improved when one can define a natural metric on the output space (Frogner et al., 2015). In language modeling, there is a well established metric space for the outputs (words in the language) based on word embeddings, with meaningful distances between words (Mikolov et al., 2013; Pennington et al., 2014). Second, in the classical framework, inputs and outputs are considered as isolated entities with no semantic link between them. This is clearly not the case for language modeling, where inputs and outputs in fact live in identical spaces. Therefore, even for models with moderately sized vocabularies, the classical framework could be a vast source of inefficiency in terms of the number of variables in the model, and in terms of utilizing the information gathered by different parts of the model (e.g. inputs and outputs).
In this work, we introduce a novel loss framework for language modeling to remedy the above two problems. Our framework is comprised of two closely linked improvements. First, we augment the classical cross-entropy loss with an additional term which minimizes the KL-divergence between the model’s prediction and an estimated target distribution based on the word embeddings space. This estimated distribution uses knowledge of word vector similarity. We then theoretically analyze this loss, and this leads to a second and synergistic improvement: tying together two large matrices by reusing the input word embedding matrix as the output classification matrix. We empirically validate our theory in a practical setting, with much milder assumptions than those in theory. We also find empirically that for large networks, most of the improvement could be achieved by only reusing the word embeddings.
We test our framework by performing extensive experiments on the Penn Treebank corpus, a dataset widely used for benchmarking language models (Mikolov et al., 2010; Merity et al., 2016). We demonstrate that models trained using our proposed framework significantly outperform models
trained using the conventional framework. We also perform experiments on the newly introduced Wikitext-2 dataset (Merity et al., 2016), and verify that the empirical performance of our proposed framework is consistent across different datasets.
2 BACKGROUND: RECURRENT NEURAL NETWORK LANGUAGE MODEL
In any variant of recurrent neural network language model (RNNLM), the goal is to predict the next word indexed by t in a sequence of one-hot word tokens (y∗1 , . . . y ∗ N ) as follows:
xt = Ly ∗ t−1, (2.1)
ht = f(xt, ht−1), (2.2) yt = softmax (Wht + b) . (2.3)
The matrix L ∈ Rdx×|V | is the word embedding matrix, where dx is the word embedding dimension and |V | is the size of the vocabulary. The function f(., .) represents the recurrent neural network which takes in the current input and the previous hidden state and produces the next hidden state. W ∈ R|V |×dh and b ∈ R|V | are the the output projection matrix and the bias, respectively, and dh is the size of the RNN hidden state. The |V | dimensional yt models the discrete probability distribution for the next word.
Note that the above formulation does not make any assumptions about the specifics of the recurrent neural units, and f could be replaced with a standard recurrent unit, a gated recurrent unit (GRU) (Cho et al., 2014), a long-short term memory (LSTM) unit (Hochreiter & Schmidhuber, 1997), etc. For our experiments, we use LSTM units with two layers.
Given yt for the tth example, a loss is calculated for that example. The loss used in the RNNLMs is almost exclusively the cross-entropy between yt and the observed one-hot word token, y∗t :
Jt = CE(y∗t ‖ yt) = − ∑ i∈|V | y∗t,i log yt,i. (2.4)
We shall refer to yt as the model prediction distribution for the tth example, and y∗t as the empirical target distribution (both are in fact conditional distributions given the history). Since crossentropy and Kullback-Leibler divergence are equivalent when the target distribution is one-hot, we can rewrite the loss for the tth example as
Jt = DKL(y∗t ‖ yt). (2.5)
Therefore, we can think of the optimization of the conventional loss in an RNNLM as trying to minimize the distance1 between the model prediction distribution (y) and the empirical target distribution (y∗), which, with many training examples, will get close to minimizing distance to the actual target distribution. In the framework which we will introduce, we utilize Kullback-Leibler divergence as opposed to cross-entropy due to its intuitive interpretation as a distance between distributions, although the two are not equivalent in our framework.
3 AUGMENTING THE CROSS-ENTROPY LOSS
We propose to augment the conventional cross-entropy loss with an additional loss term as follows:
ŷt = softmax (Wht/τ) , (3.1) Jaugt = DKL(ỹt ‖ ŷt), (3.2) J tott = Jt + αJ aug t . (3.3)
In above, α is a hyperparameter to be adjusted, and ŷt is almost identical to the regular model prediction distribution yt with the exception that the logits are divided by a temperature parameter τ . We define ỹt as some probability distribution that estimates the true data distribution (conditioned on the word history) which satisfies Eỹt = Ey∗t . The goal of this framework is to minimize the
1We note, however, that Kullback-Leibler divergence is not a valid distance metric.
distribution distance between the prediction distribution and a more accurate estimate of the true data distribution.
To understand the effect of optimizing in this setting, let’s focus on an ideal case in which we are given the true data distribution so that ỹt = Ey∗t , and we only use the augmented loss, Jaug. We will carry out our investigation through stochastic gradient descent, which is the technique dominantly used for training neural networks. The gradient of Jaugt with respect to the logits Wht is
∇ Jaugt = 1
τ (ŷt − ỹt). (3.4)
Let’s denote by ej ∈ R|V | the vector whose jth entry is 1, and others are zero. We can then rewrite (3.4) as
τ ∇ Jaugt = ŷt − [ e1, . . . , e|V | ] ỹt = ∑ i∈V ỹt,i(ŷt − ei). (3.5)
Implication of (3.5) is the following: Every time the optimizer sees one training example, it takes a step not only on account of the label seen, but it proceeds taking into account all the class labels for which the conditional probability is not zero, and the relative step size for each step is given by the conditional probability for that label, ỹt,i. Furthermore, this is a much less noisy update since the target distribution is exact and deterministic. Therefore, unless all the examples exclusively belong to a specific class with probability 1, the optimization will act much differently and train with greatly improved supervision.
The idea proposed in the recent work by Hinton et al. (2015) might be considered as an application of this framework, where they try to obtain a good set of ỹ’s by training very large models and using the model prediction distributions of those.
Although finding a good ỹ in general is rather nontrivial, in the context of language modeling we can hope to achieve this by exploiting the inherent metric space of classes encoded into the model, namely the space of word embeddings. Specifically, we propose the following for ỹ:
ut = Ly ∗ t , (3.6) ỹt = softmax ( LTut τ ) . (3.7)
In words, we first find the target word vector which corresponds to the target word token (resulting in ut), and then take the inner product of the target word vector with all the other word vectors to get an unnormalized probability distribution. We adjust this with the same temperature parameter τ used for obtaining ŷt and apply softmax. The target distribution estimate, ỹ, therefore measures the similarity between the word vectors and assigns similar probability masses to words that the language model deems close. Note that the estimation of ỹ with this procedure is iterative, and the estimates of ỹ in the initial phase of the training are not necessarily informative. However, as training procedes, we expect ỹ to capture the word statistics better and yield a consistently more accurate estimate of the true data distribution.
4 THEORETICALLY DRIVEN REUSE OF WORD EMBEDDINGS
We now theoretically motivate and introduce a second modification to improve learning in the language model. We do this by analyzing the proposed augmented loss in a particular setting, and observe an implicit core mechanism of this loss. We then make our proposition by making this mechanism explicit.
We start by introducing our setting for the analysis. We restrict our attention to the case where the input embedding dimension is equal to the dimension of the RNN hidden state, i.e. d , dx = dh. We also set b = 0 in (2.3) so that yt =Wht. We only use the augmented loss, i.e. J tot = Jaug, and we assume that we can achieve zero training loss. Finally, we set the temperature parameter τ to be large.
We first show that when the temperature parameter, τ , is high enough, Jaugt acts to match the logits of the prediction distribution to the logits of the the more informative labels, ỹ. We proceed in the same
way as was done in Hinton et al. (2015) to make an identical argument. Particularly, we consider the derivative of Jaugt with respect to the entries of the logits produced by the neural network.
Let’s denote by li the ith column of L. Using the first order approximation of exponential function around zero (exp (x) ≈ 1 + x), we can approximate ỹt (same holds for ŷt) at high temperatures as follows:
ỹt,i = exp (〈ut, li〉/τ)∑
j∈V exp (〈ut, lj〉/τ) ≈ 1 + 〈ut, li〉/τ |V |+ ∑ j∈V 〈ut, lj〉/τ . (4.1)
We can further simplify (4.1) if we assume that 〈ut, lj〉 = 0 on average:
ỹt,i ≈ 1 + 〈ut, li〉/τ
|V | . (4.2)
By replacing ỹt and ŷt in (3.4) with their simplified forms according to (4.2), we get
∂Jaugt ∂ (Wht)i → 1 τ2|V |
( Wht − LTut ) i as τ →∞, (4.3)
which is the desired result that augmented loss tries to match the logits of the model to the logits of ỹ’s. Since the training loss is zero by assumption, we necessarily have
Wht = L Tut (4.4)
for each training example, i.e., gradient contributed by each example is zero. Provided thatW and L are full rank matrices and there are more linearly independent examples of ht’s than the embedding dimension d, we get that the space spanned by the columns of LT is equivalent to that spanned by the columns of W . Let’s now introduce a square matrix A such that W = LTA. (We know A exists since LT and W span the same column space). In this case, we can rewrite
Wht = L TAht , L T h̃t. (4.5)
In other words, by reusing the embedding matrix in the output projection layer (with a transpose) and letting the neural network do the necessary linear mapping h → Ah, we get the same result as we would have in the first place.
Although the above scenario could be difficult to exactly replicate in practice, it uncovers a mechanism through which our proposed loss augmentation acts, which is trying to constrain the output (unnormalized) probability space to a small subspace governed by the embedding matrix. This suggests that we can make this mechanism explicit and constrainW = LT during training while setting the output bias, b, to zero. Doing so would not only eliminate a big matrix which dominates the network size for models with even moderately sized vocabularies, but it would also be optimal in our setting of loss augmentation as it would eliminate much work to be done by the augmented loss.
5 RELATED WORK
Since their introduction in Mikolov et al. (2010), many improvements have been proposed for RNNLMs , including different dropout methods (Zaremba et al., 2014; Gal, 2015), novel recurrent units (Zilly et al., 2016), and use of pointer networks to complement the recurrent neural network (Merity et al., 2016). However, none of the improvements dealt with the loss structure, and to the best of our knowledge, our work is the first to offer a new loss framework.
Our technique is closely related to the one in Hinton et al. (2015), where they also try to estimate a more informed data distribution and augment the conventional loss with KL divergence between model prediction distribution and the estimated data distribution. However, they estimate their data distribution by training large networks on the data and then use it to improve learning in smaller networks. This is fundamentally different from our approach, where we improve learning by transferring knowledge between different parts of the same network, in a self contained manner.
The work we present in this paper is based on a report which was made public in Inan & Khosravi (2016). We have recently come across a concurrent preprint (Press & Wolf, 2016) where the authors reuse the word embedding matrix in the output projection to improve language modeling.
However, their work is purely empirical, and they do not provide any theoretical justification for their approach. Finally, we would like to note that the idea of using the same representation for input and output words has been explored in the past, and there exists language models which could be interpreted as simple neural networks with shared input and output embeddings (Bengio et al., 2001; Mnih & Hinton, 2007). However, shared input and output representations were implicitly built into these models, rather than proposed as a supplement to a baseline. Consequently, possibility of improvement was not particularly pursued by sharing input and output representations.
6 EXPERIMENTS
In our experiments, we use the Penn Treebank corpus (PTB) (Marcus et al., 1993), and the Wikitext2 dataset (Merity et al., 2016). PTB has been a standard dataset used for benchmarking language models. It consists of 923k training, 73k validation, and 82k test words. The version of this dataset which we use is the one processed in Mikolov et al. (2010), with the most frequent 10k words selected to be in the vocabulary and rest replaced with a an <unk> token 2. Wikitext-2 is a dataset released recently as an alternative to PTB3. It contains 2, 088k training, 217k validation, and 245k test tokens, and has a vocabulary of 33, 278 words; therefore, in comparison to PTB, it is roughly 2 times larger in dataset size, and 3 times larger in vocabulary.
6.1 MODEL AND TRAINING HIGHLIGHTS
We closely follow the LSTM based language model proposed in Zaremba et al. (2014) for constructing our baseline model. Specifically, we use a 2-layer LSTM with the same number of hidden units in each layer, and we use 3 different network sizes: small (200 units), medium (650 units), and large (1500 units). We train our models using stochastic gradient descent, and we use a variant of the dropout method proposed in Gal (2015). We defer further details regarding training the models to section A of the appendix. We refer to our baseline network as variational dropout LSTM, or VD-LSTM in short.
6.2 EMPIRICAL VALIDATION FOR THE THEORY OF REUSING WORD EMBEDDINGS
In Section 4, we showed that the particular loss augmentation scheme we choose constrains the output projection matrix to be close to the input embedding matrix, without explicitly doing so by reusing the input embedding matrix. As a first experiment, we set out to validate this theoretical result. To do this, we try to simulate the setting in Section 4 by doing the following: We select a randomly chosen 20, 000 contiguous word sequence in the PTB training set, and train a 2-layer LSTM language model with 300 units in each layer with loss augmentation by minimizing the following loss:
J tot = βJaugτ2|V |+ (1− β)J. (6.1)
Here, β is the proportion of the augmented loss used in the total loss, and Jaug is scaled by τ2|V | to approximately match the magnitudes of the derivatives of J and Jaug (see (4.3)). Since we aim to achieve the minimum training loss possible, and the goal is to show a particular result rather than to achieve good generalization, we do not use any kind of regularization in the neural network (e.g. weight decay, dropout). For this set of experiments, we also constrain each row of the input embedding matrix to have a norm of 1 because training becomes difficult without this constraint when only augmented loss is used. After training, we compute a metric that measures distance between the subspace spanned by the rows of the input embedding matrix, L, and that spanned by the columns of the output projection matrix, W . For this, we use a common metric based on the relative residual norm from projection of one matrix onto another (Björck & Golub, 1973). The computed distance between the subspaces is 1 when they are orthogonal, and 0 when they are the same. Interested reader may refer to section B in the appendix for the details of this metric.
Figure 1 shows the results from two tests. In one (panel a), we test the effect of using the augmented loss by sweeping β in (6.1) from 0 to 1 at a reasonably high temperature (τ = 10). With no loss
2PTB can be downloaded at http://www.fit.vutbr.cz/ imikolov/rnnlm/simple-examples.tgz 3Wikitext-2 can be downloaded at https://s3.amazonaws.com/research.metamind.io/wikitext/wikitext-2-
v1.zip
augmentation (β = 0), the distance is almost 1, and as more and more augmented loss is used the distance decreases rapidly, and eventually reaches around 0.06 when only augmented loss is used. In the second test (panel b), we set β = 1, and try to see the effect of the temperature on the subspace distance (remember the theory predicts low distance when τ → ∞). Notably, the augmented loss causesW to approachLT sufficiently even at temperatures as low as 2, although higher temperatures still lead to smaller subspace distances.
These results confirm the mechanism through which our proposed loss pushes W to learn the same column space as LT , and it suggests that reusing the input embedding matrix by explicitly constraining W = LT is not simply a kind of regularization, but is in fact an optimal choice in our framework. What can be achieved separately with each of the two proposed improvements as well as with the two of them combined is a question of empirical nature, which we investigate in the next section.
6.3 RESULTS ON PTB AND WIKITEXT-2 DATASETS
In order to investigate the extent to which each of our proposed improvements helps with learning, we train 4 different models for each network size: (1) 2-Layer LSTM with variational dropout (VD-LSTM) (2) 2-Layer LSTM with variational dropout and augmented loss (VD-LSTM +AL) (3) 2-Layer LSTM with variational dropout and reused embeddings (VD-LSTM +RE) (4) 2-Layer LSTM with variational dropout and both RE and AL (VD-LSTM +REAL).
Figure 2 shows the validation perplexities of the four models during training on the PTB corpus for small (panel a) and large (panel b) networks. All of AL, RE, and REAL networks significantly outperform the baseline in both cases. Table 1 compares the final validation and test perplexities of the four models on both PTB and Wikitext-2 for each network size. In both datasets, both AL and RE improve upon the baseline individually, and using RE and AL together leads to the best performance. Based on performance comparisons, we make the following notes on the two proposed improvements:
• AL provides better performance gains for smaller networks. This is not surprising given the fact that small models are rather inflexible, and one would expect to see improved learning by training against a more informative data distribution (contributed by the augmented loss) (see Hinton et al. (2015)). For the smaller PTB dataset, performance with AL surpasses that with RE. In comparison, for the larger Wikitext-2 dataset, improvement by AL is more limited. This is expected given larger training sets better represent the true data distribution, mitigating the supervision problem. In fact, we set out to validate this reasoning in a direct manner, and additionally train the small networks separately on the first and second halves of the Wikitext-2 training set. This results in two distinct datasets which
are each about the same size as PTB (1044K vs 929K). As can be seen in Table 2, AL has significantly improved competitive performance against RE and REAL despite the fact that embedding size is 3 times larger compared to PTB. These results support our argument that the proposed augmented loss term acts to improve the amount of information gathered from the dataset. • RE significantly outperforms AL for larger networks. This indicates that, for large models,
the more effective mechanism of our proposed framework is the one which enforces proximity between the output projection space and the input embedding space. From a model complexity perspective, the nontrivial gains offered by RE for all network sizes and for both datasets could be largely attributed to its explicit function to reduce the model size while preserving the representational power according to our framework.
We list in Table 3 the comparison of models with and without our proposed modifications on the Penn Treebank Corpus. The best LSTM model (VD-LSTM+REAL) outperforms all previous work which uses conventional framework, including large ensembles. The recently proposed recurrent highway networks (Zilly et al., 2016) when trained with reused embeddings (VD-RHN +RE) achieves the best overall performance, improving on VD-RHN by a perplexity of 2.5.
6.4 QUALITATIVE RESULTS
One important feature of our framework that leads to better word predictions is the explicit mechanism to assign probabilities to words not merely according to the observed output statistics, but also considering the metric similarity between words. We observe direct consequences of this mechanism qualitatively in the Penn Treebank in different ways: First, we notice that the probability of generating the <unk> token with our proposed network (VD-LSTM +REAL) is significantly lower compared to the baseline network (VD-LSTM) across many words. This could be explained by noting the fact that the <unk> token is an aggregated token rather than a specific word, and it is often not expected to be close to specific words in the word embedding space. We observe the same behavior with very frequent words such as ”a”, ”an”, and ”the”, owing to the same fact that they are not correlated with particular words. Second, we not only observe better probability assignments for the target words, but we also observe relatively higher probability weights associated with the words close to the targets. Sometimes this happens in the form of predicting words semantically close together which are plausible even when the target word is not successfully captured by the model. We provide a few examples from the PTB test set which compare the prediction performance of 1500 unit VD-LSTM and 1500 unit VD-LSTM +REAL in table 4. We would like to note that prediction performance of VD-LSTM +RE is similar to VD-LSTM +REAL for the large network.
4For PTB, small models were re-trained by initializing to their final configuration from the first training session. This did not change the final perplexity for baseline, but lead to improvements for the other models.
5Large network results on Wikitext-2 are not reported since computational resources were insufficient to run some of the configurations.
6This model was developed following our work in Inan & Khosravi (2016).
7 CONCLUSION
In this work, we introduced a novel loss framework for language modeling. Particularly, we showed that the metric encoded into the space of word embeddings could be used to generate a more informed data distribution than the one-hot targets, and that additionally training against this distribution improves learning. We also showed theoretically that this approach lends itself to a second improvement, which is simply reusing the input embedding matrix in the output projection layer. This has an additional benefit of reducing the number of trainable variables in the model. We empirically validated the theoretical link, and verified that both proposed changes do in fact belong to the same framework. In our experiments on the Penn Treebank corpus and Wikitext-2, we showed that our framework outperforms the conventional one, and that even the simple modification of reusing the word embedding in the output projection layer is sufficient for large networks.
The improvements achieved by our framework are not unique to vanilla language modeling, and are readily applicable to other tasks which utilize language models such as neural machine translation, speech recognition, and text summarization. This could lead to significant improvements in such models especially with large vocabularies, with the additional benefit of greatly reducing the number of parameters to be trained.
A MODEL AND TRAINING DETAILS
We begin training with a learning rate of 1 and start decaying it with a constant rate after a certain epoch. This is 5, 10, and 1 for the small, medium, and large networks respectively. The decay rate is 0.9 for the small and medium networks, and 0.97 for the large network.
For both PTB and Wikitext-2 datasets, we unroll the network for 35 steps for backpropagation.
We use gradient clipping (Pascanu et al., 2013b); i.e. we rescale the gradients using the global norm if it exceeds a certain value. For both datasets, this is 5 for the small and the medium network, and 6 for the large network.
We use the dropout method introduced in Gal (2015); particularly, we use the same dropout mask for each example through the unrolled network. Differently from what was proposed in Gal (2015), we tie the dropout weights for hidden states further, and we use the same mask when they are propagated as states in the current layer and when they are used as inputs for the next layer. We don’t use dropout in the input embedding layer, and we use the same dropout probability for inputs and hidden states. For PTB, dropout probabilities are 0.7, 0.5 and 0.35 for small, medium and large networks respectively. For Wikitext-2, probabilities are 0.8 for the small and 0.6 for the medium networks.
When training the networks with the augmented loss (AL), we use a temperature τ = 20. We have empirically observed that setting α, the weight of the augmented loss, according to α = γτ for all the networks works satisfactorily. We set γ to values between 0.5 and 0.8 for the PTB dataset, and between 1.0 and 1.5 for the Wikitext-2 dataset. We would like to note that we have not observed sudden deteriorations in the performance with respect to moderate variations in either τ or α.
B METRIC FOR CALCULATING SUBSPACE DISTANCES
In this section, we detail the metric used for computing the subspace distance between two matrices. The computed metric is closely related with the principle angles between subspaces, first defined in Jordan (1875).
Our aim is to compute a metric distance between two given matrices, X and Y . We do this in three steps:
(1) Obtain two matrices with orthonormal columns, U and V , such that span(U)=span(X) and span(V )=span(Y ). U and V could be obtained with a QR decomposition.
(2) Calculate the projection of either one of U and V onto the other; e.g. do S = UUTV , where S is the projection of V onto U . Then calculate the residual matrix as R = V − S.
(3) Let ‖.‖Fr denote the frobenious norm, and let C be the number of columns of R. Then the distance metric is found as d where d2 = 1C ‖R‖ 2 Fr = 1 C Trace(R TR).
We note that d as calculated above is a valid metric up to the equivalence set of matrices which span the same column space, although we are not going to show it. Instead, we will mention some metric properties of d, and relate it to the principal angles between the subspaces. We first work out an expression for d:
Cd2 = Trace(RTR) = Trace ( (V − UUTV )T (V − UUTV ) ) = Trace ( V T (I − UUT )(I − UUT )V
) = Trace ( V T (I − UUT )V
) = Trace ( (I − UUT )V V T
) = Trace(V TV )− Trace ( UUTV V T
) = C − Trace ( UUTV V T
) = C − Trace ( (UTV )T (UTV )
) = C − ‖UTV ‖2Fr
= C∑ i=1 1− ρ2i , (B.1)
where ρi is the ith singular value of UTV , commonly referred to as the ith principle angle between the subspaces of X and Y , θi. In above, we used the cyclic permutation property of the trace in the third and the fourth lines.
Since d2 is 1C Trace(R TR), it is always nonnegative, and it is only zero when the residual is zero, which is the case when span(X) = span(Y). Further, it is symmetric between U and V due to the form of (B.1) (singular values of V TU and V TU are the same). Also, d2 = 1C ∑C i=1 sin
2(θi), namely the average of the sines of the principle angles, which is a quantity between 0 and 1.

This work offers a theoretical justification for reusing the input word embedding in the output projection layer. It does by proposing an additional loss that is designed to minimize the distance between the predictive distribution and an estimate of the true data distribution. This is a nice setup since it can effectively smooth over the labels given as input. However, the construction of the estimate of the true data distribution seems engineered to provide the weight tying justification in Eqs. 3.6 and 3.7.
It is not obvious why the projection matrix L in Eq 3.6 (let's rename it to L') should be the same as that in Eq. 2.1. For example, L' could be obtained through word2vec embeddings trained on a large dataset or it could be learned as an additional set of parameters. In the case that L' is a new learned matrix, it seems the result in Eq 4.5 is to use an independent matrix for the output projection layer, as is usually done.
The experimental results are good and provide support for the approximate derivation done in section 4, particularly the distance plots in figure 1.
Minor comments:
Third line in abstract: where model -> where the model
Second line in section 7: into space -> into the space
Shouldn't the RHS in Eq 3.5 be \sum \tilde{y_{t,i}}(\frac{\hat{y}_t}{\tilde{y_{t,i}}} - e_i) ?