Does Preprocessing Help Training Over-parameterized Neural Networks?
Year: 2021

Abstract:
Deep neural networks have achieved impressive performance in many areas. Designing a fast and provable method for training neural networks is a fundamental question in machine learning. The classical training method requires paying Ω(mnd) cost for both forward computation and backward computation, where m is the width of the neural network, and we are given n training points in d-dimensional space. In this paper, we propose two novel preprocessing ideas to bypass this Ω(mnd) barrier: • First, by preprocessing the initial weights of the neural networks, we can train the neural network in Õ(m1−Θ(1/d)nd) cost per iteration. • Second, by preprocessing the input data points, we can train the neural network in Õ(mnd) 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.

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

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

heading: 2 Challenges and Techniques
text: • 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).

heading: 3 Preliminaries
text: Notations For an integer n, we use [n] to denote the set {1, 2, · · · , n}. For a vector x and p ∈ {0, 1, 2,∞}, we use ‖x‖p to denote the entry-wise `p norm of a vector. We use Id to denote ddimensional identity matrix. We use N (µ, σ2) to denote Gaussian distribution with mean mu and variance σ2. We use Õ to hide the polylog factors.
This section is organized as follows. Section 3.1 introduces the neural network and present problem formulation. Section 3.2 presents the half-space report data-structure, Section 3.3 proposes our new sparsity-based Characterizations.

heading: 3.1 Problem Formulation
text: In this section, we introduce the neural network model we study in this work. Let us consider a two-layer ReLU activated neural network f that has width m and `2 loss function. 2 Definition 3.1 (Prediction function and loss function). Given b ∈ R, x ∈ Rd, W ∈ Rd×m and a ∈ Rm,
f(W,x, a) := 1√ m
m∑
r=1
arσb(〈wr, x〉),
L(W ) := 1
2
n∑
i=1
(f(W,xi, a)− yi)2.
We say function f is 2NN(m, b) for simplicity.
Here W are weights that connect input nodes with hidden nodes, a1, · · · , am ∈ R are the weights that connect hidden nodes with output node. The ReLU function σb(x) := max{x − b, 0}, where b is the threshold parameter. Following the literature, we mainly focus on optimizing W ∈ Rd×m. For weights a ∈ Rm, we will never change a during the training after we randomly choose them at the initialization.3
Definition 3.2 (Weights at initialization). We use the following initialization,
• For each r, we sample wr(0) ∼ N (0, Id) • For each r, we sample ar from {−1,+1} uniformly at random
Next, we can calculate the gradient Fact 3.3 (Gradient of the prediction function and loss function). For each r ∈ [m],
∂f(W,x, a)
∂wr = 1√ m arx1w>r x≥b. (1)
and
∂L(W )
∂wr = 1√ m
n∑
i=1
(f(W,xi, a)− yi)arxi1〈wr,xi〉≥b. (2)
To update the weights from iteration k to iteration k + 1, we follow the standard update rule of the GD algorithm,
GD: W (k + 1) = W (k)− η ·∆W (k), where ∆W (k) = ∂L(W (k)) ∂W (k) . (3)
2This is a very standard formulation in the literature, e.g., see [DZPS19, SY19, BPSW21] 3We remark, in some previous work, they do choose shift, but their shift is a random shift. In our application,
it is important that the same b is fixed for all neurons and never trained.
The ODE of the gradient flow is defined as dwr(t)
dt = −∂L(W ) ∂wr . (4)
Definition 3.4 (Error of prediction). For each t ∈ {0, 1, · · · , T}, we define err(t) ∈ Rn to be the error of prediction err(t) = y − u(t), where u(t) := f(W (t), a,X) ∈ Rn

