Title: LEARNING FROM THE EXPERIENCE OF OTHERS: APPROXIMATE EMPIRICAL BAYES IN NEURAL NETWORKS

Abstract: Learning deep neural networks could be understood as the combination of representation learning and learning halfspaces. While most previous work aims to diversify representation learning by data augmentations and regularizations, we explore the opposite direction through the lens of empirical Bayes method. Specifically, we propose a matrix-variate normal prior whose covariance matrix has a Kronecker product structure to capture the correlations in learning different neurons through backpropagation. The prior encourages neurons to learn from the experience of others, hence it provides an effective regularization when training large networks on small datasets. To optimize the model, we design an efficient block coordinate descent algorithm with analytic solutions. Empirically, we show that the proposed method helps the network converge to better local optima that also generalize better, and we verify the effectiveness of the approach on both multiclass classification and multitask regression problems with various network structures.

1. INTRODUCTION:

Empirical Bayes methods provide us a powerful tool to obtain Bayesian estimators even if we do not have complete information about prior distribution. The literature on the empirical Bayes methods and their applications are abundant (Stein, 1956; Robbins, 1956; James & Stein, 1961; Efron & Morris, 1973; 1977; Efron et al., 2001; Carlin & Louis, 2010; Efron, 2012) . Existing studies on parametric empirical Bayes methods focus on the setting where the likelihood function and the prior are assumed to have specific forms, e.g., exponential family distribution and its conjugate prior, so that marginal distribution of data has a closed form from which an estimator of the hyperparameter in the prior distribution can be obtained. While such assumption helps to simplify the setting in order to demonstrate the power of the method, it restricts us from using more expressive and rich models.

In this paper we explore extending the empirical Bayes method to expressive nonlinear models using deep neural networks. Although deep neural networks have been widely applied in various domains (Krizhevsky & Hinton, 2009; LeCun et al., 2015; He et al., 2016) , usually its parameters are learned via the principle of maximum likelihood, hence its success crucially hinges on the availability of large scale datasets. On the other hand, Bayesian modeling lends us a powerful and principled tool to prevent overfitting by incorporating prior knowledge into the design of prior distribution. To this end, we propose a regularization approach for the weight matrix in neural networks through the lens of the empirical Bayes method. We aim to address the problem of overfitting when training large networks on small dataset. Our key insight stems from the famous argument by Efron (2012) : It is beneficial to learn from the experience of others. Specifically, from an algorithmic perspective, we argue that the connection weights of neurons in the same layer (row/column vectors of the weight matrix) will be correlated with each other through the backpropagation learning. Hence by learning from other neurons in the same layer, a neuron can "borrow statistical strength" from other neurons.

As an illustrating example, consider a simple setting where the input x ∈ R d is fully connected to a hidden layer h ∈ R p , which is further fully connected to the single output ŷ ∈ R. Let σ(•) be the nonlinear activation function, e.g., ReLU (Nair & Hinton, 2010) , W ∈ R p×d be the connection matrix between the input layer and the hidden layer, and a ∈ R p be the vector connecting the output and the hidden layer. Without loss of generality, ignoring the bias term in each layer, we have: ŷ = a T h, h = σ(W x). Consider using the usual 2 loss function (ŷ, y) = 1 2 |ŷ -y| 2 and take the derivative of (ŷ, y) w.r.t. W . We obtain the update formula in backpropagation as W ← W -α(ŷ -y)(a • h ) x T , where h is the componentwise derivative of h w.r.t. its input argument, and α > 0 is the learning rate. Realize that (a • h ) x T is a rank 1 matrix, and the component of h is either 0 or 1. Hence the update for each row vector of W is linearly proportional to x. Note that the observation holds for any input pair (x, y), so the update formula implies that the row vectors of W are correlated with each other. The above observation leads us to the following question: can we define a prior distribution over W that captures the correlations between its row vectors and column vectors?

Our Contributions. To answer the above question, we develop an approximate empirical Bayes (AEB) framework to learn deep neural networks. Motivated by the example above, we propose a matrix-variate normal prior whose covariance matrix has a Kronecker product structure to capture the correlations between different neurons. The prior encourages neurons to learn from the experience of others, hence it provides an effective regularization when training networks on small datasets. Using tools from convex analysis, we design an efficient block coordinate descent algorithm with analytic solutions to optimize the model. Empirically, we show that the proposed method helps the network converge to better local optima that also generalize better, and we verify the effectiveness of the approach on both multiclass classification and multitask regression problems with various network structures.

2. PRELIMINARY:

We first introduce the notation used throughout the paper and then give a brief discussion on the empirical Bayes method (Bernardo & Smith, 2001; Gelman et al., 2013; Efron & Hastie, 2016) .

2.1. NOTATION AND SETUP:

We use lowercase letter to represent scalar and lowercase bold letter to denote vector. Capital letter, e.g., X, is reserved for matrix. Calligraphic letter, such as D, is used to denote set. We write Tr(A) as the trace of a matrix A, det(A) as the determinant of A and vec(A) as A's vectorization by column.

[n] is used to represent the set {1, . . . , n} for any integer n. Other notations will be introduced whenever needed. Suppose we have access to a training set D of n pairs of data instances (x i , y i ), i ∈ [n]. We consider the supervised learning setting where x i ∈ X ⊆ R d and y i ∈ Y. For a regression problem, Y = R; for a binary classification problem, Y = {1, -1}. Let p(y | x, w) be the conditional distribution of y given x with parameter w. The parametric form of the conditional distribution is assumed be known. In this paper, we consider a Bayesian setting where the model parameter w is sampled from a prior distribution p(w | θ) with hyperparameter θ. On the other hand, given D, the posterior distribution of w is denoted by p(w | D, θ). From a Bayesian perspective, given an unseen instance x, the goal is to infer the predictive distribution: p(y | x, D, θ) = p(y | x, w) • p(w | D, θ) dw, from which we can compute the mean, or the median, or other statistic (depending on the choice of the loss function) as our estimator of the unseen target variable y.

