Does Preprocessing Help Training Over-parameterized Neural Networks?

• First , by preprocessing the initial weights of the neural networks , we can train the neural network in Õ ( m1−Θ ( 1/d ) nd ) cost per iteration . • Second , by preprocessing the input data points , we can train the neural network in Õ ( m4/5nd ) cost per iteration . From the technical perspective , our result is a sophisticated combination of tools in different fields , greedy-type convergence analysis in optimization , sparsity observation in practical work , high-dimensional geometric search in data structure , concentration and anti-concentration in probability . Our results also provide theoretical insights for a large number of previously established fast training methods . In addition , our classical algorithm can be generalized to the Quantum computation model . Interestingly , we can get a similar sublinear cost per iteration but avoid preprocessing initial weights or input data points . 1 Introduction . Over the last decade , deep learning has achieved dominating performance over many areas , e.g. , computer vision [ LBBH98 , KSH12 , SLJ+15 , HZRS16 ] , natural language processing [ CWB+11 , DCLT18 ] , game playing [ SHM+16 , SSS+17 ] and beyond . The computational resource requirement for deep neural network training grows very quickly . Designing a fast and provable training method for neural networks is , therefore , a fundamental and demanding challenge . Almost all deep learning models are optimized by gradient descent ( or its variants ) . The total training time can be split into two components , the first one is the number of iterations and the second one is the cost per spent per iteration . Nearly all the iterative algorithms for acceleration can be viewed as two separate lines of research correspondingly , the first line is aiming for an algorithm that has as small as possible number of iterations , the second line is focusing on designing as efficient as possible data structures to improve the cost spent per iteration of the algorithm [ Vai89 , CLS19 , LSZ19 , JLSW20 , JKL+20 , JSWZ21 ] . In this paper , our major focus is on the second line . There are a number of practical works trying to use a nearest neighbor search data structure to speed up the per-step computation of the deep neural network training [ CMF+20 , LXJ+20 , CLP+21 , DMZS21 ] . However , none of the previous work is able to give a provable guarantee . In this paper , 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . our goal is to develop training algorithms that provably reduce per step time complexity . Let us consider the ReLU activation neural network and two-layer neural network1 . Let n denote the number of training data points . Let d denote the dimension of each data point . Let m denote the number of neurons . In each iteration of gradient descent ( GD ) , we need to compute prediction for each point in the neural network . Each point xi ∈ Rd , requires to compute m inner product in d dimension . Thus , Ω ( mnd ) is a natural barrier for cost per iteration in training neural networks ( in both forward computation and backward computation ) . A natural question to ask is Is it possible to improve the cost per iteration of training neural network algorithm ? E.g. , is o ( mnd ) possible ? We list our contributions as follows : • We provide a new theoretical framework for speeding up neural network training by : 1 ) adopting the shifted neural tangent kernel ; 2 ) showing that only a small fraction ( o ( m ) ) of neurons are activated for each input data in each training iteration ; 3 ) identifying the sparsely activated neurons via geometric search ; 4 ) proving that the algorithm can minimize the training loss to zero in a linear convergence rate . • We provide two theoretical results 1 ) our first result ( Theorem 6.1 ) builds a dynamic half-space report data structure for the weights of a neural network , to train neural networks in sublinear cost per iteration ; 2 ) our second result ( Theorem 6.2 ) builds a static half-space report data-structure for the input data points of the training data set for training a neural network in sublinear time . Acceleration via high-dimensional search data-structure . High-dimensional search data structures support efficiently finding points in some geometric query regions ( e.g. , half-spaces , simplices , etc ) . Currently , there are two main approaches : one is based on Locality Sensitive Hashing ( LSH ) [ IM98 ] , which aims to find the close-by points ( i.e. , small ` 2 distance [ DIIM04 , AR15 , AIL+15 , ARN17 , Raz17 , AIR18 , BIW19 , DIRW20 ] or large inner product [ SL14 , SL15b , SL15a ] ) of a query q ∈ Rd in a given set of points S ⊂ Rd . This kind of algorithms runs very fast in practice , but most of them only support approximate queries . Another approach is based on space partitioning data structures , for example , partition trees [ Mat92a , Mat92b , AEM92 , AC09 , Cha12 ] , k-d trees / range trees [ CT17 , TOG17 , Cha19 ] , Voronoi diagrams [ ADBMS98 , Cha00 ] , which can exactly search the query regions . Recent works have successfully applied high-dimensional geometric data structure to reduce the complexity of training deep learning models . SLIDE [ CMF+20 ] accelerates the forward pass by retrieving neurons with maximum inner product via an LSH-based data structure ; Reformer [ KKL20 ] similarly adopts LSH to reduce the memory usage for processing long sequence ; MONGOOSE [ CLP+21 ] accelerates the forward pass by retrieving neurons with maximum inner products via a learnable LSH-based data structure [ Cha02 ] and lazy update framework [ CLS19 ] . Despite the great empirical success , there is no theoretical understanding of such acceleration . The goal of our paper is to theoretically characterize the acceleration brought by the highdimensional geometric data structure . Specifically , our algorithm and analysis are built upon the HSR data structures [ AEM92 ] which can find all the points that have large inner products and support efficient data update . Note that HSR comes with a stronger recovery guarantee than LSH , in the sense that HSR , whereas LSH is guaranteed to find some of those points . Convergence via over-parameterization . Over the last few years , there has been a tremendous work studying the convergence result of deep neural network explicilty or implicitly based on neural tangent kernel ( NTK ) [ JGH18 ] , e.g . [ LL18 , DZPS19 , AZLS19a , AZLS19b , DLL+19 , ADH+19a , ADH+19b , SY19 , CGH+19 , ZMG19 , CG19 , ZG19 , OS20 , LSS+20 , JT20 , ZPD+20 , HLSY21 , BPSW21 ] . It has been shown that ( S ) GD can train a sufficiently wide NN with random initialization will converge to a small training error in polynomial steps . 1An alternative name of the two-layer neural network is “ one-hidden layer neural network ” . 2 Challenges and Techniques . • Empirical works combine high-dimensional search data structures ( e.g. , LSH ) with neural network training , however , they do not work theoretically due to the following reasons : – Without shifting , the number of activated ( and therefore updated ) neurons is Θ ( m ) . There is no hope to theoretically prove o ( m ) complexity ( See Challenge 1 ) . – Approximate high-dimensional search data structures might miss some important neurons , which can potentially prevent the training from converging ( see Challenge 2 ) . • Our solutions are : – We propose a shifted ReLU activation that is guaranteed to have o ( m ) number of activated neurons . Along with the shifted ReLU , we also propose a shifted NTK to rigorously provide a convergence guarantee ( see Solution 1 ) . – We adopt an exact high-dimensional search data structure that better couples with the shifted NTK . It takes o ( m ) time to identify the activated neurons and fits well with the convergence analysis as it avoids missing important neurons ( see Solution 2 ) . Challenge 1 : How to sparsify an over-parameterized neural network ? To speed up the training process , we need the neural network to be “ sparse ” , that is , for each training data x ∈ Rd , the number of activated neurons is small . Then , in the forward computation , we can just evaluate a small subset of neurons . However , in the previous NTK analysis ( e.g. , [ DZPS19 ] ) , the activation function is σ ( x ) = max { 〈wr , x〉 , 0 } , and the weights vectors wr are initially sampled from a standard ddimensional Gaussian distribution . Then , by the symmetry of Gaussian distribution , we know that for every input data x , there will be about half of the neurons being activated , which means that we can only obtain a constant-factor speedup . Solution 1 The problem actually comes from the activation function . In practice , people use a shifted ReLU function σb ( x ) = max { 〈wr , x〉 , br } to train neural networks . The main observation of our work is that threshold implies sparsity . We consider the setting where all neurons have a unified threshold parameter b . Then , by the concentration of Gaussian distribution , there will be O ( exp ( −b2 ) ·m ) activated neurons after the initialization . The next step is to show that the number of activated neurons will not blow up too much in the following training iterations . In [ DZPS19 , SY19 ] , they showed that the weights vectors are changing slowly during the training process . In our work , we open the black box of their proof and show a similar phenomenon for the shifted ReLU function . More specifically , a key component is to prove that for each training data , a large fraction of neurons will not change their status ( from non-activated to activated and vice versa ) in the next iteration with high probability . To achieve this , they showed that this is equivalent to the event that a standard Gaussian random variable in a small centered interval [ −R , R ] , and applied the anti-concentration inequality to upper-bound the probability . In our setting , we need to upper-bound the probability of z ∼ N ( 0 , 1 ) in a shifted interval [ b − R , b + R ] . On the one hand , we can still apply the anti-concentration inequality by showing that the probability is at most Pr [ z ∈ [ −R , R ] ] . On the other hand , this probability is also upper-bounded by Pr [ z > b − R ] , and for small R , we can apply the concentration inequality for a more accurate estimation . In the end , by some finer analysis of the probability , we can show that with high probability , the number of activated neurons in each iteration is also O ( exp ( −b2 ) ·m ) for each training data . If we take b = Θ ( √ logm ) , we only need to deal with truly sublinear in m of activated neurons in the forward evaluation . Challenge 2 : How to find the small subset of activated neurons ? A linear scan of the neurons will lead to a time complexity linear in m , which we hope to avoid . Randomly sampling or using LSH for searching can potentially miss important neurons which are important for a rigorous convergence analysis . Solution 2 Given the shifted ReLU function σb ( 〈wr , x〉 ) = max { 〈wr , x〉 − b , 0 } , the active neurons are those with weights wr lying in the half space of 〈wr , x〉 − b > 0 . Finding such neurons is equivalent to a computational geometry problem : given m points in Rd , in each query and a half space H , the goal is to output the points contained in H. Here we use the Half-Space Reporting ( HSR ) data structure proposed by [ AEM92 ] : after proper initialization , the HSR data structure can return all points lying in the queried half space with complexity as low as O ( log ( n ) + k ) , where k is the number of such points . Note that the HSR data structure well couples with the shifted ReLU , as the number of activated neurons k is truly sublinear in m as per the setting of b = Θ ( √ logm ) .

Summary: This paper proposed two preprocessing algorithms that can reduce the computation cost per iteration on two layer ReLU neural networks. Suppose the network has width $m$, input dimension $d$, and number of training points $n$. The first proposed algorithm preprocessed the initial weights of the neural network and can reduce the cost per iteration from $\Omega(mnd)$ to $\widetilde{O}(m^{1-\Theta(1/d)}nd)$. The second algorithm preprocessed the input data points and reduce the cost per iteration to  $\widetilde{O}(m^{4/5}nd)$. The analysis is done in the NTK regime and the authors also proved linear convergence result under the preprocessing.   The algorithm fixed the bias term to a relatively large value so that given any input the number of active neurons is only $o(m).$ The algorithm only needs to do forward and backward computation on these active neurons, which saves the computation cost. In order to efficiently identity the active neurons at each iteration, the authors adopted a dynamic half-space report data structure for the weights in one algorithm and used a static half-space report data structure for the inputs in the other.