Fractal Structure and Generalization Properties of Stochastic Optimization Algorithms

1 Introduction . In statistical learning , many problems can be naturally formulated as a risk minimization problem minw2Rd n R ( w ) : = Ez⇠⇡ [ ` ( w , z ) ] o , ( 1 ) where z 2 Z denotes a data sample coming from an unknown distribution ⇡ , and ` : Rd ⇥ Z ! R+ is the composition of a loss and a function from the hypothesis class parameterized by w 2 Rd . Since the distribution ⇡ is unknown , one needs to rely on empirical risk minimization as a surrogate to ( 1 ) , minw2Rd n R̂ ( w , Sn ) : = ( 1/n ) Xn i=1 ` ( w , zi ) o , ( 2 ) where Sn : = { z1 , . . . , zn } denotes a training set of n points that are independently and identically distributed ( i.i.d . ) and sampled from ⇡ , and model training often amounts to using an optimization algorithm to solve the above problem . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Statistical learning theory is mainly interested in understanding the behavior of the generalization error , i.e. , R̂ ( w , Sn ) R ( w ) . While classical results suggest that models with large number of parameters should suffer from poor generalization [ SSBD14 , AB09 ] , modern neural networks challenge this classical wisdom : they can fit the training data perfectly , yet manage to generalize well [ ZBH+17 , NBMS17 ] . Considering that the generalization error is influenced by many factors involved in the training process , the conventional algorithm- and data-agnostic uniform bounds are typically overly pessimistic in a deep learning setting . In order to obtain meaningful , non-vacuous bounds , the underlying data distribution and the choice of the optimization algorithm need to be incorporated in the generalization bounds [ ZBH+17 , DR17 ] . Our goal in this study is to develop novel generalization bounds that explicitly incorporate the data and the optimization dynamics , through the lens of dynamical systems theory . To motivate our approach , let us consider stochastic gradient descent ( SGD ) , which has been one of the most popular optimization algorithms for training neural networks . It is defined by the following recursion : wk = wk 1 ⌘rR̃k ( wk 1 ) , where rR̃k ( w ) : = rR̃⌦k ( w ) : = ( 1/b ) X i2⌦k r ` ( w , zi ) . Here , k represents the iteration counter , ⌘ > 0 is the step-size ( also called the learning-rate ) , rR̃k is the stochastic gradient , b is the batch-size , and ⌦k ⇢ { 1 , 2 , . . . , n } is a random subset drawn with or without replacement with cardinality |⌦k| = b for all k. Constant step-size SGD forms a Markov chain with a stationary distribution w1 ⇠ µ , which exists and is unique under mild conditions [ DDB20 , YBVE20 ] , and intuitively we can expect that the generalization performance of the trained model to be intimately related to the behavior of the risk R ( w ) over this limit distribution µ . In particular , the Markov chain defined by the SGD recursion can be written by using random functions h⌦k at each SGD iteration k , i.e. , wk = h⌦k ( wk 1 ) , with h⌦k ( w ) = w ⌘rR̃k ( w ) . ( 3 ) Here , the randomness in h⌦k is due to the selection of the subset ⌦k . In fact , such a formulation is not specific to SGD ; it can cover many other stochastic optimization algorithms if the random function h⌦k is selected accordingly , including second-order algorithms such as preconditioned SGD [ Li17 ] . Such random recursions ( 3 ) and characteristics of their stationary distribution have been studied extensively under the names of iterated random functions [ DF99 ] and iterated function system ( IFS ) [ Fal04 ] . In this paper , from a high level , we relate the ‘ complexity of the stationary distribution ’ of a particular IFS to the generalization performance of the trained model . We illustrate our context in two toy examples . In the first one , we consider a 1-dimensional quadratic problem with n = 2 and ` ( w , z1 ) = w2/2 and ` ( w , z2 ) = w2/2 w. We run SGD with constant step-size ⌘ to minimize the resulting empirical risk . We simply choose ⌦k ⇢ { 1 , 2 } uniformly random with batch-size b = 1 , and we plot the histograms of stationary distributions for different step-size choices ⌘ 2 { 1/100 , 1/3 , 2/3 } in Figure 1 . We observe that the support of the stationary distribution of SGD depends on the step-size : As the step-size increases the support becomes less dense and a fractal structure in the stationary distribution can be clearly observed . This behavior is not surprising , at least for this toy example . It is well-known that the set of points that is invariant under the resulting IFS ( termed as the attractor of the IFS ) for the specific choice of ⌘ = 2/3 is the famous ‘ middle-third Cantor set ’ [ FW09 ] , which coincides with the support of the stationary distribution of the SGD . As another example , we run SGD with constant step-size in order to train an ordinary linear regression model for a dataset of n = 5 samples and d = 2 dimensions , i.e. , a > i w ⇡ yi , where for i = 1 , . . . , 5 , yi and each coordinate of ai are drawn uniformly at random from the interval [ 1 , 1 ] . Figure 2 shows the heatmap of the resulting stationary distributions for different step-size choices ⌘ ranging from 0.1 to 0.9 ( bright colors represent higher density ) . We observe that for small step-size choices , the stationary distribution is dense , whereas a fractal structure can be clearly observed as the step-size gets larger . Fractals are complex patterns and the level of this complexity is typically measured by the Hausdorff dimension of the fractal , which is a notion of dimension that can take fractional values1 , and can be much smaller than the ambient dimension d. Recently , assuming that SGD trajectories can be well-approximated by a certain type of stochastic differential equations ( SDE ) [ ŞGN+19 , ŞSG19 , NcGR19 , ŞZTG20 ] , it is shown that the generalization error can be controlled by the Hausdorff dimension of the trajectories of the SDE , instead of their ambient dimension d [ ŞSDE20 ] . That is , the ambient dimension that appears in classical learning theory bounds is replaced with the Hausdorff dimension . The fractal geometric approach presented in [ ŞSDE20 ] can capture the ‘ low dimensional structure ’ of fractal sets and provides an alternative perspective to the compression-based approaches that aim to understand why overparametrized networks do not overfit [ AGNZ18 , SAN20 , SAM+20 , HJTW21 ] . However , SDE approximations for SGD often serve as mere heuristics , and guaranteeing a good approximation typically requires unrealistically small step-sizes [ LTE19 ] . For more realistic step- sizes , theoretical concerns have been raised about the validity of conventional SDE approximations for SGD [ LMA21 , GŞZ21 , Yai19 ] . Another drawback of the SDE approximation is that the bounds in [ ŞSDE20 ] are implicit , in the sense that they can not be related to algorithm hyperparameters , problem geometry , or data . We address these issues and present a direct , discrete-time analysis by exploiting the connections between IFSs and stochastic optimization algorithms . Our contributions are summarized as follows : • We extend [ ŞSDE20 ] and show that the generalization error can be linked to the Hausdorff dimension of invariant measures ( rather than the Hausdorff dimension of sets as in [ ŞSDE20 ] ) . More precisely , under appropriate conditions , we establish a generalization bound for the stationary distribution of IFS w1 ⇠ µ . That is , with probability at least 1 2⇣ , |R̂ ( w1 , Sn ) R ( w1 ) | . s dimHµ log 2 ( n ) n + log ( 1/⇣ ) n , ( 4 ) for n large enough , where dimHµ is the ( upper ) Hausdorff dimension of the measure µ . • By leveraging results from IFS theory , we further link dimHµ to ( i ) the form of the recursion ( e.g. , h⌦k in ( 3 ) ) , ( ii ) algorithm hyperparameters ( e.g. , ⌘ , b ) , and ( iii ) problem geometry ( e.g. , Hessian of R̃k ) , through a single term , which encapsulates all these components and their interaction . • We establish bounds on dimHµ for SGD and preconditioned SGD algorithms , when used to minimize various empirical risk minimization problems such as least squares , logistic regression , support vector machines . In all cases , we explicitly link the generalization performance of the model to the hyperparameters of the underlying training algorithm . • Finally , we numerically compute key quantities that appear in our generalization bounds , and show empirically that they have a statistically significant correlation with the generalization error . Notation and preliminaries . Bd ( x , r ) ⇢ Rd denotes the closed ball centered around x 2 Rd with radius r. A function f : Rd1 ! Rd2 is said to be ( Fréchet ) differentiable at x 2 Rd1 if there exists a d1⇥d2 matrix Jf ( x ) : Rd1 ! Rd2 such that limkhk ! 0 kf ( x+h ) f ( x ) Jf ( x ) hk/khk = 0 . The matrix Jf ( x ) is called the differential of f , also known as the Jacobian matrix at x , and determinant of Jf ( x ) is called the Jacobian determinant [ HS74 ] . For real-valued functions f , g , we define f ( n ) = ! ( g ( n ) ) if limn ! 1 |f ( n ) |/g ( n ) = 1 . For a set A , |A| denotes its cardinality . For a scalar-valued function f̃ : R ! R , we define kf̃k1 = maxr2R |f̃ ( r ) | . 1The Hausdorff dimension of the middle-third Cantor set in Figure 1 is log3 ( 2 ) ⇡ 0.63 whereas the ambient dimension is 1 [ Fal04 , Example 2.3 ] . 2 Technical Background on Fractal Geometry . Fractal sets emerge virtually in all branches of science , and fractal-based techniques have been used in machine learning [ SHTY13 , MSS19 , DSD20 , ŞSDE20 , AGZ21 ] . The inherent ‘ complexity ’ of a fractal set often plays an important role and it is typically measured by its fractal dimension , where several notions of dimension have been proposed [ Fal04 ] . In this section , we briefly mention two important notions of fractal dimension , which will be used in our theoretical development . Minkowski dimension of a set . The Minkowski dimension ( also known as the box-counting dimension [ Fal04 ] ) is defined as follows . Let F ⇢ Rd be a set and for > 0 , let N ( F ) denote a collection of sets that contains the smallest number of closed balls of diameter at most which cover F . Then the upper-Minkowski dimension of F is defined as follows : dimMF : = lim sup ! 0 h log |N ( F ) | / log ( 1/ ) i . ( 5 ) To visualize the upper-Minkowski dimension of a set F , consider the set F lying on an evenly spaced grid and count how many boxes are required to cover the set . The upper-Minkowski dimension measures how this number changes as the grid is made finer using a box-counting algorithm . Hausdorff dimension of a set . An alternative to the purely geometric Minkowski dimension , the Hausdorff dimension [ Hau18 ] is a measure theoretical notion of fractal dimension . It is based on the Hausdorff measure , which generalizes the traditional notions of area and volume to noninteger dimensions [ Rog98 ] . More precisely , for s 0 , let F ⇢ Rd and > 0 , and denote H s ( F ) : = inf P1 i=1 diam ( Ai ) s , where the infimum is taken over all the -coverings { Ai } i of F , that is , F ⇢ [ iAi with diam ( Ai ) < for every i . The s-dimensional Hausdorff measure of F is defined as the monotone limit Hs ( F ) : = lim ! 0 Hs ( F ) . When s 2 N , Hs is the s-dimensional Lebesgue measure up to a constant factor ; hence the generalization of ‘ volume ’ to fractional orders . Based on the Hausdorff measure , the Hausdorff dimension of a set F ⇢ Rd is then defined as follows : dimH F : = sup { s > 0 : H s ( F ) > 0 } = inf { s > 0 : Hs ( F ) < 1 } . In other words , the Hausdorff dimension of F is the ‘ moment ’ s when Hs ( F ) drops from 1 to 0 , that is , Hr ( F ) = 0 for all r > dimH F and Hr ( F ) = 1 for all r < dimH F . We always have 0 dimH F d , and when F is bounded , we always have 0 dimH F dimMF d [ Fal04 ] . Furthermore , the Hausdorff dimension of Rd equals d , and the Hausdorff dimension of smooth Riemannian manifolds correspond to their intrinsic dimension , e.g . dimH Sd 1 = d 1 , where Sd 1 is the unit sphere in Rd . Hausdorff dimension of a probability measure . IFSs generate invariant measures as the number of iterates goes to infinity , and random fractals arise from such invariant measures . There has been a growing literature that studies the structure of such random fractals [ Saz00 , NSB02 , MS02 , Ram06 , FST06 , JR08 ] , where the notion of fractal dimension can be extended to measures , and our theory will rely on the Hausdorff dimension of invariant measures associated with stochastic optimization algorithms . In particular , we will mainly use the upper Hausdorff dimension dimHµ of a Borel probability measure µ on Rd , which is defined as follows : dimHµ : = inf { dimH A : µ ( A ) = 1 } . In other words , dimHµ is the smallest Hausdorff dimension of all measurable sets with full measure .

Summary: The authors give a general generalization bound for stochastic optimization algorithms, depending on the Hausdorff dimension of the invariant measure. This bound takes into account the dataset and training algorithm, and draws on theory on fractal dimensions and random iterated function systems in dynamical systems. The authors link this bound to quantities that can be estimated from running the algorithm in question - namely an expected logarithm of norm of Jacobian. They specialize the theory for least squares (SGD & stochastic Newton), logistic regression, SVMs, and 1-hidden-layer neural networks. They present experiments showing that this quantity does in fact correlate with with generalization performance for fully-connected and convolutional neural networks trained on CIFAR10 and SVHN.