2.2. THE EMPIRICAL BAYES METHOD:

To compute the predictive distribution, we need access to the value of the hyperparameter θ. However, complete information about the hyperparameter θ is usually not available in practice. To this end, empirical Bayes method (Robbins, 1956; Efron & Morris, 1973) proposes to estimate θ from the data directly using the marginal distribution:

θ = arg max θ p(D | θ) = arg max θ p(D | w) • p(w | θ) dw.

(1)

Under specific choice of the likelihood function p(x, y | w) and the prior distribution p(w | θ), e.g., exponential family distribution and its corresponding conjugate prior, we can solve the integral in (1) in closed form to obtain an analytic solution of θ, which can be subsequently plugged into the prior distribution to obtain a Bayesian estimator for the model parameter w.

At a high level, by learning the hyperparameter θ in the prior distribution directly from data, the empirical Bayes method provides us a principled and convenient way to obtain a Bayesian estimator of the model parameter w. In fact, when both the prior and the likelihood functions are normal, it has been formally shown that the empirical Bayes estimators, e.g., the James-Stein estimator (James & Stein, 1961) and the Efron-Morris estimator (Efron & Morris, 1977) , dominate the classic maximum likelihood estimator (MLE) in terms of quadratic loss for every choice of the model parameter w. At a colloquial level, the success of empirical Bayes estimators can be attributed to the effect of "learning from the experience of others" (Efron, 2012) , which also makes it a powerful tool in multitask learning (Long et al., 2017; Zhao et al., 2017) and meta-learning (Grant et al., 2018) .

3 LEARNING WITH APPROXIMATE EMPIRICAL BAYES 

EQUATION

where ŵ = arg max w p(D | w) • p(w | θ) is the mode of the joint distribution. The above approximation is crude and will only be accurate if 1) the likelihood under ŵ dominates the likelihoods under other model parameters, or 2) for the fixed θ, the prior distribution p(w | θ) is sharply concentrated around its mode. When the size of the dataset D is large, the first condition is met according to the central limit theorem under some regularity conditions. We shall come back later to verify the validity of this approximation through numerical experiments in Sec. 4.

Given an estimate ŵ, by maximizing the R.H.S. of (2) w.r.t. θ, we can obtain θ as an approximation of the maximum marginal likelihood estimator. As a result, we can use θ to further refine the estimate ŵ by maximizing the posterior distribution as follows:

EQUATION

The maximizer of (3) can in turn be used to better approximate the integral in (2). Formally, we can define the following optimization problem that characterizes our framework of the approximate empirical Bayes (AEB) method:

EQUATION

It is worth to connect the optimization problem (4) to the classic maximum a posteriori (MAP) inference and also discuss their difference. If we drop the inner optimization over the hyperparameter θ in the prior distribution. Then for any fixed value θ, (4) reduces to MAP with the prior defined by the specific choice of θ, and the maximizer ŵ corresponds to the mode of the posterior distribution given by θ. From this perspective, the optimization problem in (4) actually defines a series of MAP inference problems, and the sequence { ŵj ( θj )} j defines a solution path towards the final approximate empirical Bayes estimator. On the algorithmic side, the optimization problem (4) also suggests a natural block coordinate descent algorithm where we alternatively optimize over w and θ until the convergence of the objective function. An illustration of the framework is shown in Fig. 1 . In next section, we give a specific prior distribution over the parameters of neural networks to capture the fact that neurons in the same layer are correlated with each other.

3.2. NEURAL NETWORK WITH MATRIX-NORMAL PRIOR:

Inspired by the observation from Sec. 1, we propose to define a matrix-variate normal distribution (Gupta & Nagar, 2018) over the connection weight matrix W : W ∼ MN (0 p×d , Σ r , Σ c ), where Σ r ∈ S p ++ and Σ c ∈ S d ++ are the row and column covariance matrices, respectively. 1 Equivalently, one can understand the matrix-variate normal distribution over W as a multivariate normal distribution with a Kronecker product covariance structure over vec(W ): vec(W ) ∼ N (0 p×d , Σ c ⊗ Σ r ). It 1 The probability density function is given by p

(W | Σr, Σc) = exp(-Tr(Σ -1 r W Σ -1 c W T )/2) (2π) pd/2 det(Σr ) d/2 det(Σc) p/2 .

Under review as a conference paper at ICLR 2019 is then easy to check that the marginal prior distributions over the row and column vectors of the weight matrix W are given by:

W i: ∼ N (0 d , [Σ r ] ii • Σ c ), W :j ∼ N (0 p , [Σ c ] jj • Σ r ).

We point out that the Kronecker product structure of the covariance matrix exactly captures our prior about the connection matrix W : the fan-in/fan-out of neurons in the same layer (row/column vectors of W ) are correlated with the same correlation matrix in the prior, and they only differ at the scales (variances).

For illustration purpose, let us consider the simple feed-forward network discussed in Sec. 1. Consider a reparametrization of the model by defining Ω r := Σ -1 r and Ω c := Σ -1 c to be the corresponding precision matrices and plug in the prior distribution into the general approximate empirical Bayes framework (4). After routine algebraic simplifications, we reach the following concrete optimization problem:

EQUATION

)

