Bypass Exponential Time Preprocessing: Fast Neural Network Training via Weight-Data Correlation Preprocessing
Authors: Josh Alman
Jiehao Liang
Zhao Song
Ruizhe Zhang
Danyang Zhuo
Venue: NeurIPS 2023 poster (2023)
Keywords: Training neural network
Dynamic activated neuron detection
Sparsity
Fine-grained complexity
Data structure

Abstract: Over the last decade, deep neural networks have transformed our society, and they are already widely applied in various machine learning applications. State-of-the-art deep neural networks are becoming larger in size every year to deliver increasing model accuracy, and as a result, model training consumes substantial computing resources and will only consume more in the future.
Using current training methods, in each iteration, to process a data point $x \in \mathbb{R}^d$ in a layer, we need to spend $\Theta(md)$ time to evaluate all the $m$ neurons in the layer. This means processing the entire layer takes $\Theta(nmd)$ time for $n$ data points. Recent work [Song, Yang and Zhang, NeurIPS 2021] reduces this time per iteration to $o(nmd)$, but requires exponential time to preprocess either the data or the neural network weights, making it unlikely to have practical usage. 

In this work, we present a new preprocessing method that simply stores the weight-data correlation in a tree data structure in order to quickly and dynamically detect which neurons fire at each iteration. Our method requires only $O(nmd)$ time in preprocessing and still achieves $o(nmd)$ time per iteration. We complement our new algorithm with a lower bound, proving that assuming a popular conjecture from complexity theory, one could not substantially speed up our algorithm for dynamic detection of firing neurons.

Bypass Exponential Time Preprocessing: Fast Neural Network Training via Weight-Data Correlation Preprocessing
Josh Alman'
Jiehao Liang!
Zhao Song*
Ruizhe Zhang
Danyang Zhuo
Abstract
Over the last decade, deep neural networks have transformed our society, and they are already widely applied various machine learning applications. State- of-the-art deep neural networks are becoming larger in size every year to deliver increasing model accuracy, and result, model training consumes substantial computing resources and will only consume more in the future. Using current training methods, in each iteration, to process data point x Rd in layer, we need to spend (mmd) time to evaluate all the neurons in the layer. This means processing the entire layer takes (nnmd) time for n data points. Recent work [Song, Yang and Zhang, NeurlPS 2021] reduces this time per iteration to o(nmd) but requires exponential time to preprocess either the data or the neural network weights, making it unlikely to have practical usage.
In this work, we present new preprocessing method that simply stores the weight- data correlation in a tree data structure in order to quickly, and dynamically detect which neurons fire each iteration. Our method requires only ((nmd) time in preprocessing and still achieves o(uvdd) time per iteration. We complement our new algorithm with lower bound, proving that assuming popular conjecture from complexity theory, one could not substantially speed our algorithm for dynamic detection of firing neurons.
1 Introduction
Machine learning applications are requiring larger and larger neural network size, and the computing resources required to train these large models is growing correspondingly. Determining how train these models quickly has become an important research challenge.
Training neural network is an iterative algorithm, and in each iteration, need to process each of the m neurons on each of the n data points. Assuming each data point has a length (e.g., could be the size of input image), this means the per-iteration training time of the straightforward algorithm is at least ((nmd) just to compute the activations. As we train larger neural networks on more training data, this running time can become significant obstacle.
Recent work by Song, Yang, and Zhang [SY721] gave the first training algorithm that reduces this per iteration training time to o(ndd). The high-level idea of their algorithm is to use nearest neighbor search data structure that stores the neural network weights and training data. This allows the training method to have fast access to the inner products of the training data with the current
cossicc..columiaa......oo jiiehao wwwignobeeeeeee...ooooo... edu. University California, Berkeley. 'zsong@adobe com. Adobe Research.
cow.mbia......ocoooo..o......... edu. Columbia University.
Sruizhelutexas.com edu. Simons Institute the Theory Computing. danyangecs. duke. edu. Duke University.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

weight of the iteration. However, their algorithm's 's initial preprocessing time to set up the data structure is exponential in the dimension d, making too slow in most applications. This raises natural theoretical question:
Is it possible to design an algorithm that spends polynomial time to preprocess the weights and data, and which achieves a training time of o(pmd) per iteration?
This question is important for two reasons. First, speeding up neural network training is a fundamental research challenge with real-world value. Second, dynamic data structures have been successfully used to speed up computations in many contexts throughout computer science, yet their power and limitations when applied the training of neural networks are currently poorly understood.
1.1 Our Result: An Upper Bound
Our main result answers this question in the affirmative, giving a new algorithm with efficient preprocessing and faster training in the natural wee-ppraameterization regime (which has m >> n):
Theorem 1.1 (Main result). There is a data structure which preprocesses data points in d- dimensional space, and m initialization weights points in d-dimensional space, in (mmd) prepro- cessing time and Ommn + md + nd) space, which can be used to speed up neural network training: Running the gradient descent algorithm on two-layer, m-width, over-paramin aver-parameterized ReLU neural network, which will minimize the training loss to zero, can be performed with expected running time (of the gradient descent algorithm per iteration) of
Õmm 2/5 nnn2
The following remark gives a comparison between our result and a closely related work [SYZ21]:
Remark 1.2. The prior work [SYZ21] presented two algorithms. Their first result (see Theorem 6. I and Part of Corollary B.6 in [SY/21]] has O(2 preprocessing time and uses O(m ((m2-1/dnd) cost per iteration. Their second result (see Theorem 6.2 and Part 2 of Corollary B.6 of [SYZ21]) has O(r²)) preprocessing time and uses 0mm4//md time per iteration. Our result exponentially improves the running time of the data structure in [SYZ21] in terms of the dimension d. Notably, unlike [SYZ21], J, we do not use any complicated geometric data structure in previous work, and our algorithms are much easier to implement (see Algorithms I and 2). Moreover, as we discussed in Section 5, they can be parallelized to further reduce the cost-per-iteration to (md4
Our key observation is that in each iteration of the training process, the weight updates are mostly sparse, and only small fraction of neurons are activated for each training data point. Given this observation, we construct binary search tree for each training data point (or neuron) to detect which neurons will fire. Our data structure and the corresponding algorithms are deterministic, not relying on any randomness, and solve the following dynamic algorithms problem which we prove appears as a key subroutine of the training process.
Definition 1.3 (Dynamic Detection of Firing Neurons (DDFN)). Given two set of points X (z1...., ,n,) C Zd, Y = {yi....,, ..m} C Zd and a threshold b E IR design a data structure to support the following operations:
UUDATTE[ E [m], z 2d), set yj to z
QUERY(), either output the set
Q = (ii.) E [n] x [m] (2,, 8/) 2 6}.
or report that |Q] > m 4/5
We give data structure for DDFN which takes ((mmd)-tim for preprocessing, -(nd)-time per -time per query. At high level, our data structure works as follows.
update, and (mmin (Q), m/45
We build n binary search trees to maintain the (æ., w;) pairs for i E [m] and j E [m]]. More specifically, the i-th tree has leaf nodes, storing the inner-products between u; and (リ, jemm]- Fach internal node stores the larger value of its two child nodes. The preprocessing time for the binary search trees for all the input data and neurons takes ((ndd) time and ((m)) space.
Preprocessing

---

Update Suppose we will update yy- Then, for each i E [n], we need to update a path from the leaf node corresponding to (z,, yi) to the root, which contains (log m) nodes, Hence, the running time of each update is (nd Ogdlog m) m).
Query We need to find all the leaf nodes with values greater than b. We can traverse each tree from top to bottom. At each node, if its value is most b, we will not move further. Otherwise, will try to search each of its child nodes. Note that the number of visited nodes is of the same order as the number of visited leaf nodes. And we visit leaf if and only if its value is greater than b. Hence, the total query cost is O(min [Q), m m4/5
1.2 Our Result: A Lower Bound
We complement our new algorithm with a lower bound, showing that assuming a popular conjecture from complexity theory, one could not improve much on running time for Dynamic Detection of Firing Neurons (DDFN). Prior work [SY721] got around this by using exponential preprocessing time avoid needing a dynamic algorithm. However, in our setting with polynomial preprocessing and running times, there is a limit to how quickly one can perform each iteration:
Theorem 1.4 (Lower Bound for DDFN, informal version of Theorem F.4.. Let d = 20 "), and assume the OVC or SETH conjecture from complexity theory. For every constant E > 0, there is no data structure for DDFN with O(m²/' //5--/)))))))) update time and (m query time.
Here, log" n denotes the iterated logarithm function, which grows incredibly slowly, such that the dimension 20 n) is barely larger than a constant, and one would typically pick a much larger d. The complexity-theoretie assumptions OVC and SETH are defined in Section F. We prove Theorem 1.4 by reducing the Maximum Inner Product Search problem to DDFN.
In other words, our Theorem 1.4 shows that without using large preprocessing time as in the prior work, it is impossible to substantially improve on our algorithm, no matter how sophisticated the algorithmic techniques one might use.
1.3 Related Work
Orthogonal Vector Conjecture The orthogonal vector problem (OV) is a fundamental problem in fine-grained complexity which asks, given X, Y C {0, 1]d of size |X] |Y| n, whether there are x E X and y E Y with (z. v) = 0. The state-of-the-art algorithm 14, CW16] runs in time 2221//((oge) c) in dimension d = c log n for all e 2 1; as the dimension increases, its running time approaches the trivial bound n²2 The orthogonal vector conjecture (OVC) conjectures an n 2-1o(1) lower bound for OV when d = (log n). It is also known that the popular Strong Exponential Time Hypothesis (SETH) regarding the hardness k-SAT implies OVC. This conjecture has been used to obtain conditional lower bounds for other important problems with polynomial-time algorithms in wide variety of areas, including pattern matching [AWW14, Bri14, B115, B116, BM16, BGL17, BK18, CW19], kernel sparsification [ACSS20, AA22], softmax attention computation [AS23a, AS23b, DS723], graph algorithms [RVW13, ABH 18, GIKW18, KT18, DLW22, CWX22), and computational geometry [BBK+ 16, Rub18, Will8a, Che20, KM20]; see also the survey [Will8b].
Acceleration via high-dimensional search data-structure Data structures have been designed that allow one to quickly find high-dimensional points in geometric query regions (e.g.. half-spaces, simplices, etc). Currently, there are two main approaches to designing these structures. One is based on Locality Sensitive Hashing (LSH) [IM98], which aims to find nearby points (c.g.. small &g distance [DIIM04, AR15, AIL 15, ARNI7, Raz17, AIR18. BIW19. DIRW20] or large inner product [SL14, SL15b, SL15a]) to a query E Rdd in a given set of points C Rdd LSH-based algorithms typically quickly in practice, but only support approximate nearest-neighhor queries. The other approach is based on space partitioning data structures, such as partition trees [Mat92a, Mat92b, AEM92, AC09, Chal2], k-d trees/range trees [CT17, TOG17, Cha and Voronoi diagrams [ADBMS98, Cha00], which can exactly search for points in the queried region.
There is a recent line of research that has applied high-dimensional geometric data structures to reduce deep neural networks" training time in practice. Empirically, SLIDE (CMF+20) uses -SSed methods to efficiently find neurons with maximum inner product in the forward pass; Reformer
3

---

[KKL20] also applies LSH to save the space complexity such that the neural networks are able to process very long sequences; MONGOOSE [C]P+21] combines the learnable LSH-based data structure [Cha02] with the lazy update framework [CISS19] to speedup the neural network training. Theoretically, [SY721] gives the first provable sublinear-time training algorithm for 2-layer over- parameterized neural networks using the HSR data structures in [AEM92].
The goal of our paper is design an efficient high-dimensional geometric data structure that can be embedded in the neural network training framework with provable performance guarantees. Specifically, our data structures have the same functionalities as the HSR data structures [AEM92] which can find all the points that have large inner products and support efficient data updates. However, our data structures do not have an exponential dependence on d, the dimension of the points, which appears in the time complexities of many computational geometry algorithms (including [AEM92]] due to the curse of dimensionality. Compared with the LSH-based approaches, our data structures have a stronger correctness guarantee, and will always report all the points with a sufficiently large inner product; LSH only gives approximate guarantees and might miss some of them.
Convergence via over por--aaramttrizationn are much larger than the number of training data points (i.e., m >> n). very natural and common regime in deep learning. It plays key role in explaining why deep neural networks can perform so well at many different tasks. Over the last few years, there has been a tremendous amount of work toward theoretically understanding the convergence and generalization of deep neural net- works in the over-parameterization regime, e.g., [LL.18, DZPS19, AZLS19a, AZLS119,, ADH+19; ADH 19b, SY19, CGH+ 19, ZMG19, CG19, ZG19, OS20, JT20, LSS+ HLSY21, ZPD+220, BPSW21, SZZ21, Zha22, HSWZ22, MOSW22, Y3,+ GMS23, LSY23, QSS23, QSY23J. A key observation is that when the width (m) of the neural network tends infinity, the neural network is equivalent to a neural tangent kernel (NTK) [JGHI8], and so technical tools from kernel methods can be adapted to analyze deep neural networks. In particular, it has been shown that (stochastic) gradient descent (SS)DD can train a sufficiently wide neural network with random initialization, converging to small training error in polynomial number of steps.
Over-paramettrization where the trainable parameters
This paper is organized as follows: Section 2, we formulate our problem of training neural networks. In Section 3, we develop the Correlation Tree data structure, which is the main contribution of this work. In Section 4, we state our main result for quickly training neural networks using Correlation Trees. In Section 5, we conclude and discuss some future directions. A number of our proofs are deferred to the appendix.
Roadmap
Preliminaries
Before describing our new data structure, we first present the notation we will use, and formulate the problem setting.
Basic Notation. For N., we use [n] to denote the set f1,2, ... n.] We write E[X] to denote the expected value of a random variable X, and Pr[Y] to denote the probability of a random event Y. For matrix M, write M to denote the transpose of M. We use x y to denote the inner product between vectors x and y. We use Ia to denote the d x d identity matrix. We use N ((u; o?) to denote the Gaussian distribution with mean μ and variance o².
Problem Formulation In this section, we introduce the neural network model we study in this paper. We consider two-layer ReLU-activated neural network f that has width and uses lz loss function.
Definition 2.1 (Prediction function and loss function). For threshold parameter b E R. data point I E Rdd weight matrix W E ddxm, and weights a E R.,, the prediction function functiio f(W, x, a) and loss function L(W) are given by
f(W, x) =
1
Σ
o ((w,
1 L(W) - (F(W, y.)² -2))
mm
where 0ι(xx) = max(x b, 0} the ReLU function with threshold parameter b. Following the prior work [SYZ21], we write 2N, b) to denote this function f for simplicity.

---

Remark 2.2. In the neural network, Wi,".. Wm E R are the weight vectors of the edges between input nodes and hidden neuron nodes, and 01,-. O. E P. are the weights of the edges connecting hidden neuron nodes with the output node.
Following the setups in previous work, we only train the weight parameters W Rdxm minimize the loss L(W), and will leave E IR unchanged after the initialization.
In this work, study the following training process:
Initialization:
For each hidden neuron, we sample w,,(0) ~ NO, Ta), and sample or from {-1,+1} uniformly random.
Gradient computation: For each neuron, we have
@f(,,,, Or Ow, Vm arllpa2d and
OL(W) ar Σ (f(W,
v:)x; (m)aa)2
Sw,
Weight update: We follow the standard update rule the GD algorithm from iteration to iteration & + 1:
W(k 1) ((k) (W)),
where W(k) denotes the weights at iteration k, and
(W))) = A (K))
OW(k)
Sparsity phenomenon in the training process As observed in many experimental and theoretical works [CLP*21, SY721, $7721, GS22, GSZ23, DCL+ DMS23, GSYZ23, WWD+23, GSY23, 23,+23, SWYZ23, BSZ23], for a randomly initialized over- wver-paameeterized neural network (Defi- nition 2.1), given input data I, only small fraction (o(m) of the neurons will be activated when evaluating 2NN(m, b) (Definition 2.1) in each training iteration. We refer to this phenomenon as "Sparsity".. We exploit this property to design efficient data structure to identify the sparse activated neurons, achieving sublinear training time terms of m, number of neurons the hidden layer.
To be specific, the sparsity of activated neurons during the training process is bounded by choosing proper threshold for the ReLU function. Because concentration phenomenon of the randomized initialization, we can upper-bound the number of activated neurons just after initialization, which we refer to "sparsity after initialization". Then, in subsequent training iterations, using the neural tangent kernel (NTK) property, it follows that there is only a minor increase in the number of activated neurons per iteration. Therefore, the total number of activated neurons can be bounded by small quantity.
For simplicity, define the "fire set" (tre(r)) first, which is the of neurons that is activated when the neural network's 's input is x.
Definition 2.3 (Fire set). Let the neural network be defined as in Definition 2.1. For a data point Rd. let Sace (z) denote the of neurons that are activated input x, i.e.,
(are(x)) = {i E [m] 0b((x, w.)) > 0}.
Then, we similarly define fire sets for hidden neurons and input data points for each iteration:
Definition 2.4 (Fire set per iteration). For euch data point x; E IRd with E [n] and each iteration t {0, .T),, 3, let wt(() Edd be the weight vector of the r-th neuron the t-th iteration for
r E [m]. Define
St, (t) = E [m] (nl(,, 2. (t))) > 0).
3 (t) = {i E [n]] or W. (t))) 0}.
We fürther denote the sizes of these sets by kit = (t)) and ky,t = (r.aae(t))
5