heading: 3.2 Data Structure for Half-Space Reporting
text: The half-space range reporting problem is an important problem in computational geometry, which is formally defined as following: Definition 3.5 (Half-space range reporting). Given a set S of n points in Rd. There are two operations:
• QUERY(H): given a half-space H ⊂ Rd, output all of the points in S that contain in H , i.e., S ∩H .
• UPDATE: add or delete a point in S.
– INSERT(q): insert q into S – DELETE(q): delete q from S
Let Tinit denote the pre-processing time to build the data structure, Tquery denote the time per query and Tupdate time per update. We use the data-structure proposed in [AEM92] to solve the half-space range reporting problem, which admits the interface summarized in Algorithm 1. Intuitively, the data-structure recursively partitions the set S and organizes the points in a tree data-structure. Then for a given query (a, b), all k points of S with sgn(〈a, x〉 − b) ≥ 0 are reported quickly. Note that the query (a, b) here defines the half-space H in Definition 3.5.
Algorithm 1 Half Space Report Data Structure 1: data structure HALFSPACEREPORT 2: procedures: 3: INIT(S, n, d) . Initialize the data structure with a set S of n points in Rd 4: QUERY(a, b) . a, b ∈ Rd. Output the set {x ∈ S : sgn(〈a, x〉 − b) ≥ 0} 5: ADD(x) . Add point x ∈ Rd to S 6: DELETE(x) . Delete point x ∈ Rd from S 7: end data structure
Adapted from [AEM92], the algorithm comes with the following complexity: Corollary 3.6 ([AEM92]). Given a set of n points in Rd, the half-space reporting problem can be solved with the following performances:
• Part 1. Tquery(n, d, k) = Od(n1−1/bd/2c + k), amortized Tupdate = Od(log2(n)).
• Part 2. Tquery(n, d, k) = Od(log(n) + k), amortized Tupdate = Od(nbd/2c−1). We remark that Part 1 will be used in Theorem 6.1 and Part 2 will be used in Theorem 6.2.

heading: 3.3 Sparsity-based Characterizations
text: In this section, we consider the ReLU function with a nonzero threshold: σb(x) = max{0, x − b}, which is commonly seen in practise, and also has been considered in theoretical work [ZPD+20].
We first define the set of neurons that are firing at time t. Definition 3.7 (fire set). For each i ∈ [n], for each t ∈ {0, 1, · · · , T}, let Si,fire(t) ⊂ [m] denote the set of neurons that are “fire” at time t, i.e.,
Si,fire(t) := {r ∈ [m] : 〈wr(t), xi〉 > b}. We define ki,t := |Si,fire(t)|, for all t in {0, 1, · · · , T}.
We propose a new “sparsity” lemma in this work. It shows that σb gives the desired sparsity.
Lemma 3.8 (Sparsity after initialization). Let b > 0 be a tunable parameter. If we use the σb as the activation function, then after the initialization, with probability at least 1 − n · exp(−Ω(m · exp(−b2/2))), it holds that for each input data xi, the number of activated neurons ki,0 is at most O(m · exp(−b2/2)), where m is the total number of neurons.
Proof. By the concentration of Gaussian distribution, the initial fire probability of a single neuron is
Pr[σb(〈wr(0), xi〉) > 0] = Pr z∼N (0,1) [z > b] ≤ exp(−b2/2).
Hence, for the indicator variable 1r∈Si,fire(0), we have
E[1r∈Si,fire(0)] ≤ exp(−b2/2). By standard concentration inequality (Lemma B.1),
Pr [|Si,fire(0)| > k0 + t] ≤ exp ( − t 2/2
k0 + t/3
) ,∀t > 0 (5)
where k0 := m · exp(−b2/2). If we choose t = k0, then we have: Pr [|Si,fire(0)| > 2k0] ≤ exp (−3k0/8)
Then, by union bound over all i ∈ [n], we have that with high probability
1− n · exp(−Ω(m · exp(−b2/2))),
the number of initial fire neurons for the sample xi is bounded by ki,0 ≤ 2m · exp(−b2/2).
The following remark gives an example of setting the threshold b, and will be useful for showing the sublinear complexity in the next section. Remark 3.9. If we choose b = √ 0.4 logm then k0 = m4/5. For t = m4/5, Eq. (5) implies that
Pr [ |Si,fire(0)| > 2m4/5 ] ≤ exp ( −min{mR,O(m4/5)} ) .

heading: 4 Training Neural Network with Half-Space Reporting Data Structure
text: In this section, we present two sublinear time algorithms for training over-parameterized neural networks. The first algorithm (Section 4.1) relies on building a high-dimensional search data-structure for the weights of neural network. The second algorithm (Section 4.2) is based on building a data structure for the input data points of the training set. Both of the algorithms use the HSR to quickly identify the fired neurons to avoid unnecessary calculation. The time complexity and the sketch of the proof are provided after each of the algorithms.

