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

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

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

heading: 1 Introduction
text: 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.

heading: 2 Proxy Convexity in Comparison to Other Non-convex Optimization Frameworks
text: 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.

heading: 3 Proxy PL Inequality Implies Proxy Objective Guarantees
text: In this section, we show that for loss functions satisfying a proxy PL inequality, SGD2 efficiently minimizes the proxy. We leave the proofs for Section 5. Theorem 3.1. Suppose F (w) = Ez∼Df(w; z) where f(· ; z) satisfies the (g(· ; z), ξ(z), α, µ)-proxy PL inequality for some function g(· ; z) : Rp → R for each z. Denote by G(w) := Ez∼Dg(w; z). Assume that f is non-negative and has L2-Lipschitz gradients. Then for any ε > 0, provided η < 1/L2, online SGD with fixed step size η and run for T = 2η−1(µε/2)−2/αf(w0; z0) iterations results in the following guarantee in expectation over z0, . . . , zT−1 ∼ DT ,
min t<T
G(wt) ≤ Ez∼Dξ(z) + ε. (3.1)
To get a feel for how a proxy PL inequality might be useful for learning neural networks, consider a classification problem with labels y ∈ {±1}, and suppose that N(W ;x) is a neural network function. A standard approach for learning neural networks is to minimize the cross-entropy loss `(yN(W ;x)) = log ( 1 + exp(−yN(W ;x)) ) using gradient descent. Using the variational form of the norm, we have
‖∇`(yN(W ;x))‖ = sup ‖U‖=1 〈∇`(yN(W ;x)), U〉
≥ −`′(yN(W ;x)) · y〈∇N(W ;x), V 〉, (3.2)
where V is any matrix satisfying ‖V ‖ = 1. Now, although the function −`′ is not an upper bound for ` (indeed, −`′ < `), it is an upper bound for a constant multiple of the zero-one loss, and can thus serve as a proxy for the classification error.3 This is because for convex and decreasing losses `, the function −`′ is non-negative and decreasing, and hence by Markov’s inequality,
P(y 6= sgn(N(W ;x))) = P(y ·N(W ;x) < 0) = P(−`′(yN(w;x)) > −`′(0))
≤ 1 −`′(0) E[−`′(yN(W ;x))].
Thus, if one can bound the population risk under −`′, one has a bound for the classification error. Indeed, this property has been used in a number of recent works on neural networks [6, 13, 22, 16]. This lets the −`′ term in (3.2) represent the desired proxy g in the definition of the (g, ξ, α, µ)-proxy PL inequality. Thus, for neural network classification problems, the problem of showing the neural network has small classification error is reduced to constructing a matrix V that allows for the quantity y〈∇N(W ;x), V 〉 to be large and non-negative. The quantity y〈∇N(W ;x), V 〉 can be thought of
2We focus on online SGD in this paper for simplicity. Analogous optimization guarantees would hold for other variants of gradient descent that utilize samples in batches.
3In fact, the function z 7→ [`′(z)]2 is also a proxy for the classification error; see [15, Appendix A].
as a margin function that is large when the gradient of the neural network loss points in a good direction. Although we shall see below that in some instances one can derive a lower bound for y〈∇N(W ;x), V 〉 that holds for all W , x, and y, a more general approach would be to show that along the gradient descent trajectory W (t), a lower bound for y〈∇N(W (t);x), V 〉 holds.4
In the remainder of this section, we will show that a number of recent works on learning neural networks with gradient descent utilized proxy PL inequalities. In our first example, we consider recent work by Charles and Papailiopoulos [7] that directly used a (standard) PL inequality. Example 3.2 (Standard PL inequality for single leaky ReLU neurons and deep linear networks). Charles and Papailiopoulos [7] showed that the standard PL inequality holds in two distinct settings. The first is that of a single leaky ReLU neuron x 7→ σ(〈w, x〉), where σ(z) = max(cσz, z) for cσ 6= 0. They showed that if smin(X) is the smallest singular value of the matrix X ∈ Rn×d of n samples, then for a λ-strongly convex loss `, the loss f(w) = `(σ(〈w, x〉)) satisfies the standard µ-PL inequality, i.e., the (f, f∗, 2, µ)-proxy PL inequality for µ = λsmin(X)2c2σ (Charles and Papailiopoulos [7, Theorem 4.1]).
The same authors also showed that under certain conditions the standard PL inequality holds when the neural network takes the form N(W ;x) = WL · · ·W1x and the loss is the squared loss, f(W ) = 1/2(y − N(W ;x))2. In particular, they showed that if smin(Wi) ≥ τ > 0 throughout the gradient descent trajectory, then f satisfies the standard µ-PL inequality for µ = Lτ2L−2/ ∥∥(XX>)−1X∥∥2 F (Charles and Papailiopoulos [7, Theorem 4.5]).
The standard PL inequality has been used by a number of other authors in the deep learning theory literature, see e.g. Xie et al. [36, Theorem 1], Hardt and Ma [20, Eq. 2.3], Zhou and Liang [40, Theorem 1], Shamir [35, Theorem 3].
In our next example, we show that a proxy PL inequality holds for deep neural networks in the neural tangent kernel (NTK) regime. Example 3.3 (Proxy PL inequality for deep neural networks in NTK regime). Consider the class of deep, L-hidden-layer ReLU networks, either with or without residual connections:
N1(W ;x) = σ(W (1)x), Nl(W ;x) = slNl−1(W ;x) + σ(W (l)Nl−1(W ;x)), l = 2, . . . , L,
N(W ;x) = m∑ j=1 aj [NL(W ;x)]j ,
where sl = 0 for fully-connected networks and sl = 1 for residual networks. Cao and Gu [6, Theorem 4.2], Frei, Cao, and Gu [13, Lemma 4.3], and Zou et al. [41, Lemma B.5] have shown that under certain distributional assumptions and provided the iterates of gradient descent stay close to their intialization, one can guarantee that for the cross-entropy loss f(W ; (x, y)) = `(yN(W ;x)),
‖∇f(W ; (x, y))‖ ≥ C1 · −`′(yN(W ;x)). (3.3) By defining g(W ; z) = −`′(yN(W ;x)), the loss f satisfies the (g, 0, 1, 2C1)-proxy PL inequality. Since the ReLU is not smooth, the loss f will not have Lipschitz gradients, and thus a direct application of Theorem 3.1 is not possible. Instead, the authors show that in the NTK regime, the loss obeys a type of semi-smoothness that still allows for an analysis simliar to that of Theorem 3.1. Example 3.4 (Proxy PL inequality for one-hidden-layer networks outside NTK regime). Consider a one-hidden-layer network with activation function σ,
N(W ; (x, y)) = m∑ j=1 ajσ(〈wj , x〉), (3.4)
where the second layer weights {aj}mj=1 are randomly initialized and fixed at initialization, but the {wj}mj=1 are trained. Assume σ satisfies σ′(z) ≥ cσ > 0 for all z (e.g., the leaky ReLU activation). Frei, Cao, and Gu have shown [16, Lemma 3.1] that there exists a matrix V ∈ Rm×d with ‖V ‖F = 1 such that for distributions satisfying anti-concentration, for any x, y and W ,
y〈∇N(W ;x), V 〉 ≥ C1[cσ − ξ(x, y)], 4Although our results as stated would not immediately apply in this setting, the proof would be the same up
to trivial modifications.
where E[ξ(x, y)] = O( √ OPT) where OPT is the best classification error achieved by a halfspace over D. Thus, when f(W ; (x, y)) = `(yN(W ; (x, y)) is the cross-entropy loss, ‖∇f(W ; (x, y))‖ = sup
‖Z‖F=1 〈∇f(W ; (x, y)), Z〉
≥ −`′(yN(W ;x)) · y〈∇N(W ;x), V 〉 ≥ C1cσ · [−`′(yN(W ;x))− c−1σ ξ(x, y)].
As in Example 3.3, by defining g(W ; z) = −`′(yN(W ;x)), the loss f satisfies the (g, c−1σ ξ(x, y), 1, 2C1cσ)-proxy PL inequality. Thus, provided we can show that f(·; z) has L2Lipschitz gradients for some constant L2 > 0, Theorem 3.1 shows that
min t<T P(x,y)(y 6= sgn(N(W (t);x))) ≤ min t<T E(x,y)∼D[−`′(yN(W (t);x)] |`′(0)|
≤ 2C2c−1σ E(x,y)∼Dξ(x, y) + ε = O( √ OPT) + ε.
Provided σ is such that σ′ is continuous and differentiable, then f has L2-Lipschitz gradients and thus the guarantees will follow. In particular, this analysis follows if σ is any smoothed approximation to the leaky ReLU which satisfies σ′(z) ≥ cσ > 0. Note that the above optimization analysis is an original contribution of this work as we utilize a completely different proof technique than that of [16]. In that paper, the authors utilize a Perceptronstyle proof technique that analyzes the correlation 〈W (t), V 〉 of the weights found by gradient descent and a reference matrix V . Their proof relies crucially on the homogeneity of the (non-smooth) leaky ReLU activation, namely that zσ′(z) = σ(z) for z ∈ R, and cannot accommodate more general smooth activations. By contrast, the proxy PL inequality proof technique in this example relies upon the smoothness of the activation function and is more similar to smoothness-based analyses of gradient descent.

heading: 4 Proxy Convexity Implies Proxy Objective Guarantees
text: In this section, we show that if f satisfies (g, h)-proxy convexity, we can guarantee that by minimizing f with gradient descent, we find a hypothesis for which g(w) is at least as small as the best normbounded predictor as measured by the loss h. We present two versions of our result: one that relies upon fewer assumptions on the loss f but needs a small step size, and another that requires a proxy smoothness assumption on f but allows for a constant step size. The proofs for the theorem are given in Section 5. Theorem 4.1. Suppose that F (w) := Ez∼Df(w; z) and f(·; z) is (g(·; z), h(·; z))-proxy convex for each z. Denote H(w) := Ez∼Dh(w; z) and G(w) := Ez∼Dg(w; z).
(a) Assume there exists L1 > 0 such that for all w, Ez∼D[‖∇f(w; z)‖2] ≤ L21. Then for any v ∈ Rp and any ε > 0, performing online SGD on F (w) from an arbitrary initialization w0 with fixed step size η ≤ εL−21 for T = η−1ε−1 ‖w0 − v‖
2 iterations implies that, in expectation over (z0, . . . , zT−1) ∼ DT ,
min t<T
G(wt) ≤ H(v) + ε.
(b) Assume there exists L2 > 0 such that for all w, Ez∼D[‖∇f(w; z)‖2] ≤ 2L2Ez∼Dg(w; z). Then for any v ∈ Rp and any ε > 0, performing online SGD on F (w) from an arbitrary initialization with fixed step size η ≤ L−12 /2 for T = η−1ε−1 ‖w0 − v‖
2 implies that, in expectation over (z0, . . . , zT−1) ∼ DT ,
min t<T
G(wt) ≤ (1 + 2ηL2)H(v) + ε.
In order for (g, h)-proxy convexity to be useful, there must be a way to relate guarantees for g into guarantees for the desired objective function f . In the remainder of this section, we will discuss two neural network learning problems which are non-convex and yet satisfy proxy convexity which leads to generalization guarantees. Our first example is the problem of learning a neural network with a single nonlinear unit.
Example 4.2 (Single neuron satisfies proxy convexity). Consider the problem of learning a single neuron x 7→ σ(〈w, x〉) under the squared loss, where σ is the ReLU activation. The objective function of interest is F (w) = E(x,y)∼D1/2(σ(〈w, x〉)− y)2. Denote
F ∗ := min ‖w‖≤1 F (w).
It is known that F is non-convex [37]. Under the assumption that learning sparse parities with noise is computationally hard, it is known that no polynomial time algorithm can achieve risk F ∗ exactly; moreover, it is known that (unconditionally) the standard gradient descent algorithm cannot achieve risk F ∗ [18].5 However, Frei, Cao, and Gu [14] showed that although F is non-convex and no algorithm can achieve risk F ∗, F does satisfy a form of proxy convexity that allows for gradient descent to achieve risk O( √ F ∗). They showed that the loss function
f(w; (x, y)) = 1/2(σ(〈w, x〉)− y)2
satisfies (g, h)-proxy convexity, where
g(w; (x, y)) = 2 [σ(〈w, x〉)− σ(〈v∗, x〉)]2 σ′(〈w, x〉), h(v; (x, y)) = |σ(〈v, x〉)− y| = √ 2f(v; (x, y),
where v∗ is the population risk minimizer of F (w) (see their Eq. (3.13)). Moreover, they showed (see their Eq. (3.9)) ‖∇f(w; z)‖2 ≤ 8g(w; z). Thus Theorem 4.1(b) implies that SGD with step size η ≤ 1/8 and T = 2η−1ε−1 ‖w0 − v∗‖2 iterations will find a point wt satisfying
G(wt) = 2E(x,y) [ (σ(〈wt, x〉)− σ(〈v∗, x〉))2 σ′(〈wt, x〉) ] ≤ (1 + 8η)H(v∗) + ε ≤ (1 + 8η)E|σ(〈v∗, x〉)− y|+ ε
≤ (1 + 8η) √
E[(σ(〈v∗, x〉)− y)2] + ε = O( √ F ∗).
The authors then show that under some distributional assumptions on D, G(wt) = O( √ F ∗) implies
F (wt) = O( √ F ∗) [14, Lemma 3.5]. Thus, the optimization problem for F induces a proxy convex optimization problem defined in terms of G which yields guarantees for G in terms of H , and this in turn leads to approximate optimality guarantees for the original objective F .
In our next example, we show that a number of works on learning one-hidden-layer ReLU networks in the neural tangent kernel regime [21] can be cast as problems satisfying proxy convexity. Example 4.3 (Proxy convexity for one-hidden-layer ReLU neural networks in the NTK regime). Consider the class of one-hidden-layer ReLU networks consisting of m neurons,
N(W ; (x, y)) = m∑ j=1 ajσ(〈wj , x〉),
where the {aj}mj=1 are randomly initialized and fixed at initialization, but the {wj}mj=1 are trained. Suppose we consider a binary classification problem, where y ∈ {±1} and we minimize the crossentropy loss,
F (W ) = E(x,y)∼Df(W ; (x, y)), f(W ; (x, y)) = ` ( yN(W ; (x, y)) ) , `(z) = log(1+exp(−z)).
Ji and Telgarsky [22, Proof of Lemma 2.6] showed that there exists a function h̃(a,W, V ; (x, y)) such that the iterates of gradient descent satisfy
〈∇f(W ; (x, y)), w − v〉 ≥ f(W ; (x, y))− h̃(a,W, V ; (x, y)). 5This stands in contrast to learning a single leaky ReLU neuron x 7→ max(αx, x) for α 6= 0, which as we
showed in Example 3.2 can be solved using much simpler techniques.
Under the assumption that the iterates of gradient descent stay close to the initialization (i.e., the neural tangent kernel regime), they show that h̃(a,W, V ; (x, y)) ≤ ε under distributional assumptions, and thus F (w) will satisfy (f, h̃ ≡ ε)-proxy convexity. The cross entropy loss satisfies [`′(z)]2 ≤ `(z), and thus we can apply Theorem 4.1(b) to get guarantees of the form mint<T F (wt) ≤ ε for the cross-entropy loss F (W ) of SGD-trained neural networks in the NTK regime.
In another problem of learning one-hidden-layer networks, Allen-Zhu, Li, and Liang [3, Proof of Lemma B.4] show that there exists a proxy loss function g(a,W ; (x, y)) such that provided the neural network weights stay close to their initialized values, f(a,W ; (x, y)) satisfies (g, g+ ε) proxy convexity. Again using that the cross-entropy loss satisfies [`′(z)]2 ≤ `(z), Theorem 4.1(b) shows that SGD-trained neural networks in the NTK regime satisfy mint<T G(Wt) ≤ minV G(V ) + ε. They further show that the proxy loss g is close to the cross entropy loss, i.e. |E[g(a,W ; (x, y))]− E[f(a,W ; (x, y))]| ≤ ε, implying a bound of the form mint<T F (Wt) ≤ minV F (V ) + ε.

heading: 5 Proof of the Main Results
text: In this section we provide the proofs of the theorems given in Sections 3 and 4.
We first give the proof of Theorem 3.1 which provides guarantees for learning with objectives satisfying proxy PL inequalities.
Proof of Theorem 3.1. Since f has L2-Lipschitz gradients, we have for any w,w′ and fixed z,
f(w; z) ≤ f(w′; z) + 〈∇f(w′; z), w − w′〉+ L2 2 ‖w − w′‖2 .
Taking w = wt+1, w′ = wt, and z = zt,
f(wt+1; zt) ≤ f(wt; zt)− η ‖∇f(wt; zt)‖2 + η2L2
2 ‖∇f(wt; zt)‖2
= f(wt; zt)− η [1− ηL2/2] ‖∇f(wt; zt)‖2 . (5.1)
Since η < 1/L2, we have (1− ηL2/2)−1 ≤ 2, and thus we can rearrange the above to get
‖∇f(wt; zt)‖2 ≤ 1
η(1− ηL2/2) [f(wt; zt)− f(wt+1; zt)]
≤ 2 η [f(wt; zt)− f(wt+1; zt)]. (5.2)
Summing the above from t = 0 to t = T − 1 and using that f is non-negative, we get
1
T T−1∑ t=0 ‖∇f(wt; zt)‖2 ≤ 2f(w0; z0) ηT .
Using the definition of proxy PL inequality, this implies
1
T T−1∑ t=0 (µ/2)2/α(g(wt; zt)− ξ(zt))2/α ≤ 2f(w0; z0) ηT .
Taking the minimum over t < T and re-arranging terms, this means
min t<T
(g(wt; zt)− ξ(zt))2/α ≤ 2f(w0; z0)
ηT (µ/2)2/α .
Therefore, we have
min t<T
g(wt; zt) ≤ ξ(zt) + 2 µ · ( 2f(w0; z0) ηT )α/2 .
Taking T = 2η−1f(w0; z0)(µε/2)−2/α and taking expectations over z0, . . . , zT−1, we get (3.1).
We next prove guarantees for SGD when the objective satisfies proxy convexity.
Proof of Theorem 4.1. By the definition of proxy convexity,
‖wt − v‖2 − ‖wt+1 − v‖2 = 2η〈∇f(wt; zt), wt − v〉 − η2 ‖∇f(wt; zt)‖2
≥ 2η[g(wt; zt)− h(v; zt)]− η2 ‖∇f(wt; zt)‖2 . Conditional on the values of z0, . . . , zt−1 (and hence on the value of wt), taking expectations of both sides with respect to zt ∼ D results in ‖wt − v‖2 − Ezt∼D ‖wt+1 − v‖ 2 ≥ 2η[G(wt)−H(v)]− η2Ezt∼D ‖∇f(wt; zt)‖ 2 . (5.3)
For case (a), this results in
‖wt − v‖2 − Ezt∼D ‖wt+1 − v‖ 2 ≥ 2η[G(wt)−H(v)− η/2L21].
Dividing both sides by 2ηT and summing from t = 0, . . . , T − 1, we get
1
T T−1∑ t=0 G(wt) ≤ 1 T T−1∑ t=0 H(v) + ηL21 2 + ‖w0 − v‖2 − EzT−1∼D ‖wT − v‖ 2 2ηT .
Taking expectations over z0:t = (z0, . . . , zt−1) ∼ Dt, we get
min t<T
Ez0:t∼DtG(wt) ≤ 1
T T−1∑ t=0 Ez0:t∼DtG(wt) ≤ H(v) + ηL21 2 + ‖w0 − v‖2 2ηT .
In particular, for η ≤ εL−21 and T = η−1ε−1 ‖w0 − v‖ 2, we get
min t<T Ez0:t∼DtG(wt) ≤ H(v) + ε.
For case (b), (5.3) becomes
‖wt − v‖2 − Ezt∼D ‖wt+1 − v‖ 2 ≥ 2η[G(wt)−H(v)− ηL2G(wt)]
= 2η [(1− ηL2)G(wt)−H(v)] .
Dividing both sides by 2ηT (1− ηL2) and summing from t = 0, . . . , T − 1, we get
1
T T−1∑ t=0 G(wt) ≤ 1 1− ηL2 H(v) + ‖w0 − v‖2 − EzT−1∼D ‖wT − v‖ 2 2ηT (1− ηL2)
≤ (1 + 2ηL2)H(v) + (1 + 2ηL2) ‖w0 − v‖2
2ηT ,
where in the last line we have used that η ≤ L−12 /2 and that 1/(1− x) ≤ 1 + 2x on [0, 1/2]. Taking expectations over z0:t = (z0, . . . , zt−1) ∼ Dt, we get
min t<T
Ez0:t∼DtG(wt) ≤ 1
T T−1∑ t=0 Ez0:t∼DtG(wt)
≤ (1 + 2ηL2)H(v) + 2 ‖w0 − v‖2
2ηT .
In particular, for T = η−1ε−1 ‖w0 − v‖2, we get min t<T Ez0:t∼DtG(wt) ≤ (1 + 2ηL2)H(v) + ε.
We note that under additional assumptions on D and the loss function, we could improve the results from holding in expectation over the draws of the sample to high probability guarantees by using standard concentration arguments. This is easily done when the objective function satisfies proxy convexity: we can make a slight modification to the proof of Theorem 4.1 to argue inductively that until we reach a point withG(wt) ≤ H(v)+ηL21+ε, we have that that ‖wt − v‖
2−‖wt+1 − v‖2 ≥ ε. This implies that the norm of the predictors remain bounded throughout the trajectory of gradient descent until we reach the desired point with G(wt) ≤ H(v) + ηL21 + ε, which can then be used in Rademacher complexity-type arguments [4]. This type of argument was previously used by e.g., Frei et al. [14].

heading: 6 Additional Related Work
text: The Polyak–Lojasiewicz inequality can be dated back to the original works of Polyak [33] and Lojasiewicz [27]. Recent work by Karimi et al. [23] proved linear convergence under the PL condition and showed that the PL condition is one of the weakest assumptions under which linear convergence is possible. In particular, they showed that the error bound inequality [28], essential strong convexity [26], weak strong convexity [31], and the restricted secant inequality [38] are all assumptions under which linear convergence is possible and that each of these assumptions implies the PL inequality.
As we described in Section 2, the standard PL condition was shown to hold under certain assumptions for neural network objective functions [20, 36, 40, 7]. In addition to those covered in this paper, there are a number of other provable guarantees for generalization of SGD-trained networks which rely on a variety of different techniques, such as tensor methods [24] and utilizing connections with partial differential equations by way of mean field approximations [29, 9, 30, 8].
In the optimization literature, recent work has shown that SGD can efficiently find stationary points and can escape saddle points [17, 12]. As the proxy PL inequality implies guarantees for the proxy objective function at stationary points of the original optimization objective, our framework can naturally be used for other optimization algorithms that are known to efficiently find stationary points, such as SVRG [2, 34], Natasha2 [1], SARAH/SPIDER [32, 11], and SNVRG [39].

heading: 7 Conclusion
text: In this paper we have introduced the notion of proxy convexity and proxy PL inequality and developed guarantees for learning with stochastic gradient descent under these conditions. We demonstrated that many recent works in the learning of neural networks with gradient descent can be framed in terms of optimization problems that satisfy either proxy convexity or a proxy PL inequality. While the proxy convexity framework cannot unify all existing analyses of learning neural networks, we hope that it can reveal some of the principles underlying the success of SGD-trained neural networks.
Acknowledgments and Disclosure of Funding
QG is partially supported by the National Science Foundation CAREER Award 1906169, IIS1855099 and IIS-2008981. SF acknowledges the support of the NSF and the Simons Foundation for the Collaboration on the Theoretical Foundations of Deep Learning through awards DMS-2031883 and #814639. The views and conclusions contained in this paper are those of the authors and should not be interpreted as representing any funding agencies.

Review summaries:
This papers introduces two notions, proxy convexity and proxy PL condition, in order to prove convergence results for non-convex functions, in particular for the analysis of neural network training.  1. A function $f$ is $(g,h)$-proxy convex if $\langle \nabla f(w),w-v\rangle \geq g(w)-h(v)$ for all $w,v$. Under this condition, it's proved that the function value of $g$ will be decreased to the smallest value of $h$ after polynomial number of SGD updates.  2. A function $f$ satisfies $(g,\xi,\alpha,\mu)$-proxy PL condition if $\lVert \nabla f(w)\rVert^\alpha \geq \frac{1}{2}\mu(g(w)-\xi).$ Under this condition, it's proved that the function value of $g$ will be decreased to $\xi.$  Then the authors showed that many existing analyses of neural network training can be proved through proxy convexity or proxy PL condition.  

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. 

The Polyak-Lojasiewicz (PL) inequality is a condition on a function that the gradient is lower bounded by the in terms of the gap between the functions value and the optimal value.  This paper presents a generalized conditioned with additional parameters, including a proxy function $g$ and a notion of proxy convexity using a similar proxy function $g$. Unlike the PL inequality, the new condition can be satisfied by functions having non-optimal local minima.  A theorem bounds the the error of gradient descent (with respect to the proxy convexity function $g$).  They give some examples where guarantees with respect to the desired function can be obtained from the guarantees on the proxy function $g$.