DEEP MULTI-TASK REPRESENTATION LEARNING: A TENSOR FACTORISATION APPROACH
Most contemporary multi-task learning methods assume linear models. This setting is considered shallow in the era of deep learning. In this paper, we present a new deep multi-task representation learning framework that learns cross-task sharing structure at every layer in a deep network. Our approach is based on generalising the matrix factorisation techniques explicitly or implicitly used by many conventional MTL algorithms to tensor factorisation, to realise automatic learning of end-to-end knowledge sharing in deep networks. This is in contrast to existing deep learning approaches that need a user-defined multi-task sharing strategy. Our approach applies to both homogeneous and heterogeneous MTL. Experiments demonstrate the efficacy of our deep multi-task representation learning in terms of both higher accuracy and fewer design choices.
1 INTRODUCTION
The paradigm of multi-task learning is to learn multiple related tasks simultaneously so that knowledge obtained from each task can be re-used by the others. Early work in this area focused on neural network models (Caruana, 1997), while more recent methods have shifted focus to kernel methods, sparsity and low-dimensional task representations of linear models (Evgeniou & Pontil, 2004; Argyriou et al., 2008; Kumar & Daumé III, 2012). Nevertheless given the impressive practical efficacy of contemporary deep neural networks (DNN)s in many important applications, we are motivated to revisit MTL from a deep learning perspective.
While the machine learning community has focused on MTL for shallow linear models recently, applications have continued to exploit neural network MTL (Zhang et al., 2014; Liu et al., 2015). The typical design pattern dates back at least 20 years (Caruana, 1997): define a DNN with shared lower representation layers, which then forks into separate layers and losses for each task. The sharing structure is defined manually: full-sharing up to the fork, and full separation after the fork. However this complicates DNN architecture design because the user must specify the sharing structure: How many task specific layers? How many task independent layers? How to structure sharing if there are many tasks of varying relatedness?
In this paper we present a method for end-to-end multi-task learning in DNNs. This contribution can be seen as generalising shallow MTL methods (Evgeniou & Pontil, 2004; Argyriou et al., 2008; Kumar & Daumé III, 2012) to learning how to share at every layer of a deep network; or as learning the sharing structure for deep MTL (Caruana, 1997; Zhang et al., 2014; Spieckermann et al., 2014; Liu et al., 2015) which currently must be defined manually on a problem-by-problem basis.
Before proceeding it is worth explicitly distinguishing some different problem settings, which have all been loosely referred to as MTL in the literature. Homogeneous MTL: Each task corresponds to a single output. For example, MNIST digit recognition is commonly used to evaluate MTL algorithms by casting it as 10 binary classification tasks (Kumar & Daumé III, 2012). Heterogeneous MTL: Each task corresponds to a unique set of output(s) (Zhang et al., 2014). For example, one may want simultaneously predict a person’s age (task one: multi-class classification or regression) as well as identify their gender (task two: binary classification) from a face image.
In this paper, we propose a multi-task learning method that works on all these settings. The key idea is to use tensor factorisation to divide each set of model parameters (i.e., both FC weight matrices,
and convolutional kernel tensors) into shared and task-specific parts. It is a natural generalisation of shallow MTL methods that explicitly or implicitly are based on matrix factorisation (Evgeniou & Pontil, 2004; Argyriou et al., 2008; Kumar & Daumé III, 2012; Daumé III, 2007). As linear methods, these typically require pre-engineered features. In contrast, as a deep network, our generalisation can learn directly from raw image data, determining sharing structure in a layer-wise fashion. For the simplest NN architecture – no hidden layer, single output – our method reduces to matrix-based ones, therefore matrix-based methods including (Evgeniou & Pontil, 2004; Argyriou et al., 2008; Kumar & Daumé III, 2012; Daumé III, 2007) are special cases of ours.
2 RELATED WORK
Multi-Task Learning Most contemporary MTL algorithms assume that the input and model are both D-dimensional vectors. The models of T tasks can then be stacked into a D × T sized matrix W . Despite different motivations and implementations, many matrix-based MTL methods work by placing constrains on W . For example, posing an `2,1 norm on W to encourage low-rank W (Argyriou et al., 2008). Similarly, (Kumar & Daumé III, 2012) factorises W as W = LS, i.e., it assigns a lower rank as a hyper-parameter. An earlier work (Evgeniou & Pontil, 2004) proposes that the linear model for each task t can be written as wt = ŵt + ŵ0. This is the factorisation L = [ŵ0, ŵ1, . . . , ŵT ] and S = [11×T ; IT ]. In fact, such matrix factorisation encompasses many MTL methods. E.g., (Xue et al., 2007) assumes S·,i (the ith column of S) is a unit vector generated by a Dirichlet Process and (Passos et al., 2012) models W using linear factor analysis with Indian Buffet Process (Griffiths & Ghahramani, 2011) prior on S.
Tensor Factorisation In deep learning, tensor factorisation has been used to exploit factorised tensors’ fewer parameters than the original (e.g., 4-way convolutional kernel) tensor, and thus compress and/or speed up the model, e.g., (Lebedev et al., 2015; Novikov et al., 2015). For shallow linear MTL, tensor factorisation has been used to address problems where tasks are described by multiple independent factors rather than merely indexed by a single factor (Yang & Hospedales, 2015). Here the D-dimensional linear models for all unique tasks stack into a tensor W , of e.g. D × T1 × T2 in the case of two task factors. Knowledge sharing is then achieved by imposing tensor norms on W (Romera-paredes et al., 2013; Wimalawarne et al., 2014). Our framework factors tensors for the different reason that for DNN models, parameters include convolutional kernels (N -way tensors) or D1 × D2 FC layer weight matrices (2-way tensors). Stacking up these parameters for many tasks results in D1 × · · · ×DN × T tensors within which we share knowledge through factorisation. Heterogeneous MTL and DNNs Some studies consider heterogeneous MTL, where tasks may have different numbers of outputs (Caruana, 1997). This differs from the previously discussed studies (Evgeniou & Pontil, 2004; Argyriou et al., 2008; Bonilla et al., 2007; Jacob et al., 2009; Kumar & Daumé III, 2012; Romera-paredes et al., 2013; Wimalawarne et al., 2014) which implicitly assume that each task has a single output. Heterogeneous MTL typically uses neural networks with multiple sets of outputs and losses. E.g., Huang et al. (2013) proposes a shared-hidden-layer DNN model for multilingual speech processing, where each task corresponds to an individual language. Zhang et al. (2014) uses a DNN to find facial landmarks (regression) as well as recognise facial attributes (classification); while Liu et al. (2015) proposes a DNN for query classification and information retrieval (ranking for web search). A key commonality of these studies is that they all require a user-defined parameter sharing strategy. A typical design pattern is to use shared layers (same parameters) for lower layers of the DNN and then split (independent parameters) for the top layers. However, there is no systematic way to make such design choices, so researchers usually rely on trial-and-error, further complicating the already somewhat dark art of DNN design. In contrast, our method learns where and how much to share representation parameters across the tasks, hence significantly reducing the space of DNN design choices.
Parametrised DNNs Our MTL approach is a parameterised DNN (Sigaud et al., 2015), in that DNN weights are dynamically generated given some side information – in the case of MTL, given the task identity. In a related example of speaker-adaptive speech recognition (Tan et al., 2016) there may be several clusters in the data (e.g., gender, acoustic conditions), and each speaker’s model could be a linear combination of these latent task/clusters’ models. They model each speaker i’s weight matrix W (i) as a sum of K base models W̃ , i.e., W (i) = ∑K k=1 λ (i) p W̃ (p). The difference between speakers/tasks comes from λ and the base models are shared. An advantage of this is that,
when new data come, one can choose to re-train λ parameters only, and keep W̃ fixed. This will significantly reduce the number of parameters to learn, and consequently the required training data. Beyond this, Yang & Hospedales (2015) show that it is possible to train another neural network to predict those λ values from some abstract metadata. Thus a model for an unseen task can be generated on-the-fly with no training instances given an abstract description of the task. The techniques developed here are compatible with both these ideas of generating models with minimal or no effort.
3 METHODOLOGY
3.1 PRELIMINARIES
We first recap some tensor factorisation basics before explaining how to factorise DNN weight tensors for multi-task representation learning. An N -way tensorW with shape D1 ×D2 × · · ·DN is an N -dimensional array containing ∏N n=1Dn elements. Scalars, vectors, and matrices can be seen as 0, 1, and 2-way tensors respectively, although the term tensor is usually used for 3-way or higher. A mode-n fibre ofW is a Dn-dimensional vector obtained by fixing all but the nth index. The mode-n flatteningW(n) ofW is the matrix of sizeDn× ∏ i¬nDi constructed by concatenating
all of the ∏
i¬nDi mode-n fibres along columns.
The dot product of two tensors is a natural extension of matrix dot product, e.g., if we have a tensor A of sizeM1×M2×· · ·P and a tensor B of size P ×N1×N2 . . . , the tensor dot productA•B will be a tensor of size M1 ×M2 × · · ·N1 × N2 · · · by matrix dot product AT(−1)B(1) and reshaping
1. More generally, tensor dot product can be performed along specified axes, A • B(i,j) = AT(i)B(j) and reshaping. Here the subscripts indicate the axes of A and B at which dot product is performed. E.g., when A is of size M1 × P ×M3 × · · ·MI and B is of size N1 × N2 × P × · · ·NJ , then A • B(2,3) is a tensor of size M1 ×M3 × · · ·MI ×N1 ×N2 × · · ·NJ .
Matrix-based Knowledge Sharing Assume we have T linear models (tasks) parametrised by Ddimensional weight vectors, so the collection of all models forms a size D × T matrix W . One commonly used MTL approach (Kumar & Daumé III, 2012) is to place a structure constraint on W , e.g., W = LS, where L is a D ×K matrix and S is a K × T matrix. This factorisation recovers a shared factor L and a task-specific factor S. One can see the columns of L as latent basis tasks, and the model w(i) for the ith task is the linear combination of those latent basis tasks with task-specific information S·,i.
w(i) := W·,i = LS·,i = K∑ k=1 L·,kSk,i (1)
From Single to Multiple Outputs Consider extending this matrix factorisation approach to the case of multiple outputs. The model for each task is then a D1 × D2 matrix, for D1 input and D2 output dimensions. The collection of all those matrices constructs a D1 × D2 × T tensor. A straightforward extension of Eq. 1 to this case is
W (i) :=W·,·,i = K∑
k=1
L·,·,kSk,i (2)
This is equivalent to imposing the same structural constraint on WT(3) (transposed mode-3 flattening ofW). It is important to note that this allows knowledge sharing across the tasks only. I.e., knowledge sharing is only across-tasks not across dimensions within a task. However it may be that the knowledge learned in the mapping to one output dimension may be useful to the others within one task. E.g., consider recognising photos of handwritten and print digits – it may be useful to share across handwritten-print; as well as across different digits within each. In order to support general knowledge sharing across both tasks and outputs within tasks, we propose to use more general tensor factorisation techniques. Unlike for matrices, there are multiple definitions of tensor factorisation, and we use Tucker (Tucker, 1966) and Tensor Train (TT) (Oseledets, 2011) decompositions.
1We slightly abuse ‘-1’ referring to the last axis of the tensor.
3.2 TENSOR FACTORISATION FOR KNOWLEDGE SHARING
Tucker Decomposition Given anN -way tensor of sizeD1×D2 · · ·×DN , Tucker decomposition outputs a core tensor S of size K1 × K2 · · · × KN , and N matrices U (n) of size Dn × Kn, such that,
Wd1,d2,...,dN = K1∑
k1=1 K2∑ k2=1 · · · KN∑ kN=1 Sk1,k2,...,kNU (1) d1,k1 U (2) d2,k2 · · ·U (N)dN ,kN (3)
W = S • U(1)(1,2) • U (2) (1,2) · · · • U (N) (1,2) (4)
Tucker decomposition is usually implemented by an alternating least squares (ALS) method (Kolda & Bader, 2009). However (Lathauwer et al., 2000) treat it as a higher-order singular value decomposition (HOSVD), which is more efficient to solve: U (n) is exactly the U matrix from the SVD of mode-n flattening W(n) ofW , and the core tensor S is obtained by,
S =W • U(1)(1,1) • U (2) (1,1) · · · • U (N) (1,1) (5)
Tensor Train Decomposition Tensor Train (TT) Decomposition outputs 2 matrices U (1) and U (N) of size D1 × K1 and KN−1 × DN respectively, and (N − 2) 3-way tensors U (n) of size Kn−1 ×Dn ×Kn. The elements ofW can be computed by,
Wd1,d2,...,dN = K1∑
k1=1 K2∑ k2=1 · · · KN−1∑ kN−1=1 U (1) d1,k1 U (2)k1,d2,k2U (3) k2,d3,k3 · · ·U (N)kN−1,dN (6)
= U (1) d1,·U (2) ·,d2,·U (3) ·,d3,· · · ·U (d) ·,dN (7)
W = U (1) • U (2) · · · • U (N) (8)
where U (n)·,dn,· is a matrix of size Kn−1 ×Kn sliced from U (n) with the second axis fixed at dn. The TT decomposition is typically realised with a recursive SVD-based solution (Oseledets, 2011).
Knowledge Sharing If the final axis of the input tensor above indexes tasks, i.e. if DN = T then the last factor U (N) in both decompositions encodes a matrix of task specific knowledge, and the other factors encode shared knowledge.
3.3 DEEP MULTI-TASK REPRESENTATION LEARNING
To realise deep multi-task representation learning (DMTRL), we learn one DNN per-task each with the same architecture2. However each corresponding layer’s weights are generated with one of the knowledge sharing structures in Eq. 2, Eq. 4 or Eq. 8. It is important to note that we apply these ‘right-to-left’ in order to generate weight tensors with the specified sharing structure, rather than actually applying Tucker or TT to decompose an input tensor. In the forward pass, we synthesise weight tensorsW and perform inference as usual, so the method can be thought of as tensor composition rather than decomposition.
Our weight generation (construct tensors from smaller pieces) does not introduce non-differentiable terms, so our deep multi-task representation learner is trainable via standard backpropagation. Specifically, in the backward pass over FC layers, rather than directly learning the 3-way tensor W , our methods learn either {S, U1, U2, U3} (DMTRL-Tucker, Eq. 4), {U1,U2, U3} (DMTRL-TT, Eq. 8), or in the simplest case {L, S} (DMTRL-LAF3, Eq. 2). Besides FC layers, contemporary
2Except heterogeneous MTL, where the output layer is necessarily unshared due to different dimensionality. 3LAF refers to Last Axis Flattening.

DNN designs often exploit convolutional layers. Those layers usually contain kernel filter parameters that are 3-way tensors of sizeH×W ×C, (whereH is height, W is width, and C is the number of input channels) or 4-way tensors of sizeH×W ×C×M , whereM is the number of filters in this layer (i.e., the number of output channels). The proposed methods naturally extend to convolution layers as convolution just adds more axes on the left-hand side. E.g., the collection of parameters from a given convolutional layer of T neural networks forms a tensor of shapeH×W×C×M×T . These knowledge sharing strategies provide a way to softly share parameters across the corresponding layers of each task’s DNN: where, what, and how much to share are learned from data. This is in contrast to the conventional Deep-MTL approach of manually selecting a set of layers to undergo hard parameter sharing: by tying weights so each task uses exactly the same weight matrix/tensor for the corresponding layer (Zhang et al., 2014; Liu et al., 2015); and a set of layers to be completely separate: by using independent weight matrices/tensors. In contrast our approach benefits from: (i) automatically learning this sharing structure from data rather than requiring user trial and error, and (ii) smoothly interpolating between fully shared and fully segregated layers, rather than a hard switching between these states. An illustration of the proposed framework for different problem settings can be found in Fig. 1.
4 EXPERIMENTS
Implementation Details Our method is implemented with TensorFlow (Abadi et al., 2015). The code is released on GitHub4. For DMTRL-Tucker, DMTRL-TT, and DMTRL-LAF, we need to assign the rank of each weight tensor. The DNN architecture itself may be complicated and so can benefit from different ranks at different layers, but grid-search is impractical. However, since
4https://github.com/wOOL/DMTRL
both Tucker and TT decomposition methods have SVD-based solutions, and vanilla SVD is directly applicable to DMTRL-LAF, we can initialise the model and set the ranks as follows: First train the DNNs independently in single task learning mode. Then pack the layer-wise parameters as the input for tensor decomposition. When SVD is applied, set a threshold for relative error so SVD will pick the appropriate rank. Thus our method needs only a single hyper parameter of max reconstruction error (we set to = 10% throughout) that indirectly specifies the ranks of every layer. Note that training from random initialisation also works, but the STL-based initialisation makes rank selection easy and transparent. Nevertheless, like (Kumar & Daumé III, 2012) the framework is not sensitive to rank choice so long as they are big enough. If random initialisation is desired to eliminate the pre-training requirement, good practice is to initialise parameter tensors by a suitable random weight distribution first, then do decomposition, and use the decomposed values for initialising the factors (the real learnable parameters in our framework). In this way, the resulting re-composed tensors will have approximately the intended distribution. Our sharing is applied to weight parameters only, bias terms are not shared. Apart from initialisation, decomposition is not used anywhere.
4.1 HOMOGENEOUS MTL
Dataset, Settings and Baselines We use MNIST handwritten digits. The task is to recognise digit images zero to nine. When this dataset is used for the evaluation of MTL methods, ten 1-vs-all binary classification problems usually define ten tasks (Kumar & Daumé III, 2012). The dataset has a given train (60,000 images) and test (10,000 images) split. Each instance is a monochrome image of size 28× 28× 1. We use a modified LeNet (LeCun et al., 1998) as the CNN architecture. The first convolutional layer has 32 filters of size 5 × 5, followed by 2 × 2 max pooling. The second convolutional layer has 64 filters of size 4× 4, and again a 2× 2 max pooling. After these two convolutional layers, two fully connected layers with 512 and 1 output(s) are placed sequentially. The convolutional and first FC layer use RELU f(x) = max(x, 0) activation function. We use hinge loss, `(y) = max(0, 1− ŷ ·y), where y ∈ ±1 is the true label and ŷ is the output of each task’s neural network. Conventional matrix-based MTL methods (Evgeniou & Pontil, 2004; Argyriou et al., 2008; Kumar & Daumé III, 2012; Romera-paredes et al., 2013; Wimalawarne et al., 2014) are linear models taking vector input only, so they need a preprocessing that flattens the image into a vector, and typically reduce dimension by PCA. As per our motivation for studying Deep MTL, our methods decisively outperform such shallow linear baselines. Thus to find a stronger MTL competitor, we instead search user defined architectures for Deep-MTL parameter sharing (cf (Zhang et al., 2014; Liu et al., 2015; Caruana, 1997)). In all of the four parametrised layers (pooling has no parameters), we set the first N (1 ≤ N ≤ 3) to be hard shared5. We then use cross-validation to select among the three userdefined MTL architectures and the best option is N = 3, i.e., the first three layers are fully shared (we denote this model UD-MTL). For our methods, all four parametrised layers are softly shared with the different factorisation approaches. To evaluate different MTL methods and a baseline of single task learning (STL), we take ten different fractions of the given 60K training split, train the model, and test on the 10K testing split. For each fraction, we repeat the experiment 5 times with randomly sampled training data. We report two performance metrics: (1) the mean error rate of the ten binary classification problems and (2) the error rate of recognising a digit by ranking each task’s 1-vs-all output (multi-class classification error).
Results As we can see in Fig. 2, all MTL approaches outperform STL, and the advantage is more significant when the training data is small. The proposed methods, DMTRL-TT and DMTRLTucker outperform the best user-defined MTL when the training data is very small, and their performance is comparable when the training data is large.
Further Discussion For a slightly unfair comparison, in the case of binary classification with 1000 training data, shallow matrix-based MTL methods with PCA feature (Kang et al., 2011; Kumar & Daumé III, 2012) reported 14.0% / 13.4% error rate. With the same amount of data, our methods
5This is not strictly all possible user-defined sharing options. For example, another possibility is the first convolutional layer and the first FC layer could be fully shared, with the second convolutional layer being independent (task specific). However, this is against the intuition that lower/earlier layers are more task agnostic, and later layers more task specific. Note that sharing the last layer is technically possible but not intuitive, and in any case not meaningful unless at least one early layer is unshared, as the tasks are different.
have error rate below 6%. This shows the importance of our deep end-to-end multi-task representation learning contribution versus conventional shallow MTL. Since the error rates in (Kang et al., 2011; Kumar & Daumé III, 2012) were produced on a private subset of MNIST dataset with PCA representations only, to ensure a direct comparison, we implement several classic MTL methods and compare them in Appendix A.
For readers interested in the connection to model capacity (number of parameters), we present further analysis in Appendix B.
4.2 HETEROGENEOUS MTL: FACE ANALYSIS
Dataset, Settings and Baselines The AdienceFaces (Eidinger et al., 2014) is a large-scale face images dataset with the labels of each person’s gender and age group. We use this dataset for the evaluation of heterogeneous MTL with two tasks: (i) gender classification (two classes) and (ii) age group classification (eight classes). Two independent CNN models for this benchmark are introduced in (Levi & Hassncer, 2015). The two CNNs have the same architecture except for the last fully-connected layer, since the heterogeneous tasks have different number of outputs (two / eight). We take these CNNs from (Levi & Hassncer, 2015) as the STL baseline. We again search for the best possible user-defined MTL architecture as a strong competitor: the proposed CNN has six layers – three convolutional and three fully-connected layers. The last fully-connected layer has non-shareable parameters because they are of different size. To search the MTL design-space, we try setting the first N (1 ≤ N ≤ 5) layers to be hard shared between the tasks. Running 5-fold cross-validation on the train set to evaluate the architectures, we find the best choice is N = 5 (i.e., all layers fully shared before the final heterogeneous outputs). For our proposed methods, all the layers before the last heterogeneous dimensionality FC layers are softly shared.
We select increasing fractions of the AdienceFaces train split randomly, train the model, and evaluate on the same test set. For reference, there are 12245 images with gender labelled for training, 4007 ones for testing, and 11823 images with age group labelled for training, and 4316 ones for testing.
Results Fig. 3 shows the error rate for each task. For the gender recognition task, we find that: (i) User-defined MTL is not consistently better than STL, but (ii) our methods, esp., DMTRLTucker, consistently outperform both STL and the best user-defined MTL. For the harder age group classification task, our methods generally improve on STL. However UD-MTL does not consistently improve on STL, and even reduces performance when the training set is bigger. This is the negative transfer phenomenon (Rosenstein et al., 2005), where using a transfer learning algorithm is worse than not using it. This difference in outcomes is attributed to sufficient data eventually providing some effective task-specific representation. Our methods can discover and exploit this, but UDMTL’s hard switch between sharing and not sharing can not represent or exploit such increasing task-specificity of representation.
4.3 HETEROGENEOUS MTL: MULTI-ALPHABET RECOGNITION
Dataset, Settings and Baselines We next consider the task of learning to recognise handwritten letters in multiple languages using the Omniglot (Lake et al., 2015) dataset. Omniglot contains handwritten characters in 50 different alphabets (e.g., Cyrillic, Korean, Tengwar), each with its own number of unique characters (14 ∼ 55). In total, there are 1623 unique characters, and each has exactly 20 instances. Here each task corresponds to an alphabet, and the goal is to recognise its characters. MTL has a clear motivation here, as cross-alphabet knowledge sharing is likely to be useful as one is unlikely to have extensive training data for a wide variety of less common alphabets.
The images are monochrome of size 105 × 105. We design a CNN with 3 convolutional and 2 FC layers. The first conv layer has 8 filters of size 5 × 5; the second conv layer has 12 filters of size 3 × 3, and the third convolutional layer has 16 filters of size 3 × 3. Each convolutional layer is followed by a 2 × 2 max-pooling. The first FC layer has 64 neurons, and the second FC layer has size corresponding to the number of unique classes in the alphabet. The activation function is tanh.
We use a similar strategy to find the best user-defined MTL model: the CNN has 5 parametrised layers, of which 4 layers are potentially shareable. So we tried hard-sharing the firstN (1 ≤ N ≤ 4) layers. Evaluating these options by 5-fold cross-validation, the best option turned out to be N = 3, i.e., the first three layers are hard shared. For our methods, all four shareable layers are softly shared.
Since there is no standard train/test split for this dataset, we use the following setting: We repeatedly pick at random 5, . . . 90% of images per class for training. Note that 5% is the minimum, corresponding to one-shot learning. The remaining data are used for evaluation.
Results Fig. 4 reports the average error rate across all 50 tasks (alphabets). Our proposed MTL methods surpass the STL baseline in all cases. User-defined MTL does not work well when the training data is very small, but does help when training fraction is larger than 50%.
Measuring the Learned Sharing Compared to the conventional user-defined sharing architectures, our method learns how to share from data. We next try to quantify the amount of sharing estimated by our model on the Omniglot data. Returning to the key factorisation W = LS, we can find that S-like matrix appears in all variants of proposed method. It is S in DMTRL-LAF, the transposed U (N) in DMTRL-Tucker, and U (N) in DMTRL-TT (N is the last axis of W). S is a K × T size matrix, where T is the number of tasks, and K is the number of latent tasks (Kumar & Daumé III, 2012) or the dimension of task coding (Yang & Hospedales, 2015). Each column of S is a set of coefficients that produce the final weight matrix/tensor by linear combination. If we put STL and user-defined MTL (for a certain shared layer) in this framework, we see that STL is to assign (rather than learn) S to be an identity matrix IT . Similarly, user-defined MTL (for a certain shared layer) is to assign S to be a matrix with all zeros but one particular row is all ones, e.g., S = [11×T ;0]. Between these two extremes, our method learns the sharing structure in S. We propose the following equation to measure the learned sharing strength:
ρ = 1( T 2 ) ∑ i<j Ω(S·,i, S·,j) = 2 T (T − 1) ∑ i<j Ω(S·,i, S·,j) (9)
Here Ω(a, b) is a similarity measure for two vectors a and b and we use cosine similarity. ρ is the average on all combinations of column-wise similarity. So ρmeasures how much sharing is encoded by S between ρ = 0 for STL (nothing to share) and ρ = 1 for user-defined MTL (completely shared). Since S is a real-valued matrix in our scenario, we normalise it before applying Eq. 9: First we take absolute values, because large either positive or negative value suggests a significant coefficient. Second we normalise each column of S by applying a softmax function, so the sum of every column is 1. The motivation behind the second step is to make a matched range of our S with S = IT or S = [11×T ;0], as for those two cases, the sum of each column is 1 and the range is [0, 1].
For the Omniglot experiment, we plot the measured sharing amount for training fraction 10%. Fig. 4 reveals that three proposed methods tend to share more for bottom layers (‘Conv1’, ‘Conv2’, and ‘Conv3’) and share less for top layer (‘FC1’). This is qualitatively similar to the best user-defined MTL, where the first three layers are fully shared (ρ = 1) and the 4th layer is completely not shared (ρ = 0). However, our methods: (i) learn this structure in a purely data-driven way and (ii) benefits from the ability to smoothly interpolate between high and low degrees of sharing as depth increases. As an illustration, Fig. 4 also shows example text from the most and least similar language pairs as estimated at our multilingual character recogniser’s FC1 layer (the result can vary across layers).
5 CONCLUSION
In this paper, we propose a novel framework for end-to-end multi-task representation learning in contemporary deep neural networks. The key idea is to generalise matrix factorisation-based multitask ideas to tensor factorisation, in order to flexibly share knowledge in fully connected and convolutional DNN layers. Our method provides consistently better performance than single task learning and comparable or better performance than the best results from exhaustive search of user-defined MTL architectures. It reduces the design choices and architectural search space that must be explored in the workflow of Deep MTL architecture design (Caruana, 1997; Zhang et al., 2014; Liu et al., 2015), relieving researchers of the need to decide how to structure layer sharing/segregation. Instead sharing structure is determined in a data-driven way on a layer-by-layer basis that moreover allows a smooth interpolation between sharing and not sharing in progressively deeper layers.
Acknowledgements This work was supported by EPSRC (EP/L023385/1), and the European Union’s Horizon 2020 research and innovation program under grant agreement No 640891.
A COMPARISON WITH CLASSIC (SHALLOW) MTL METHODS
We provide a comparison with classic (shallow, matrix-based) MTL methods for the first experiment (MNIST, binary one-vs-rest classification, 1% training data, mean of error rates for 10-fold CV). A subtlety in making this comparison is what feature should the classic methods use? Conventionally they use a PCA feature (obtained by flattening the image, then dimension reduction by PCA). However for visual recognition tasks, performance is better with deep features – a key motivation for our focus on deep approaches to MTL. We therefore also compare the classic methods when using a feature extracted from the penultimate layer of the CNN network used in our experiment.
As expected, the classic methods improve on STL, and they perform significantly better with CNN than PCA features. However, our DMTRL methods still outperform the best classic methods, even when they are enhanced by CNN features. This is due to soft (cf hard) sharing of the feature extraction layers and the ability of end-to-end training of both the classifier and feature extractor. Finally, we note that more fundamentally, the classic methods are restricted to binary problems (due to their matrix-based nature) and so, unlike our tensor-based approach, they are unsuitable for multi-class problems like omniglot and age-group classification.
B MODEL CAPACITY AND PERFORMANCE
We list the number of parameters for each model in the first experiment (MNIST, binary one-vs-rest classification) and the performance (1% training data, mean of error rate for 10-fold CV).
The conventional hard-sharing method (UD-MTL) design is to share all layers except the top layer. Its number of parameter is roughly 10% of the single task learning method (STL), as most parameters are shared across the 10 tasks corresponding to 10 digits. Our soft-sharing methods also significantly reduce the number of parameters compared to STL, but are larger than UD-MTL’s hard sharing.
To compare our method to UD-MTL, while controlling for network capacity, we expanded UDMDL by adding more hidden neurons so its number of parameter is close to our methods (denoted UD-MTL-Large). However UD-MDL performance does not increase. This is evidence that our model’s good performance is not simply due to greater capacity than UD-MTL.