heading: 4.1 Weights Preprocessing
text: We first introduce the algorithm that preprocesses the weights wr for r ∈ [m], which is commonly used in practice [CLP+21, CMF+20, KKL20]. Recall 2NN(m, b) is f(W,x, a) := 1√ m ∑m r=1 arσb(〈wr, x〉). By constructing a HSR data-structure for wr’s, we can quickly find the set of active neurons Si,fire for each of the training sample xi. See pseudo-code in Algorithm 2.
In the remaining part of this section, we focus on the time complexity analysis of Algorithm 2. The convergence proof will be given in Section 5.
Lemma 4.1 (Running time part of Theorem 6.1). Given n data points in d-dimensional space. Running gradient descent algorithm (Algorithm 2) on 2NN(m, b = √ 0.4 logm) (Definition 3.1) the expected cost per-iteration of the gradient descent algorithm is
Õ(m1−Θ(1/d)nd).
Algorithm 2 Training Neural Network via building a data structure of weights of the neural network 1: procedure TRAININGWITHPREPROCESSWEIGHTS({(xi, yi)}i∈[n],n,m,d) . Theorem 6.1 2: Initialize wr, ar for r ∈ [m] and b according to Definition 3.2 and Remark 3.9 3: HALFSPACEREPORT HSR.INIT({wr(0)}r∈[m],m, d) . Algorithm 1 4: for t = 1→ T do 5: Si,fire ← HSR.QUERY(xi, b) for i ∈ [n] 6: Forward pass for xi only on neurons in Si,fire for i ∈ [n] 7: Calculate gradient for xi only on neurons in Si,fire for i ∈ [n] 8: Gradient update for the neurons in ∪i∈[n]Si,fire 9: HSR.DELETE(wr(t)) for r ∈ ∪i∈[n]Si,fire 10: HSR.ADD(wr(t+ 1)) for r ∈ ∪i∈[n]Si,fire 11: end for 12: return Trained weights wr(T + 1) for r ∈ [m] 13: end procedure
Algorithm 3 Training Neural Network via building a data-structure of the input data points 1: procedure TRAININGWITHPROCESSDATA({(xi, yi)}i∈[n],n,m,d) . Theorem 6.2 2: Initialize wr, ar for r ∈ [m] and b according to Definition 3.2 and Remark 3.9 3: HALFSPACEREPORT HSR.INIT({xi}i∈[n], n, d) . Algorithm 1 4: S̃r,fire ← HSR.QUERY(wr(0), b) for r ∈ [m] . S̃r,fire are samples which neuron r fires for 5: Si,fire ← {r | i ∈ S̃r,fire} . Si,fire is the set of neurons, which fire for xi 6: for t = 1→ T do 7: Forward pass for xi only on neurons in Si,fire for i ∈ [n] 8: Calculate gradient for xi only on neurons in Si,fire for i ∈ [n] 9: Gradient update for the neurons in ∪i∈[n]Si,fire 10: for r ∈ ∪i∈[n]Si,fire do 11: Si,fire.DEL(r) for i ∈ S̃r,fire 12: S̃r,fire ← HSR.QUERY(wr(t+ 1), b) 13: Si,fire.ADD(r) for i ∈ S̃r,fire 14: end for 15: end for 16: return Trained weights wr(T + 1) for r ∈ [m] 17: end procedure
Proof. The per-step time complexity is
n∑
i=1
TQUERY(m, d, ki,t) + (TDELETE + TINSERT) · | ∪i∈[n] Si,fire(t)|+ d ∑
i∈[n] ki,t
The first term ∑n i=1 TQUERY(m, d, ki,t) corresponds to the running time of querying the active neuron set Si,fire(t) for all training samples i ∈ [n]. With the first result in Corollary 3.6, the complexity is bounded by Õ(m1−Θ(1/d)nd).
The second term (TDELETE + TINSERT) · | ∪i∈[n] Si,fire(t)| corresponds to updating wr in the highdimensional search data-structure (Lines 9 and 10). Again with the first result in Corollary 3.6, we have TDELETE + TINSERT = O(log2m). Combining with the fact that | ∪i∈[n] Si,fire(t)| ≤ | ∪i∈[n] Si,fire(0)| ≤ O(nm4/5), the second term is bounded by O(nm4/5 log2m). The third term is the time complexity of gradient calculation restricted to the set Si,fire(t). With the bound on ∑ i∈[n] ki,t (Lemma C.10), we have d ∑ i∈[n] ki,t ≤ O(m4/5nd).
Putting them together completes the proof.