where λ is a constant that only depends on p and d, 0 < u ≤ v and uv = 1. Note that the constraint is necessary to guarantee the feasible set to be compact so that the optimization problem is well formulated and a minimum is attainable.2 It is not hard to show that in general the optimization problem (5) is not jointly convex in terms of {a, W, Ω r , Ω c }, and this holds even if the activation function is linear and we do not have the hidden layer. However, as we will show later, for any fixed a, W , the reparametrization makes the partial optimization over Ω r and Ω c bi-convex. More importantly, we can derive an efficient algorithm that finds the optimal Ω r (Ω c ) for any fixed a, W, Ω c (Ω r ) in O(max{d 3 , p 3 }) time with closed form solutions. This allows us to apply our algorithm to networks of large sizes, where a typical hidden layer can contain thousands of nodes. Before we delve into the details on solving (5), it is instructive to discuss some of its connections and differences to other learning paradigms.

Maximum-A-Posteriori Estimation. Essentially, for model parameter W , (5) defines a sequence of MAP problems where each MAP is indexed by the pair of precision matrices (Ω

EQUATION

c ) at iteration t. Equivalently, at each stage of the optimization, we can interpret (5) as placing a matrix variate normal prior on W where the precision matrix in the prior is given by Ω ∀t, then (5) naturally reduces to learning with 2 regularization, or weight decay (Krogh & Hertz, 1992) . More generally, for non-diagonal precision matrices, the regularization term for W becomes:

(t) r ⊗ Ω (t) c . From this perspective, if we fix Ω (t) r = I p and Ω (t) c = I d ,

EQUATION

)

and this is exactly the Tikhonov regularization (Golub et al., 1979) imposed on W where the Tikhonov matrix Γ is given by

Γ := Ω 1/2 c ⊗ Ω 1/2

r . But instead of manually designing the regularization matrix Γ to improve the conditioning of the estimation problem, under the principle of empirical Bayes we propose to also learn both precision matrices (so Γ as well) from data.

Approximate Volume Minimization. Let us consider the log det(•) function over the positive definite cone. It is well known that the log-determinant function is concave (Boyd & Vandenberghe, 2004) . Hence for any pair of matrices A 1 , A 2 ∈ S m ++ , the following inequality holds:

log det(A 1 ) ≤ log det(A 2 ) + ∇ log det(A 2 ), A 1 -A 2 = log det(A 2 ) + Tr(A -1 2 A 1 ) -m (7) Applying the above inequality twice by fixing A 1 = W Ω c W T /2d, A 2 = Σ r and A 1 = W T Ω r W/2p, A 2 = Σ c

respectively leads to the following inequalities:

d log det(W Ω c W T /2d) ≤ -d log det(Ω r ) + 1 2 Tr(Ω r W Ω c W T ) -dp p log det(W T Ω r W/2p) ≤ -p log det(Ω c ) + 1 2 Tr(Ω r W Ω c W T ) -dp Using the fact that Tr(Ω r W Ω c W T ) = ||Ω 1/2 r W Ω 1/2 c || 2

F , we immediately have:

d log det(W ΩcW T ) + p log det(W T ΩrW ) ≤ ||Ω 1/2 r W Ω 1/2 c || 2 F -(d log det(Ωr) + p log det(Ωc)) + c (8)

where c is a constant that only depends on d and p. Recall that | det(A T A)| computes the squared volume of the parallelepiped spanned by the column vectors of A. Hence (8) gives us a natural interpretation of the objective function in (5): the regularizer essentially upper bounds the log-volume of the two parallelpipeds spanned by the row and column vectors of W . But instead of measuring the volume using standard Euclidean inner product, it also takes into account the local curvatures defined by Σ r and Σ c , respectively. For vectors with fixed lengths, the volume of the parallelepiped spanned by them becomes smaller when they are more linearly correlated, either positively or negatively. At a colloquial level, this means that the regularizer in (5) forces fan-in/fan-out of neurons at the same layer to be either positively or negatively correlated with each other, and this corresponds exactly to the effect of learning from the experience of others.

3.3. THE ALGORITHMS:

Algorithm 1 Block Coordinate Descent for Approximate Empirical Bayes

Input: Initial value φ (0) := {a (0) , W (0) }, Ω (0) r ∈ S p ++ and Ω (0) c ∈ S d ++ , first-order optimization algorithm A, constants 0 < u ≤ v. 1: for t = 1, . . . , ∞ until convergence do 2:

Fix Ω (t-1) r

, Ω

(t-1) c

, optimize φ (t) by backpropagation and algorithm A 3:

Ω (t) r ← INVTHRESHOLDING(W (t) Ω (t-1) c W (t)T , d, u, v) 4: Ω (t) c ← INVTHRESHOLDING(W (t)T Ω (t)

r W (t) , p, u, v) 5: end for 6: 7: procedure INVTHRESHOLDING(∆, m, u, v) 8:

Compute SVD: Qdiag(r)Q T = SVD(∆) 9:

Hard thresholding r ← T [u,v] (m/r) 10:

return Qdiag(r )Q T 11: end procedure

In this section we describe a block coordinate descent algorithm to optimize the objective function in (5) and detail how to efficiently solve the matrix optimization subproblems in closed form using tools from convex analysis. Due to space limit, we defer all the proofs to appendix. Given a pair of constants 0 < u ≤ v, we define the following thresholding function T [u,v] (x):

T [u,v] (x) := max{u, min{v, x}}.

(9) We summarize our block coordinate descent algorithm to solve (5) in Alg. 1. In each iteration, Alg. 1 takes a first-order algorithm A, e.g., the stochastic gradient descent, to optimize the parameters of the neural network by backpropagation. It then proceeds to compute the optimal solutions for Ω r and Ω c using INVTHRESHOLDING as a sub-procedure. Alg. 1 terminates when a stationary point is found.

We now proceed to show that the procedure INVTHRESHOLDING finds the optimal solution given all the other variables fixed. Due to the symmetry between Ω r and Ω c in (5), we will only prove this for Ω r , and similar arguments can be applied to Ω c as well. Fix both W , Ω c and ignore all the terms that do not depend on Ω r , the sub-problem on optimizing Ω r becomes:

EQUATION

)

Proposition 1. The optimization problem (10) is convex.

Define the constraint set C := {A ∈ S p ++ | uI p A vI p } and the indicator function

I C (A) = 0 iff A ∈ C else ∞.

Given the convexity of (10), we can use the indicator function to first transform (10) into an unconstrained one and use the first-order optimality condition to characterize the optimal solution:

EQUATION

where

N C (A) := {B ∈ S p | Tr(B T (Z -A)) ≤ 0, ∀Z ∈ C} is the normal cone w.r.t. C at A. Equivalently, we have Ω -1 r -W Ω c W T /d ∈ N C (Ω r ).

The following key lemma characterizes the structure of the normal cone w.r.t. C: 

Lemma 1. Let Ω r ∈ C, then N C (Ω r ) = -N C (Ω -1 r ). Lemma 1 implies W Ω c W T /d -Ω -1 r ∈ N C (Ω -1 r ).

( Ω r ) = QT [u,v] (Λ)Q T .

Corollary 1. Let W Ω c W T be eigendecomposed as Qdiag(r)Q T , then the optimal solution to ( 10) is given by QT

[u,v] (d/r)Q T .

Similar arguments can be made to derive the solution for Ω c in (5). The final algorithm is very simple as it only contains one SVD, hence its time complexity is O(max{d 3 , p 3 }). Note that the total number of parameters in the network is at least dp, hence the algorithm is efficient as it scales sub-quadratically in terms of number of parameters in the network.

4. EXPERIMENTS:

So far we develop our model and algorithms based on a simple neural network with one hidden layer and a single output. However, it is straightforward to extend the AEB framework to more sophisticated models with various structures. In this section we demonstrate the effect of our AEB method on learning practical deep neural networks.

4.1. EXPERIMENTAL SETUP:



Multiclass Classification (MNIST & CIFAR10)::

In the experiments, we show that AEB provides an effective regularization on the network parameters. To this end, we use a convolutional neural network as our baseline model. MNIST considers the following structure: CONV 5×5×1×10 -CONV 5×5×10×20 -FC 320×50 -FC 50×10 . The notation CONV 5×5×1×10 denotes the convolutional layer with kernel size 5 × 5 from depth 1 to 10; the notation FC 320×50 denotes the fully connected layer with size 320 × 50. Similarly, CIFAR10 considers the structure: CONV 5×5×3×10 -CONV 5×5×10×20 -FC 500×500 -FC 500×500 -FC 500×10 . To show the effect of regularization, we gradually increase the training set size. MNIST considers the step from 60 to 60,000 (11 different experiments) and CIFAR10 considers the step from 5,000 to 50,000 (10 different experiments). For each training set size, we repeat the experiments for 10 times. The mean along with its standard deviation are shown as the statistics. Moreover, since both the optimization and generalization of neural networks are sensitive to the size of minibatches (Keskar et al., 2016; Goyal et al., 2017) , we study two minibatch settings for 256 and 2048, respectively. In our AEB model, we place a matrix-variate normal prior over the weight matrix of the last softmax layer, and we use Alg. 1 to optimize both the model weights of the convolutional network and two covariance matrices of the weight matrix in the last layer.

Multitask Regression (SARCOS): SARCOS relates to an inverse dynamics problem for a seven degree-of-freedom (DOF) SARCOS anthropomorphic robot arm (Vijayakumar & Schaal, 2000) . The goal of this task is to map from a 21-dimensional input space (7 joint positions, 7 joint velocities, 7 joint accelerations) to the corresponding 7 joint torques. Hence there are 7 tasks and the inputs are shared among all the tasks. The training set and test set contain 44,484 and 4,449 examples, respectively. The network structure is given by FC 21×256 -FC 256×100 -FC 100×7 . Again, we apply our AEB method on the last layer weight matrix, where each row corresponds to a separate task vector.

We compare our AEB method with classic regularization methods in the literature, including weight decay, dropout (Srivastava et al., 2014) , batch normalization (BN) (Ioffe & Szegedy, 2015) and the DeCov method (Cogswell et al., 2015) . We also note that we fix all the hyperparameters such as learning rate to be the same for all the methods. To better understand the working mechanism of the proposed method, we report evaluation metrics on test set as a measure of generalization, the trajectory of the loss function during training, and the correlation of the weight matrix. (2017) . DeCov, on the other hand, improves the generalization in MNIST with batch size 256 (see Fig. 2 (a )), while it demonstrates only comparable or even worse performance in other settings. To conclude, as training set size grows, AEB consistently performs better generalization as comparing to other regularization methods. We also note that AEB is not sensitive to the size of minibatches while most of the methods suffer from large minibatches. In appendix, we show the combination of AEB with other generalization methods can usually lead to even better results.

4.2. RESULTS:



Multiclass:



Multitask Regression (SARCOS)::

In this experiment we are interested in investigating whether our AEB method can lead to better generalization for multiple related regression problems. To do so, we report the explained variance as a normalized metric, e.g., one minus the ratio between mean squared error and the variance of different methods in Table 1 . The larger the explained variance, the better the predictive performance. In this case we observe a consistent improvement of AEB over other competitors on all the 7 regression tasks. We would like to emphasize that all the experiments share exactly the same experimental protocol, including network structure, optimization algorithm, training iteration, etc, so that the performance differences can only be explained by different ways of regularizations. For better visualization, we also plot the result in appendix.