---

The following lemma upper bounds the sparsity after initialization.
Lemma 2.5 (Sparsity after initialization, informal version of Lemma B.3, [SYZ21]. Let > 0 be tunable parameter. If we setup the neural network as Definition 2.1. then after the randomized initialization, with probability at least 1 - exp(-52(m) exp(-b² (2)()bb))))..... it holds that for any input data æ, the number of activated neurons is at most (m -xp( (2))) where m is the total number of neurons.
Remark 2.6. This suggests that take b v0.4 log m, we achieve sublinear number, Omm (m3) of activated neurons.
We can similarly control the sparsity in each iteration, and not just the first iteration; we defer the details to Section B.2.
In the next section, we will show how our weight-tree correlation data structure can take advantage of this sparsity phenomenon.
3 Correlation Tree Data Structure
In this section, we consider neural network 2NN(m, b) (Definition 2.1) with data points. We let (Wi,-.. Wm3 C Fd be the weights, {Æ1,-.. ,' C IRd be the data points, and [(wr, C. gm+" be the weight-data pairs.
We propose two data structures: Correlation DTTree and Correlation WTrre. The DTTree data structure has trees, and its i-th tree has m leaf nodes corresponding to the set of inner-products between x; and all hidden neurons, i.e., { [we, 2z) fre[m]. Similarly, the WTree data structure consists of m trees, and its r-th tree has leaf nodes corresponding the set of inner-products between the r-th neuron and all data points, i.e., { {W,, ar) hejm]-
The Correlation Tree is simple binary tree data structure. At high level, it works as follows:
Tree construction We first calculate the inner products of all weight-data pairs (w:, each representing the evaluation of neuron a data point. To search activated neurons efficiently, we create tree structure in the following way (taking the Correlation DTree an example): we first build m leaf nodes, where the r-th leaf stores (wy, xx)) for E [m]. Then, recursively construct binary tree such that parent node takes the larger value from two child nodes, Finally, we obtain a tree with root having value max,s im] { (wy, xx) Moreover, the value of each internal node equals to the maximum value among the leaf nodes in this subtree.
Efficient search Given threshold b, the data structure find all the pairs of vectors whose inner product is greater than b. Take the Correlation DTTre as an example. It outputs the indices of those activated neurons (i.e., (w,,,,) b) in recursive way: starting from the root, it checks whether it is "activated" (i.e., with value > b). If not, the search ends. Otherwise, moves each of the child nodes and repeats this searching process until stops. This is typical depth-first search strategy. Its running time is determined by how many nodes it visits during searching. The number of visited nodes has the same magnitude the number of visited leaf nodes, i.e., the number of activated neurons. Hence, the efficiency of our data structures relies on the sparsity phenomenon of the training process.
Relation between DTree and WTree In the Correlation DTree, each weight vector appears only different trees. In the Correlation WTree, each weight vector w, appears only in one of the m trees. When is updated, D'Tree will change the nodes along root-to-leat path in n trees, whereas WTree only changes such paths the r-th tree.
3.1 Correlation DTree data structure
We now state our main theorem summarizing the correlation DTtrree data structure. Its pseudocode is given Algorithms 1 and 2 below. Its proof are deferred to Section D.I.
Theorem 3.1 (Correlation DTTree data structure). There exists data structure with the following procedures:

---

INIT([W1, W2. um C R f21, ,In C E E N,d N). Given series of weights w1, u2.-.- Wm and data 21, 21,22, In in d-dimensional space, performs preprocessing in time (
UPDATE(2 E R2, r E [m]). Given a weight 2 and an index r, updates weight Wy to 2 in time (n (d log m.).
QUERY(i E |n],T E R). Given an index indicating data point æi and a threshold T, finds all indices r E [mi] such that (w,, mi) > T in time O(1((T) log m), where
S(7) = fr: (wi, x2) > r}.
Algorithm Correlation DTree data structure
data structure CORRELATIOND
Theorem 3.1
2: members
3: W E Rm (ra weight vectors
4: X E IRaxd (n data points)
5: Binary tree T1, T2. Tn
binary search trees
6: end members
7: procedure INIT(W1, W2: ... wm E Rd m, 21, 22, :. In 2
E Rd n, m, d)
Lemma D.2
8: for i 1 do
9: 2; 6 I;;
10: end for
11: for j = 1 - m do
12:
W W
13:
end for
14: for
do
for data point, create tree
15:
for j =1 + m do
Uj + (2i, wi)
end for Ti < MAKEMAXTREE(U) um two children, Algorithm 7 end for 20: end procedure
17:
Each node stores the maximum value for his
procedure UPDATE(2 IRd, r [m])
Lemma D.3
22:
<-
for 1 do
24:
1 + the 1-th leaf of tree T;
25:
Lvvauee (z,x;)
26:
while / is not root do
27:
p f parent o
28:
valuu maxx(p.vluee I. value)
29:
1 t-p p
30:
end while
31: end for
32:
end procedure
33: end data structure
3.2 Correlation WTrre data structure
We next state the main theorem summarizing our similar Correlation WTree data structure. Both the Correlation DTree and Correlation WTree have query time that roughly equal to the output size, but since they have different outputs, each can be faster than the other depending on the setting. The pseudocode and proof for Correlation WTrre are deferred to Section D.3.
Theorem 3.2 (Correlation WTrre data structure). There exists a data structure with the following procedures:
7

---

Algorithm 2 Correlation DTrees
1: data structure DORRELATIONDTREE
Theorem 3,
2: procedure QUERY(i E [n],T R2o)
Lemma D.4
3: return FIND(T, (7))
4: end procedure
5: procedure FIND(T R20, 6 T)
6: if' r is leaf then
7: return r
8: else
9:
< left child of r, 12 < right child of
if valuue 2 T then
11:
S FIND((7,))))
12:
end
13:
if r2. value 2 r then
S2 FiND((r r2)
15: end if
16: end
17: return S₁ U.2
18:
end procedure
19:
end data structure
INIT((W1, W2. Wmm C Rd f21, series of weights wi, w2, Wm and data "2,"""" Fn. in d-dimensional space, performs preprocessing in time O(nmd).
... Fnn C ERRn E N, N). Given
UPDATE(2 E Rd, E [m]). Given weight z and index it updates weight wy z in time (
QUERY(r E [m]. R). Given an index indicating weight wr and a threshold it finds all indices E [n] such that (w,, 24) T in time O([S(T) logmm) m), where S()) = (w,,,;) T).
4 Running Time of Our Algorithm
In this section, we show how apply the Correlation Tree data structures developed in Section to speed up neural network training.
4.1 Weights Preprocessing
Algorithm 3 Training Neural Network based on Correlation DTrre
1: procedure TRAINING WITHDTREE((( (It; 31)) mejnj.nmm.d
Theorem 4.1
2:
Initialize Wy. a, for r E [m] and according to Section 2 (1 for tor T do
3:
m.d)
Algorithm 10
4:
5:
Si DTREE. QUBRY for E [n]]
6:
Forward pass for æ; only on neurons in Si.ir for i E [n]] Calculate gradient for I; only on neurons in St for E [n]
7:
8: Gradient update for the neurons in St.rre
9:
D'REE.UUPDTER (t 1),r)
10:
end for
11: return Trained weights w, (T + 1) for r E [m]
end procedure
In Algorithm 8, we use DTree structure to speed up the training process. We preprocess weights Wr, E [m]] for each data point H,, e [n]] by constructing weight-data correlation trees. In each

---

iteration, QUERY finds the set of activated neurons St,r (Definition 2.4) efficiently for each data point 2, and UPDATE helps change the weights in backward propagation.
Our main result for weight preprocessing is as follows.
Theorem 4.1 (Running time part, informal version of Theorem E.I). Given a data points in TR: gradient descent using the DTree data structure (Algorithm 8) for the neural network 2NN b 1o m. (Definition 2.1) takes O (m=/6n2dd) time per iteration in expectation.
4.2 Data Preprocessing
Algorithm 4 Training Neural Network based on Correlation WTree
1:
procedure TRAINING WATWWTER (æ;, y/ (e() n.m,d))
> Theorem 4.2
2: Initialize wy. a, for r E [m]] and according to Section 2
3: WTREE.INIT( (e(n d)
Algorithm 13 b. Data points fire set Hidden neurons fire set
4: Sp.fire +- WT. (VT.QUER((W<(O), b) for E [m]]
5: Si,tre (r 1 Sr,
6:
for 1 do
7:
Forward pass for only neurons in Sy,ir for i [n] Calculate gradient for x; only on neurons in Si. for E [n]]
8:
9:
for r Usejn] do
10:
Sy.ire + WTREE. QUERY 1),6)
11:
Update Si,fre for each E Sr,tr
12: end for
end for
14: return Trained weights wo (T' + 1) for r E [m]
end procedure
Preprocessing weights based on data points is common practice for neural networks. Here consider its dual form: preprocessing input data I, E [m] based on neural network weights Wr,r E [m)]. This can be easily done due the symmetric property of the inner product that we used in the correlation tree structure.
Given a weight vector wy, we can quickly find Strr (Definition 2.4) which contains the indices of data points that "fire" for weight w,. By the dual relationship between Sr and S.,rr we can recover St,r easily.
One advantage of the data preprocessing approach is that the data structure only depends on the training dataset, instead of the neural network architecture. Therefore, the data structure could be pre-computed and stored in cloud platforms.
The performance guarantee of our data preprocessing training algorithm is shown as follows:
Theorem 4.2 (Running time part, informal version of Theorem E.2). Given data points in TR. gradient descent algorithm using the WTree data structure (Algorithm 9) for the neural network 2N b log m.) takes (m log n) -time per iteration to initialize Sy,rre Ste for r E [m], E [n). and the total running time per iteration is
0m5m
in expectation.
Conclusion
Deep neural networks are becoming larger every year offer improved model accuracy. Training these models consumes substantial resources, and resource consumption will only increase as these models grow. In traditional training methods, for each iteration, we need to spend ((mmd) time to evaluate the neurons on data points with dimension d. Recent work [SYZ21] reduced the per-iteration cost to onnmd), but required exponential time to preprocess either the data or the neural weights. We develop new method that reduces the preprocessing cost to O(rmd) while keeping the

---

per-iteration running time at o(nmd). One limitation of our algorithm is that it has an n²² dependence in the cost-per-ierration. However, for very wide neural networks (with m > n), the runtime is still sublinear. More importantly, we design a simple binary tree-based dynamic geometric data structure that can efficiently identify all the activated neurons in each training iteration and bypass the high-dimensional barrier of the prior approach. We further remark that the Update procedure of DTree/Wtree structure (Algorithm and 14) can be parallelized, where we can update all the correlation trees using distributed computing simultaneously. It will improve the running time from ((nd) to (d), resulting in total running time Omm*// per iteration.
Our work naturally raises some open questions for future study:
First, can we apply our data structure, together with an analysis of the sparsity in training over-parameterized neural networks [SYZ21, SZZ21J, to speed up neural network training with more than two layers? Giving a provable, theoretical backing for quickly training multi-layer networks remains open, difficult challenge.
Second, many empirical results (e.g., [CMF+ 20, 2LP+21J) indicate that only approximately identifying the activated neurons (i.e., neurons with top-k inner products) in each itera- tion may still be enough to train a neural network. Can we provide a more theoretical understanding of these approaches?
Third, our current algorithms use more memory (i.e., O(mn) space) to store the correlation tree data structure. Is it possible to reduce the space complexity of the algorithms?
Fourth, we think it is possible that data structures will work for more general activation functions. Roughly speaking, as long as the activated neurons are sparse or approximately sparse, data structures will be able theoretically reduce the postiter-itrration. However, we need to re-prove the sparsification results [SY721] for the activation function other than ReLU,
The authors would like to thank Lichen Zhang for his helpful discussions. JA was partly supported by grant from the Simons Foundation (Grant Number 825870 JA). Part of this research was performed while RZ was visiting the Institute for Pure and Applied Mathematics (IPAM), which is supported by the National Science Foundation (Grant No.
Acknowledgements
References
[AA22] Amol Aggarwal and Josh Alman. Optimal-degree polynomial approximations for exponentials and gaussian kernel density estimation. In 37th Computational Complexity Conference (CCC 2022). Schloss Pogstuhl-Liinii--eeettu für Informatik, 2022.
[ABH+18] 18] Amir Abboud, Arturs Backurs, Thomas Dueholm Hansen, Virginia Vas- silevska Williams, and Or Zamir. Subtree isomorphism revisited. ACM Transactions on Algorithms (TALG), 14(3):1-23, 2018.
[ABW15] Amir Abboud, Arturs Backurs, and Virginia Vassilevska Williams. Tight hardness results for Ics and other sequence similarity measures. In 2015 IEEE 56th Annual Symposium on Foundations of Computer Science, pages 59-78. IEEE, 2015.
[AC09] Peyman Afshani and Timothy M Chan. Optimal halfspace range reporting in three dimensions. In Proceedings of the twentieth annual ACM-SIAM symposium on Discrete algorithms, pages 180-186. SIAM, 2009.
[ACSS20] Josh Alman, Timothy Chu, Aaron Schild, and Zhao Song, Algorithms and hardness for linear algebra on geometric graphs. In 2020 IEEE 61 Annual Symposium on Foundations of Computer Science (FOCS), pages 541-552, IEEE, 2020.
[ADBMS98] Pankaj K. Agarwal, Mark De Berg, Jiri Matousek, and Otfried Schwarzkopf. Con- structing levels in arrangements and higher order voronoi diagrams. SIAM journal on computing, 27(3):654 667, 1998.
[ADH+19a] Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks. In International Conference on Machine Learning, pages 322-332, 2019.
10

---

[ADH* 19b] Sanjeev Arora, Simon S Du, Wei Hu, Zhiyuan Li, Ruslan Salakhutdinov, and Ruosong Wang. On exact computation with an infinitely wide neural net. In NeurlPS, 2019.
[AEM92] Pankaj Agarwal, David Eppstein, and Jiri Matousek. Dynamic half-space reporting, geometric optimization, and minimum spanning trees. In Annual Symposium on Foundations of Computer Science (FOCS), volume 33, pages 80-80, 1992.
[5]]+15] Alexandr Andoni, Piotr Indyk, TMM Laarhoven, Ilya Razenshteyn, and Ludwig Schmidt. Practical and optimal Ish for angular distance. In Advances in Neural Information Processing Systems (NIPS), pages 1225-1233, Curran Associates, 2015.
[AIR18] Alexandr Andoni, Piotr Indyk, and Ilya Razenshteyn. Approximate nearest neighbor search in high dimensions. arXiv preprint /vxiv:/806.09822 7, 2018.
[AR15] Alexandr Andoni and Ilya Razenshteyn. Optimal data-dependent hashing for approxi- mate near neighbors. In Proceedings of the serty-seventh annual ACM symposium Theory of computing (STOC), pages 793-801, 2015.
[ARN17] Alexandr Andoni, Ily Razenshteyn, and Negev Shekel Nosatzki. Lsh forest: Practical algorithms made theoretical. In Proceedings of the Twenty-Fighth Annual ACM-SIAM Symposium on Discrete Algorithms (SODA), pages 67-78. SIAM, 2017.
[AS23a] Josh Alman and Zhao Song. Fast attention requires bounded entries. In NeurIPS. arXiv preprint arXiv:2302.13214, 2023.
[AS23b] Josh Alman and Zhao Song. How to capture higher-order correlations? generalizing matrix softmax attention to kronecker computation. arXiv preprint arXiv:2310.04064, 2023.
[AWW14] Amir Abboud, Virginia Vassilevska Williams, and Oren Weimann. Consequences of faster alignment of sequences. In International Colloquium Automata, Languages, and Programming, pages 39-51. Springer, 2014.
[AWY14] Amir Abboud, Ryan Williams, and Huacheng Yu. More applications of the polynomial method to algorithm design. In Proceedings of the weenty-sixth annual ACM-SIAM symposium on Discrete algorithms, pages 218-230. SIAM, 2014.
[AZLS19a] Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via over-parameteriizaion. ICML. 2019.
[AZLS19b] Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. On the convergence rate of training recurrent neural networks. In NeurlPS, 2019.
[BBK 16] Kevin Buchin, Maike Buchin, Maximilian Konzack, Wolfgang Mulzer, and André Schulz. Fine-grained analysis of problems on curves. EuroCG, Lugano, Switzerland, 2016.
[Be224] Sergei Bernstein. On a modification of chebyshev's inequality and of the error formula of laplace. Ann. Sci. Inst. Sav. Ukraine, Sect. Math, 144::38 49, 1924.
[BGLI7] Karl Bringmann, Allan Grønlund, and Kasper Green Larsen. A dichotomy for regular expression membership testing. In 2017 IEEE 58th Annual Symposium on Foundations of Computer Science (FOCS), pages 307-318. IEEE, 2017.
[B115] Arturs Backurs and Piotr Indyk. Edit distance cannot be computed in strongly sub- quadratic time (unless seth is false). In Proceedings of the forry-seventh annual ACM symposium Theory of computing, pages 51-58, 2015.
[BI16] Arturs Backurs and Piotr Indyk. Which regular expression patterns are hard match? In 2016 IEEE 57th Annual Symposium on Foundations of Computer Science (FOCS), pages 457-466, IEEE, 2016.
[BIW19] Arturs Backurs, Piotr Indyk, and Tal Wagner. Space and time efficient kernel density estimation in high dimensions. In NewrIPS, pages 15773-15782, 2019.
11

---

[BK18] Karl Bringman and Marvin Künnemann. Multivariate fine- fine-grained complexity of longest common subsequence. In Proceedings of the Twenty-Ning -Nith Annual ACM-SIAM Symposium on Discrete Algorithms, pages 1216-1235. SIAM, 2.18.
[BM16] Karl Bringmann and Wolfgang Mulzer. Approximability of the discrete fréchet distance. Journal of Computational Geometry, 7(2):46-76, 2016.
[BPSW21] Jan van den Brand, Binghui Peng, Zhao Song, and Omri Weinstein. Training (over- parametrized) neural networks in near-linear time. In /2th Innovations Theoretical Computer Science Conference (ITCS), 2021.
[Bri14] Karl Bringmann. Why walking the dog takes time: Frechet distance has no strongly subquadratic algorithms unless seth fails. In 2014 IEEE 55th Annual Symposium Foundations of Computer Science, pages 661-670. IEEE, 2014.
[BSZ23] Jan van den Brand, Zhao Song, and Tianyi Zhou. Algorithm and hardness for dynamic attention maintenance in large language models. arXiv preprint arXiv:2304.02200 2023.
[CG19] Yuan Cao and Quanquan Gu. Generalization bounds of stochastic gradient descent for wide and deep neural networks. In NewrIPS, pages 10835-10845, 2019.
[CGH+ 19] Tianle Cai, Ruiqi Gao, Jikai Hou, Siyu Chen, Dong Wang. Di He, Zhihua Zhang, and Liwei Wang. Gram Gram-gauss gauss-newonn method: Learning overparameterized neural networks for regression problems. arXiv preprint 1r0iv 1675, 2019.
[Cha00] Timothy M Chan. Random sampling, halfspace range reporting, and construction of ( ky-levels three dimensions. SIAM Journal on Computing, 30(2):561-575, 2000.
[Cha02] Moses S Charikar. Similarity estimation techniques from rounding algorithms. In Proceedings of the thiry-fourth annual ACM symposium on Theory of computing (STOC), pages 380-388, 2002.
[Cha12] Timothy M Chan. Optimal partition trees. Discrete & Computational Geometry, 47(4):661-690, 2012.
[Cha19] Timothy M Chan. Orthogonal range searching in moderate dimensions: kd trees and range trees strike back. Discrete & Computational Geometry, 64488999,,, 2019.
[Che20] Lijie Chen. On the hardness approximate and exact (bichromatic) maximum inner product. Theory OF Computing, 16(4):1-50, 2020.
[CIP09] Chris Calabro, Russell Impagliazzo, and Ramamohan Paturi. The complexity of satisfiability of small depth circuits. In International Workshop on Parameterized and Exact Computation, pages 75-85. Springer, 2009.
[CLP+21] Beidi Chen, Zichang Liu, Binghui Peng. Zhaozhuo Xu, Jonathan Lingjie Li, Tri Dao, Zhao Song, Anshumali Shrivastava, and Christopher Re. Mongoose: A learnable Ish framework for efficient neural network training. In ICLR oral, 2021.
[CLS19] Michael B Cohen, Yin Tat Lee, and Zhao Song. Solving linear programs in the current matrix multiplication time. In STOC, 2019.
[CMF+ 20] Beidi Chen, Tharun Medini, James Farwell, Sameh Gobriel, Charlie Tai, and Anshu- mali Shrivastava. Slide: In defense of smart algorithms over hardware acceleration for large-scale deep learning systems. In In Proceedings of the 3rd Conference on Machine Learning and Systems (MLSys), 2020.
[CT17] Timothy M Chan and Konstantinos Tsakalidis. Dynamic orthogonal range searching on the ram, revisited. Leibniz International Proceedings in Informatics, LJPlcs, 77:281- 13, 2017.
12

---

[CW16] Timothy M Chan and Ryan Williams. Deterministic apsp. orthogonal vectors, and more: Quickly derandomizing razborov-smolensky. In Proceedings of the rwenty- seventh annual ACM-SIAM symposium on Discrete algorithms, pages 1246-1255. SIAM, 2016.
[CW19] Lijie Chen and Ryan Williams. An equivalence class for orthogonal vectors. In Proceedings of the Thirtieth Annual ACM-SIAM Symposium on Discrete Algorithms, pages 21-40. SIAM, 2019.
[CWX22] Timothy Chan, Virginia Vassilevska Williams, and Yinzhan Xu. Hardness for triangle problems under even more believable hypotheses: Reductions from real apsp. real 3sum, and ov. arXiv preprint 22X3::2003008355 2022.
[DCL+22] Tri Dao, Beidi Chen, Kaizhao Liang. Jiaming Yang, Zhao Song, Atri Rudra, and Christopher Re. Pixelated butterfly: Simple and efficient sparse training for neural network models. In ICLR. arXiv preprint arXiv: 122.009,,, 2022.
[DIM044] Mayur Datar, Nicole Immorlica, Piotr Indyk, and Vahab S Mirrokni. Locality-sensitive hashing scheme based on p-stable distributions. In Proceedings of the twentieth annual symposium on Computational geometry (SoCG), pages 253-262, 2004.
[DIRW20] Yihe Dong, Piotr Indyk, Ilya Razenshteyn, and Tal Wagner. Learning space partitions for nearest neighbor search. In ICLR. arXiv preprint
2020.
[DLW22] Mina Dalirrooyfard, Ray Li, and Virginia Vassilevska Williams. Hardness of approxi- mate diameter: Now for undirected graphs. In 202/ IEEE 62nd Annual Symposium on Foundations of Computer Science (FOCS), pages 1021-1032 IEEE, 2022.
[DMS23] Yichuan Deng, Sridhar Mahadevan, and Zhao Song. Randomized and deterministic attention sparsification algorithms for over-parameterized feature dimension. arXiv preprint arxiv:2304.04399) 2023.
[DSZ23] Yichuan Deng, Zhao Song, and Tianyi Zhou. Superiority of softmax: Unveiling the performance edge over linear attention. arXiv preprint arXiv:2310. 1685, 2023.
[DZPS19] Simon S Du, Xiyu Zhai, Barnabas Poezos, and Aarti Singh. Gradient descent provably optimizes wee-pparaetterized neural networks. In ICLR, 2019.
[GIKW18] Jiawei Gao, Russell Impagliazzo, Antonina Kolokolova, and Ryan Williams. Com- pleteness for first-order properties sparse structures with algorithmic applications. ACM Transactions on Algorithms (TALG), 15(2):1-35, 2018.
[GMS23] Yeqi Gao,, Sridhar Mahadevan, and Zhao Song. An over-paramettrrized exponential regression. arXiv preprint arXiv:2303. 16504, 2023.
[GS22] Yuzhou Gu and Zhao Song. A faster small treewidth sdp solver. arXiv preprint arXiv:2211.06033 2022.
[GSY23] Yeqi Gao, Zhao Song, and Junze Yin. Gradientcoin: A peer-to-peer decentralized large language models. arXiv preprint 0rXiv:2308.1050/ 2023.
[GSYZ23] Yeqi Gao, Zhao Song, Xin Yang, and Ruizhe Zhang. Fast quantum algorithm for attention computation. arXiv preprint arXiv:2307.0804. 08045, 2023.
[GSZ23] Yuzhou Gu. Zhao Song, and Lichen Zhang. A nearly-linear time algorithm for structured support vector machines. arXiv preprint arXiv:2307.07735, 2023.
[HI.SY2]] Baihe Huang, Xiaoxiao Li, Zhao Song, and Xin Yang. Fl-ntk: A neural tangent kernel-based framework for federated learning convergence analysis. In ICML, 2021.
[HSW722] Hang Hu, Zhao Song, Omri Weinstein, and Danyang Zhuo. Training overparametrized neural networks in sublinear time. arXiv preprint arXiv:2208.04455 2022.
13

---

[IM98] Piotr Indyk and Rajeev Motwani. Approximate nearest neighbors: towards removing the curse of dimensionality. In Proceedings of the thirtieth annual ACM symposium on Theory of computing (STOC), pages 604-613, 613, 1998.
[IPO1] Russell Impagliazzo and Ramamohan Paturi. On the complexity of k-sat. Journal of Computer and System Sciences, 62(2):367-377 2001.
[JGHI8] Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Con- vergence and generalization in neural networks. In Advances in neural information processing systems, pages 8571-8580, 2018.
[JT20] Ziwei Ji and Matus Telgarsky. Polylogarithmic width suffices for gradient descent to achieve arbitrarily small test error with shallow relu networks. In ICLR, 2020.
[KKI,20] Nikita Kitaev, łukasz Kaiser, and Anselm Levskaya. Reformer: The efficient trans- former. arXiv preprint wwXiv:2001.0445) 2020.
[KM20] CS Karthik and Pasin Manurangsi. On closest pair in euclidean metrie: Monochromatic is as hard as bichromatic. Combinatorica, 40(4):539-577 2020.
[KT18] Robert Krauthgamer and Ohad Trabelsi. Conditional lower bounds for all all-pairs max- flow. ACM Transactions on Algorithms (TALG), 14(4):1-15, 2018.
[LL18] Yuanzhi Li and Yingyu Liang. Learning overparameterized neural networks via stochastic gradient descent on structured data. In NewrlPS, 2018.
[LSO1] W.V. Li and Q.-M. Shao. Gaussian processes: Inequalities, small ball probabilities and applications. In Stochastic Processes: Theory and Methods, volume 19 of Handbook of Statistics, pages 533-597. Elsevier, 2001.
[LS Jason D Lee, Ruoqi Shen, Zhao Song, Mengdi Wang, and Zheng Yu. Generalized leverage score sampling for neural networks. In NeurlPS, 2020.
[LSY23] Xiaoxiao Li, Zhao Song, and Jiaming Yang. Federated adversarial learning: frame- work with convergence analysis. In International Conference on Machine Learning, pages 19932-19959, PMLR, 2023.
[LDD*23] Zichang Liu, Jue Wang. Tri Dao, Tianyi Zhou, Binhuang Yuan, Zhao Song, Anshumali Shrivastava, Yuandong Tian Ce Zhang, Christopher Re, and Beidi Chen. Deja vu: Contextual sparsity for efficient Ilms at inference time. In ICML, 2023.
[Mat92a] Jiri Matousek. Efficient partition trees. Discrete & Computational Geomelry, 8(3):315 334, 1992.
[Ma192b] Jiri Matousek. Reporting points halfspaces. Computational Geometry, 2(3):169-186 1992.
[MOSW22] Alexander Munteanu, Simon Omlor, Zhao Song, and David Woodruff. Bounding the width of neural networks via coupled initialization a worst case analysis. In International Conference Machine Learning (ICML), pages 16083-16122. PMLR, 2022.
[OS20] Samet Oymak and Mahdi Soltanolkotabi. Toward moderate overparameteriizatio:: Global convergence guarantees for training shallow neural networks. JEEE Journal on Selected Areas in Information Theory, 1(1):84-105, 2020.
[QSS23] Lianke Qin, Zhao Song, and Baocheng Sun. Is solving graph neural tangent kernel equivalent to training graph neural network? arXiv preprint arXiv:2309.07452 2023.
[QSY23] Lianke Qin, Zhao Song, and Yuanyuan Yang. Efficient sgd neural network training via sublinear activated neuron identification. arXiv preprint urXiv:2307.06566 2023.
[Raz17] Ilya Razenshteyn. High-dimensional similarity search and sketching: algorithms and hardness. PhD thesis, Massachusetts Institute of Technology, 2017.
14

---

[Rub18] Aviad Rubinstein. Hardness of approximate nearest neighbor search. In Proceedings of the 50th annual ACM SIGACT symposium on theory of computing, pages 1260-1268, 2018.
[RVW13] Liam Roditty and Virginia Vassilevska Williams. Fast approximation algorithms for the diameter and radius of sparse graphs. In Proceedings of the forty-fisth annual ACM symposium on Theory of computing, pages 515-524, 2013.
[SL14] Anshumali Shrivastava and Ping Li. Asymmetric Ish (alsh) for sublinear time maximum inner product search (mips). Advances Neural Information Processing Systems (NIPS), pages 2321-2329, 2014.
[SL15a] Anshumali Shrivastava and Ping Li. Asymmetric minwise hashing for indexing binary inner products and set containment. In Proceedings of the 24th international conference on world wide web (WWW), pages 98 981-991, 2015.
[SL15b] Anshumali Shrivastava and Ping Li. Improved asymmetric locality sensitive hashing (alsh) for maximum inner product search (mips). Proceedings of the Thirty- -irst Conference on Uncertainty in Artificial Intelligence (UAI), pages 812 821, 2015.
[SWYZ23] Zhao Song, Yitan Wang, Zheng Yu, and Lichen Zhang. Sketching for first order method: Efficient algorithm for low-bandwidth channel and vulnerability. In International Conference on Machine Learning, pages
PMLR, 2023.
[SY19] Zhao Song and Xin Yang. Quadratic suffices for over-parametrriztioo via matrix chernoff bound. arXiv preprint 1rXiv:1906.03359, 2019.
[SYZ21] Zhao Song, Shuo Yang, and Ruizhe Zhang. Does preprocessing help training over- parameterized neural networks? Advances in Neural Information Processing Systems, 34, 2021.
[SZZ21] Zhao Song, Lichen Zhang, and Ruizhe Zhang. Training multi-layer over wer-parametrized neural network in subquadratic time. arXiv preprint arXiv:21 2021.
[TOG17] Csaba D Toth, Joseph O' Rourke, and Jacob E Goodman. Handbook of discrete and computational geometry. CRC press, 2017.
[Wi105] Ryan Williams. A new algorithm for optimal 2-constraint satisfaction and its implica- tions. Theoretical Computer Science, 348(2-3):357-365, 2005.
[Wil18a] Ryan Williams. On the difference between closest, furthest, and orthogonal pairs: Nearly-linear vs barely-subquadratic complexity. In Proceedings of the Twenty-innth Annual ACM-SIAM Symposium on Discrete Algorithms, pages 1207-1215. SIAM, 2018.
[Wil18b] Virginia Vassilevska Williams. On some fine-grained questions in algorithms and complexity. In Proceedings of the International Congress of Mathematicians: Rio de Janeiro 2018, pages 3447-3487. World Scientific, 2018
[Y] Hongru Yang, Ziyu Jiang, Ruizhe Zhang, Zhangyang Wang. and Yingbin Liang. Convergence and generalization of wide neural networks with large bias. CoRR, abs/2301.00% 100027...007,... 2023.
[ZG19] Difan Zou and Quanquan Gu. An improved analysis of training over-parameterrized deep neural networks. NeurlPS, pages 2053-2062, 2019.
[Zha22] Lichen Zhang. Speeding up optimizations via data structures: Faster search, sample and maintenance. Master's thesis, Carnegie Mellon University, 2022.
[ZMG19] Guodong Zhang. James Martens, and Roger B Grosse. Fast convergence of natural gradient descent for wver-parameterized neural networks. In Advances in Neural Information Processing Systems (NeurIPS), 2019,
15

---

[OPD*20] Yi Zhang, Orestis Plevrakis, Simon S Du, Xingguo Li, Zhao Song, and Sanjeev Arora. Over-parameterized adversarial training: An analysis overeoming the curse of dimensionality. In NeurIPS. arXiv preprint arXiv:2002.06666, 2020.
[ZSZ*23] Zhenyu Zhang, Ying Sheng. Tianyi Zhou, Tianlong Chen, Lianmin Zheng. Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Ré, Clark Barrett, Zhangyang Wang, and Beidi Chen. H2O Heavy-hitter oracle for efficient generative inference of large language models. In NeurlPS. arXiv preprint arXiv:2306. 14048, 2023.
16

---

Appendix
Roadmap. We restate notation and provide additional tools about probability in section A. Then we present the results about sparsity in section B. In section C, we demonstrate the idea of two different correlation trees, DTree and WTree, and present the full version of training algorithms using our data structure. In section D, we provide detailed implementation and analysis of running time for our data structure. Section E presents the proof of running time for training 2NN b) using DTree and WTree. Section F shows formal version of lower bound for dynamic detection of firing neurons.
A Preliminary
A.1
Basic Notation
For any positive integer nI, use [n]] to denote the set 21, ,n]. We [[X] denote the expected value of a random variable X, and Pr[Y] to denote the probability of a random event Y. For a matrix M, we write M T to denote the transpose of M. We use x T y to denote the inner product between vectors z and y. We use La to denote d-dimensional identity matrix. We use (u; to denote the Gaussian distribution with mean μ. and variance x²
A.2 Upper bound on the movement of weights per iteration
The following Claim is quite standard in the literature, we omit the details.
Claim A.1 (Corollary 4.1 in [D]PS19], Lemma 3.8 in [SY19]). Lererr(i) be defined as Definition A.2. If E [t), | & (1 - 7//2)' ||er((0)]l,, then
|| (t + 1) - w, I 4X m vn ||er((0)]22 = D.
This claim shows uniform bound on the movement weights.
Next, we introduce the definition of error of prediction.
Definition A.2 (Error of prediction). For each t {0, 1,... ,T}, we define err(t) e R" be the error of prediction err(t) = y u(t), where u(t) = ((W(t), IR
A.3 Probabilities
We introduce the classical Bernstein inequality here.
Lemma A.3 (Bernstein inequality [Ber24]). Assume Z,,... Zn are n i.idd random variables. Vi [n). E[Z,] = 0 and [z| < M almost surely. Let 7 = =n -i=1 Zz. Then,
Pr Z > t < cxp
Vt 0.
Σ
Σド= E E[Z3] + Mt/3
Next, we show an inequality on shifted small ball with a Gaussian distribution.
Claim A.4 (Theorem 3.1 in [LSO1]). Let b 0 and 0. Then.
exp(-b2// x>>))..1)
Pr [|||
r| v
Pr 20
[ bl srls
Pr 2
[|0| S sr).
We state the anti-concentration inequality here.
Lemma A.5 (Anti-concentration for Gaussian distribution). Let Z ~ NN0, o2). Then, for > O,
Pr| & <
V
B Sparsity
In this section, we start by restating the result about sparsity after initialization in Section 2. Then we show how bound the number of fire neurons per iteration in Section B.2.
17

---

B.1 Sparsity after initialization
The goal of this section to prove the sparsity after the initialization of the neural network. We start by defining the "fire" set. Definition B.1 (fire set, Definition 3.7 in [SYZ21]). Fix query point x E Rd, let (are(x) denote the set of neurons that are "fire", i.e.
Stre (2) = {i E [m] (c, ws) > bj
Next, we introduce the fire for each training iteration.
Definition B.2 (fire set per iteration, Definition 3.7 in [SYZ21]. For each data point 2; E R [n]. weight w, RRd r [m]] and each iteration E {0, 1, T we define
(,,r == r E [m] (ar, (t)) (y.ire := i E [n]] (24, (x)(t)
Also. define ks,t (Se, and ky, := [Sr,are (t)l
With the above definitions, we can state the sparsity after initialization.
Lemma B.3 (Sparsity after initialization, formal version of Lemma 2.5, Lemma 3.8 in [SYZ21]). Let b > O be a tunable parameter. If we use the 4h as the activation function, then after the initialization. with probability at least 1 -xp( e exp(-b² 112 (2))), it holds that for input data x, the number of activated neurons b at most Om cxp(-1? /2)), where m the total number of neurons.
B.2 Bounding the number of fired neuron per iteration per level
In this section, will show that for I....... T, 0, log m, the number of fire neurons ki, = (irre.free() (t)) is small with high probability.
We define the set of neurons that are flipping at time t:
Definition B.4 (flip set, Definition C.8 in [SYZ21] For each [n]], for each time t [T] let St,ip (t) C [m]] denote the set of neurons that are never flipped during the entire training process,
(t. = fr E [m] (g xx)) )) b) # sgn((w, (1-1),,;))) 6)}.
Over all the iterations of training algorithm, there are some neurons that never flip states. We provide mathematical formulation of that set,
Definition B.5 (noflip set, Definition C.9 in [SY721]). For each [n], let S [m]] denote the set of neurons that are never flipped during the entire training process,
S fr e [m] Vt E [T] sgn((w,(t), x.) -b) sgn((w; (ww,(0)) (0), xz)) - b)}.
(1)
In Lemma 2.5, we already show that kv O(m exp(- (2) for all E [n] with high probability. We can show that it also holds for > 0.
Lemma B.6 (Bounding the number of fired neuron per iteration, Lemma C.10 [SYZ21]). Let b 2 0 be a parameter, and let ov(x) = max(æ, b} be the activation function. funnciion For each E [ni, [T]. hi, is the number of activated neurons at the t-th iteration. For 0 < S T, with probability at least 1 n exp (-1(mm)- min cxp(-b? b2/2))), ki, most Omm cxp(-1? /2)) for all i [n].
C Algorithm
In this section, we explain the procedures for two correlation tree data structure. We use the same setting Section
For the DTrre data structure, contains n binary trees indexed by data points and supports the following operations:
18

---

Initialize It takes data points fæ1,-.. In C. R and weights {wi. Wm } C R as imput and compute inner products of all weight-data pairs (w,, zt). It uses these inner products to create different trees. For the i-th tree based on data point it is constructed from m leaf nodes (Wy, æt), E [m] and satisfies the property that the value of parent node is the maximum value of its child nodes.
Update It takes calculate the new inner product (2, æi) and stores the value into the r-th leaf node. Then compares the new value with its parent node. It replaces parent node with new value if is larger and continue this comparing process. Otherwise it stops. Repeat the same operation for all trees.
weight 2 Rd and index E [m] as input. For the i-th tree, it
Query takes threshold E R20 and an index E [n] as input. Starting from the root of the i-th tree, it checks if its value is greater than threshold b. If no, search ends. If yes, it treats the child nodes as the root of a new subtree and repeat this searching process until stop. Then finds all indices r [m] that satisfy {w, sgug( (w,, 24) 6) > 0}.
For the WTree data structure, it contains m binary trees indexed by m weights and supports the following operations:
Initialize Similar to DTree, it takes data points {I1,-.- Inn C Rd and weights f....... Wm C IRd input and compute inner products of all weight-data pairs (w,, J. It uses these inner products to create nm different trees. For the r-th tree based on weight wy, it is constructed from n leaf nodes (Wr, E [n] and satisfies the property that the value of parent node is the maximum value of child nodes.
Update It takes a new weight z E TRd and an index r E [m] as input. Then it re-constructs the r-th tree with weight
Query takes a threshold b Rgo and an index r [m]] as input. Starting from the root of the r-th tree, it checks if its value is greater than threshold b. If no, search ends. If yes, treats the child nodes as the root of new subtree and repeat this searching process until stop. Then it finds all indices i E [n] that satisfy {w, sgu( ((w,, Hi) b) 2 0j.
Algorithm 5 Correlation DTree Data Structure
I:
data structure CORRELATIONDDREE
procedures:
3:
INIT(S C Rd, ", ".,WC C R" m. d) QUERY(i, b)
b Initialize the data structure via building n trees
4:
DiE [n]. R. Output the set (r [m]] sgn((w;, (x) - b) 2 0}
UPDATE(Z, i)
b Update I'th point IR with
6: end data structure
Algorithm 6 Correlation WTree Data Structure
I:
data structure CORRELATIONWTREE
2:
procedures:
3:
INIT(S Rd, W. C Cd m,
Initialize the data structure via building m trees
4: QUERY(r, b)
DrE (m],
Output the set (i [m] : sgn ((w,(( 6) 0} > Update the r'th point in Bd with w
5: UPDATE(W, r)
6:
end data structure
We present MAKEMAXTREN algorithm (Algorithm 7) which shows how to construct tree satisfying the property that the value of parent node is the max value of its child node.
We then give two training algorithms (Algorithm 8 and Algorithm 9) to show how DTree and WTree help in training neural network efficiently.
D Correlation Tree Data Structure
In this section, we demonstrate detailed results for DTree and WTree data structures.
19

---

Algorithm 7 Make MaxTree
1:
procedure MATEAXTREEINBITT
2: if then
3: return V1
4: else
5:
for i [n/2]] do
6:
Create node r
7:
if 1.lue. > P2i value then
8:
r/ < r2--1 121-1
9:
else
r r2 end if
11:
Insert 72:-1 as left child
13:
Insert right child
end for
15:
D n is odd, create parent node for the last node. return MAKEMAXTREEINEER(1)
16:
17:
end if
18:
end procedure
19: procedure MAKEMAXTRHE(C))
u,))
20:
for E [n] do
21:
Create nodes r; r:.value +- V4g end for
22:
23:
24:
return MATEMAXTREENNNRRTT.
end procedure
D.1
Correlation DTreee data structure
start by stating the main theorem of correlation DTree data structure.
Theorem D.1 (Correlation DTree data structure). There exists a data structure with the following procedures:
INIT((W), .2...... Wm Rdd, J2....,.. ,.,) C Rdd N,m E N. N). Given series of weights w1, u2,"." ; and datas 21; U2, Fn in d-dimensional space, it preprocesses time ((ndd)
UPDATE(2 E Rd,r S. [ml). Given a weight and index r, it updates weight w, with z in time O(n (d log
(d m))
QUERY(i [m], IR). Given an index indicating data point x; and threshold T, it finds all index [m] such that (wr, T time O([S(T) log m). where 5(7) = (ww, xz) > +}
D.2 Running time for CORRELATIONDTREB
The goal of this secion to prove the running time of INIT, UPDATE and QUERY. We start by showing the running time of INIT.
Lemma D.2 (Running time of INIT), Given series of weights {w1, w2, .. .m C. Rd and datas {21, F2,
Fn C Rdd, it preprocesses in time ((nmd)
Proof. The INIT consists two independent forloop and two recursive forloops. The first forloop (start from line 8) has n interations, which takes O(n) time. The second forloop (start from line 11) has m iterations, which takes (m) time. Now we consider the recursive forloop. The outer loop (line 14) has iterations. In inner loop has m iterations. In each iteration of the inner loop, line 16 takes (d) time. Line 18 takes 0(m) time. Putting it all together, the running time INIT
On + + n(md + m))
20

---

Algorithm 8 Training Neural Network via building trees, where each tree is the correlation between one data point and all the weights
1: procedure WAAAANNwWTTIPPRRRSSSSSSEEEECCTT
Theorem 4.1
PlInitiaizzation step*/
3: Sample (0) and according to Definition
4: b - v04 log m.
5: I"A dynamic /
CORRELATIONDTREE CDT CDT (DT. Netw) {ur n,, |Tterative step*/
Theorem b. takes Time m, d) time. Alg. 10
8:
9:
for = 0 - 7 do
10:
/ooowar computation step*/ for
do
12:
St,ir +- CDT. QUE.QU)RY(i.))
It takes Tauary (m, ks.s) time. Alg. 12 takes O(d: ks,t) time
13:
u(t), 1
mm Enesi, ar (t)
14:
end for
BBBckwardd computation step*/ + x
16:
b P E 2
17:
for i = do
18:
for E St do
19:
Por -- Vm dr (t)
20:
end for end for
21:
22:
M + X diag(y u(t))
M E Rd it takes O d) time D DAW E dd takes Old mnz(P)) time, unz( (P) = O(mm
23:
AW + M P
dxn
24:
W(t + + ((t) AW.
25:
U*Update data structure*/
26:
Let Q c [m]] where for each r E Q, the AW., not all zeros
|Q] S 0(mm
27:
for E Q do
28:
CDT.UPDATE(W; (1 r),r)
Alg. 11
29:
end for end for
30:
31:
return W
W E Rdxm
32:
end procedure
(
Thus, we complete the proof.
Next, we analyze the running time of UPDATE.
Lemma D.3 (Running time of UPDATE), Given weight z E Rd and index E [m], updates weight wj with in time O(n (d ( log logm)) m))
Proof. The running time of UPDATE mainly comes from the forloop (line 23), which consists of iterations. each iteration, line 24 takes Ollog m) time, line 25 takes O(d) time and the while loop takes Ollog m) time since it go through path bottom up. Putting it together, the running time of UPDATE (n + log m)).
Finally, we state the running time for QUERY procedure.
Lemma D.4 (Running time of QUERY). Given a query q E Id and a threshold F > 0, it finds all index E [n) such that (wo, q) > in time O(IS(T)) log m)) where S(T) := {i (wi, o) > r}
Proof. The rumning time comes from FIND with input 7 and root(T;). In FIND, we start from the root node r and find indices in recursive way. The INIT guarantees that for a node r satisfying r.value > the sub-tree with root must contains a leaf whose value is greater than If not satisfied, all the values the nodes in the sub-tree with root r less thanr. This guarantees that all the paths it search does not have any branches that leads to the leaf we don't 't want and it will report all the
21

---

Algorithm 9 Training Neural Network via building m trees, where each tree is the correlation between one weight and all the data points
1: procedure TRAININGWITHHPP TTAINNG
Theorem 4.2
Plnitialization step*/
3: Sample (0) and according to Definition
4: - v0.4 log m.
5: I"A static data-structure"/
CORRELATIONW CWT
Algorithm 10, of Theorem takes Tine m. d) time
(=T (w-((0) (R), MIiitillize Sp and Si,a, */
8:
It takes Em (nu kr.s) O(m 4/= log n) time Sy.fre is the set of samples, for which neuron r fires D. Si,fir is the set of neurons, which fire for æ;
10: Sy.e + for r E [m]]
11:
Si.tre < 1 for E [n].
12:
for r-
m do Sr.fire <- (UETYQURRY(
13:
14:
for ioriE E SM.rc do ADD(r.-DD(r)) for end for
15:
16:
17:
18: [Iterative step*/
19: for t ''' do
20:
/*Forward computation step*/
21: for i n do
22:
u(t), + E , or - (t)
b It takes O(d ks.)) time
23:
end for
24:
/BBacwward computation step*/ P x=
25:
P E P R
26:
for + do
27:
for E St do Per - ar (w)((t) (x))
28:
29:
end for for
30:
31:
M + diag(u u(t))
M E RRdxn, takes O d) time
32:
AW M
AAW E Rd takes Old mnz(P)) time, mnz(P) O(mm 4/5)
dxn axm W(L t ((L) AW.
33:
34:
U'pdtate Soo Si step*/
35:
It takes O Ki, Eres Tauery (n,, d. kr))+1)) O - logn 4/5)
36:
Sin) + U for r E Spn],are do
37:
38:
for E Sr. do
b Removing old fired neuron indices. It takes (kr.,) time
39:
40:
end for CWT. WWT.UUDAAT(W.-(1+ 1),r) Sr & CWT.QUERY(1;)) i do ADD(rc.AD((r)
41:
D takes Tupdata d) time It takes Tawery (n,, d, kr,4+1) time Adding fired neuron indices. It takes (kr, time
42:
43:
44:
45:
end for end for
46:
47:
end for
48:
return W
W Rdxm
49:
end procedure
indices i satisfying (wi, q) 0. Note that the depth i' O n), the running time of QUERY is O([((7)) logn).
D.3 Correlation WTrre data structure
In this section, we state the main theorem of correlation wtree data structure.
22

---

Algorithm 10 Correlation DTTree data structure
1: data structure CORRELATIOND TREE
>> Theorem D.1
2: members
3: W E IRxd (m weight vectors
4: X E Rixd (n data points)
5: Binary tree 12. Tn
> binary search trees
6: end members
7:
8: public:
9: procedure INIT(W1, W2.
Wm E Rd m, 21, 22,"" En E R d)
Lemma D.2
for 1 do
11: Ti <-
end for 13: for
- do
w <- W
15:
end for
16:
for = + TI do
> for data point, we create a tree
17:
for j =
do
18:
Uj + (mi, wi)
19:
end for
20:
T, < MAKETREE(U ... 22
D Each node stores the maximum value for his two
children
21:
end for end procedure
22:
23: end data structure
Algorithm 11 Correlation DTrees
1: data structure CORRELATIONTREE
Theorem D.1
2: public:
3: procedure UPDATE(2 IRd, IRd, r
[m])
Lemma D.3
4: <- 2
5: for
do 6: [ + the 1-th leaf of tree 7;
7:
1.value (z, x:))
8:
while is not root do
9:
p parent !
10:
v.luuue + (nalue.aaue,,.vval v.aaue)
1tp p
12:
end while
end for
14:
end procedure
15:
end data structure
Theorem D.5 (Correlation WTrre data structure). There exists a data structure with the following procedures:
INIT((W), W2..... Wm) Rdd, (21, (E1,, ,n) C Rdd N,m E N, N). Given a series of weights w1, W2,...., Wu and datas 21, D2. Fn in d-dimensional space, it preprocesses time O(nmd)
UpDATE(Z R²,re [m]). Given weight and index r, it updates weight w, with in time ((nd)
(U [m], E IR), Given an index indicating weight w, and a threshold T, finds all index E [n]] such that (w,, Izi > T in time S{S((T) logm), m), where S(T) {i (w) r)
23

---

Algorithm 12 Correlation DTTree
1: data structure CORRELATIONDTREN
Theorem D.1
2: public:
3: procedure QUERY(Í E [n], E Rgo)
Lemma D.4 >
4: return FIND(T, (oot
5: end procedure
6:
7: private:
8: procedure FIND(T Rzu,r E T)
9: if r is leaf then
return
11: else
r1 < left child of r, 12 < right child of
13:
if value 2 r then
S FFNND(((,)
15:
end if
16:
value 2 r then S2 FINDDNT r2)
17:
18:
end 19: end if
20: return S1 U.2
21:
end procedure
22:
end data structure
Algorithm 13 Correlation WTTree data structure
1: data structure CORRELATION
Theorem D.5
2: members
3: W E mg (m weight vectors
4: X Raxd (n data points)
5: Binary tree T2. ... TM
b m binary search trees
6: end members
7:
8: public:
9: procedure INIT(W1, W2, Wm E Rdd m, 21, 02,
Fn E Rdd n, m, d)
b Lemma D.6
10: io 1
do
11:
2; - Iz end for
13: for j = + m do
Wiy < W
15:
end for
for
+ m do
for weight, we create tree
17:
for j = + do
U +- (Mi, wi)
19:
end for
Each node stores the maximum value for his two T, <- MAKETREE(U ,.,)
children
21:
end for end procedure
22:
23:
end data structure
D.4 Running time for Correlation WTrre
The goal of this secion is to prove the running time of INIT, UPDATE and QUERY.
As in DTree, first show the running time for INIT.
Lemma D.6 (Running time of INIT). Given series of weights {w1, w2, Wm C Bd and datas {21, D2, Inn C. Rdd preprocesses in time O (nndd)

---

end data structure Algorithm 14 Correlation WTTrees
1: data structure CORRELATIONWTRES
Theorem D.5
2: public:
3: procedure UPDATE(2 E IR= E [m))
Lemma D.7
4: W/ <-
5: for j = 11110 do
6:
"j (x)
7: T; < MAKETREE(4, "1 u2)
>> Each node stores the maximum value for his two
children
8: end for
9:
end procedure
10:
Algorithm 15 Correlation WTrre
1: data structure CORRELATIONWTRES
2: public:
3: procedure QUERY(r E [m]. r E Rz0)
Lemma D.8
4: return FIND(T,
5: end procedure
6:
7: private:
8: procedure FIND(T Rzu,r T)
9: if is leaf then return
11: else
12:
F1 < left child of r2 right child of if value 2 then
13:
S FIND((TT
15:
end
r2. value 2 r then S2 + r2)
17:
end 19: end if
return Si U S2
21:
end procedure
end data structure
Proof. The INIT consists of two independent forloop and two recursive forloops. The first forloop (start from line 10) has n interations, which takes Onn) time. The second forloop (start from line 13) has m iterations, which takes (mm) time. Now we consider the recursive forloop. The outer loop (line 16) has m iterations. In inner loop has iterations. each iteration of the inner loop. line 18 takes (d) time. Line 20 takes O(n) time. Putting it all together, the running time of INIT is
Onn + m + m(nd + n2)
(nmd)
Thus, we complete the proof.
Next, we turn to the running time for UPDATE.
Lemma D.7 (Running time UPDATE), Given a weight z E gd and index r E [m], updates weight wj with in time O(nd)
Proof. In this procedure, it generates a new tree for weight w. with n leaves, which takes (nd) time. Thus, we complete the proof.
Finally, present the running time of QUERY.
25

---

Lemma D.8 (Running time of QUERY). Given a query q E I and a threshold T > 0. it finds all index [n] such that (w,, q) > T in time O[[((T) log m)), where S(T) := {i (w., q) > 7}
Proof. The running time comes from FIND with input and (T:). In FIND, we start from the root node r and find indices in a recursive way. The INIT guarantees that for a node r satisfying valuu > r, the sub-tree with root r must contains a leaf whose value is greater than If not satisfied, all the values of the nodes in the sub-tree with root r is less thanT. This guarantees that all the paths search does not have any branches that leads the leaf we don't want and it will report all the indiex i satisfying (w;, q) > 0. Note that the depth of T is Ollog n.), the running time of QUERY is ((((())] l)
E
More Details of Our Training Algorithms
E.1 Weights Preprocessing
In this section, we present the formal version of our training algorithm using DTree, which prepro- cessing weights for each data point.
Theorem E.1 (Running time part, formal version of Theorem 4.1). Given data points in IR Running gradient descent algorithm (Algorithum 8) on 2NN v0.4 log m) (Definition 2.1) the expected cost per-iteration of the gradient descent algorithm is
O(mA/On22
Proof. The per-step time complexity is
T Ti + T
Σ TQUERY (m, d, ki,e) + TUPDATE Uie(n) Se,tre (t) + +d
Σ kj,t
sE(m)
The first term Ti = =E= TQUERY (m, d, ks,c) corresponds to the running time of querying the active neuron set Si,r (t) for all training samples i [m]. With the first result in Theorem 3.1, the complexity is bounded by (mn2/5 log log
The second term To TUPPATe: - [uejn) (t)) corresponds to updating w, in the high-dimensional search data-structure (Line 28). Again with the first result in Theorem 3.1, we have TUPDATE O(n(d + log m)). Combining with the fact that [[e[m) (t) I |Uig(m) Si (0) S Omm4/n)) the second term is bounded by (m))/
The third term is the time complexity of gradient calculation restricted to the set St, With the bound on Eve()) ky (Lemma B.6), we have d d' k S (0m4/5nd)
Putting them together, we have
T s 1/5 n m) 0mm4/5n2d + Omm
0(m 4/5m4/,
Thus, we complete the proof.
E.2 Data Preprocessing
In this section, we describe a similar version of training algorithm aforementioned but it uses WTree to preprocess data points based on weights.
Theorem E.2 (Running time part, formal version of Theorem 4.2), Given data points in Rd. Running gradient descent algorithi (Algorithm 9) on 2NN(m, b = v0.4 m), the expected per- iteration running time of initializing Sre St.rr for E [m], E [n] is 0(m4/5 log The cost per-ite per-iteration of the training algorithm is 0(m 1 ?dd).
26

---

Proof. We analyze the initialization and training parts separately.
Initialization From Line 10 to Line 17, the sets Sr.free Sire for E [m]] [m] are initialized. For each r E [m]] we need to query the data structure the set of data points x's such that (s (0) x)) > 0.
Hence the running time of this step
m
TQUERY (n, kr,Kr.))= =02
ky,0 log
0 Σ
ks, log
==1
O logn)
where the second step follows from Σ kv.0
Σi= Ks,0-
Training Consider training the neural networkfor T steps. For each step, first notice that the forward
and backward computation parts (Line 21 Line 33) the same as previous algorithm. The time complexity is O(m4/5n)
We next show that maintaining Sr,ire, r E [m]] and Si,tro; i [n)] (Line 36 Line 45) takes O(m4/5 time. For each fired neuron E [m]. we first remove the indices of data in the sets Si,ire which takes time
0(1)
Σ
h = 0(1)
Σ hins = (m
rEULE(N)
Then, we find the new set of z's such that oz (,, 1), x)) > 0 by querying the correlation tree data structure. The total ruunning time for all fired neurons
Σ
TUrDNEE( d) + TQUERY (n, kr,01) m/ log
Σ
ky,441 log
Se,
(EUicta) St
=m
Then, we update the index sets Si,rr in time n/m Therefore, each training step takes ( time, which completes the proof.
Lower Bound for Dynamic Detection of Firing Neurons
The goal of this section to prove the lower bound for Dynamic Detection of Firing Neurons.
We start by introducing the strong exponential time hypothesis, SETH in abbreviation.
Definition F.1 (Strong exponential time hypothesis, SETH, [IP01, CIPO9]). For every e 0, there exists a k k(()) N such that algorithm can solve k-SAT (i.e., satisfiability on CNF of width k)) in (((())))) time where is the number of variables.
We present another relative concept called orthogonal vector conjecture, OVC in abbreviation.
Definition F2 (Orthogonal vector conjecture, OVC, [Wil05, AWW14, B115, ABW15]). For every O, there exists a 1 such that the orthogonal vector problem of size in d-dimension requires 2 -tme when d = c log n.
refer a theorem about maximum bichromatic inner product lower bound in [Che20].
Theorem F.3 (Maximum bichromatic inner product lower bound. [Che20]). Assuming SETH (Definition F.I) or OVC (Definition F.2), there is constant such that any exact algorithm for E Z-Max-IP, n.a in dimension d = log" requires n 2-o((1) time, with vectors of O -bi entries.
Putting things together, we state the main result for the lower bound for Dynamic Detection of Firing Neurons.
27

---

Theorem F.4 (Lower Bound for Dynamic Detection of Firing Neurons, Formal version of Theo- rem 1.4,, Let = 20("o"" n)_ Unless OVC or SETH fails, for any constants e E (0, 1), no data
structure can solve DDFN with less than m 2 o(1) -time per update and m m ( - per query.
Proof. Without loss of generality, assume that Let d -olog" m where is defined in Theorem F.3.
Suppose there exists data structure that for (m, d + 1)-sized instance, perform updates in m f and answer queries m=2-en -e-------ttem for some c (0, 1) and E (0, c).
Let X ( Fm Zd Y (.n,.... V/ zd be hard instance of Z.- Z--ax-IP m.d problem constructed in Theorem F.3. For each vector 2; (or Vj), we construct a new vector I. (or yj) in (d + di-meennoo such that (F:)a+1 -1 and (ū) d++1 w, where w is parameter be chosen later.
Then, construct k = [m/n] instances of the DDFN problem in Definition 1.3 as follows:
X = ( J.h, y = . }m
and O.
We show that the data structures for these instances ((x(1,, ((), b)]s(K)))))))))) can be used to solve MaxxIIx,,aXXXY)
We perform a binary search for the value of Z-Max-IP,, m.d (X, Y). Note that at most O(logn) iterations suffice to find the exact answer.
Suppose the current value in the binary search t E Consider the i-th instance (X ( b) for any i [k]. We first call UPDATE( to set Wj]a+1 1 for each j [m]] By the data structure's 's guarantee, this step takes O(m m ( = 0 (m 2 c) that
time. Then, we call QUERY(). Notice
( F1) (Ia, Mz) t2
(x Ub) 2t. t.
Hence, QUERY( will return all pairs of (i,j) such that (I,, 2 t. This step runs in O(m! ---mm- We repeat this process for all h instances. And based whether the outputs of the QUERY() are empty or not, we know the direction of the binary search for the next iteration.
Hence, the total running time of each iteration
O(k 2m² + m on (t0-0-)))
Om c-ce-1 m
which follows from the assumption of m Thus, can solve Z-Max-IPn.dde (X, Y) in Ö(m²-s) < mmm which contradicts to the lower bound for Z-Max-IP,, a. Theorem F.3. Therefore, no such data structure can exist.
28


Reviews:
[{"review_id": "j7jMcj2u8Q", "reviewer": "Reviewer_7T2Q", "summary": "This paper investigate the efficient training methods than the usual training protocol which requires the complexity $O(nmd)$ for 2-Layer ReLU networks. The authors improve the complexity in the previous study [SYZ21] by proposing the preprocessing method utilizing the tree data structure for both data and weights. Moreover, the authors successfully provide the upper bound/lower bound (for lb, the authors assume some conjecture) for their proposed preprocessing complexity. Unlike in previous study [SYZ21], this paper is the first presentation for the lower bound.", "questions": "Please see the weaknesses part.", "limitations": "The authors have adequately addressed the limitations and potential negative societal impact of their work.", "rating": 6, "confidence": 2, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1. The authors efficiently improve the preprocessing time from exponential complexity $O(2^d)$ (in previous research [SYZ21]) to polynomial $O(nmd)$ (in this paper) for both data and weight parameters using the tree data structures which is popular in general computer science.\n\n2. For the tree data structure preprocessing, both the upper bound/lower bound are provided in a solid theory under the NTK regime and the theory basically depends on the previous study [SYZ21].\n\n3. Unlike in previous study [SYZ21], the authors also provide the lower bound for their proposed method (although assume some conjecture), so the tree-based preprocessing method is nearly optimal.", "weaknesses": "Actually, I'm not an expert in this field, but I have carefully read this paper along with the previous study [SYZ21].\n\nHere are my main concerns:\n\n1. Based on my understanding, the authors have successfully reduced the complexity of preprocessing from $O(2^d)$ to $O(nmd)$ and from $O(n^d)$ to $O(nmd)$. However, it seems that the per-iteration time in this paper (for example, in Theorem 4.1, the time per-iteration is $O(m^{4/5}n^2d)$, but it is $O(m^{4/5} n d)$ in previous study [SYZ21]) has actually increased compared to the previous study [SYZ21]. In fact, if preprocessing time constitutes a significant portion of the neural network's training process, this research would have more significance. Therefore, it seems that an (empirical) analysis of the portion that preprocessing takes in neural network training, in computational terms, is necessary.\n\n2. As this study presents more advanced preprocessing techniques compared to the previous research [SYZ21], there should be experimental analysis on how much the actual preprocessing time is reduced and its impact on per-iteration time (under quite theoretical settings or even for synthetic data/architecture). This analysis seems to be necessary, even for simple neural network models and simple datasets, to further validate the improvements made by the proposed preprocessing methods.\n\n3. In the theoretical perspective, it is necessary to provide remarks on what is the challenging points of the theory in this paper compared to the previous study [SYZ21]. When I was examining the supplementary material, it is not clear which aspects have significantly changed compared to the theoretical analysis in the previous study [SYZ21].", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks so much for taking the time to read and understand our paper and for your helpful suggestions. Please refer to our general response for your concerns about the empirical analysis and the comparison to [SYZ21]. In the final version, we will add a remark to compare our techniques to [SYZ21].   \n"}, {"review_id": "hMTWHRxOje", "reviewer": "Reviewer_GDL4", "summary": "In the paper, the authors proposed fast optimization algorithm for over-parameterized two-layer networks. They proved that by using the sparsity firing feature from the neural network, the proposed method requires only O(nmd) time in preprocessing and still achieves o(nmd) time per iteration.", "questions": "1. The space required by the proposed method is high, \n2. Is there any possibility to do a toy experiment?", "limitations": "1. The space required by the proposed method is high. \n2. There are few works training a two-layer network and showed competitive performance with deep NN. It is non-trivial to simplify the problem and prove the convergence of NN. However, why we need to do training based on an un-empirical structure? I hope at least a tory experiment should be provided to show that this training method is viable in deep NN.", "rating": 6, "confidence": 1, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. A thorough theoretical analysis is provided and prove the statement of the paper.", "weaknesses": "1. The proposed method require O(nm) space for storing, which is a lot when n and m is big. \n2. No experimental results shown in the paper. A toy experiment can show empirical impact of this work.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks so much for your great questions. Please refer to our general response for your concern about the experiments. \n\nAnd we agree with the referee that our data structure requires $O(mn)$ space. However, we would like to highlight that the primary objective of our research is to study the neural network training problem from a theoretical perspective. And in this work, we focus on optimizing the time complexity while preserving the convergence guarantee. As a trade-off, we need to use more memory. We believe it is a very interesting open question to develop a training algorithm that theoretically reduces space complexity. It is also interesting to investigate the time-space trade-off for neural network training algorithms. We leave these as open questions for future study.    \n"}, {"review_id": "ANzaOUrNJA", "reviewer": "Reviewer_kiFp", "summary": "This paper analyzes a specific neural network setting: two-layer neural networks with $m$ neurons in the hidden layer and a ReLU activation. Given input data of dimension $d$ and $n$ training examples, it normally requires O(mnd) operations to compute the hidden activations. This paper follows prior work in showing that this can be asymptotically reduced to $O(m^{4/5}n^2d)$, which is better in the overparameterized (large m) regime. Improving on the previous work, the proposed algorithm requires polynomial instead of exponential time preprocessing. The core technical contribution is the Correlation Tree data structure, which is a collection of binary trees that store inner products between datapoints and neurons, and allows efficiently updating based on the sparsity of activated neurons.", "questions": "N/A", "limitations": "The authors discuss most of the prominent limitations, but I think it is not completely clear how practical the algorithm is. It seems like it should be relatively straightforward to implement, which could strengthen the paper.", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 2, "strengths": "- The paper analyzes an interesting theoretical setting, and provides an original data structure and algorithm for this problem of efficient computation. The main result that training with sublinear cost per iteration is novel.\n- The paper is well-structured and written well. The problem setting is clear and the contributions are clearly decribed. Overall it is high quality.\n", "weaknesses": "The main weakness of this submission is the empirical practicality of the method. The authors are transparent about these weaknesses, and it is not the intended focus of this direction of research, so I think this weakness does not significantly detract from the paper. It might be interesting to comment more on the technical challenges behind extending it to more complex settings (e.g. beyond 2 layer, or other activation functions).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks so much for your helpful suggestions! Please refer to our general response for your concern about the practicality of the algorithm. \n\nRegarding the generalization to more complex settings, we think it is possible that our data structure will work for other activation functions. Roughly speaking, as long as the activated neurons are sparse or approximately sparse, our data structures will be able to theoretically reduce the cost-per-iteration. However, we need to re-prove the sparsification results in [SYZ21] for the activation function other than ReLU. We believe this is a very interesting direction for further research. For multi-layer neural network training, a natural barrier is $O(m^2nd)$-per iteration (assuming each layer has $m$ neurons and the number of layers is constant). [SZZ21] proposed an algorithm that runs in sub-quadratic time with a provable convergence guarantee. Their algorithm does not use the half-space reporting or tree-based data structure like ours or [SYZ21] but employs some other techniques like tensor sketching to speed up matrix-related computations, which are bottlenecks only appearing in the multi-layer setting.  Since they showed that the ReLU sparsifier will also work even beyond 2-layer networks, our tree-based data structure could be applied to this setting. However, our lower bound for Dynamic Detection of Firing Neurons (Theorem 1.4) with $n=m$ may prevent us from surpassing the quadratic barrier.  We will add some discussions about this point in the final version. \n\nReference:\n\n[SZZ21] Song, Zhao, Lichen Zhang, and Ruizhe Zhang. \"Training multi-layer over-parametrized neural network in subquadratic time.\" arXiv preprint arXiv:2112.07628 (2021).\n"}, {"review_id": "UDZblzq4ah", "reviewer": "Reviewer_oEJX", "summary": "The paper presents a new preprocessing method for training shallow overparametrized sparse neural networks. It significantly improves the preprocessing time yet achieves same performance on query time. They also show that their algorithm is very close to optimal.", "questions": "1. Does your work have applications in domains other than deep learning?", "limitations": "1. The setting is limited. Overparametrization is currently largely a theoretical assumption for the convenience of proving things, rarely used in practice. Shallow network is also too limited for characterizing deep learning. Also neural networks are commonly run on parallel machines like GPU with some complications for sparsity processing, so it's hard to say whether data structures would be really helpful. But it's very common, due to the nature of things, to have these limitations in a theoretical work, so it's not a serious flaw. Maybe it's better to design the algorithm for a general problem and put deep learning as a possible use case.\n\nThey have addressed the limitation of their work in terms of shallowness.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. Clearly written. Easy to understand. Well structured.\n2. Story is rather complete. Lower bound is included showing that their result is close to optimal.\n3. Connect ideas from applied algorithms (LSH etc.) to deep learning, insightful", "weaknesses": "1. Not all assumptions appear in the statement of theorem 1.1. Expect a more detailed description of the sparsity assumption.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks so much for your great efforts and helpful comments. Please refer to the general response for the practicality concerns.\n\nRegarding the sparsity assumption, we note that we do not assume the activated neurons are sparse. But we use a result proved in [SYZ21] that as long as setting a unified threshold ($b=\\sqrt{0.4\\log m}$) to the ReLU functions, after the standard randomized initialization, the number of activated neurons in each training iteration will be $O(m^{4/5})$ with high probability (see Remark 2.5). We will state it more explicitly in Theorem 1.1 in the final version. \n\nRegarding other potential applications, our data structure actually solves the problem of dynamically finding pairs of points with large inner products, which is a very useful sub-routine for applications such as clustering, database searching, etc. And compared to the classical computational geometric approach (e.g. [AEM92]), our data structure does not suffer from the curse-of-dimensionality problem, which is a big advantage for dealing with high-dimensional data.   \n"}]

Author rebuttal:
## General response:

We thank all the referees for the valuable comments. Here, we give general responses to some common questions.

First, regarding the concerns on empirical practicality, we want to kindly emphasize that our purpose is to present a training algorithm that is both efficient and accurate, supported by solid theoretical guarantees. For comparison, numerous recent works focused on empirically accelerating the training procedure (e.g., MONGOOSE [CLP+20], etc.).  They used some locality-sensitive hashing (LSH)-based data structure to *approximately* find the neurons with large outputs. The impressive experiments in those works show that the neural network training time can be significantly reduced while achieving similar accuracies. However, it remains to be a big challenge to offer any performance guarantee for those methods. Our work and [SYZ21] seek to bridge the gap and develop more theoretical insights into fast neural network training. The algorithm proposed in [SYZ21] provides a provable convergence guarantee but falls short in terms of efficiency due to the slow preprocessing. In contrast, our algorithm not only provides a theoretical performance guarantee but also mitigates the curse-of-dimensionality issue in [SYZ21]. Thus, the contribution of this paper is more theoretical, and we acknowledge that our algorithm will not be as fast as the previous empirical methods. Furthermore, we note that it is quite difficult to empirically gauge the preprocessing time in [SYZ21] since the half-space reporting (HSR)-data structure in [SYZ21] is an extremely complicated computational geometric data structure developed by [AEM92]. Even if we can implement this data structure, the exponential dependence in $d$ restricts its usage to very tiny datasets. 

Second, we would like to draw a comparison between our techniques to those developed in [SYZ21]. Our work consists of three technical components: a) designing the correlation-tree data structure to retrieve the activated neurons, b) proving the complexity and correctness of our algorithms, and c) showing the nearly-optimality of our data structure. The most important contribution of this work is in Part a), which exponentially improves the running time of the data structure in [SYZ21] in terms of $d$. Notably, unlike [SYZ21], we do not use any “black-magic” in previous work, and our data structures are simple and clean (see Algorithms 1 and 2 in our paper). As we discussed in Section 5, they can be easily parallelized to reduce the cost-per-iteration to $O(m^{4/5}nd)$. Part b) relies on the ReLU-sparsifier developed in [SYZ21], which demonstrated that by adding a unified threshold to the ReLU functions, the activated neurons in each iteration will be sparse while the training dynamics will still converge to the zero training loss. This insight helps us establish the performance guarantees for our training algorithm. Part c) is another novel contribution of this work. We are the first to develop a *non-trivial* fine-grained complexity result for a key sub-problem in neural network training, namely the dynamic detection of activated neurons.  This result shows that even for a nearly-constant dimension, it is still hard to construct a data structure with sublinear-time per update and subquadratic-time per query, implying that our data structure is nearly-optimal in the computational-theoretic sense. 

Reference: 

[AEM92] Agarwal, Pankaj K., David Eppstein, and Jirí Matousek. "Dynamic half-space reporting, geometric optimization, and minimum spanning trees." Annual Symposium on Foundations of Computer Science. Vol. 33. IEEE Computer Society Press, 1992.

[CLP+20] Chen, Beidi, et al. "Mongoose: A learnable lsh framework for efficient neural network training." International Conference on Learning Representations. 2020.

[SYZ21] Song, Zhao, Shuo Yang, and Ruizhe Zhang. "Does preprocessing help training over-parameterized neural networks?." Advances in Neural Information Processing Systems 34 (2021): 22890-22904.


Meta-review: The paper proposes a new preprocessing procedure to improve the runtime of training neural networks assuming sequential computation. Compared to prior work, they improve the preprocessing time by an exponential factor at the cost of an extra 
 (number of samples) factor in the training time. Underlying their result is a new data structure that can dynamically find pairs of inputs with large inner-products efficiently.

Technically the data structure is novel, and the optimality result of the data structure is a plus. However, I share the concern with the reviewers about the empirical practicality of the method, and the scope of the current setting of heavily overparameterized depth-2 networks trained in the NTK regime (width). Furthermore, this work assumes sequential computation for each neuron in a layer, which is not a practical concern. I am on the fence about this, but I am leaning towards accepting based on the reviewers unanimous agreement to accept the paper. I encourage the authors to address these concerns with additional discussion in the final version of the paper, and add more discussion on the limitations of their techniques.
Decision: Accept (poster)