heading: 4.2 Data Preprocessing
text: While the weights preprcessing algorithm is inspired by the common practise, the dual relationship between the input xi and model weights wr inspires us to preprocess the dataset before training (i.e., building HSR data-structure for xi). This largely improves the per-iteration complexity and avoids the frequent updates of the data structure since the training data is fixed. More importantly, once the training dataset is preprocessed, it can be reused for different models or tasks, thus one does not need to perform the expensive preprocessing for each training.
The corresponding pseudocode is presented in Algorithm 3. With xi preprocessed, we can query HSR with weights wr and the result S̃r,fire is the set of training samples xi for which wr fires for. Given S̃r,fire for r ∈ [m], we can easily reconstruct the set Si,fire, which is the set of neurons fired for sample xi. The forward and backward pass can then proceed similar to Algorithm 2.
At the end of each iteration, we will update S̃r,fire based on the new wr estimation and update Si,fire accordingly. For Algorithm 3, the HSR data-structure is static for the entire training process. This is the main difference from Algorithm 2, where the HSR needs to be updated every time step to account for the changing weights wr.
We defer the convergence analysis to Section 5 and focus on the time complexity analysis of Algorithm 2 in the rest of this section. We consider d being a constant for the rest of this subsection. Lemma 4.2 (Running time part of Theorem 6.2). Given n data points in d-dimensional space. Running gradient descent algorithm (Algorithm 2) on 2NN(m, b = √ 0.4 logm) (Definition 3.1), the
expected per-iteration running time of initializing S̃r,fire, Si,fire for r ∈ [m], i ∈ [n] is O(m log n+ m4/5n). The cost per iteration of the training algorithm is O(m4/5n log n).
Proof. We analyze the initialization and training parts separately.
Initialization In Lines 4 and 5, the sets S̃r,fire, Si,fire for r ∈ [m], i ∈ [n] are initialized. For each r ∈ [m], we need to query the data structure the set of data points x’s such that σb(wr(0)>x) > 0. Hence, the running time of this step is
m∑
r=1
Tquery(n, d, k̃r,0) = O(m log n+ m∑
r=1
k̃r,0)
= O(m log n+
n∑
i=1
ki,0)
= O(m log n+m4/5n).
where the second step follows from ∑m r=1 k̃r,0 = ∑n i=1 ki,0.
Training Consider training the neural network for T steps. For each step, first notice that the forward and backward computation parts (Line 7 - 9) are the same as previous algorithm. The time complexity is O(m4/5n log n).
We next show that maintaining S̃r,fire, r ∈ [m] and Si,fire, i ∈ [n] (Line 10 - 14) takes O(m4/5n log n) time. For each fired neuron r ∈ [m], we first remove the indices of data in the sets Si,fire, which takes time
O(1) · ∑
r∈∪i∈[n]Si,fire k̃r,t = O(1) ·
m∑
r=1
k̃r,t = O(m 4/5n).
Then, we find the new set of x’s such that σb(〈wr(t + 1), x〉) > 0 by querying the half-space reporting data structure. The total running time for all fired neurons is
∑
r∈∪i∈[n]Si,fire Tquery(n, d, k̃r,t+1) . m4/5n log n+
∑
r∈∪i∈[n]Si,fire k̃r,t+1 = O(m
4/5n log n)
Then, we update the index sets Si,fire in time O(m4/5n). Therefore, each training step takes O(m4/5n log n) time, which completes the proof.

heading: 5 Convergence of Our Algorithm
text: We state the result of our training neural network algorithms (Lemma 5.2) can converge in certain steps. An important component in our proof is to find out a lower bound on minimum eigenvalue of the continuous Hessian matrix λmin(Hcts). It turns out to be an anti-concentration problem of the Gaussian random matrix. In [OS20], they gave a lower bound on λmin(Hcts) for ReLU function with b = 0, assuming the input data are separable. One of our major technical contribution is generalizing it to arbitrary b ≥ 0. Proposition 5.1 (Informal version of Theorem F.1). Given n (normalized) input data points {x1, x2, · · · , xn} ⊆ Rd such that ∀i ∈ [n], ‖xi‖2 = 1. Let parameter δ := mini 6=j{‖xi − xj‖2, ‖xi + xj‖2} denote the data separability. For any shift parameter b ≥ 0, we define shifted NTK Hcts ∈ Rn×n as follows
Hctsi,j := E w∼N (0,Id)
[ 〈xi, xj〉 · 1〈w,xi〉≥b · 1〈w,xj〉≥b ] ,∀i ∈ [n], j ∈ [n].
Then
λmin(H cts) ≥ 0.01e−b2/2δ/n2.
With proposition 5.1, we are ready to show the convergence rate of training an over-parameterized neural network with shifted ReLU function.
Lemma 5.2 (Convergence part of Theorem 6.1 and Theorem 6.2). Suppose input data-points are δ-separable, i.e., δ := mini 6=j{‖xi − xj‖2, ‖xi + xj‖2}. Let m = poly(n, 1/δ, log(n/ρ)) and η = O(λ/n2). Let b = Θ( √ logm). Then
Pr [ ‖err(k)‖22 ≤ (1− ηλ/2)k · ‖err(0)‖22, ∀k ∈ {0, 1, · · · , T} ] ≥ 1− ρ.
Note that the randomness is over initialization. Eventually, we choose T = λ−2n2 log(n/ ) where is the final accuracy.
This result shows that despite the shifted ReLU and sparsely activated neurons, we can still retain the linear convergence. Combined with the results on per-step complexity in the previous section, it gives our main theoretical results of training deep learning models with sublinear time complexity (Theorem 6.1 and Theorem 6.2).