Optimization: It has recently been empirically shown that BN helps optimization not by reducing internal covariate shift, but instead by smoothing the landscape of the loss function (Santurkar et al., 2018) . To understand how AEB improves generalization, in Fig. 3 , we plot the values of the cross entropy loss function on both the training and test sets during optimization using Alg. 1. The experiment is performed in MNIST with batch size 2048. In this experiment, we fix the number of outer loop to be 2 and each block optimization over network weights contains 50 epochs. Because of the stochastic optimization over model weights, we can see several unstable peaks in function value around iteration 50 when trained with AEB, which corresponds to the transition phase between two consecutive outer loops with different row/column covariance matrices. In both cases AEB converges to better local optima of the loss landscape, which lead to better generalization on the test set as well because they have smaller loss values on the test set when compared with training without AEB. Stable rank and spectral norm: Given a matrix W , the stable rank of W , denoted as s-rank(W ), is defined as s-rank(W

) := ||W || 2 F /||W || 2 2 .

In other words, the stable rank of a matrix is the ratio of its squared Frobenius norm and squared spectral norm. Clearly, for any matrix W , we have 1 ≤ s-rank(W ) ≤ rank(W ). As its name suggests, the stable rank is more stable than the rank because it is largely unaffected by tiny singular values. Let r = rank(W ), then s-rank(W ) = rank(W ) when σ 1 (W ) = • • • = σ r (W ) and s-rank(W ) = 1 when rank(W ) = 1. It has recently been shown (Neyshabur et al., 2017 , Theorem 1) that the generalization error of neural networks crucially depends on both the stable ranks and the spectral norms of connection matrices in the network. Specifically, it can be shown that the generalization error is upper bounded by O L j=1 ||W j || 2 2 L j=1 s-rank(W j )/n , where we omit constant terms that do not depend on weight matrices and L is the number of layers in the network. Essentially, this upper bound suggests that smaller spectral norm (smoother function mapping) and stable rank (skewed spectrum) leads to better generalization.

To understand why AEB improves generalization, in Fig. 5 and Fig. 6 , we plot both the stable rank and the spectral norm of the weight matrix in the last layer of the CNNs used in our MNIST and CIFAR10 experiments. We compare 3 methods: CNN without any regularization, CNN trained with weight decay and CNN with AEB. For each setting we repeat the experiments for 5 times, and we plot the mean along with its standard deviation. From Fig. 5a and Fig. 6a it is clear that AEB leads to a significant reduction in terms of the stable rank when compared with weight decay, and this effect is consistent in all the experiments with different training size. Similarly, in Fig. 5b and Fig. 6b we plot the spectral norm of the weight matrix. Again, both weight decay and AEB help reduce the spectral norm in all settings, but AEB plays a more significant role than the usual weight decay.

Combining the experiments with the generalization upper bound introduced above, we can see that training with AEB leads to an estimator of W that has lower stable rank and smaller spectral norm, which explains why it achieves a better generalization performance. Furthermore, this observation holds on the SARCOS datasets as well, and we show the results in the appendix.

Ablations: In all the experiments, the AEB algorithm is performed on the softmax layer. Here, we study the effects of applying AEB algorithm in all CONV/FC layers, all CONV layers, all FC layers, and the last FC layer (i.e., softmax layer). We first discuss how we handle the convolutions in our AEB algorithm. Consider a convolutional layer with {input channel, output channel, kernel width, kernel height} being {a, b, k w , k h }, we vectorize the original 4-D tensor to be a 2-D matrix of size ak w k h × b. The AEB algorithm can therefore be directly applied on this transformed matrix. Next, we perform the experiment on MNIST with batch size 2048 in Fig. 4 . The training set size here is chosen as {128, 256, 512, 1024, 2048, 4096, 8192, 16384, 32768, 60000} . We find that simply applying the AEB algorithm in the softmax layer reaches best generalization as comparing to applying AEB on more layers. The improvement is more obvious when the training set size is small. We argue that neural networks can be realized as a combination of a complex nonlinear transformation (i.e., feature extraction) and a linear model (i.e., softmax layer). Since AEB represents a correlation learning in the weight matrix, it implies that implicit correlations of neurons can also be discovered. In the real world setting, different tasks should be correlated. Therefore, applying AEB in the linear model shall improve the model performance by discovering these tasks correlations. On the contrary, the nonlinear features should be decorrelated for the purpose of generalization (Cogswell et al., 2015) . Hence, applying AEB in previous layers may lead to adversarial effect.

Correlation Matrix: To verify that AEB imposes the effect of "learning from the experience of others" during training, we visualize the weight matrix of the softmax layer by computing the corresponding correlation matrix, as shown in Fig. 7 . In Fig. 7 , darker color means stronger correlation. We conduct two experiments with training size 600 and 60,000 respectively. As we can observe, training with AEB leads to weight matrix with stronger correlations, and this effect is more evident when the training set is large. This is consistent with our analysis that AEB encourages learning from the experience of others. As a sanity check, from Fig. 7 we can also see that similar digits, e.g., 1 and 7, share a positive correlation while dissimilar ones, e.g., 1 and 8, share a negative correlation.

5. RELATED WORK:

Different kinds of regularization approaches have been studied and designed for neural networks, e.g., weight decay (Krogh & Hertz, 1992) , early stopping (Caruana et al., 2001) , Dropout (Srivastava et al., 2014) and the more recent DeCov (Cogswell et al., 2015) method. BN was proposed to reduce the internal covariate shift during training, but recently it has been empirically shown to actually smooth the landscape of the loss function (Santurkar et al., 2018) . As a comparison, we propose AEB through the lens of the empirical Bayes method, with the aim to reduce overfitting by allowing neurons to learn from each other. From the optimization perspective, learning the row and column covariance matrices help to converge to better local optimum that also generalizes better.

