Learning From the Experience of Others: Approximate Empirical Bayes in Neural Networks

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 ∈ 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 . 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 ( 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 . 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 . 3.1 APPROXIMATE EMPIRICAL BAYES . When the likelihood function p ( D | w ) is implemented as a neural network , the marginalization in ( 1 ) over model parameter w can not 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 .

Summary: 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.