heading: 6 Main Classical Results
text: We present two theorems (under classical computation model) of our work, showing the sublinear running time and linear convergence rate of our two algorithms. We leave the quantum application into Appendix G. The first algorithm is relying on building a high-dimensional geometric search data-structure for the weights of a neural network.
Theorem 6.1 (Main result I, informal of Theorem E.2). Given n data points in d-dimensional space. We preprocess the initialization weights of the neural network. Running gradient descent algorithm (Algorithm 2) on a two-layer, m-width, over-parameterized ReLU neural network will minimize the training loss to zero, and the expected running time of gradient descent algorithm (per iteration) is
Õ(m1−Θ(1/d)nd).
The second algorithm is based on building a data structure for the input data points of the training set. Our second algorithm can further reduce the cost per iteration from m1−1/d to truly sublinear in m, e.g. m4/5.
Theorem 6.2 (Main result II, informal of Theorem E.2). Given n data points in d-dimensional space. We preprocess all the data points. Running gradient descent algorithm (Algorithm 3) on a two-layer, m-width, over-parameterized ReLU neural network will minimize the training loss to zero, and the expected running time of gradient descent algorithm (per iteration) is
Õ(m4/5nd).

heading: 7 Discussion and Limitations
text: In this paper, we propose two sublinear algorithms to train neural networks. By preprocessing the weights of the neuron networks or preprocessing the training data, we rigorously prove that it is possible to train a neuron network with sublinear complexity, which overcomes the Ω(mnd) barrier in classical training methods. Our results also offer theoretical insights for many previously established fast training methods.
Our algorithm is intuitively related to the lottery tickets hypothesis [FC18]. However, our theoretical results can not be applied to explain lottery tickets immediately for two reasons: 1) the lottery ticket hypothesis focuses on pruning weights; while our results identify the important neurons. 2) the lottery ticket hypothesis identifies the weights that need to be pruned after training (by examining their magnitude), while our algorithms accelerate the training via preprocessing. It would be interesting to see how our theory can be extended to the lottery ticket hypothesis.
One limitation of our work is that the current analysis framework does not provide a convergence guarantee for combining LSH with gradient descent, which is commonly seen in many empirical works. Our proof breaks as LSH might miss important neurons which potentially ruins the convergence analysis. Instead, we refer to the HSR data structure, which provides a stronger theoretical guarantee of successfully finding all fired neurons.
Acknowledgments and Disclosure of Funding
SY’s research is supported by NSF grants 1564000 and 1934932. RZ’s research is supported by NSF Grant CCF-1648712 and Scott Aaronson’s Vannevar Bush Faculty Fellowship from the US Department of Defense.

Review summaries:
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.  

This paper studied the problem of training 2-layer over-parameterized neural networks. The commonly used backpropagation algorithm has a running time of O(mnd) per iteration, which is also the running time for neural network evaluation. In this paper, the authors study if the traning could be done in o(mnd). They give affirmative answers and show two main contributions: 1. They proposed a framework for speeding up the training algorithm, where they observed that after the randomized initialization, the activated neurons are sparse. They also gave a convergence analysis for their framework. 2. They provide an instance of the framework by using half-space report data structures to select and update the weights of neurons during the training process. Their first algorithm uses a dynamic data structure to maintain the weights and achieves sublinear cost per iteration. By the dual relation between input data and neuron weights, their second algorithm uses a static data structure to store the training data and achieves a truly sublinear cost. 

The authors propose a variant of the over-parameterized two-layer RELU network where the original RELU activation is shifted so that fewer neurons are activated. The advantage of such modification is that since fewer neurons are activated one may carefully design techniques to only update and carry out computation on the weights related to the activated neurons. In particular there exist algorithms/data structures such as HSR to efficiently search for the activated neurons. One could either build HSR data structures on the data and query the sample or vice versa, producing two variants of the algorithms. Both variants are able to improve the computational dependency on m, the number of hidden neurons, to sub-linear. Convergence guarantees are also provided for the proposed algorithm.