Despite the name, empirical Bayes method is in fact a frequentist approach to obtain estimator with favorable properties. On the other hand, truly Bayesian inference would instead put a (approximate/variational) posterior distribution over model weights to characterize the uncertainty during training (MacKay, 1992; Hernández-Lobato & Adams, 2015; Blundell et al., 2015) . However, due to the complexity of nonlinear neural networks, analytic posterior is not available, hence strong independent assumptions over model weight have to be made in order to achieve computationally tractable variational solution. Typically, both the prior and the variational posterior are assumed to fully factorize over model weights. As a comparison, in AEB we characterize the correlation between model weights via learned covariance matrix with a Kronecker product structure, which is more flexible and realistic. The application of empirical Bayes in the machine learning community is not new. McInerney (2017) proposed Bayes Empirical Bayes method for choosing the hyper-hyperparameters of models, a.k.a., the Type-III MLE method. It proposes to use Monte Carlo sampling to approximate both the marginal likelihood and the prior distribution, in order to optimize the hyper-prior distribution over the hyper-hyperparameters. As a comparison, in our work we study the matrix-variate prior with Kronecker factorization structure and we propose an analytic solution that could be computed in closed form without any sampling.

Orthogonal to our work, the Kronecker factorization assumption has also been applied in the literature of neural networks to approximate the Hessian matrix, or the Fisher information matrix, in secondorder optimization methods (Martens & Grosse, 2015; Zhang et al., 2017) . The main idea here is to approximate the curvature of the loss function's landscape, in order to achieve better convergence speed compared with first-order method while maintaining the tractability of such computation.

Determinantal point process (DPP) has been previously applied to compress neural networks (Mariet & Sra, 2015) . Specifically, a DPP kernel is placed over the activations of neurons from the same layer, and then neurons with similar activations over a fixed dataset are merged into a single one. However, it is well known that DPPs can capture only negative correlations (Kulesza & Taskar, 2011; Kulesza et al., 2012) , and as a result they do not stimulate neurons to learn from the experience of other neurons. As a comparison, by explicitly learning both precision (covariance) matrices, our framework can account for both positive and negative correlations among fan-in/fan-out of neurons from the same layer.

A PROOFS Proposition 1. The optimization problem (10) is convex.

Proof. It is clear that the objective function is convex: the trace term is linear in Ω r and it is wellknown that the log det(•) is concave in the positive definite cone (Boyd & Vandenberghe, 2004) , hence it trivially follows that

Tr(Ω r W Ω c W T ) -d log det(Ω r ) is convex in Ω r .

It remains to show that the constraint set is also convex. Let Ω 1 , Ω 2 be any feasible points, i.e., uI p Ω 1 vI p and uI p Ω 2 vI p . Let ∀t ∈ (0, 1), we have:

||tΩ 1 + (1 -t)Ω 2 || 2 ≤ t||Ω 1 || 2 + (1 -t)||Ω 2 || 2 ≤ tv + (1 -t)v = v,

where we use || • || 2 to denote the spectral norm of a matrix. Now since both Ω 1 and Ω 2 are positive definite, the spectral norm is also the largest eigenvalue, hence this shows that tΩ 1 + (1 -t)Ω 2 vI p .

To show the other direction, we use the Courant-Fischer characterization of eigenvalues. Let λ min (A) denote the minimum eigenvalue of a real symmetric matrix A, then by the Courant-Fischer min-max theorem, we have:

λ min (A) := min x =0,||x||2=1

||Ax|| 2 .

For the matrix tΩ 1 + (1 -t)Ω 2 , let x * be the vector corresponding to the minimum eigenvalue, hence we have:

λ min (tΩ 1 + (1 -t)Ω 2 ) = min x =0,||x||2=1 ||(tΩ 1 + (1 -t)Ω 2 )x|| 2 = (tΩ 1 + (1 -t)Ω 2 )x * ≥ tλ min (Ω 1 ) + (1 -t)λ min (Ω 2 ) ≥ tu + (1 -t)u = u,

which also means that tΩ 1 + (1 -t)Ω 2 uI p , and this completes the proof.

Lemma 1. Let Ω r ∈ C, then N C (Ω r ) = -N C (Ω -1 r ).

Proof. Let S ∈ N C (Ω r ). We want to show -S ∈ N C (Ω -1 r ). By definition of the normal cone, since S ∈ N C (Ω r ), we have:

Tr(SZ) ≤ Tr(SΩ r ), ∀Z ∈ C

Now realize that Ω r ∈ C and C is a compact set, it follows Ω r is the solution of the following linear program: max Tr(SZ), subject to Z ∈ C

Since both S and Z are real symmetric matrix, we can decompose them as

Z := Q Z Λ Z Q T Z and S := Q S Λ S Q T S

, where both Q Z , Q S are orthogonal matrices and Λ Z , Λ S are diagonal matrices with the corresponding eigenvalues in decreasing order. Plug them into the objective function, we have:

Tr(SZ) = Tr(Q S Λ S Q T S Q Z Λ Z Q T Z ) = Tr(Λ S Q T S Q Z Λ Z Q T Z Q S ). Define K := Q T S Q Z and D = K • K,

where we use • to denote the Hadamard product between two matrices. Since both Q S and Q Z are orthogonal matrices, we know that K is also orthogonal, which implies:

p j=1 D ij = 1, ∀i ∈ [p], and p i=1 D ij = 1, ∀j ∈ [p].

As a result, D is a doubly stochastic matrix and we can further simplify the objective function as:

Tr(Λ S Q T S Q Z Λ Z Q T Z Q S ) = Tr(Λ S KΛ Z K T ) = λ T S Dλ Z = p i,j=1 λ S,i D ij λ Z,j ,

where λ S and λ Z are p dimensional vectors that contain the eigenvalues of S and Z in decreasing order, respectively. Now for any λ S and λ Z in decreasing order, we have:

EQUATION

