Efficient Algorithms for Learning Depth-2 Neural Networks with General ReLU Activations

1 Introduction . The empirical success of deep learning in recent years has led to a flurry of recent works exploring various theoretical aspects of deep learning such as learning , optimization and generalization . A fundamental question in the theory of deep learning is to identify conditions under which one can design provably time-efficient and sample-efficient learning algorithms for neural networks . Perhaps surprisingly , even for the simplest case of a depth-2 feedforward neural network , the learning question remains unresolved . In this work we make progress on this front by studying the problem of learning an unknown neural network of the form y = f ( x ) = aT ( WTx+ b ) . ( 1 ) We are given access to a finite amount of samples of the form ( xi , yi ) drawn i.i.d . from the data distribution , where each xi is comes from the standard Gaussian distribution N ( 0 , I ) , and yi = f ( xi ) . The goal is to design an algorithm that outputs an approximation of the function f up to an arbitrary error measured in the expected mean squared metric ( squared ` 2 loss ) . An efficient learning algorithm has running time and sample complexity that are polynomial in the different problem parameters such as the input dimensionality , number of hidden units , and the desired error . Without any further assumptions on the depth-2 network , efficient learning algorithms are unlikely . The recent work of [ DKKZ20 ] provides evidence by proving exponential statistical query lower bounds ( even when x is Gaussian ) that rule out a broad class of algorithms . Several recent works have designed efficient algorithms for depth-2 neural networks in the special setting when the bias term b = 0 . One prominent line of work [ GLM18 , BJW19 ] give polynomial time algorithms under the non-degeneracy assumption that the matrix W has full-rank . Another body of work relaxes the full-rank assumption by designing algorithms that incur an exponential dependence on the number of hidden units [ DKKZ20 , CKM20 , GK19 ] , or a quasipolynomial dependence when the coefficients { ai : i 2 [ m ] } are all non-negative [ DK20 ] . There is little existing literature on learning neural networks in the presence of the bias term . A notable exception is an approach based on computing “ score functions ” [ JSA15 ] that applies to certain activations with bias and requires 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . various assumptions that are not satisfied by the ReLU function . The diminished expressivity of neural networks without the bias terms leads to the following compelling question : Can we design polynomial time algorithms even in the presence of bias terms in the ReLU units ? We answer the question affirmatively by designing efficient algorithms for learning depth-2 neural networks with general ReLU activations , under the assumption that Wd⇥m has linearly independent columns ( hence m d ) . In fact , our algorithms can be extended to work under much weaker assumptions on W , that allow for m d ( m O ( d ` ) for any constant ` 1 ) in a natural smoothed analysis setting considered in prior works [ GKLW19 ] ( see Theorem 3.2 and Corollary 3.3 ) . An important consequence of our techniques is the fact that the network parameters are identifiable up to signs , as long as no two columns of W are parallel , and all the { ai : i 2 [ m ] } are nonzero . Furthermore we show that this ambiguity in recovering the signs is inherent unless stronger assumptions are made . Conceptual and technical challenges with bias terms . Similar to prior works [ GKLW19 , JSA15 ] , our techniques rely on the use of tensor decomposition algorithms to recover the parameters of the network . In the absence of any bias , it can be shown that the 4th Hermite coefficient of the function f ( x ) takes the form f̂4 = P i=1 aiw ⌦4 i where wi are the columns of W and a = ( a1 , a2 , . . . , ) . When the columns of W are linearly independent , existing algorithms for tensor decompositions in the full-rank setting can be used to recover the parameters [ Har70 ] .1 However , when bias terms are present , there are several challenges that we highlight below . In the presence of biases the kth Hermite expansion of f ( x ) takes the form f̂k = P m i=1 aigk ( bi ) w ⌦k i , where gk is a function that may vanish on some of the unknown bi parameters . This creates a hurdle in recovering the corresponding wi . A simple example where the above approach fails is when some of the bi = ±1 , since the corresponding rank-1 terms vanish from the decomposition of f̂4 . To overcome this obstacle , we first give a precise expression for the function gk ( bi ) involving the ( k 2 ) th Hermite polynomial ( see Lemma 3.5 ) . We then design an algorithm that decomposes multiple tensors obtained from Hermite coefficients to recover the parameters . We use various properties of Hermite polynomials to analyze the algorithm e.g. , the separation of roots of consecutive Hermite polynomials is used to argue that each wi is recovered from decomposing at least one of the tensors . Secondly , in the presence of the bias terms , recovery of all the parameters ( even up to sign ambiguities ) may not even be possible from polynomially many samples . For instance , consider a particular hidden node with output ( w > i x+ bi ) . If bi is a large positive number then it behaves like a linear function ( always active ) . Hence if multiple bis are large positive numbers then one can only hope to recover a linear combination of their corresponding weights and biases . On the other hand if bi is a large negative constant then the activation is 0 except with probability exponentially small in |bi| . We can not afford a sample complexity that is exponential in the magnitude of the parameters . Furthermore , when the columns of W are not linearly independent , the tensor decomposition based method will only recover good approximations up to a sign ambiguity for the terms whose bias does not have very large magnitude i.e. , we recover ± ( w > i x+ bi ) if |bi| is not very large . To handle the above issue we proceed in two stages . In the first stage we recover the network parameters ( up to signs ) of all the “ good ” terms , i.e. , hidden units with biases of small magnitude . To handle the “ bad ” terms ( large magnitude bias ) we show that a linear functions is a good approximation to the residual function comprising of the bad terms . Based on the above , we show that one can solve a truncated linear regression problem to learn a function g ( x ) that achieves low mean squared error with respect to the target f ( x ) . The output function g ( x ) is also a depth-2 ReLU network with at most two additional hidden units than the target network . There are several other technical challenges that arise in the analysis sketched above , when there are sampling errors due to using only a polynomial amount of data ( for example , the tensors obtained from f̂k may have some rank-1 terms that are small but not negligible , that may affect the robust recovery guarantees for tensor decompositions ) . We obtain our robust guarantees by leveraging many useful properties of Hermite polynomials , and a careful analysis of how the errors propagate . 1Tensor decompositions will in fact recover each ReLU activation up to an ambiguity in the sign . However , in the full-rank setting , the correct sign can also be recovered ( as we demonstrate later in Theorem 3.1 ) . The rest of the paper is organized as follows . We present preliminaries in Section 1.1 followed by related work in Section 2 . We then formally present and discuss our main results in Section 3 . In Section 4 we present our main algorithm and analysis in the population setting , i.e. , under the assumption that one has access to infinite data from the distribution . We then present the finite sample extension of our algorithm in Section 5 that achieves polynomial runtime and sample complexity . 1.1 Model Setup and Preliminaries . We consider the supervised learning problem with input x 2 Rd drawn from a standard ddimensional Gaussian distribution N ( 0 , Id⇥d ) and labels y generated by a neural network y = f ( x ) = aT ( WTx+ b ) , where a , b 2 Rm , W 2 Rd⇥m and is the element-wise ReLU activation function , i.e. , ( t ) = max ( t , 0 ) . We denote the column vectors of W as wi 2 Rd with i 2 [ m ] and ai as the i ’ th element of vector a , similarly for b and x . We pose a constraint on magnitudes of a , b , W such that they are all B-bounded for some 1 B poly ( m , d ) , i.e . kak1 , kbk1 , kWk1 B , and mini2 [ m ] |ai| 1/B . Furthermore , we assume kwik2 = 1 without loss of generality . If wi are not unit vectors , we can always scale ai and bi to kwikai and bikwik respectively so that wi are normalized . We will denote by ( · ) the cumulative density function ( CDF ) of the standard Gaussian distribution . Finally , for a matrix M , we will use sk ( M ) to denote the kth largest singular value of M . For introduction and further preliminaries regarding poly ( · ) notation , basics of Hermite polynomials and tensor decomposition , please refer to Appendix A . 2 Related Work . By now there is a vast literature exploring various aspects of deep learning from a theoretical perspective . Here we discuss the works most relevant in the context of our results . As discussed earlier , the recent works of [ GLM18 , BJW19 , GKLW19 ] provide polynomial time algorithms for learning depth-2 feedforward ReLU networks under the assumption that the input distribution is Gaussian and that the matrix W is full rank . Some of these works consider a setting where the output is also a high dimensional vector [ BJW19 , GKLW19 ] , and also consider learning beyond the Gaussian distribution . However , these works do not extend to the case of non-zero bias . The work of [ JSA15 ] proposed a general approach based on tensor decompositions for learning an unknown depth-2 neural network that could also handle the presence of the bias terms . The tensor used in the work of [ JSA15 ] is formed by taking the weighted average of a “ score ” function evaluated on each data point . In this way their approach generalizes to a large class of distributions provided one has access to the score function . However , for most data distributions computing the score function is a hard task itself . When the input distribution is Gaussian , then the score functions correspond to the Hermite coefficients of the target function f ( x ) and can be evaluated efficiently . However , the analysis in [ JSA15 ] does not extend to the case of ReLU activations for several reasons . Their technique needs certain smoothness and symmetry assumptions on the activations that do not hold for ReLU . These assumptions also ensure that all the terms in the appropriate tensor are non-zero . We do not make such assumptions , and tackle one of the main challenges by showing that one can indeed recover a good approximation to the network by analyzing multiple higher order tensors . Furthermore , the authors in [ JSA15 ] assume that the biases , and the spectral norm of W are both bounded by a constant . As a result they do not handle the case of biases of large magnitude where some of the ReLU units mostly function as linear functions ( with high probability ) . There have also been works on designing learning algorithms for neural networks without assumptions on the linear independence of columns of W . These results incur an exponential dependence on either the input dimensionality or the number of parameters in the unknown network [ DKKZ20 , CKM20 ] , or quasipolynomial dependence when the coefficients { ai : i 2 [ m ] } are all non-negative [ DK20 ] . In particular , the result of [ CKM20 ] provides a learning algorithm for arbitrary depth neural networks under the Gaussian distribution with an “ FPT ” guarantee ; its running time is polynomial in the dimension , but exponential in the number of ReLU units . Given the recent correlational statistical query lower bounds on learning deep neural networks [ DKKZ20 , GGJ+20 , DV21 ] , getting a fully polynomial time algorithm without any assumptions is a challenging open problem , even under Gaussian marginals . Polynomial time algorithms with fewer assumptions and beyond depth-2 can be designed if the activation functions in the first hidden layer are sigmoid functions [ GK19 ] . Finally , there is also extensive literature on analyzing the convergence properties of gradient descent and stochastic gradient descent for neural networks . The results in this setting implicitly or explicitly assume that the target function is well approximated in the Neural Tangent Kernel ( NTK ) space of an unknown network . Under this assumption these results show that gradient descent on massively overparameterized neural networks can learn the target [ APVZ14 , Dan17 , DFS16 , AZLS19 , DZPS19 , ADH+19 , LXS+19 , CB18 , JGH21 ] .

Summary: The paper gives a polynomial-time algorithm for learning depth-2 neural networks with the ReLU activation and bias terms in the realizable setting w.r.t. the square loss. The result assumes Gaussian marginals and certain assumptions on the weights matrix W. Assumptions of this nature are required due to known harness results. Prior works on efficient learning of depth-2 networks either assume that there is no bias or require some assumptions on the activation function that are not satisfied by the ReLU function. Thus, the algorithm shown in this paper is the first efficient algorithms for learning depth-2 ReLU networks with bias.  The main results are: - Assume that the input distribution is Gaussian, the parameters of the target network are bounded, and the smallest singular value of W^{\odot \ell} for some constant \ell is at least inverse polynomial, then the algorithm runs in polynomial time in all relevant parameters (and in exponential time in \ell) and finds w.h.p. a depth-2 network that approximates the target network. If the bias terms of the target network are not too large, then the algorithm also recovers the parameters of the target network. If the bias terms are large then recovering (some of) the parameters is not possible. Where \ell=1 the assumption on W^{\odot \ell} implies that W has full rank. Where \ell>1 the requirement may allow the number of hidden neurons m to be significantly larger than the input dimension d.  - The authors show that the above assumption on W^{\odot \ell} holds w.h.p. in the smoothed-analysis setting, namely, where the target weights matrix is slightly perturbed with a random i.i.d. gaussian noise. Hence, in the smoothed-analysis setting the assumption on W^{\odot \ell} is not required. In this case, the number of hidden neurons can be m=O(d^\ell) and the running time of the algorithm is exponential in \ell (and polynomial in the other parameters). - Another consequence of the analysis, is that if the weights in the second layer are non-zero, and there are no two parallel weight vectors in the fist layer, then the parameters of the target network are uniquely identifiable (not necessarily in polynomial time).  The algorithm is obtained in two steps. First, assuming that the Hermite coefficients are already known, the authors give an algorithm for recovering the parameters. This algorithm is non-robust since in reality the coefficients are not known, and we have to estimate them (with some errors) from the training set. In the second step the authors give the robust algorithm, that uses an unbiased estimator for the Hermite coefficients obtained from the training set. The algorithm relies on Jennrich’s algorithm for tensor decomposition.