VARIATIONAL SGD: DROPOUT , GENERALIZATION AND CRITICAL POINT AT THE END OF CONVEXITY

1 INTRODUCTION . Finding a generalizable solution is a critical problem for any machine learning task . The ultimate goal of learning from the available ground truths is to make a good prediction for new data . The Bayesian method is a very powerful approach that gives a probabilistic measure of the ability of a proposed model to predict by estimating how well the model predicts known data . The accuracy of the predictions depends on how the found solution is able to overcome a sampling bias to avoid overfitting for given particular samples of training data . Specifically , in Bayesian method predictions of labels y for an input x are made by using a probabilistic model , for certainty a neural network , which defines a function parametrized by weights w that allows computing probabilities P ( y|x , w ) for each weight point . Each weight point contributes to predicted probabilities of labels Prob ( y|x ) in accordance with probability distribution of weights . The distribution of weights is learned from a known training data set { xn , yn ; n = 1 .. N } and its prior probability distribution P0 ( w ) in the following way : Prob ( y|x ) = ∫ w P ( y|x , w ) P0 ( w ) N∏ n=1 P ( yn|xn , w ) / ∫ w P0 ( w ) N∏ n=1 P ( yn|xn , w ) ( 1 ) Here the predicted probability Prob ( y|x ) is an average of the model probability P ( y|x , w ) at a weight w over the learned weight distribution . To make predictions we are only interested in a method that allows to find averages in eq . ( 1 ) and not absolute values of the integrals . According to mean value theorem ( Cauchy ( 1813 ) , also in Encyclopedia of Mathematics Hazewinkel ( 1994 ) ) values of the averages can be represented by a single point , which in our case means that there is a single point in the weight space w0 that represents a result of computing the integrals , so Prob ( y|x ) = P ( y|x , w0 ) . That point w0 is a solution of the training of the neural network . A standard approach to get the solution is a maximum likelihood method that finds a maximum of the integrand . However , there are some cases when the maximum likelihood fails to represent main contribution to the integral by weights . Consider this example : if log-likelihood for N data samples has a maximum at some weight point w1 , then in general its first derivative by weights is zero , second derivative is negative and proportional to N , so corresponding Gaussian integral by the weights is proportional to N−d/2 , where d is number of weights . This will change if there is a flat maximum , which has not only first but also second and third derivatives equal to zero . In this case the integral is proportional to N−d/4 . For large number of samples the flat maximum makes the most significant contribution to the integral by weights : I1 = P N 1 ∫ w e−NC2w 2 ∝ PN1 ×N−d/2 , I2 = PN2 ∫ w e−NC4w 4 ∝ PN2 ×N−d/4 , and I2/I1 ∝ ( P2/P1 ) N ×Nd/4 . For a typical case when the number of weights d ∼ N and average sample probabilities at maxima are comparable O ( P1 ) ∼ O ( P2 ) the integral around flat maximum I2 is always bigger than the integral around narrow maximum I1 , unless P2 is zero . While in general a likelihood has a number of regular local maxima and no flat maximum the effect of integration over multiple frequent local maxima can result in an effective flat maximum that defines a solution . We argue that any local or global maximum of likelihood gives a wrong solution that is not generalizable and so makes inaccurate predictions , because the locations for the global maximum and local maxima depend on specific samples in the training data and any modification of it by adding or removing samples will change the solution ( Zhang et al. , 2017 ) . Instead we will show that there is another solution that more associated with properties of the distribution of training data and less with particular samples . The purpose of this paper is to show that the effective flat maximum always exists for specific parameters of prior weight distribution P0 ( w ) ( regularization parameters ) and corresponding solution is the most generalizable solution that can be found in training . We show that the solution is a critical point in an effective loss that represents the result of integration over the weights . In the next sections we derive the algorithm for the optimizer for finding the critical point solution and analyze properties of the solutions . The empirical study is outside of the scope of the paper and will be presented separately . For simplicity we use same notations for a vector of weights and its components , as well as corresponding parameters of distributions of weights because all weight components are independent in our consideration and it is clear from context when it is a vector or its component . 2 METHOD FOR COMPUTING THE INTEGRALS . We use a recurrent approach for estimating the integrals . First , we represent a probability of each training sample as a product of factors close to one , P ( y|x , w ) = ( 1 + 1/T lnP ( y|x , w ) ) T , where free parameter T 1 is a number of epochs : ∫ w P0 ( w ) N∏ n=1 P ( yn|xn , w ) = ∫ w P0 ( w ) N∏ n=1 ( 1 + 1 T lnP ( yn|xn , w ) ) T , then model each factor as a product of Gaussian distributions Q ( w|µ , σ ) one for each component of weight vector . For each iteration a running prior distribution of weights is updated by absorbing a single factor to produce a new prior Qt ( w ) ( 1 + 1 T lnP ( y|x , w ) ) → Qt+1 ( w ) , withQ0 ( w ) = P0 ( w ) andQ ( w|µ , σ ) = e− ( w−µ ) 2 2σ2 √ 2πσ2 . Specifically , we do the following : First , let ’ s enumerate all factors for all data samples N∏ n=1 ( 1 + 1 T lnP ( yn|xn , w ) ) T = N×T∏ t=1 ( 1 + 1 T lnP ( yt|xt , w ) ) Then , under the integral by weights we use an identical re-writing for a product of a prior Qt ( w ) and one of the factors and a new prior Qt+1 ( w ) and normalization factor Nt∫ w Qt ( w ) ( 1 + 1 T lnP ( yt|xt , w ) ) ( . . . ) = Nt ∫ w [ Rt ( w ) Qt+1 ( w ) ] Qt+1 ( w ) ( . . . ) , where normalization factor Nt = ∫ w Qt ( w ) ( 1 + 1 T lnP ( yt|xt , w ) ) and distribution Rt ( w ) = Qt ( w ) ( 1 + 1T lnP ( yt|xt , w ) ) Nt . Finally , we make an approximation by replacing ratio of distributions Rt ( w ) /Qt+1 ( w ) by its mean for distribution Qt+1 which is equal to 1 . Then the iterations are repeated until all factors from probabilities of data are replaced by a single final Gaussian distribution and some normalization factor . To minimize the introduced error on each iteration we select the means and variances of the new prior Qt+1 to minimize variance of the ratio Rt/Qt+1 . The variance is equal to VarQt+1 [ Rt/Qt+1 ] = ∫ w [ Rt ( w ) Qt+1 ( w ) ] 2 Qt+1 ( w ) − 1 = ERt [ Rt ( w ) Qt+1 ( w ) ] − 1 . The lower bound of the variance is expressible via KL divergence ( Kullback & Leibler ( 1951 ) ) of Rt and Qt+1 VarQt+1 [ Rt/Qt+1 ] ≥ exp ( ∫ w Rt ( w ) log Rt ( w ) Qt+1 ( w ) ) − 1 = exp ( DKL ( Rt‖Qt+1 ) ) − 1 . Finding the minimum of KL divergence is equivalent to minimizing the lower bound of the variance of the ratio of Rt/Qt+1 which leads to equations ∂ ∂µt+1 , ∂ ∂σt+1 ∫ w Q ( w|µt , σt ) ( 1 + 1 T lnP ( y|x , w ) ) lnQ ( w|µt+1 , σt+1 ) = 0 . Then solving the above equations gives the update rules for means and variances of each weight component µt+1 = µt + σ2t NT ∂ ∂µt 〈 lnP ( w ) 〉 µt , σt , 1 σ2t+1 = 1 σ2t − 1 NT ∂2 ∂µ2t 〈 lnP ( w ) 〉 µt , σt , ( 2 ) where averages are defined as 〈A ( w ) 〉µ , σ = ∫ w Q ( w|µ , σ ) A ( w ) and averages of gradients by weights are equal to gradients of averages by means 〈 ∂ ∂wA ( w ) 〉 µ , σ = ∂∂µ 〈 A ( w ) 〉 µ , σ . Another useful identity allows to replace second gradient of ( µ , σ ) -average w.r.t mean on first gradient by variance ∂ ∂σ2 〈A ( w ) 〉µ , σ = 1 2 ∂2 ∂µ2 〈A ( w ) 〉µ , σ . ( 3 ) For a full batch the log of probability of data in eqs . ( 2 ) lnP ( w ) = ∑N n lnP ( yn|xn , w ) , while for a minibatch the sum goes over the size of the minibatch . In eqs . ( 2 ) we used rescaled variances σ2 → σ2/N so all gradients are normalized per data sample . With the rescaled variances the prior distribution of weights is a product over all weight dimensions d P0 ( w|µ0 , σ0 ) = Nd/2 d∏ i=1 e−N ( wi−µi,0 ) 2 2σ2 i,0√ 2πσ2i,0 . The index t enumerates iterations over all minibatches in all epochs . Then for a minibatch size one the total number of iterations equals to a number of epochs T times number of data samples N in a training set . By recursively applying the update rules in eqs . ( 2 ) over N samples and T epochs we obtain the approximation of Bayesian integrals that allows to compute averages∫ w P0 ( w ) N∏ n P ( yn|xn , w ) ( . . . ) ≈ ∫ w Qtmax ( w ) ( . . . ) × exp ( 1 T tmax∑ t ∫ w Qt ( w ) lnPt ( w ) ) . It is important to emphasize that the averages are defined by distribution Qtmax ( w ) after a finite number of iterations tmax , which for the minibatch one is equal to a product N × T and not by distribution at the infinite number of iterations Q∞ ( w ) . Number of epochs T controls the accuracy of the factorized representation and in practical computing any T larger than 10 or 100 is good enough . The prediction probability in eq . ( 1 ) is defined by weight w0 that is a mean µtmax of distribution Qtmax ( w ) Prob ( y|x ) ≈ P ( y|x , µtmax ) . That mean µtmax is a final point of iterations in eqs. ( 2 ) . The trajectory in mean-variance space is defined by starting point ( µ0 , σ0 ) , the mean and variance of a prior distribution of weights in eq . ( 1 ) which are regularization parameters : P0 ( w ) = Q ( w|µ0 , σ0 ) .

Summary: Non-convex learning problems can have multiple solutions with different generalization properties, thus it is important to find solutions that generalize well. The goal of this paper is to derive an algorithm for finding a solution to the learning problem with the best possible generalization properties. This is achieved by using a Bayesian approach in which the parameters w (e.g., weights of a network) are random variables and the effective loss (integral wrt to w) is minimized in lieu of the usual loss. The paper assumes that each component of the weight vector w is Gaussian and derives a formula for updating the mean and covariance of said Gaussians (this is an SGD method). The paper claims that the resulting effective loss is convex for large variances (sigma > threshold), nonconvex for small variances (sigma < threshold), and converges to original loss as sigma goes to zero. The paper also claims that when sigma=0 there are trivial solutions that are unstable as data changes, but that when sigma=threshold (assuming this is what is meant by end of convexity) there are non-trivial solutions that are less sensitive to data changes and hence the most generalizable.