From ( 13), in order for Ω r to maximize the linear program, it must hold that D = K = I p and all the eigenvalues of Ω r are v. But due to the assumption that uv = 1, in this case we also know that all the eigenvalues of Ω -1 r are 1/v = u, hence Ω -1 r also minimizes the above linear program, which implies:

Tr(SΩ -1 r ) ≤ Tr(SZ), ∀Z ∈ C ⇔ Tr(-S(Z -Ω -1 r )) ≤ 0 ∀Z ∈ C.

In other words, we have -S ∈ N C (Ω -1 r ). Using exactly the same arguments it is clear to see that the other direction also holds, hence we have

N C (Ω r ) = -N C (Ω -1 r ).

Theorem 1. Let Ω r ∈ S p with eigendecomposition as Ω r = QΛQ T and Proj C (•) be the Euclidean projection operator onto C, then Proj

C ( Ω r ) = QT [u,v] (Λ)Q T .

Proof. Since Ω r ∈ C is real and symmetric, we can reparametrize Ω r as Ω r := U Λ Ωr U T where U is an orthogonal matrix and Λ Ωr is a diagonal matrix whose entries corresponds to the eigenvalues of Ω r . Recall that U corresponds to a rigid transformation that preserves length, so we have:

EQUATION

)

Define B := U T Ω r U . Now by using the fact that Ω r can be eigendecomposed as Ω r = QΛQ T , we can further simplify ( 14) as:

||Λ Ωr -U T Ω r U || 2 F = i∈[p] (Λ Ωr,ii -B ii ) 2 + i =j B 2 ij ≥ i∈[p] (Λ Ωr,ii -B ii ) 2 ≥ i∈[p] (T [u,v] (B ii )-B ii ) 2 ,

where the last inequality holds because u ≤ Λ Ωr,ii ≤ v, ∀i ∈ [p]. In order to achieve the first equality, B = U T Ω r U should be a diagonal matrix, which means U T Q = I p ⇔ U = Q. In this case, diag(B) = Λ. To achieve the second equality, simply let

Λ Ωr = T [u,v] (diag(B)) = T [u,v] (Λ),

which completes the proof.

B MORE EXPERIMENTS:

In this section we present more experimental results.

Combination. As discussed in the main text, combining the proposed AEB with BN can further improve the generalization performance, due to the complementary effects between these two approaches: BN helps smoothing the landscape of the loss function while AEB also changes the curvature via the row and column covariance matrices (see Fig. 8 ).

On the other hand, we do not observe significant difference when combining AEB with Dropout on this dataset. While we are not clear what is the exact reason for this effect, we conjecture this is due to the fact that Dropout works as a regularizer that prevents coadaptation while AEB instead encourages neurons to learn from each other.

Covariance matrices in the prior. One byproduct that AEB brings to us is the learned row and column covariance matrices, which can be used in exploratory data analysis to understand the correlations between learned features and different output tasks. To this end, we visualize both the row and column covariance matrices in Fig. 9 . The two covariance matrices on the first row correspond to the ones learned on a training set with 600 instances while the two on the second row are trained with the full dataset on MNIST.

From Fig. 9 we can make the following observations: the structure of both covariance matrices become more evident when trained with larger dataset, and this is consistent with the Bayesian principle because more data provide more evidence. Second, we observe in our experiments that the variances of both matrices are small. In fact, the variance of the row covariance matrix Σ r achieves the lower bound limit u at convergence, which also validates our approximation in computing the marginal distribution in (2). Lastly, comparing the row covariance matrix Σ r in Fig. 9 with the one computed from model weights in Fig. 7 , we can see that both matrices exhibit the same correlation patterns, except that the one obtained from model weights are more evident, which is due to the fact that model weights are closer to data evidence than the row covariance matrix in our Bayesian hierarchy.

On the other hand, the column covariance matrix in Fig. 9 also exhibit rich correlations between the learned features, e.g., the neurons in the penultimate layer. Again, with more data, these patterns become more evident.

Stable rank and spectral norm. In this paragraph we show the experimental results of stable ranks and spectral norms on the SARCOS dataset. For the SARCOS dataset, the weight matrix being regularized is of dimension 100 × 7. Again, we compare the results using three methods: MTL, MTL-WeightDecay and MTL-AEB. All the network architectures used in this experiment are consistent with the ones used before. As can be observed from Table 2 , compared with the weight decay regularization, our AEB regularization greatly reduces both the stable rank and the spectral norm of learned weight matrix, which also helps to explain why MTL-AEB generalizes better compared with MTL and MTL-WeightDecay.

To conclude this section, we plot a bar chart in Fig. 10 to show the explained variance achieved by different methods on 7 regression tasks from the SARCOS dataset for better visualization. Again, we can see that AEB improves uniformly over all the other methods on all the 7 related regression tasks. 

:

The constraint uv = 1 is only for the ease of presentation in the following part and can be readily removed.

CONCLUSIONIn this paper we propose an approximate empirical Bayes method with matrix-variate normal prior to learn the model parameters of deep neural networks. The prior encourages neurons to borrow statistical strength from other neurons during the learning process, and it provides an effective regularization when training large networks on small datasets. We connect our approach with sequential MAP inference and volume minimization, and we also design an efficient block coordinate descent algorithm to optimize the model. Empirically, on three datasets we demonstrate that our AEB method improves generalization by finding better local optima that does not overfit. One future direction is to develop a better approximate solution to optimize the two covariance matrices from the marginal log-likelihood function.

Figures and Tables:

Figure 1: Illustration for Bayes/ Empirical Bayes, and Approximate Empirical Bayes.