This paper proposed a deep multi-task representation learning framework that learns cross-task sharing structure at every layer in a deep network with tensor factorization and end-to-end knowledge sharing. This approach removed the requirement of a user-deﬁned multi-task sharing strategy in conventional approach. Their experimental results indicate that their approach can achieve higher accuracy with fewer design choices.
Although factorization ideas have been exploited in the past for other tasks I think applying it to MTL is interesting. The only thing I want to point out is that the saving of parameter is from the low-rank factorization. In the conventional MTL each layer's weight size can also be reduced if SVD is used.
BTW, recent neural network MTL was explored first (earlier than 2014, 2015 work cited) in speech recognition community. see, e.g.,
Huang, J.T., Li, J., Yu, D., Deng, L. and Gong, Y., 2013, May. Cross-language knowledge transfer using multilingual deep neural network with shared hidden layers. In 2013 IEEE International Conference on Acoustics, Speech and Signal Processing (pp. 7304-7308). IEEE.


ENTROPY-SGD: BIASING GRADIENT DESCENT INTO WIDE VALLEYS
This paper proposes a new optimization algorithm called Entropy-SGD for training deep neural networks that is motivated by the local geometry of the energy landscape. Local extrema with low generalization error have a large proportion of almost-zero eigenvalues in the Hessian with very few positive or negative eigenvalues. We leverage upon this observation to construct a local-entropy-based objective function that favors well-generalizable solutions lying in large flat regions of the energy landscape, while avoiding poorly-generalizable solutions located in the sharp valleys. Conceptually, our algorithm resembles two nested loops of SGD where we use Langevin dynamics in the inner loop to compute the gradient of the local entropy before each update of the weights. We show that the new objective has a smoother energy landscape and show improved generalization over SGD using uniform stability, under certain assumptions. Our experiments on convolutional and recurrent neural networks demonstrate that Entropy-SGD compares favorably to state-of-the-art techniques in terms of generalization error and training time.
1 INTRODUCTION
This paper presents a new optimization tool for deep learning designed to exploit the local geometric properties of the objective function. Consider the histogram we obtained in Fig. 1 showing the spectrum of the Hessian at an extremum discovered by Adam (Kingma & Ba, 2014) for a convolutional neural network on MNIST (LeCun et al., 1998) (≈ 47,000 weights, cf. Sec. 5.1). It is evident that:
(i) a large number of directions (≈ 94%) have near-zero eigenvalues (magnitude less than 10−4), (ii) positive eigenvalues (right inset) have a long tail with the largest one being almost 40,
(iii) negative eigenvalues (left inset), which are directions of descent that the optimizer missed, have a much faster decay (the largest negative eigenvalue is only −0.46).
Interestingly, this trend is not unique to this particular network. Rather, its qualitative properties are shared across a variety of network architectures, network sizes, datasets or optimization algorithms (refer to Sec. 5 for more experiments). Local minima that generalize well and are discovered by gradient descent lie in “wide valleys” of the energy landscape, rather than in sharp, isolated minima. For an intuitive understanding of this phenomenon, imagine a Bayesian prior concentrated at the minimizer of the expected loss, the marginal likelihood of wide valleys under this prior is much
higher than narrow, sharp valleys even if the latter are close to the global minimum in training loss. Almost-flat regions of the energy landscape are robust to data perturbations, noise in the activations, as well as perturbations of the parameters, all of which are widely-used techniques to achieve good generalization. This suggests that wide valleys should result in better generalization and, indeed, standard optimization algorithms in deep learning seem to discover exactly that — without being explicitly tailored to do so. For another recent analysis of the Hessian, see the parallel work of Sagun et al. (2016).
Based on this understanding of how the local geometry looks at the end of optimization, can we modify SGD to actively seek such regions? Motivated by the work of Baldassi et al. (2015) on shallow networks, instead of minimizing the original loss f (x), we propose to maximize
F(x,γ) = log ∫ x′∈Rn exp ( − f (x′)− γ 2 ‖x− x′‖22 ) dx′.
The above is a log-partition function that measures both the depth of a valley at a location x ∈ Rn, and its flatness through the entropy of f (x′); we call it “local entropy” in analogy to the free entropy used in statistical physics. The Entropy-SGD algorithm presented in this paper employs stochastic gradient Langevin dynamics (SGLD) to approximate the gradient of local entropy. Our algorithm resembles two nested loops of SGD: the inner loop consists of SGLD iterations while the outer loop updates the parameters. We show that the above modified loss function results in a smoother energy landscape defined by the hyper-parameter γ which we can think of as a “scope” that seeks out valleys of specific widths. Actively biasing the optimization towards wide valleys in the energy landscape results in better generalization error. We present experimental results on fullyconnected and convolutional neural networks (CNNs) on the MNIST and CIFAR-10 (Krizhevsky, 2009) datasets and recurrent neural networks (RNNs) on the Penn Tree Bank dataset (PTB) (Marcus et al., 1993) and character-level text prediction. Our experiments show that Entropy-SGD scales to deep networks used in practice, obtains comparable generalization error as competitive baselines and also trains much more quickly than SGD (we get a 2x speed-up over SGD on RNNs).
2 RELATED WORK
Our above observation about the spectrum of Hessian (further discussed in Sec. 5) is similar to results on a perceptron model in Dauphin et al. (2014) where the authors connect the loss function of a deep network to a high-dimensional Gaussian random field. They also relate to earlier studies such as Baldi & Hornik (1989); Fyodorov & Williams (2007); Bray & Dean (2007) which show that critical points with high training error are exponentially likely to be saddle points with many negative directions and all local minima are likely to have error that is very close to that of the global minimum. The authors also argue that convergence of gradient descent is affected by the proliferation of saddle points surrounded by high error plateaus — as opposed to multiple local minima. One can also see this via an application of Kramer’s law: the time spent by diffusion is
inversely proportional to the smallest negative eigenvalue of the Hessian at a saddle point (Bovier & den Hollander, 2006).
The existence of multiple — almost equivalent — local minima in deep networks has been predicted using a wide variety of theoretical analyses and empirical observations, e.g., papers such as Choromanska et al. (2015a;b); Chaudhari & Soatto (2015) that build upon results from statistical physics as also others such as Haeffele & Vidal (2015) and Janzamin et al. (2015) that obtain similar results for matrix and tensor factorization problems. Although assumptions in these works are somewhat unrealistic in the context of deep networks used in practice, similar results are also true for linear networks which afford a more thorough analytical treatment (Saxe et al., 2014). For instance, Soudry & Carmon (2016) show that with mild over-parameterization and dropout-like noise, training error for a neural network with one hidden layer and piece-wise linear activation is zero at every local minimum. All these results suggest that the energy landscape of deep neural networks should be easy to optimize and they more or less hold in practice — it is easy to optimize a prototypical deep network to near-zero loss on the training set (Hardt et al., 2015; Goodfellow & Vinyals, 2015).
Obtaining good generalization error, however, is challenging: complex architectures are sensitive to initial conditions and learning rates (Sutskever et al., 2013) and even linear networks (Kawaguchi, 2016) may have degenerate and hard to escape saddle points (Ge et al., 2015; Anandkumar & Ge, 2016). Techniques such as adaptive (Duchi et al., 2011) and annealed learning rates, momentum (Tieleman & Hinton, 2012), as well as architectural modifications like dropout (Srivastava et al., 2014), batch-normalization (Ioffe & Szegedy, 2015; Cooijmans et al., 2016), weight scaling (Salimans & Kingma, 2016) etc. are different ways of tackling this issue by making the underlying landscape more amenable to first-order algorithms. However, the training process often requires a combination of such techniques and it is unclear beforehand to what extent each one of them helps.
Closer to the subject of this paper are results by Baldassi et al. (2015; 2016a;b) who show that the energy landscape of shallow networks with discrete weights is characterized by an exponential number of isolated minima and few very dense regions with lots of local minima close to each other. These dense local minima can be shown to generalize well for random input data; more importantly, they are also accessible by efficient algorithms using a novel measure called “robust ensemble” that amplifies the weight of such dense regions. The authors use belief propagation to estimate local entropy for simpler models such as committee machines considered there. A related work in this context is EASGD (Zhang et al., 2015) which trains multiple deep networks in parallel and modulates the distance of each worker from the ensemble average. Such an ensemble training procedure enables improved generalization by ensuring that different workers land in the same wide valley and indeed, it turns out to be closely related to the replica theoretic analysis of Baldassi et al. (2016a).
Our work generalizes the local entropy approach above to modern deep networks with continuous weights. It exploits the same phenomenon of wide valleys in the energy landscape but does so without incurring the hardware and communication complexity of replicated training or being limited to models where one can estimate local entropy using belief propagation. The enabling technique in our case is using Langevin dynamics for estimating the gradient of local entropy, which can be done efficiently even for large deep networks using mini-batch updates.
Motivated by the same final goal, viz. flat local minima, the authors in Hochreiter & Schmidhuber (1997b) introduce hard constraints on the training loss and the width of local minima and show using the Gibbs formalism (Haussler & Opper, 1997) that this leads to improved generalization. As the authors discuss, the effect of hyper-parameters for the constraints is intricately tied together and they are difficult to choose even for small problems. Our local entropy based objective instead naturally balances the energetic term (training loss) and the entropic term (width of the valley). The role of γ is clear as a focusing parameter (cf. Sec. 4.3) and effectively exploiting this provides significant computational advantages. Among other conceptual similarities with our work, let us note that local entropy in a flat valley is a direct measure of the width of the valley which is similar to their usage of Hessian, while the Gibbs variant to averaging in weight space (Eqn. 33 of Hochreiter & Schmidhuber (1997b)) is similar to our Eqn. (7). Indeed, Gibbs formalism used in their analysis is a very promising direction to understanding generalization in deep networks (Zhang et al., 2016).
Homotopy continuation methods convolve the loss function to solve sequentially refined optimization problems (Allgower & Georg, 2012; Mobahi & Fisher III, 2015), similarly, methods that perturb the weights or activations to average the gradient (Gulcehre et al., 2016) do so with an aim to smooth the rugged energy landscape. Such smoothing is however very different from local entropy. For instance, the latter places more weight on wide local minima even if they are much shallower than the global minimum (cf. Fig. 2); this effect cannot be obtained by smoothing. In fact, smoothing can introduce an artificial minimum between two nearby sharp valleys which is detrimental to generalization. In order to be effective, continuation techniques also require that minimizers of successively smaller convolutions of the loss function lie close to each other (Hazan et al., 2016); it is not clear whether this is true for deep networks. Local entropy, on the other hand, exploits wide minima which have been shown to exist in a variety of learning problems (Monasson & Zecchina, 1995; Cocco et al., 1996). Please refer to Appendix C for a more elaborate discussion as well as possible connections to stochastic variational inference (Blei et al., 2016).
3 LOCAL ENTROPY
We first provide a simple intuition for the concept of local entropy of an energy landscape. The discussion in this section builds upon the results of Baldassi et al. (2016a) and extends it for the case of continuous variables. Consider a cartoon energy landscape in Fig. 2 where the x-axis denotes the configuration space of the parameters. We have constructed two local minima: a shallower although wider one at xrobust and a very sharp global minimum at xnon−robust. Under a Bayesian prior on the parameters, say a Gaussian of a fixed variance at locations xrobust and xnon−robust respectively, the wider local minimum has a higher marginalized likelihood than the sharp valley on the right.
The above discussion suggests that parameters that lie in wider local minima like xrobust, which may possibly have a higher loss than the global minimum, should generalize better than the ones that are simply at the global minimum. In a neighborhood of xrobust, “local entropy” as introduced in Sec. 1 is large because it includes the contributions from a large region of good parameters; conversely, near xnon−robust, there are fewer such contributions and the resulting local entropy is low. The local entropy thus provides a way of picking large, approximately flat, regions of the landscape over sharp, narrow valleys in spite of the latter possibly having a lower loss. Quite conveniently, the local entropy is also computed from the partition function with a local re-weighting term.
Formally, for a parameter vector x ∈Rn, consider a Gibbs distribution corresponding to a given energy landscape f (x):
P(x; β ) = Z−1β exp (−β f (x)) ; (1) where β is known as the inverse temperature and Zβ is a normalizing constant, also known as the partition function. As β→∞, the probability distribution above concentrates on the global minimum of f (x) (assuming it is unique) given as:
x∗ = argmin x f (x), (2)
which establishes the link between the Gibbs distribution and a generic optimization problem (2). We would instead like the probability distribution — and therefore the underlying optimization problem — to focus on flat regions such as xrobust in Fig. 2. With this in mind, let us construct a modified Gibbs distribution:
P(x′; x,β ,γ) = Z−1x,β , γ exp ( −β f (x′)−β γ 2 ‖x− x′‖22 ) . (3)
The distribution in (3) is a function of a dummy variable x′ and is parameterized by the original location x. The parameter γ biases the modified distribution (3) towards x; a large γ results in a
P(x′; x,β ,γ) with all its mass near x irrespective of the energy term f (x′). For small values of γ , the term f (x′) in the exponent dominates and the modified distribution is similar to the original Gibbs distribution in (1). We will set the inverse temperature β to 1 because γ affords us similar control on the Gibbs distribution.
Definition 1 (Local entropy). The local free entropy of the Gibbs distribution in (1), colloquially called “local entropy” in the sequel and denoted by F(x,γ), is defined as the log-partition function of modified Gibbs distribution (3), i.e.,
F(x,γ) = logZx,1,γ = log ∫ x′ exp ( − f (x′)− γ 2 ‖x− x′‖22 ) dx′. (4)
The parameter γ is used to focus the modified Gibbs distribution upon a local neighborhood of x and we call it a “scope” henceforth.
Effect on the energy landscape: Fig. 2 shows the negative local entropy −F(x,γ) for two different values of γ . We expect xrobust to be more robust than xnon-robust to perturbations of data or parameters and thus generalize well and indeed, the negative local entropy in Fig. 2 has a global minimum near xrobust. For low values of γ , the energy landscape is significantly smoother than the original landscape and still maintains our desired characteristic, viz. global minimum at a wide valley. As γ increases, the local entropy energy landscape gets closer to the original energy landscape and they become equivalent at γ → ∞. On the other hand, for very small values of γ , the local entropy energy landscape is almost uniform.
Connection to classical entropy: The quantity we have defined as local entropy in Def. 1 is different from classical entropy which counts the number of likely configurations under a given distribution. For a continuous parameter space, this is given by
S(x,β ,γ) =− ∫
x′ logP(x′; x,β ,γ) dP(x′; x,β ,γ)
for the Gibbs distribution in (3). Minimizing classical entropy however does not differentiate between flat regions that have very high loss versus flat regions that lie deeper in the energy landscape. For instance in Fig. 2, classical entropy is smallest in the neighborhood of xcandidate which is a large region with very high loss on the training dataset and is unlikely to generalize well.
4 ENTROPY-GUIDED SGD
Simply speaking, our Entropy-SGD algorithm minimizes the negative local entropy from Sec. 3. This section discusses how the gradient of local entropy can be computed via Langevin dynamics. The reader will see that the resulting algorithm has a strong flavor of “SGD-inside-SGD”: the outer SGD updates the parameters, while an inner SGD estimates the gradient of local entropy.
Consider a typical classification setting, let x ∈ Rn be the weights of a deep neural network and ξk ∈ Ξ be samples from a dataset Ξ of size N. Let f (x;ξk) be the loss function, e.g., cross-entropy of the classifier on a sample ξk. The original optimization problem is:
x∗ = argmin x 1 N
N
∑ k=1 f (x; ξk); (5)
where the objective f (x,ξk) is typically a non-convex function in both the weights x and the samples ξk. The Entropy-SGD algorithm instead solves the problem
x∗Entropy-SGD = argmin x −F(x,γ; Ξ); (6)
where we have made the dependence of local entropy F(x,γ) on the dataset Ξ explicit.
4.1 GRADIENT OF LOCAL ENTROPY
The gradient of local entropy over a randomly sampled mini-batch of m samples denoted by ξ`i ∈ Ξ` for i≤ m is easy to derive and is given by
−∇xF ( x,γ; Ξ` ) = γ ( x− 〈 x′; Ξ` 〉) ; (7)
where the notation 〈·〉 denotes an expectation of its arguments (we have again made the dependence on the data explicit) over a Gibbs distribution of the original optimization problem modified to focus on the neighborhood of x; this is given by
P(x′; x,γ) ∝ exp [ − (
1 m
m
∑ i=1
f ( x′; ξ`i )) − γ
2 ‖x− x′‖22
] . (8)
Computationally, the gradient in (7) involves estimating 〈 x′; Ξ` 〉 with the current weights fixed to x. This is an expectation over a Gibbs distribution and is hard to compute. We can however approximate it using Markov chain Monte-Carlo (MCMC) techniques. In this paper, we use stochastic gradient Langevin dynamics (SGLD) (Welling & Teh, 2011) that is an MCMC algorithm for drawing samples from a Bayesian posterior and scales to large datasets using mini-batch updates. Please see Appendix A for a brief overview of SGLD. For our application, as lines 3-6 of Alg. 1 show, SGLD resembles a few iterations of SGD with a forcing term −γ(x− x′) and additive gradient noise. We can obtain some intuition on how Entropy-SGD works using the expression for the gradient: the term 〈x′; ·〉 is the average over a locally focused Gibbs distribution and for two local minima in the neighborhood of x roughly equivalent in loss, this term points towards the wider one because 〈x′; ·〉 is closer to it. This results in a net gradient that takes SGD towards wider valleys. Moreover, if we unroll the SGLD steps used to compute (x−〈x′; ·〉) (cf. line 5 in Alg. 1), it resembles one large step in the direction of the (noisy) average gradient around the current weights x and Entropy-SGD thus looks similar to averaged SGD in the literature (Polyak & Juditsky, 1992; Bottou, 2012). These two phenomena intuitively explain the improved generalization performance of Entropy-SGD.
4.2 ALGORITHM AND IMPLEMENTATION DETAILS
Alg. 1 provides the pseudo-code for one iteration of the Entropy-SGD algorithm. At each iteration, lines 3-6 perform L iterations of Langevin dynamics to estimate µ = 〈 x′;Ξ` 〉 . The weights x are updated with the modified gradient on line 7.
Algorithm 1: Entropy-SGD algorithm Input : current weights x, Langevin iterations L Hyper-parameters: scope γ , learning rate η , SGLD step size η ′
// SGLD iterations; 1 x′,µ ← x; 2 for `≤ L do 3 Ξ` ← sample mini-batch; 4 dx′ ← 1m ∑mi=1 ∇x′ f ( x′; ξ`i ) − γ (x− x′);
5 x′ ← x′−η ′ dx′+ √
η ′ ε N(0, I); 6 µ ← (1−α)µ +α x′;
// Update weights; 7 x ← x−η γ (x−µ)
Let us now discuss a few implementation details. Although we have written Alg. 1 in the classical SGD setup, we can easily modify it to include techniques such as momentum and gradient preconditioning (Duchi et al., 2011) by changing lines 5 and 7. In our experiments, we have used SGD with Nesterov’s momentum (Sutskever et al., 2013) and Adam for outer and inner loops with similar qualitative results. We use exponential averaging to estimate µ in the SGLD loop (line 6)
with α = 0.75 so as to put more weight on the later samples, this is akin to a burn-in in standard SGLD.
We set the number of SGLD iterations to L = [5,20] depending upon the complexity of the dataset. The learning rate η ′ is fixed for all our experiments to values between η ′ ∈ [0.1,1]. We found that annealing η ′ (for instance, setting it to be the same as the outer learning rate η) is detrimental; indeed a small learning rate leads to poor estimates of local entropy towards the end of training where they are most needed. The parameter ε in SGLD on line 5 is the thermal noise and we fix this to ε ∈ [10−4,10−3]. Having thus fixed L,ε and η ′, an effective heuristic to tune the remaining parameter γ is to match the magnitude of the gradient of the local entropy term, viz. γ (x− µ), to the gradient for vanilla SGD, viz. m−1 ∑mi=1 ∇x f (x; ξ`i).
4.3 SCOPING OF γ
The scope γ is fixed in Alg. 1. For large values of γ , the SGLD updates happen in a small neighborhood of the current parameters x while low values of γ allow the inner SGLD to explore further away from x. In the context of the discussion in Sec. 3, a “reverse-annealing” of the scope γ , i.e. increasing γ as training progresses has the effect of exploring the energy landscape on progressively finer scales. We call this process “scoping” which is similar to that of Baldassi et al. (2016a) and use a simple exponential schedule given by
γ(t) = γ0 (1+ γ1)t ;
for the t th parameter update. We have experimented with linear, quadratic and bounded exponential (γ0 (1− e−τt)) scoping schedules and obtained qualitatively similar results. Scoping of γ unfortunately interferes with the learning rate annealing that is popular in deep learning, this is a direct consequence of the update step on line 7 of Alg. 1. In practice, we therefore scale down the local entropy gradient by γ before the weight update and modify the line to read
x← x−η(x−µ). Our general strategy during hyper-parameter tuning is to set the initial scope γ0 to be very small, pick a large value of η and set γ1 to be such that the magnitude of the local entropy gradient is comparable to that of SGD. We can use much larger learning rates than SGD in our experiments because the local entropy gradient is less noisy than the original back-propagated gradient. This also enables very fast progress in the beginning with a smooth landscape of a small γ .
4.4 THEORETICAL PROPERTIES
We can show that Entropy-SGD results in a smoother loss function and obtains better generalization error than the original objective (5). With some overload of notation, we assume that the original loss f (x) is β -smooth, i.e., for all x,y∈Rn, we have ‖∇ f (x)−∇ f (y)‖ ≤ β ‖x−y‖. We additionally assume for the purpose of analysis that no eigenvalue of the Hessian ∇2 f (x) lies in the set [−2γ−c,c] for some small c > 0. Lemma 2. The objective F(x,γ; Ξ) in (6) is α1+γ−1 c -Lipschitz and β 1+γ−1 c -smooth.
Please see Appendix B for the proof. The local entropy objective is thus smoother than the original objective. Let us now obtain a bound on the improvement in generalization error. We denote an optimization algorithm, viz., SGD or Entropy-SGD by A(Ξ), it is a function of the dataset Ξ and outputs the parameters x∗ upon termination. Stability of the algorithm (Bousquet & Elisseeff, 2002) is then a notion of how much its output differs in loss upon being presented with two datasets, Ξ and Ξ′, that differ in at most one sample:
sup ξ ∈ Ξ ∪ Ξ′
[ f (A(Ξ),ξ )− f ( A(Ξ′),ξ )] ≤ ε.
Hardt et al. (2015) connect uniform stability to generalization error and show that an ε-stable algorithm A(Ξ) has generalization error bounded by ε , i.e., if A(Ξ) terminates with parameters x∗,
|EΞ (RΞ(x∗)−R(x∗))| ≤ ε;
where the left hand side is the generalization error: it is the difference between the empirical loss RΞ(x) := N−1 ∑Nk=1 f (x,ξk) and the population loss R(x) := Eξ f (x,ξ ). We now employ the following theorem that bounds the stability of an optimization algorithm through the smoothness of its loss function and the number of iterations on the training set.
Theorem 3 (Hardt et al. (2015)). For an α-Lipschitz and β -smooth loss function, if SGD converges in T iterations on N samples with decreasing learning rate ηt ≤ 1/t the stability is bounded by
ε / 1 N α1/(1+β ) T 1−1/(1+β ).
Using Lemma 2 and Theorem 3 we have
ε Entropy-SGD / ( α T−1 )(1− 1 1+γ−1c ) β ε SGD, (9)
which shows that Entropy-SGD generalizes better than SGD for all T > α if they both converge after T passes over the samples.
Let us note that while the number of passes over the dataset for Entropy-SGD and SGD are similar for our experiments on CNNs, Entropy-SGD makes only half as many passes as SGD for our experiments on RNNs. As an aside, it is easy to see from the proof of Lemma 2 that for a convex loss function f (x), the local entropy objective does not change the minimizer of the original problem.
Remark 4. The above analysis hinges upon an assumption that the Hessian ∇2 f (x) does not have eigenvalues in the set [−2γ − c,c] for a constant c > 0. This is admittedly unrealistic, for instance, the eigenspectrum of the Hessian at a local minimum in Fig. 1 has a large fraction of its eigenvalues almost zero. Let us however remark that the result in Thm. 3 by Hardt et al. (2015) assumes global conditions on the smoothness of the loss function; one imagines that Eqn. 9 remains qualitatively the same (with respect to T in particular) even if this assumption is violated to an extent before convergence happens. Obtaining a rigorous generalization bound without this assumption would require a dynamical analysis of SGD and seems out of reach currently.
5 EXPERIMENTS
In Sec. 5.1, we discuss experiments that suggest that the characteristics of the energy landscape around local minimal accessible by SGD are universal to deep architectures. We then present experimental results on two standard image classification datasets, viz. MNIST and CIFAR-10 and two datasets for text prediction, viz. PTB and the text of War and Peace. Table 1 summarizes the results of these experiments on deep networks.
5.1 UNIVERSALITY OF THE HESSIAN AT LOCAL MINIMA
We use automatic differentiation1 to compute the Hessian at a local minimum obtained at the end of training for the following networks:
(i) small-LeNet on MNIST: This network has 47,658 parameters and is similar to LeNet but with 10 and 20 channels respectively in the first two convolutional layers and 128 hidden units in the fully-connected layer. We train this with Adam to obtain a test error of 2.4%.
(ii) small-mnistfc on MNIST: A fully-connected network (50,890 parameters) with one layer of 32 hidden units, ReLU non-linearities and cross-entropy loss; it converges to a test error of 2.5% with momentum-based SGD.
(iii) char-lstm for text generation: This is a recurrent network with 48 hidden units and Long Short-Term Memory (LSTM) architecture (Hochreiter & Schmidhuber, 1997a). It has 32,640 parameters and we train it with Adam to re-generate a small piece of text consisting of 256 lines of length 32 each and 96-bit one-hot encoded characters.
1https://github.com/HIPS/autograd
(iv) All-CNN-BN on CIFAR-10: This is similar to the All-CNN-C network (Springenberg et al., 2014) with ≈ 1.6 million weights (cf. Sec. 5.3) which we train using Adam to obtain an error of 11.2%. Exact Hessian computation is in this case expensive and thus we instead compute the diagonal of the Fisher information matrix (Wasserman, 2013) using the element-wise first and second moments of the gradients that Adam maintains, i.e., diag(I) = E(g2)− (E g)2 where g is the back-propagated gradient. Fisher information measures the sensitivity of the log-likelihood of data given parameters in a neighborhood of a local minimum and thus is exactly equal to the Hessian of the negative log-likelihood. We will consider the diagonal of the empirical Fisher information matrix as a proxy for the eigenvalues of the Hessian, as is common in the literature.
We choose to compute the exact Hessian and to keep the computational and memory requirements manageable, the first three networks considered above are smaller than standard deep networks used in practice. For the last network, we sacrifice the exact computation and instead approximate the Hessian of a large deep network. We note that recovering an approximate Hessian from Hessianvector products (Pearlmutter, 1994) could be a viable strategy for large networks.
Fig. 1 in the introductory Sec. 1 shows the eigenspectrum of the Hessian for small-LeNet while Fig. 3 shows the eigenspectra for the other three networks. A large proportion of eigenvalues of the Hessian are very close to zero or positive with a very small (relative) magnitude. This suggests that the local geometry of the energy landscape is almost flat at local minima discovered by gradient descent. This agrees with theoretical results such as Baldassi et al. (2016c) where the authors predict that flat regions of the landscape generalize better. Standard regularizers in deep learning such as convolutions, max-pooling and dropout seem to bias SGD towards flatter regions in the energy landscape. Away from the origin, the right tails of the eigenspectra are much longer than the left tails. Indeed, as discussed in numerous places in literature (Bray & Dean, 2007; Dauphin et al., 2014; Choromanska et al., 2015a), SGD finds low-index critical points, i.e., optimizers with few negative eigenvalues of the Hessian. What is interesting and novel is that the directions of descent that SGD misses do not have a large curvature.
5.2 MNIST
We consider two prototypical networks: the first, called mnistfc, is a fully-connected network with two hidden layers of 1024 units each and the second is a convolutional neural network with the same size as LeNet but with batch-normalization (Ioffe & Szegedy, 2015); both use a dropout of probability 0.5 after each layer. As a baseline, we train for 100 epochs with Adam and a learning rate of 10−3 that drops by a factor of 5 after every 30 epochs to obtain an average error of 1.39±0.03% and 0.51±0.01% for mnistfc and LeNet respectively, over 5 independent runs. For both these networks, we train Entropy-SGD for 5 epochs with L = 20 and reduce the dropout probability (0.15 for mnistfc and 0.25 for LeNet). The learning rate of the SGLD updates is fixed to η ′ = 0.1 while the outer loop’s learning rate is set to η = 1 and drops by a factor of 10 after the second epoch; we use Nesterov’s momentum for both loops. The thermal noise in SGLD updates (line 5 of Alg. 1) is set to 10−3. We use an exponentially increasing value of γ for scoping, the initial value of the scope is set to γ = 10−4 and this increases by a factor of 1.001 after each parameter update. The results in Fig. 4a and Fig. 4b show that Entropy-SGD obtains a comparable generalization error: 1.37± 0.03% and 0.50± 0.01%, for mnistfc and LeNet respectively. While Entropy-SGD trains slightly faster in wall-clock time for LeNet; it is marginally slower for mnistfc which is a small network and trains in about two minutes anyway.
Remark on the computational complexity: Since Entropy-SGD runs L steps of SGLD before each parameter update, the effective number of passes over the dataset is L times that of SGD or Adam for the same number of parameter updates. We therefore plot the error curves of Entropy-SGD in Figs. 4, 5, and 6 against the “effective number of epochs”, i.e. by multiplying the abscissae by a factor of L. (we define L = 1 for SGD or Adam). Modulo the time required for the actual parameter updates (which are fewer for Entropy-SGD), this is a direct measure of wall-clock time, agnostic to the underlying hardware and implementation.
5.3 CIFAR-10
We train on CIFAR-10 without data augmentation after performing global contrast normalization and ZCA whitening (Goodfellow et al., 2013). We consider the All-CNN-C network of Springenberg et al. (2014) with the only difference that a batch normalization layer is added after each convolutional layer; all other architecture and hyper-parameters are kept the same. We train for 200 epochs with SGD and Nesterov’s momentum during which the initial learning rate of 0.1 decreases by a factor of 5 after every 60 epochs. We obtain an average error of 7.71± 0.19% in 200 epochs vs. 9.08% error in 350 epochs that the authors in Springenberg et al. (2014) report and this is thus a very competitive baseline for this network. Let us note that the best result in the literature on non-augmented CIFAR-10 is the ELU-network by Clevert et al. (2015) with 6.55% test error.
We train Entropy-SGD with L = 20 for 10 epochs with the original dropout of 0.5. The initial learning rate of the outer loop is set to η = 1 and drops by a factor of 5 every 4 epochs, while the learning rate of the SGLD updates is fixed to η ′ = 0.1 with thermal noise ε = 10−4. As the scoping scheme, we set the initial value of the scope to γ0 = 0.03 which increases by a factor of 1.001 after each parameter update. Fig. 5 shows the training and validation error curves for Entropy-SGD compared with SGD. It shows that local entropy performs as well as SGD on a large CNN; we obtain a validation error of 7.81±0.09% in about 160 effective epochs. We see almost no plateauing of training loss or validation error for Entropy-SGD in Fig. 5a; this trait is shared across different networks and datasets in our experiments and is an indicator of the additional smoothness of the local entropy landscape coupled with a good scoping schedule for γ .
5.4 RECURRENT NEURAL NETWORKS
5.4.1 PENN TREE BANK
We train an LSTM network on the Penn Tree Bank (PTB) dataset for word-level text prediction. This dataset contains about one million words divided into a training set of about 930,000 words, a validation set of 74,000 words and 82,000 words with a vocabulary of size 10,000. Our network called PTB-LSTM consists of two layers with 1500 hidden units, each unrolled for 35 time steps; note that this is a large network with about 66 million weights. We recreated the training pipeline of Zaremba et al. (2014) for this network (SGD without momentum) and obtained a word perplexity of 81.43± 0.2 on the validation set and 78.6± 0.26 on the test set with this setup; these numbers closely match the results of the original authors.
We run Entropy-SGD on PTB-LSTM for 5 epochs with L = 5, note that this results in only 25 effective epochs. We do not use scoping for this network and instead fix γ = 10−3. The initial learning rate of the outer loop is η = 1 which reduces by a factor of 10 at each epoch after the
third epoch. The SGLD learning rate is fixed to η ′ = 1 with ε = 10−4. We obtain a word perplexity of 80.116± 0.069 on the validation set and 77.656± 0.171 on the test set. As Fig. 6a shows, Entropy-SGD trains significantly faster than SGD (25 effective epochs vs. 55 epochs of SGD) and also achieves a slightly better generalization perplexity.
5.4.2 CHAR-LSTM ON WAR AND PEACE
Next, we train an LSTM to perform character-level text-prediction. As a dataset, following the experiments of Karpathy et al. (2015), we use the text of War and Peace by Leo Tolstoy which contains about 3 million characters divided into training (80%), validation (10%) and test (10%) sets. We use an LSTM consisting of two layers of 128 hidden units unrolled for 50 time steps and a vocabulary of size 87. We train the baseline with Adam for 50 epochs with an initial learning rate of 0.002 that decays by a factor of 2 after every 5 epochs to obtain a validation perplexity of 1.224±0.008 and a test perplexity of 1.226±0.01. As noted in Sec. 4.2, we can easily wrap Alg. 1 inside other variants of SGD such as Adam; this involves simply substituting the local entropy gradient in place of the usual back-propagated gradient. For this experiment, we constructed Entropy-Adam which is equivalent to Adam with the local entropy gradient (which is computed via SGLD). We run Entropy-Adam for 5 epochs with L = 5 and a fixed γ = 0.01 with an initial learning rate of 0.01 that decreases by a factor of 2 at each epoch. Note that this again results in only 25 effective epochs, i.e. half as much wall-clock time as SGD or Adam. We obtain a validation perplexity of 1.213±0.007 and a test perplexity of 1.217±0.005 over 4 independent runs which is better than the baseline. Fig. 6b shows the error curves for this experiment.
Tuning the momentum in Entropy-SGD was crucial to getting good results on RNNs. While the SGD baseline on PTB-LSTM does not use momentum (and in fact, does not train well with momentum) we used a momentum of 0.5 for Entropy-SGD. On the other hand, the baseline for char-LSTM was trained with Adam with β1 = 0.9 (β1 in Adam controls the moving average of the gradient) while we set β1 = 0.5 for Entropy-Adam. In contrast to this observation about RNNs, all our experiments on CNNs used a momentum of 0.9. In order to investigate this difficulty, we monitored the angle between the local entropy gradient and the vanilla SGD gradient during training. This angle changes much more rapidly for RNNs than for CNNs which suggests a more rugged energy landscape for the former; local entropy gradient is highly uncorrelated with the SGD gradient in this case.
6 DISCUSSION
In our experiments, Entropy-SGD results in generalization error comparable to SGD, but always with lower cross-entropy loss on the training set. This suggests the following in the context of energy landscapes of deep networks. Roughly, wide valleys favored by Entropy-SGD are located deeper in the landscape with a lower empirical loss than local minima discovered by SGD where it presumably gets stuck. Such an interpretation is in contrast to theoretical models of deep networks (cf. Sec. 2) which predict multiple equivalent local minima with the same loss. Our work suggests that geometric properties of the energy landscape are crucial to generalize well and provides algorithmic approaches to exploit them. However, the literature lacks general results about the geometry of the loss functions of deep networks — convolutional neural networks in particular — and this is a promising direction for future work.
If we focus on the inner loop of the algorithm, SGLD updates compute the average gradient (with Langevin noise) in a neighborhood of the parameters while maintaining the Polyak average of the new parameters. Such an interpretation is very close to averaged SGD of Polyak & Juditsky (1992) and Bottou (2012) and worth further study. Our experiments show that while Entropy-SGD trains significantly faster than SGD for recurrent networks, it gets relatively minor gains in terms of wallclock time for CNNs. Estimating the gradient of local entropy cheaply with few SGLD iterations, or by using a smaller network to estimate it in a teacher-student framework (Balan et al., 2015) is another avenue for extensions to this work.
7 CONCLUSIONS
We introduced an algorithm named Entropy-SGD for optimization of deep networks. This was motivated from the observation that the energy landscape near a local minimum discovered by SGD is almost flat for a wide variety of deep networks irrespective of their architecture, input data or training methods. We connected this observation to the concept of local entropy which we used to bias the optimization towards flat regions that have low generalization error. Our experiments showed that Entropy-SGD is applicable to large convolutional and recurrent deep networks used in practice.
8 ACKNOWLEDGMENTS
This work was supported by ONR N00014-13-1-034, AFOSR F9550-15-1-0229 and ARO W911NF-15-1-0564/66731-CS.
A STOCHASTIC GRADIENT LANGEVIN DYNAMICS (SGLD)
Local entropy in Def. (1) is an expectation over the entire configuration space x ∈ Rn and is hard to compute; we can however approximate its gradient using Markov chain Monte-Carlo (MCMC) techniques. In this section, we briefly review stochastic gradient Langevin dynamics (Welling & Teh, 2011) that is an MCMC algorithm designed to draw samples from a Bayesian posterior and scales to large datasets using mini-batch updates.
For a parameter vector x ∈ Rn with a prior distribution p(x) and if the probability of generating a data item ξk given a model parameterized by x is p(ξk |x), the posterior distribution of the parameters based on N data items can be written as
p(x |ξ k≤N) ∝ p(x) N
∏ k=1 p(ξk |x) . (10)
Langevin dynamics (Neal, 2011) injects Gaussian noise into maximum-a-posteriori (MAP) updates to prevent over-fitting the solution x∗ of the above equation. The updates can be written as
∆xt = η 2
( ∇ log p(xt)+ N
∑ k=1
∇p(ξk |xt) ) + √ η εt ; (11)
where εt ∼ N(0,ε2) is Gaussian noise and η is the learning rate. In this form, Langevin dynamics faces two major hurdles for applications to large datasets. First, computing the gradient ∑Nk=1 ∇p(ξk |xt) over all samples for each update ∆xt becomes prohibitive. However, as Welling & Teh (2011) show, one can instead simply use the average gradient over m data samples (mini-batch) as follows:
∆xt = ηt 2
( ∇ log p(xt)+
N m
m
∑ k=1
∇p(ξk |xt) ) + √ ηt εt . (12)
Secondly, Langevin dynamics in (11) is the discrete-time approximation of a continuous-time stochastic differential equation (Mandt et al., 2016) thereby necessitating a Metropolis-Hastings (MH) rejection step (Roberts & Stramer, 2002) which again requires computing p(ξk |x) over the entire dataset. However, if the learning rate ηt → 0, we can also forgo the MH step (Chen et al., 2014). Welling & Teh (2011) also argue that the sequence of samples xt generated by updating (12) converges to the correct posterior (10) and one can hence compute the statistics of any function g(x) of the parameters using these samples. Concretely, the posterior expectation E [g(x)] is given by E [g(x)]≈ ∑
t s=1 ηt g(xt ) ∑ts=1 ηt
; which is the average computed by weighing each sample by the corresponding learning rate in (12). In this paper, we will consider a uniform prior on the parameters x and hence the first term in (12), viz., ∇ log p(xt) vanishes.
Let us note that there is a variety of increasingly sophisticated MCMC algorithms applicable to our problem, e.g., Stochastic Gradient Hamiltonian Monte Carlo (SGHMC) by Chen et al. (2014) based on volume preserving flows in the “parameter-momentum” space, stochastic annealing thermostats (Santa) by Chen et al. (2015) etc. We can also employ these techniques, although we use SGLD for ease of implementation; the authors in Ma et al. (2015) provide an elaborate overview.
B PROOFS
Proof of Lemma 2. The gradient −∇F(x) is computed in Sec. 4.1 to be γ ( x− 〈 x′; Ξ` 〉) . Consider the term
x− 〈 x′; x 〉 = x−Z−1x,γ ∫ x′ x′ e− f (x ′)− γ2 ‖x−x′‖2 dx′
≈ x−Z−1x,γ ∫
s (x+ s) e− f (x)−∇ f (x) >s− 12 s>(γ+∇2 f (x))s ds
= x ( 1−Z−1x,γ ∫
s e− f (x)−∇ f (x)
>s− 12 s>(γ+∇2 f (x))s ds ) −Z−1x,γ ∫ s s e− f (x)−∇ f (x) >s− 12 s>(γ+∇2 f (x))s ds
=−Z−1x,γ e− f (x) ∫
s s e−∇ f (x) >s− 12 s>(γ+∇2 f (x))s ds.
The above expression is the mean of a distribution ∝ e−∇ f (x) >s− 12 s>(γ+∇2 f (x))s. We can approximate it using the saddle point method as the value of s that minimizes the exponent to get
x− 〈 x′; x 〉 ≈ ( ∇2 f (x)+ γ I )−1 ∇ f (x). Let us denote A(x) := ( I + γ−1 ∇2 f (x)
)−1. Plugging this into the condition for smoothness, we have ‖∇F(x,γ)−∇F(y,γ)‖= ‖A(x) ∇ f (x)−A(y) ∇ f (y)‖
≤ (
sup x ‖A(x)‖
) β ‖x− y‖.
Unfortunately, we can only get a uniform bound if we assume that for a small constant c > 0, no eigenvalue of ∇2 f (x) lies in the set [−2γ− c,c]. This gives(
sup x ‖A(x)‖
) ≤ 1
1+ γ−1 c .
This shows that a smaller value of γ results in a smoother energy landscape, except at places with very flat directions. The Lipschitz constant also decreases by the same factor.
C CONNECTION TO VARIATIONAL INFERENCE
The fundamental motivations of (stochastic) variational inference (SVI) and local entropy are similar: they both aim to generalize well by constructing a distribution on the weight space. In this section, we explore whether they are related and how one might reconcile the theoretical and algorithmic implications of the local entropy objective with that of SVI.
Let Ξ denote the entire dataset, z denote the weights of a deep neural network and x be the parameters of a variational distribution qx(z). The Evidence Lower Bound (ELBO) can be then be written as
log p(Ξ)≥ Ez∼qx(z) [log p(Ξ | z)]−KL(qx(z) || p(z)) ; (13) where p(z) denotes a parameter-free prior on the weights and controls, through their KL-divergence, how well the posited posterior qx(z) fits the data. Stochastic variational inference involves maximizing the right hand side of the above equation with respect to x after choosing a suitable prior p(z) and a family of distributions qx(z). These choices are typically dictated by the ease of sampling z ∼ qx(z), e.g. a mean-field model where qx(z) factorizes over z, and being able to compute the KL-divergence term, e.g. a mixture of Gaussians.
On the other hand, if we define the loss as the log-likelihood of data, viz. f (z) := − log p(Ξ|z), we can write the logarithm of the local entropy in Eqn. (4) as
logF(x,γ) = log ∫ z ∈ Rn exp [ − f (z; Ξ)− γ 2 ‖x− z‖2 ] dz,
≥ ∫
z ∈ Rn
[ log p(Ξ | z)− γ
2 ‖x− z‖2
] dz; (14)
by an application of Jensen’s inequality. It is thus clear that Eqn. (13) and (14) are very different in general and one cannot choose a prior, or a variational family, that makes them equivalent and interpret local entropy as ELBO.
Eschewing rigor, formally, if we modify Eqn. (13) to allow the prior p(z) to depend upon x, we can see that the two lower bounds above are equivalent iff qx(z) belongs to a “flat variational family”, i.e. uniform distributions with x as the mean and px(z) ∝ exp ( − γ2 ‖x− z‖2 ) . We emphasize that the distribution px(z) depends on the parameters x themselves and is thus, not really a prior, or one that can be derived using the ELBO.
This “moving prior” is absent in variational inference and indeed, a crucial feature of the local entropy objective. The gradient of local entropy in Eqn. (7) clarifies this point:
∇F(x,γ) =−γ (x−〈z; Ξ〉) =−γ Ez ∼r(z;x) [z] ;
where the distribution r(z;x) is given by r(z; x) ∝ p(Ξ | z) exp ( −γ
2 ‖x− z‖2
) ;
it thus contains a data likelihood term along with a prior that “moves” along with the current iterate x.
Let us remark that methods in the deep learning literature that average the gradient through perturbations in the neighborhood of x (Mobahi, 2016) or noisy activation functions (Gulcehre et al., 2016) can be interpreted as computing the data likelihood in ELBO (without the KL-term); such an averaging is thus different from local entropy.
C.1 COMPARISON WITH SGLD
We use stochastic gradient Langevin dynamics (cf. Appendix A) to estimate the gradient of local entropy in Alg. 1. It is natural then, to ask the question whether vanilla SGLD performs as well as local entropy. To this end, we compare the performance of SGLD on two prototypical networks: LeNet on MNIST and All-CNN-BN on CIFAR-10. We follow the experiments in Welling & Teh (2011) and Chen et al. (2015) and set the learning rate schedule to be η/(1+ t)b where the initial learning rate η and b are hyper-parameters. We make sure that other architectural aspects (dropout, batch-normalization) and regularization (weight decay) are consistent with the experiments in Sec. 5.
After a hyper-parameter search, we obtained a test error on LeNet of 0.63±0.1% after 300 epochs and 9.89± 0.11% on All-CNN-BN after 500 epochs. Even if one were to disregard the slow convergence of SGLD, its generalization error is much worse than our experimental results; we get 0.50± 0.01% on LeNet and 7.81± 0.09% on All-CNN-BN with Entropy-SGD. For comparison, the authors in Chen et al. (2015) report 0.71% error on MNIST on a slightly larger network. Our results with local entropy on RNNs are much better than those reported in Gan et al. (2016) for SGLD. On the PTB dataset, we obtain a test perplexity of 77.656± 0.171 vs. 94.03 for the same model whereas we obtain a test perplexity of 1.213±0.007 vs. 1.3375 for char-LSTM on the War and Peace dataset.
Training deep networks with SGLD, or other more sophisticated MCMC algorithms such as SGHMC, SGNHT etc. (Chen et al., 2014; Neal, 2011) to errors similar to those of SGD is difficult, and the lack of such results in the literature corroborates our experimental experience. Roughly speaking, local entropy is so effective because it operates on a transformation of the energy landscape that exploits entropic effects. Conventional MCMC techniques such as SGLD or Nose’-Hoover thermostats (Ding et al., 2014) can only trade energy for entropy via the temperature parameter which does not allow the direct use of the geometric information of the energy landscape and does not help with narrow minima.

Overview:
This paper introduces a biasing term for SGD that, in theoretical results and a toy example, yields solutions with an approximately equal or lower generalization error. This comes at a computational cost of estimating the gradient of the biasing term for each iteration through stochastic gradient Langevin dynamics, approximating an MCMC sample of the log partition function of a modified Gibbs distribution. The cost is equivalent to adding an inner for-loop to the standard SGD algorithm for each minibatch.
Pros:
- Reviews and distills many results and theorems from past 2 decades that suggest a promising way forward for increasing the generalizability of deep neural networks
- Generally very well written and well presented results, with interesting discussion of eigenvalues of Hessian as a way to characterize “flat” minima
- Promising mathematical arguments suggest that E-SGD has generalization error bounded below by SGD, motivating further research in the area
Cons / points suggested for a rebuttal:
(1) One claim of the paper given in the abstract is ”experiments on competitive baselines demonstrate that Entropy-SGD leads to improved generalization and has the potential to accelerate training.“ This does not appear to be supported by the current set of experiments. As the authors comment in the discussion section, “In our experiments, Entropy-SGD results in a comparable generalization error as SGD, but always has a lower cross-entropy loss.” It's not clear to me how to reconcile those two claims.
(2) Similarly, the claim of accelerated training is not convincingly supported in the present version of the paper. Vanilla SGD requires a single forward pass through all M minibatches during one epoch for a parameter update, but the new method, E-SGD requires, L*M forward passes during one epoch where L is the number of Langevin updates, which require a minibatch sample each. This could in fact mean that E-SGD has worse computational complexity to reach the same point. In a remark on p.9, the authors note that a single epoch is defined to be “the number of parameter updates required to run through the dataset once.” It’s not clear to me how this answers the objection to a factor of L additional computations required for the inner-loop SGLD iterations. SGLD appears to introduces a potentially costly tradeoff that must be carefully managed by a user of E-SGD.
(3) As the previous two points suggest, the paper could use some attention to the magnitude of the claims. For example, the introduction reads “Actively biasing towards wide valleys aids generalization, in fact, we can optimize solely the free energy term to obtain similar generalization error as SGD on the original loss function.“ According the the values reported on pp.9-10, only on MNIST is the generalization error, using only the free energy term (the log partition function of the modified Gibbs distribution), equivalent to using only the SGD loss function. This corresponds to setting rho to 0 in equation (6). On CIFAR-10, rho = 0.01 is used.
(4) Another contribution of this paper, the characterization of the optimization landscape in terms of the eigenvalues of the Hessian and low generalization error being associated with flat local extrema, is helpful and interesting. I found the plots clear and useful. As another reviewer has already pointed out, there are high-level similarities to “Flat Minima” by Hochreiter and Schmidhuber (1997). The authors have responded already by adding a paragraph that helpfully explores some differences with H&S 1997. However, the similarities should also be carefully identified and mentioned. H&S 1997 includes detailed theoretical analysis that could be helpful for future work in this area, and has independently discovered a similar approach to training generalizable networks.
(5) It's not clear how the assumption about the eigenvalues that were made in section 4.4 / Appendix B affect the application of this result to real-world problems. What magnitude of c>0 needs to be chosen? Does this correspond to a measurable characteristic of the dataset? It's a little mysterious in the current version of the paper.