Year: 2018

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.

heading: 1 INTRODUCTION
text: 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 ∈ Rd is fully connected to a hidden layer h ∈ Rp, which is further fully connected to the single output ŷ ∈ R. Let σ(·) be the nonlinear activation function, e.g., ReLU (Nair & Hinton, 2010), W ∈ Rp×d be the connection matrix between the input layer and the hidden layer, and a ∈ Rp be the vector connecting the output and the hidden layer. Without loss of generality, ignoring the bias term in each layer, we have: ŷ = aTh,h = σ(Wx). Consider using the usual `2 loss function `(ŷ, y) = 12 |ŷ − y| 2 and take the derivative of `(ŷ, y) w.r.t. W . We obtain the update formula in backpropagation as W ← W − α(ŷ − y)(a ◦ h′) xT , where h′ is the componentwise derivative of h w.r.t. its input
argument, and α > 0 is the learning rate. Realize that (a ◦ h′) xT 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.

heading: 2 PRELIMINARY
text: 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).

heading: 2.1 NOTATION AND SETUP
text: 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 (xi, yi), i ∈ [n]. We consider the supervised learning setting where xi ∈ X ⊆ Rd and yi ∈ 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.

heading: 2.2 THE EMPIRICAL BAYES METHOD
text: 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).

heading: 3 LEARNING WITH APPROXIMATE EMPIRICAL BAYES

heading: 3.1 APPROXIMATE EMPIRICAL BAYES
text: When the likelihood function p(D | w) is implemented as a neural network, the marginalization in (1) over model parameter w cannot be computed exactly. Nevertheless, instead of performing expensive Monte-Carlo simulation, we can use point estimate of w to approximate the integral as follows:∫
p(D | w) · p(w | θ) dw ≈ p(D | ŵ) · p(ŵ | θ), (2)
where ŵ = arg maxw 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 ŵ 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 ŵ, 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 ŵ by maximizing the posterior distribution as follows:
ŵ← max w p(w | D) = max w p(D | w) · p(w | θ̂). (3)
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:
max w max θ
log p(D | w) + log p(w | θ) (4)
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 ŵ 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 {ŵ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.

heading: 3.2 NEURAL NETWORK WITH MATRIX-NORMAL PRIOR
text: 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 (0p×d,Σr,Σc), where Σr ∈ Sp++ and Σc ∈ Sd++ 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 (0p×d,Σc ⊗ Σr). It
1The probability density function is given by p(W | Σr,Σc) = exp(−Tr(Σ−1r WΣ −1 c W T )/2) (2π)pd/2 det(Σr)d/2 det(Σc)p/2 .
is then easy to check that the marginal prior distributions over the row and column vectors of the weight matrix W are given by:
Wi: ∼ N (0d, [Σr]ii · Σc), W:j ∼ N (0p, [Σ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 := Σ−1r 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:
min W,a min Ωr,Ωc
1
2n ∑ i∈[n] (ŷ(xi;W,a)− yi)2 + λ||Ω1/2r WΩ1/2c ||2F − λ (d log det(Ωr) + p log det(Ωc))
subject to uIp Ωr vIp, uId Ωc vId (5)
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{d3, p3}) 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 (Ω(t)r ,Ω (t) 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)r ⊗ Ω(t)c . From this perspective, if we fix Ω(t)r = Ip and Ω (t) c = Id, ∀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:
||Ω1/2r WΩ1/2c ||2F = ||vec(Ω1/2r WΩ1/2c )||22 = ||(Ω1/2c ⊗ Ω1/2r )vec(W )||22, (6)
and this is exactly the Tikhonov regularization (Golub et al., 1979) imposed onW where the Tikhonov matrix Γ is given by Γ := Ω1/2c ⊗ Ω1/2r . 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.
2The constraint uv = 1 is only for the ease of presentation in the following part and can be readily removed.
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 A1, A2 ∈ Sm++, the following inequality holds:
log det(A1) ≤ log det(A2) + 〈∇ log det(A2), A1 −A2〉 = log det(A2) + Tr(A−12 A1)−m (7) Applying the above inequality twice by fixing A1 = WΩcWT /2d,A2 = Σr and A1 = WTΩrW/2p,A2 = Σc respectively leads to the following inequalities:{
d log det(WΩcW T /2d) ≤ −d log det(Ωr) + 12 Tr(ΩrWΩcW T )− dp p log det(WTΩrW/2p) ≤ −p log det(Ωc) + 12 Tr(ΩrWΩcW T )− dp
Using the fact that Tr(ΩrWΩcWT ) = ||Ω1/2r WΩ1/2c ||2F , we immediately have: d log det(WΩcW
T )+p log det(WTΩrW ) ≤ ||Ω1/2r WΩ1/2c ||2F −(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(ATA)| 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.

heading: 3.3 THE ALGORITHMS
text: Algorithm 1 Block Coordinate Descent for Approximate Empirical Bayes
Input: Initial value φ(0) := {a(0),W (0)}, Ω(0)r ∈ Sp++ and Ω (0) c ∈ Sd++, 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)QT = SVD(∆) 9: Hard thresholding r′ ← T[u,v](m/r)
10: return Qdiag(r′)QT 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:
min Ωr
Tr(ΩrWΩcW T )− d log det(Ωr), subject to uIp Ωr vIp. (10)
Proposition 1. The optimization problem (10) is convex.
Define the constraint set C := {A ∈ Sp++ | uIp A vIp} and the indicator function IC(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:
0 ∈ ∂ ( 1
d Tr(ΩrWΩcW
T )− log det(Ωr) + IC(Ωr) ) = WΩcW T /d− Ω−1r +NC(Ωr), (11)
where NC(A) := {B ∈ Sp | Tr(BT (Z − A)) ≤ 0,∀Z ∈ C} is the normal cone w.r.t. C at A. Equivalently, we have Ω−1r −WΩcWT /d ∈ NC(Ωr). The following key lemma characterizes the structure of the normal cone w.r.t. C: Lemma 1. Let Ωr ∈ C, then NC(Ωr) = −NC(Ω−1r ). Lemma 1 implies WΩcWT /d−Ω−1r ∈ NC(Ω−1r ). Geometrically, this means that the optimum Ω−1r is the Euclidean projection of WΩcWT /d onto C. Hence it suffices if we can solve the following Euclidean projection problem efficiently, where Ω̃r ∈ Sp is a fixed real symmetric matrix:
min Ωr
||Ωr − Ω̃r||2F , subject to uIp Ωr vIp (12)
Theorem 1. Let Ω̃r ∈ Sp with eigendecomposition as Ω̃r = QΛQT and ProjC(·) be the Euclidean projection operator onto C, then ProjC(Ω̃r) = QT[u,v](Λ)QT .
Corollary 1. Let WΩcWT be eigendecomposed as Qdiag(r)QT , then the optimal solution to (10) is given by QT[u,v](d/r)QT .
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{d3, p3}). 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.

heading: 4 EXPERIMENTS
text: 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.

heading: 4.1 EXPERIMENTAL SETUP
text: 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: CONV5×5×1×10-CONV5×5×10×20FC320×50-FC50×10. The notation CONV5×5×1×10 denotes the convolutional layer with kernel size 5× 5 from depth 1 to 10; the notation FC320×50 denotes the fully connected layer with size 320× 50. Similarly, CIFAR10 considers the structure: CONV5×5×3×10-CONV5×5×10×20-FC500×500FC500×500-FC500×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 FC21×256-FC256×100-FC100×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.

heading: 4.2 RESULTS
text: Multiclass 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. (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.
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.
Method 1st 2nd 3rd 4th 5th 6th 7th
MTL 0.4418 0.3472 0.5222 0.5036 0.6024 0.4727 0.5298 MTL-Dropout 0.4413 0.3271 0.5202 0.5063 0.6036 0.4711 0.5345 MTL-BN 0.4768 0.3770 0.5396 0.5216 0.6117 0.4936 0.5479 MTL-DeCoV 0.4027 0.3137 0.4703 0.4515 0.5229 0.4224 0.4716 MTL-AEB 0.4769 0.3969 0.5485 0.5308 0.6202 0.5085 0.5561
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 ||2F /||W ||22. 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 ||Wj ||22 ∑L j=1 s-rank(Wj)/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, kw, kh}, we vectorize the original 4-D tensor to be a 2-D matrix of size akwkh × 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.

heading: 5 RELATED WORK
text: 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.

heading: 6 CONCLUSION
text: In 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.

heading: A PROOFS
text: 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(ΩrWΩcWT )− 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., uIp Ω1 vIp and uIp Ω2 vIp. 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 vIp. 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 x6=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 x6=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 uIp, and this completes the proof.
Lemma 1. Let Ωr ∈ C, then NC(Ωr) = −NC(Ω−1r ).
Proof. Let S ∈ NC(Ωr). We want to show −S ∈ NC(Ω−1r ). By definition of the normal cone, since S ∈ NC(Ω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 := QZΛZQTZ and S := QSΛSQ T S , where both QZ , QS 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(QSΛSQ T SQZΛZQ T Z) = Tr(ΛSQ T SQZΛZQ T ZQS).
Define K := QTSQZ and D = K ◦K, where we use ◦ to denote the Hadamard product between two matrices. Since both QS and QZ are orthogonal matrices, we know that K is also orthogonal, which implies:
p∑ j=1 Dij = 1,∀i ∈ [p], and p∑ i=1 Dij = 1,∀j ∈ [p].
As a result, D is a doubly stochastic matrix and we can further simplify the objective function as:
Tr(ΛSQ T SQZΛZQ T ZQS) = Tr(ΛSKΛZK T ) = λTSDλZ = p∑ i,j=1 λS,iDijλ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:
u p∑ i=1 λS,i ≤ p∑ i=1 λS,iλZ,1+p−i ≤ p∑ i,j=1 λS,iDijλZ,j ≤ p∑ i=1 λS,iλZ,i ≤ v p∑ i=1 λS,i (13)
From (13), in order for Ωr to maximize the linear program, it must hold that D = K = Ip 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 Ω−1r are 1/v = u, hence Ω −1 r also minimizes the above linear program, which implies: Tr(SΩ−1r ) ≤ Tr(SZ), ∀Z ∈ C ⇔ Tr(−S(Z − Ω−1r )) ≤ 0 ∀Z ∈ C. In other words, we have −S ∈ NC(Ω−1r ). Using exactly the same arguments it is clear to see that the other direction also holds, hence we have NC(Ωr) = −NC(Ω−1r ).
Theorem 1. Let Ω̃r ∈ Sp with eigendecomposition as Ω̃r = QΛQT and ProjC(·) be the Euclidean projection operator onto C, then ProjC(Ω̃r) = QT[u,v](Λ)QT .
Proof. Since Ωr ∈ C is real and symmetric, we can reparametrize Ωr as Ωr := UΛΩrUT 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:
||Ωr − Ω̃r||2F = ||UΛΩrUT − UUT Ω̃rUUT ||2F = ||ΛΩr − UT Ω̃rU ||2F (14)
Define B := UT Ω̃rU . Now by using the fact that Ω̃r can be eigendecomposed as Ω̃r = QΛQT , we can further simplify (14) as:
||ΛΩr−UT Ω̃rU ||2F = ∑ i∈[p] (ΛΩr,ii−Bii)2+ ∑ i 6=j B2ij ≥ ∑ i∈[p] (ΛΩr,ii−Bii)2 ≥ ∑ i∈[p] (T[u,v](Bii)−Bii)2,
where the last inequality holds because u ≤ ΛΩr,ii ≤ v,∀i ∈ [p]. In order to achieve the first equality, B = UT Ω̃rU should be a diagonal matrix, which means UTQ = Ip ⇔ 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.

heading: B MORE EXPERIMENTS
text: 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.

Review summaries:
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.