Proxy Convexity: A Unified Framework for the Analysis of Neural Networks Trained by Gradient Descent

Although the optimization objectives for learning neural networks are highly nonconvex , gradient-based methods have been wildly successful at learning neural networks in practice . This juxtaposition has led to a number of recent studies on provable guarantees for neural networks trained by gradient descent . Unfortunately , the techniques in these works are often highly specific to the particular setup in each problem , making it difficult to generalize across different settings . To address this drawback in the literature , we propose a unified non-convex optimization framework for the analysis of neural network training . We introduce the notions of proxy convexity and proxy Polyak-Lojasiewicz ( PL ) inequalities , which are satisfied if the original objective function induces a proxy objective function that is implicitly minimized when using gradient methods . We show that stochastic gradient descent ( SGD ) on objectives satisfying proxy convexity or the proxy PL inequality leads to efficient guarantees for proxy objective functions . We further show that many existing guarantees for neural networks trained by gradient descent can be unified through proxy convexity and proxy PL inequalities . 1 Introduction . Understanding the ability of gradient-based stochastic optimization algorithms to find good minima of non-convex objective functions has become an especially important problem due to the success of stochastic gradient descent ( SGD ) in learning deep neural networks . Although there exist non-convex objective functions and domains for which SGD will necessarily lead to sub-optimal local minima , it appears that for many problems of interest in deep learning , across domains as varied as natural language and images , these worst-case situations do not arise . Indeed , a number of recent works have developed provable guarantees for GD and SGD when used for objective functions defined in terms of neural networks over certain distributions , despite the non-convexity of the underlying optimization problem [ 5 , 3 , 6 , 22 , 14 , 16 ] . To date , however , there has not been a framework which could unify the variegated approaches for guarantees in these settings . In this work , we introduce the notion of proxy convexity and demonstrate that many existing provable guarantees for learning with neural networks trained by gradient-based optimization fall into a problem satisfying proxy convexity . First , let us define the learning problem over a distribution D , where the goal is to minimize the expected loss min w∈W F ( w ) : = Ez∼Df ( w ; z ) , ( 1.1 ) 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . whereW is a parameter domain and f ( · ; z ) is a loss function . We are interested in guarantees using the online SGD algorithm , which uses a set of i.i.d . samples zt i.i.d.∼ D and has updates given by wt+1 = wt − η∇f ( wt ; zt ) , where η > 0 is a fixed learning rate . We now introduce the first notion of proxy convexity we will consider in the paper . Definition 1.1 ( Proxy convexity ) . We say that a function f : Rp → R satisfies ( g , h ) -proxy convexity if there exist functions g , h : Rp → R such that for all w , v ∈ Rp , 〈∇f ( w ) , w − v〉 ≥ g ( w ) − h ( v ) . Clearly , every convex function f satisfies ( f , f ) -proxy convexity . We next introduce the analogy of proxy convexity for the Polyak–Łojasiewicz ( PL ) inequality [ 23 ] . Definition 1.2 ( g-proxy , ξ-optimal PL inequality ) . We say that a function f : Rp → R satisfies a g-proxy , ξ-optimal Polyak–Łojasiewicz inequality with parameters α ∈ ( 0 , 2 ] and µ > 0 ( in short , f satisfies the ( g , ξ , α , µ ) -PL inequality ) if there exists a function g : Rp → R and scalars ξ ∈ R , µ > 0 such that for all w ∈ Rp , ‖∇f ( w ) ‖α ≥ 1 2 µ ( g ( w ) − ξ ) . As we shall see below , the proxy PL inequality is a natural extension of the standard PL inequality . Our main contributions are as follows . 1 . When f satisfies ( g , h ) -proxy convexity , and f is either Lipschitz or satisfies a particular smoothness assumption , then for any norm bound R > 0 , the online SGD algorithm run for polynomial ( in 1/ε and R ) number of iterations satisfies the following in expectation over z1 , . . . , zT ∼ DT , min t < T Ez∼Dg ( wt ; z ) ≤ min ‖w‖≤R Ez∼Dh ( w ; z ) + ε . 2 . When f satisfies a ( g , ξ , α , µ ) -proxy PL inequality and has Lipschitz gradients , SGD run for a polynomial ( in 1/ε ) number of iterations satisfies the following in expectation over z1 , . . . , zT ∼ DT , min t < T Ez∼Dg ( wt ; z ) ≤ ξ + ε . 3 . We demonstrate that many previous guarantees for neural networks trained by gradient descent can be unified in the framework of proxy convexity . As we will describe in more detail below , if a loss function ` is ( g , h ) -proxy convex or satisfies a g-proxy PL inequality , then the optimization problem is straightforward and the crux of the problem then becomes connecting guarantees for the proxy g with approximate guarantees for f . Notation . We use uppercase letters to refer to matrices , and lowercase letters will either refer to vectors or scalars depending on the context . For vectors w , we use ‖w‖ to refer to the Euclidean norm , and for matrices W we use ‖W‖ to refer to the Frobenius norm . We use the standard O ( · ) , Ω ( · ) notations to hide universal constants , with Õ ( · ) and Ω̃ ( · ) additionally hiding logarithmic factors . 2 Proxy Convexity in Comparison to Other Non-convex Optimization Frameworks . In this section , we describe how proxy convexity and proxy PL-inequalities relate to other notions in non-convex optimization . In Section 6 , we will discuss additional related work . First , recall that a function f is ( g , h ) -proxy convex if there exist functions g and h such that for all w , v , 〈∇f ( w ) , w − v〉 ≥ g ( w ) − h ( v ) . One notion from the non-convex optimization literature that is related to our notion of proxy convexity is that of invexity [ 19 ] . A function f is invex if it is differentiable and there exists a vector-valued function k ( w , v ) such that for any w , v , 〈∇f ( w ) , k ( w , v ) 〉 ≥ f ( w ) − f ( v ) . It has been shown that a smooth function f is invex if and only if every stationary point of f is a global minimum [ 10 ] . However , for many problems of interest involving neural networks , it is not the case that every stationary point will be a global optimum , which makes invexity a less appealing framework for understanding neural networks . Indeed , we shall see in Example 4.2 below that if one considers the problem of learning a single ReLU neuron x 7→ max ( 0 , 〈w , x〉 ) under the squared loss , it is not hard to see that there exist stationary points which are not global minima ( e.g. , w = 0 assuming the convention σ′ ( 0 ) = 0 ) . By contrast , we shall see that the single ReLU neuron does satisfy a form of proxy convexity that enables SGD to find approximately ( but not globally ) optimal minima . Thus even the simplest neural networks induce objective functions which are proxy convex and non-invex . We shall see in Example 3.3 that proxy convexity appears in the objective functions induced by wide and deep neural networks as well . To understand how the proxy PL inequality compares to other notions in the optimization literature , recall that an objective function f satisfies the standard PL inequality [ 33 , 27 ] if there exists µ > 0 such that ‖∇f ( w ) ‖2 ≥ µ 2 [ f ( w ) − f∗ ] ∀w , where f∗ = minw f ( w ) . Clearly , any stationary point of an objective satisfying the standard PL inequality is globally optimal . Thus , the presence of local minima among stationary points in neural network objectives makes the standard PL inequality suffer from the same drawbacks that invexity does for understanding neural networks trained by gradient descent . This further applies to any of the conditions which are known to imply the PL inequality , like weak strong convexity , the restricted secant inequality , and the error bound condition ( Karimi et al . [ 23 ] ) .1 In comparison , the ( g , ξ , α , µ ) -proxy PL inequality is satisfied if there exists a function g and constants ξ > 0 , α ∈ ( 0 , 2 ] and µ > 0 such that ‖∇f ( w ) ‖α ≥ µ 2 [ g ( w ) − ξ ] ∀w . It is clear that if a function f satisfies the standard PL inequality , then it satisfies the ( f , f∗ , 2 , µ ) proxy PL inequality . Stationary points w0 of objective functions satisfying the proxy PL inequality have ‖∇f ( w0 ) ‖ = 0 which imply g ( w0 ) ≤ ξ . In the case that g = f , the slack error term ξ allows for the proxy PL inequality framework to accommodate the possibility that stationary points may not be globally optimal ( i.e . have objective value f∗ = minw f ( w ) ) , but could be approximately optimal by , for example , having objective value at most ξ = C · f∗ or ξ = C · √ f∗ for some constant C ≥ 1 . When g 6= f , the proxy PL inequality allows for the possibility of analyzing a proxy loss function g which is implicitly minimized when using gradient-based optimization of the objective f . At a high level , proxy convexity and the proxy PL inequality are well-suited to situations where stationary points may not be globally optimal , but may be approximately optimal with respect to a related optimization objective . The proxy convexity framework allows for one to realize this through developing problem-specific analyses that connect the proxy objective g to the original objective f . As we shall see below , rich function classes like neural networks are often more easily analyzed by considering a proxy objective function that naturally appears when one analyzes the gradient of the loss . Finally , we note that [ 25 ] introduced a different generalization of the PL inequality , namely the PL∗ and PL∗ε inequality , which relaxes the standard PL inequality definition so that the PL condition only needs to hold on a subset of the domain . In particular , a function f : Rp → R satisfies the PL∗ inequality on a set S ⊂ Rp if there exists µ > 0 such that ‖∇f ( w ) ‖2 ≥ µf ( w ) ∀w ∈ S. Likewise , f satisfies the PL∗ε inequality on S if there exists a set S and ε ≥ 0 such that the PL∗ inequality holds on the set Sε = { w ∈ S : f ( w ) ≥ ε } . One can see that if f satisfies the PL∗ε inequality on S , then the function g ( w ) : = f ( w ) + ε satisfies the g-proxy , ε-optimal PL inequality on S. We wish to emphasize the differences in the framing and motivation of the PL∗ε inequality by [ 25 ] and that of proxy convexity and the proxy PL inequality in this paper . [ 25 ] focus on the geometry of 1Karimi et al . [ 23 ] shows that these conditions imply the PL inequality under the assumption that the objective function has Lipschitz-continuous gradients . optimization in the overparameterized setting where one has a fixed set of samples { ( xi , yi ) } ni=1 and a parametric model class g ( x ; w ) ( for w ∈ Rp , p > n ) and the goal is to solve g ( xi ; w ) = yi for all i ∈ [ n ] . In this setting one can view the optimization problem as a nonlinear least squares system with p unknowns and n equations , and [ 25 ] use geometric arguments to show that when p > n the PL∗ condition is satisfied throughout most of the domain . They extend the PL∗ condition to the PL∗ε condition with the motivation that in underparameterized settings , or when performing early stopping , there may not exist interpolating solutions . By contrast , we focus on the stochastic optimization setting , where the goal is to minimize the expected loss over some distribution ; it is unclear how geometric arguments for minimizing the training loss can lead to generalization guarantees in the overparameterized setting , especially when considering online SGD where samples are observed one-by-one and ‘ overparameterization ’ has a less clear meaning . Furthermore , in this work we are not primarily interested in understanding how overparameterization affects optimization . Rather , our aim is to develop a framework that allows for formal characterizations of optimization problems where stationary points are not globally optimal with respect to the original objective but are approximately optimal with respect to proxy objective functions . We will demonstrate below that such a framework can help unify a number of works on learning with neural networks trained by gradient descent .

Summary: The authors propose a generalization of the well established notions of convexity and PL inequalities, which they term proxy convexity and proxy PL inequalities correspondingly.  These generalizations are particularly well suited for studying non convex models like neural networks since they do not require all stationary points to be global optima. In fact the authors find that inequalities used in several prior works to prove bounds for neural nets trained with SGD can be unified under proxy convexity and proxy PL inequalities.