Classification (MNIST & CIFAR10): Results on the multiclass classification for different training sizes are show in Fig. 2. For both MNIST and CIFAR10, we find AEB, Weight Decay, and Dropout are the effective regularization methods, while Batch Normalization and DeCov vary in different settings. Batch Normalization suffers from large batch size in CIFAR10 (comparing Fig. 2 (c) and (d)) but is not sensitive to batch size in MNIST (comparing Fig. 2 (a) and (b)). The performance deterioration in large batch size of Batch Normalization is also observed by Hoffer et al.

Figure 2: Generalization of AEB on MNIST and CIFAR10. AEB improves generalization under both minibatch settings and is most beneficial when training set is small.

Spectral norm vs training size.

Figure 5: Comparisons of stable ranks and spectral norms from different methods on MNIST.

Figure 6: Comparisons of stable ranks and spectral norms from different methods on CIFAR10.

Figure 7: Correlation matrix of the weight matrix in the softmax layer. All the networks share the same structure and optimization protocol. The left two correspond to dataset with training size 600 and the right two with size 60,000. Acc means the test set classification accuracy.

Figure 8: Combine AEB with BN and Dropout on MNIST.

Cov. matrix trained on 60,000 instances.

Figure 9: Recovered row covariance matrix Σ r and column covariance matrix Σ c in the prior distribution on MNIST.

Figure 10: Explained variance of different methods on 7 regression tasks from the SARCOS dataset.

Geometrically, this means that the optimum Ω -1 r is the Euclidean projection of W Ω c W T /d onto C. Hence it suffices if we can solve the following Euclidean projection problem efficiently, where Ω r ∈ S p is a fixed real symmetric matrix:

Table content: <table><tr><td>min</td></tr></table>

Explained variance of different methods on 7 regression tasks from the SARCOS dataset.

Table content: <table><tr><td>Method</td><td>1st</td><td>2nd</td><td>3rd</td><td>4th</td><td/><td>5th</td><td colspan="2">6th</td><td>7th</td></tr><tr><td>MTL</td><td colspan="8">0.4418 0.3472 0.5222 0.5036 0.6024 0.4727 0.5298</td></tr><tr><td colspan="9">MTL-Dropout 0.4413 0.3271 0.5202 0.5063 0.6036 0.4711 0.5345</td></tr><tr><td>MTL-BN</td><td colspan="8">0.4768 0.3770 0.5396 0.5216 0.6117 0.4936 0.5479</td></tr><tr><td>MTL-DeCoV</td><td colspan="8">0.4027 0.3137 0.4703 0.4515 0.5229 0.4224 0.4716</td></tr><tr><td>MTL-AEB</td><td colspan="8">0.4769 0.3969 0.5485 0.5308 0.6202 0.5085 0.5561</td></tr><tr><td/><td/><td/><td/><td/><td/><td>95</td><td/></tr><tr><td/><td/><td/><td/><td/><td/><td>90</td><td/></tr><tr><td/><td/><td/><td/><td/><td>Accuracy</td><td>85</td><td/></tr><tr><td/><td/><td/><td/><td/><td/><td>75 80</td><td/><td>ALL CONV FC LAST</td></tr><tr><td>(a) Training Set Size = 600</td><td/><td colspan="2">(b) Training Set Size = 6000</td><td/><td/><td>10 2</td><td>10 3</td><td>Train Size</td><td>10 4</td></tr><tr><td colspan="5">Figure 3: Optimization of AEB on MNIST with batch size 2048.</td><td colspan="4">Figure 4: Applying AEB on different lay-</td></tr><tr><td colspan="5">Sequential Tikhonov regularizations in AEB help to converge to</td><td colspan="4">ers in neural networks for MNIST with</td></tr><tr><td colspan="2">better local optima during training.</td><td/><td/><td/><td colspan="3">batch size 2048.</td></tr></table>

(d) Column Cov. matrix trained on 60,000 instances.

Table content: <table><tr><td>0 2 4 6 14 16 24 26 34 36 44 46 48 42 40 38 32 30 28 22 20 18 12 10 8</td><td>0.8 0.4 0.0 0.4 0.8</td></tr><tr><td>0 2 4 6 8 10 12 14 16 18 20 22 24 26 28 30 32 34 36 38 40 42 44 46 48</td><td/></tr></table>

Stable rank and spectral norm on the SARCOS dataset.

Table content: <table><tr><td/><td colspan="2">Stable Rank Spectral Norm</td></tr><tr><td>MTL</td><td>4.48</td><td>0.96</td></tr><tr><td>MTL-WeightDecay</td><td>4.83</td><td>0.92</td></tr><tr><td>MTL-AEB</td><td>2.88</td><td>0.70</td></tr></table>

Peer review (iclr2019):
 Starting from a simple neural network with only one hidden layer and a single output, the basic idea of approximate empirical Bayes (AEB) method is proposed, defining a matrix-variate normal prior distribution with a Kronecker product structure, so as to capture correlations between the row and column vectors of the weight matrix. Then, a block coordinate descent algorithm for solving the optimization problem is proposed. It consists of alternating three steps to obtain the optimal solutions of model parameters, row and column covariance matrices. The current method is investigated and tested on three data sets for both classification (MNIST & CIFAR10) and regression (SARCOS) tasks. Encouraging experimental results demonstrate that the correlation learning in the weight matrix significantly improves performance when the training set size is relatively small. It is also shown that the proposed AEB method does not seem sensitive to the size of mini-batches and its combination with other generalization methods can lead to better results in some cases. Strengths: This paper is mostly well written and overall is easy to follow. It clearly reveals that correlation in the weight matrix plays a crucial role in better generalizing on small training sets. Minor comments: * The authors state that it is straightforward to extend the proposed method to more sophisticated models with various structures, such as CNN. Perhaps a bit more detail should be given in the main text. * Fig. 6 on page 12 is not explicitly mentioned in the main text. It seems a bit confusing.