A Theoretical and Empirical Model of the Generalization Error under Time-Varying Learning Rate

1 INTRODUCTION . While deep learning addresses a variety of practical tasks , the reason deep learning generalizes well is not fully understood . It is hoped that the performance of deep learning methods can be improved by understanding deep learning . Consequently , deep learning will become a cornerstone for the next generation of technology . Stochastic gradient descent ( SGD ) is a principled learning algorithm for deep learning , and because most deep learning algorithms are developed based on SGD , its theoretical analysis is crucial . He et al . ( 2019 ) modeled the generalization bounds when both the batch size and learning rate are constant in SGD ; however , they did not model the generalization error . Furthermore , it is crucial to decrease the learning rate as the training progresses to make SGD converge ( Robbins & Monro , 1951 ) , which results in good performance ( You et al. , 2020 ) , but He et al . ( 2019 ) did not consider the case in which these parameters are time-varying . We analyze the time-varying case in SGD and propose a generalization error model for both constant and time-varying cases . Based on these models , we also propose a novel hyperparameter search method . In this study , we utilized a PAC-Bayes-based generalization bound to generate a model of the generalization error . To validate the proposed model , we trained approximately 300 networks on both practical networks ( VGG16 ( Simonyan & Zisserman , 2015 ) , Wide-ResNet 28-10 ( Zagoruyko & Komodakis , 2016 ) ) and data ( CIFAR10 , CIFAR100 ( Krizhevsky et al. , 2009 ) ) . We also showed that the proposed model can approximate the experimental results well . We also show that the proposed model can be used for hyperparameter search via Bayesian optimization and can yield better results , compared with the previous method for searching in the uniform distribution , and comparable results for searching in the logarithmic uniform distribution . The contributions of the present study are summarized as follows : • We analyzed the generalization bound using PAC-Bayes with the timevarying batch size and learning rate . • We directly modeled the generalization error on the batch size and learning rate and proposed its functional form . • We experimentally showed that it is possible to model the generalization error with high accuracy for practical data . • Our experimental results demonstrated that the hyperparameter search method built on the proposed model can outperform current libraries when searching over uniform distributions and has comparable performance when searching over logarithmic uniform distributions . 2 RELATED WORK . In this section , we present related work on SGD modeling , generalization bounds , prediction of model performance , and hyperparameter optimization . Theoretical analysis of SGD : SGD is often treated as a continuous-time stochastic differential equation ( SDE ) by modeling the noise of the gradient and the loss function and applying the Euler– Maruyama approximation ( Lin et al. , 2018 ; Mandt et al. , 2017 ; Li et al. , 2017 ; Jastrzębski et al. , 2017 ) . This approximation can be solved analytically as an Ornstein–Uhlenbeck process ( Uhlenbeck & Ornstein , 1930 ) . In this study , we also analyzed SGD using the same process as the previously mentioned studies . However , in most previous studies , the batch size and learning rate in SGD were treated as constants . Hence , in this study , we treat them as time-varying and seek an analytic solution . Theoretical error bounds : There is a large body of research on theoretical explanations for the generalization capability of deep neural networks ( Negrea et al. , 2019 ; Arora et al. , 2018 ; Cao & Gu , 2019 ; Deng et al. , 2021 ) . He et al . ( 2019 ) , building on their analysis of SGD , provided a theoretical explanation of the effects of the batch size and learning rate on the generalization error . However , this research only handles the case of constant batch size and learning rate and does not model the generalization error in the function form . In this study , we derive a generalization bound based on an analytic solution for the time-varying case and derive a model of the generalization error . Predicting model performance : Most studies are based on the results from training networks , assuming hyperparameters as inputs , and searching for the best parameters without training the networks . In neural architecture search , the performance of a model was predicted for a given model on a given dataset ( Istrate et al. , 2019 ) , small data ( Klein et al. , 2017 ) , and model size ( Real et al. , 2019 ) , respectively . Several studies predict the performance of a model in a functional form by treating both the dataset and model sizes as input ( Rosenfeld et al. , 2020 ) , or the pruning scale as input ( Rosenfeld et al. , 2021 ) . In this study , we exploit the findings of the above studies to predict the performance of a model in a functional form , based on a theoretically derived generalization bound to the batch size and learning rate . Hyperparameter optimization : Hyperparameter optimization algorithms are mainly categorized into two types : nonadaptive and adaptive . Nonadaptive algorithms search for the best hyperparameters without using any past search results and , thus , have the advantage of being easily parallelizable . Grid search and random search fall into this category , and many studies have shown that random search is more useful ( Bergstra & Bengio , 2012 ; Li & Talwalkar , 2020 ) . Adaptive algorithms search for the best hyperparameters by exploiting past search results and can obtain better hyperparameters with fewer searches . Adaptive algorithms include specific algorithms , such as the evolutionary algorithm ( Young et al. , 2015 ) and Bayesian optimization ( Wu et al. , 2019 ) , and practical tools , such as Hyperopt ( Bergstra et al. , 2015 ) and Optuna ( Akiba et al. , 2019 ) . These methods are flexible enough to be optimized for arbitrary functions . However , there is still room for more efficient search when searching only for specific hyperparameters . In this study , we aim to further improve the search efficiency of both the batch size and learning rate by employing a hyperparameter search based on the proposed generalization error model . In the following sections , we will analyze the influence of the batch size and learning rate on the generalization error in SGD , based on assumptions about the distribution of the gradients and the shape of the risk function near minima , an approximation to SDE , and the PAC-Bayes method . 3 PRELIMINARIES . In this section , we describe the preliminary knowledge used to model the learning process of SGD . Generalization error of stochastic algorithms . In general , machine learning algorithms are designed to select a hypothesis function Fθ from a hypothesis class { Fθ|θ ∈ Θ ⊂ Rd } such that the expected value of the risk function R computed from the loss function l is minimized . In a proba- bilistic algorithm such as SGD , the parameter θ after learning is considered to lie on the distribution Q . The risk function for parameter θ is defined as follows : R ( θ ) = E ( X , Y ) ∼D [ l ( Fθ ( X ) , Y ) ] . ( 1 ) Therefore , R ( Q ) = Eθ∼Q [ R ( θ ) ] . However , one can not compute the risk function , and in practice , the empirical risk R̂ ( θ ) from the data is approximated as the risk functionR ( θ ) . R̂ ( θ ) = 1 N N∑ i=1 l ( Fθ ( Xi ) , Yi ) , ( 2 ) where ( Xi , Yi ) is the N training data consisting of the input data and label pairs . Thus , R̂ ( Q ) = Eθ∼Q [ R̂ ( θ ) ] . Stochastic gradient descent . Gradient descent is an algorithm that computes the gradient of the risk function defined in Eq.1 and updates the parameter θ based on the gradient . The gradient g ( θ ( t ) ) is defined as g ( θ ( t ) ) , ∇θ ( t ) R ( θ ( t ) ) = ∇θ ( t ) E ( X , Y ) [ l ( Fθ ( t ) ( X ) , Y ) ] . We update the parameter θ by θ ( t + 1 ) = θ ( t ) − ηg ( θ ( t ) ) , where η is the learning rate . In practice , the true gradient of g ( θ ( t ) ) is approximated by the gradient of ĝS ( θ ( t ) ) computed from the empirical risk R̂ ( θ ) defined in Eq.2 , and the parameter θ is updated with ĝS ( θ ( t ) ) . Thus , the gradient ĝS ( θ ( t ) ) is defined as ĝS ( θ ( t ) ) = ∇θ ( t ) R̂ ( θ ( t ) ) = 1|S| ∑ i∈S ∇θ ( t ) l ( Fθ ( t ) ( Xi ) , Yi ) , and the parameter θ is updated by θ ( t+1 ) = θ ( t ) −ηĝS ( θ ( t ) ) . We use a mini-batch S to train the model in parallel , and |S| represents its size . Then , we treat li ( θ ) = l ( Fθ ( t ) ( Xi ) , Yi ) . Analysis of the learning process . The loss function li ( θ ) and its mean R̂ ( θ ) is an unbiased estimator of the risk functionR ( θ ) , and the gradient of the loss function ∇θli ( θ ) and its mean ĝS ( θ ) is an unbiased estimator of the gradient g ( θ ) = ∇θR ( θ ) . Hence , we have E [ li ( θ ) ] = E [ R̂ ( θ ) ] = R ( θ ) and E [ ∇θli ( θ ) ] = E [ ĝS ( θ ) ] = g ( θ ) = ∇θR ( θ ) , where the expectations are computed from the data distribution ( X , Y ) . Here , as in Mandt et al . ( 2017 ) , we assume that { ∇li ( θ ) } is independently and identically sampled from a Gaussian distribution with mean g ( θ ) = ∇θR ( θ ) . We can express ∇θli ( θ ) ∼ N ( g ( θ ) , C ) , where C is the covariance matrix , which is assumed to be a constant matrix for all θ . Then the gradient of ĝS ( θ ) computed for a mini-batch follows the following Gaussian distribution : ĝS ( θ ) = 1 |S| ∑ i∈S ∇θli ( θ ) ∼ N ( g ( θ ) , 1 |S| C ) . ( 3 ) Because the covariance matrix is symmetric and ( semi ) positive definite , we suppose that C can be factorized as C = BB > . In an actual SGD algorithm , the parameters are iteratively updated using Eq.3 to minimize the risk functionR ( θ ) : ∆θ ( t ) = θ ( t+ 1 ) − θ ( t ) = −ηĝS ( θ ( t ) ) = −ηg ( θ ) + η√ |S| B∆W , ∆W ∼ N ( 0 , I ) . ( 4 ) Analytical solution of the stochastic differential equation . When the learning rate η is small , Eq.4 can be approximated by SDE , which is known as Ornstein–Uhlenbeck process ( Uhlenbeck & Ornstein , 1930 ) . Here , we assume that the risk function is convex and second-order differentiable in the neighborhood of the minimum ( He et al. , 2019 ) . In addition , when the risk function R ( θ ) is minimized at θ = θ∗ , it is shifted so that it is minimized at θ = 0 . Thus , the risk function R ( θ ) can be approximated by second-order statistic on the distribution Q as follows : R ( θ ) = 1 2 θ > Aθ , ( 5 ) whereA is the Hessian matrix around the minimum of the risk function and is assumed to be positive definite . This assumption was experimentally demonstrated by Li et al . ( 2018 ) , and other studies have made similar assumptions ( Jastrzębski et al. , 2017 ; Poggio et al. , 2017 ) . Because E [ R̂ ( θ ) ] = R ( θ ) , we further assume that the empirical risk R̂ ( θ ) is in the following quadratic form : R̂ ( θ ) = 1 2 ( θ − θb ) > Â ( θ − θb ) ( 6 ) where θb is the bias term , Â is the Hessian matrix , and E [ θb ] = 0 , E [ Â ] = A . It is well known that Eq.4 has an analytic solution ( Gardiner , 2004 ) , θ ( t ) = e−Atθ ( 0 ) + √ η |S| ∫ t 0 e−A ( t−t ′ ) BdW ( t′ ) , ( 7 ) where W ( t′ ) is white noise and followsN ( 0 , I ) and Q is the distribution of θ in Eq.7 and Gaussian .

Summary: This paper approximates the learning dynamics by an Ornstein-Uhlenbeck process and uses some PAC-Bayes results to obtain a functional form for the generalisation behaviour as a function of the ration of the learning rate to the batch size.  It obtains results for both constant and time-varying parameters.  It then uses these to develop a kernel function for hyper-parameter optimisation using Gaussian Processes modelling.