SlowGuess/NeurIPS_2023_OCR huggingface dataset:
- Tree construction We first calculate the inner products of all weight-data pairs  $\langle w_i, x_j \rangle$ , each representing the evaluation of a neuron at a data point. To search activated neurons efficiently, we create a tree structure in the following way (taking the Correlation DTree as an example): we first build  $m$  leaf nodes, where the  $r$ -th leaf stores  $\langle w_r, x_i \rangle$  for  $r \in [m]$ . Then, we recursively construct a binary tree such that a parent node takes the larger value from its two child nodes. Finally, we obtain a tree with root having value  $\max_{r \in [m]} \{\langle w_r, x_i \rangle\}$ . Moreover, the value of each internal node equals to the maximum value among the leaf nodes in this subtree.
- Efficient search Given a threshold  $b$ , the data structure can find all the pairs of vectors whose inner product is greater than  $b$ . Take the Correlation DTree as an example. It outputs the indices of those activated neurons (i.e.,  $\langle w_r, x_i \rangle > b$ ) in a recursive way: starting from the root, it checks whether it is "activated" (i.e., with value  $> b$ ). If not, the search ends. Otherwise, it moves to each of the child nodes and repeats this searching process until stops. This is a 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 as 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  $w_{r}$  appears only in  $n$  different trees. In the Correlation WTree, each weight vector  $w_{r}$  appears only in one of the  $m$  trees. When  $w_{r}$  is updated, DTree will change the nodes along a root-to-leaf path in  $n$  trees, whereas WTree only changes such paths in the  $r$ -th tree.
3.1 Correlation DTree data structure
We now state our main theorem summarizing the correlation DTtree data structure. Its pseudocode is given in Algorithms 1 and 2 below. Its proof are deferred to Section D.1.
Theorem 3.1 (Correlation DTree data structure). There exists a data structure with the following procedures:
- INIT  $\left(\left\{w_{1}, w_{2}, \cdots, w_{m}\right\} \subset \mathbb{R}^{d}, \left\{x_{1}, x_{2}, \cdots, x_{n}\right\} \subset \mathbb{R}^{d}, n \in \mathbb{N}, m \in \mathbb{N}, d \in \mathbb{N}\right)$ . Given a series of weights  $w_{1}, w_{2}, \cdots, w_{m}$  and data  $x_{1}, x_{2}, \cdots, x_{n}$  in  $d$ -dimensional space, it performs preprocessing in time  $O(nmd)$ .
- UPDATE  $(z \in \mathbb{R}^d, r \in [m])$ . Given a weight  $z$  and an index  $r$ , it updates weight  $w_r$  to  $z$  in time  $O(n \cdot (d + \log m))$ .
- QUERY  $(i \in [n], \tau \in \mathbb{R})$ . Given an index  $i$  indicating data point  $x_i$  and a threshold  $\tau$ , it finds all indices  $r \in [m]$  such that  $\langle w_r, x_i \rangle > \tau$  in time  $O(|\widetilde{S}(\tau)| \cdot \log m)$ , where
$$
\widetilde {S} (\tau) := \{r: \langle w _ {r}, x _ {i} \rangle > \tau \}.
$$
Algorithm 1 Correlation DTree data structure
1: data structure CORRELATIONDTREE  $\triangleright$  Theorem 3.1  
2: members  
3:  $W\in \mathbb{R}^{m\times d}$  (m weight vectors)  
4:  $X\in \mathbb{R}^{n\times d}$  (n data points)  
5: Binary tree  $T_{1},T_{2},\dots ,T_{n}$ $\triangleright n$  binary search trees  
6: end members  
7: procedure  $\mathrm{INIT}(w_1,w_2,\dots ,w_m\in \mathbb{R}^d,m,x_1,x_2,\dots ,x_n\in \mathbb{R}^d,n,m,d)$ $\triangleright$  Lemma D.2  
8: for  $i = 1\rightarrow n$  do  
9:  $x_{i}\gets x_{i}$   
10: end for  
11: for  $j = 1\rightarrow m$  do  
12:  $w_{j}\gets w_{j}$   
13: end for  
14: for  $i = 1\rightarrow n$  do  $\triangleright$  for data point, we create a tree  
15: for  $j = 1\rightarrow m$  do  
16:  $u_{j}\gets \langle x_{i},w_{j}\rangle$   
17: end for  
18:  $T_{i}\gets \mathrm{MAKEMAXTREE}(u_{1},\dots ,u_{m})$ $\triangleright$  Each node stores the maximum value for his two children, Algorithm 7  
19: end for  
20: end procedure  
21: procedure UPDATE  $(z\in \mathbb{R}^d,r\in [m])$ $\triangleright$  Lemma D.3  
22:  $w_{r}\gets z$   
23: for  $i = 1\rightarrow n$  do  
24:  $l\gets$  the l-th leaf of tree  $T_{i}$   
25:  $l.\text{value} = \langle z,x_i\rangle$   
26: while  $l$  is not root do  
27:  $p\gets$  parent of  $l$   
28:  $p.\text{value}\gets \max \{p.\text{value},l.\text{value}\}$   
29:  $l\gets p$   
30: end while  
31: end for  
32: end procedure  
33: end data structure
3.2 Correlation WTree data structure
We next state the main theorem summarizing our similar Correlation WTree data structure. Both the Correlation DTree and Correlation WTree have a query time that is 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 WTree are deferred to Section D.3.
Theorem 3.2 (Correlation WTree data structure). There exists a data structure with the following procedures:
Algorithm 2 Correlation DTrees
1: data structure CORRELATIONDTREE  $\triangleright$  Theorem 3.1  
2: procedure QUERY  $(i\in [n],\tau \in \mathbb{R}_{\geq 0})$ $\triangleright$  Lemma D.4  
3: return FIND(τ, root(Ti))  
4: end procedure  
5: procedure FIND(τ ∈ R≥0, r ∈ T)  
6: if r is leaf then  
7: return r  
8: else  
9:  $r_1\gets$  left child of r,  $r_2\gets$  right child of r  
10: if  $r_1.\mathrm{value}\geq \tau$  then  
11:  $S_{1}\gets \mathrm{FIND}(\tau ,r_{1})$   
12: end if  
13: if  $r_2.\mathrm{value}\geq \tau$  then  
14:  $S_{2}\gets \mathrm{FIND}(\tau ,r_{2})$   
15: end if  
16: end if  
17: return  $S_{1}\cup S_{2}$   
18: end procedure  
19: end data structure
- INIT( $\{w_1, w_2, \dots, w_m\} \subset \mathbb{R}^d$ ,  $\{x_1, x_2, \dots, x_n\} \subset \mathbb{R}^d$ ,  $n \in \mathbb{N}$ ,  $m \in \mathbb{N}$ ,  $d \in \mathbb{N}$ ). Given a series of weights  $w_1, w_2, \dots, w_m$  and data  $x_1, x_2, \dots, x_n$  in  $d$ -dimensional space, it performs preprocessing in time  $O(nmd)$ .
- UPDATE  $(z \in \mathbb{R}^d, r \in [m])$ . Given a weight  $z$  and index  $r$ , it updates weight  $w_r$  to  $z$  in time  $O(nd)$ .
- QUERY  $(r \in [m], \tau \in \mathbb{R})$ . Given an index  $r$  indicating weight  $w_{r}$  and a threshold  $\tau$ , it finds all indices  $i \in [n]$  such that  $\langle w_{r}, x_{i} \rangle > \tau$  in time  $O(|S(\tau)| \cdot \log m)$ , where  $S(\tau) := \{i : \langle w_{r}, x_{i} \rangle > \tau\}$ .
4 Running Time of Our Algorithm
In this section, we show how to apply the Correlation Tree data structures developed in Section 3 to speed up neural network training.
4.1 Weights Preprocessing
Algorithm 3 Training Neural Network based on Correlation DTree
1: procedure TRAININGWITHDTREE((  $\{ (x_{i},y_{i})\}_{i\in [n],n,m,d}$  ） ▷ Theorem 4.1   
2: Initialize  $w_{r},a_{r}$  for  $r\in [m]$  and  $b$  according to Section 2   
3: DTREE.INIT(  $\{w_{r}(0)\}_{r\in [m]},m,d)$  ） ▷ Algorithm 10   
4: for  $t = 1\to T$  do   
5:  $S_{i,\mathrm{fire}}\gets \mathrm{DTree.SEQRY}(x_i,b)$  for  $i\in [n]$    
6: Forward pass for  $x_{i}$  only on neurons in  $S_{i,\mathrm{fire}}$  for  $i\in [n]$    
7: Calculate gradient for  $x_{i}$  only on neurons in  $S_{i,\mathrm{fire}}$  for  $i\in [n]$    
8: Gradient update for the neurons in  $\cup_{i\in [n]}S_{i,\mathrm{fire}}$    
9: DTREE.UPDATE(wr(t+1),r)   
10: end for   
11: return Trained weights  $w_{r}(T + 1)$  for  $r\in [m]$    
12: end procedure
In Algorithm 8, we use DTree structure to speed up the training process. We preprocess weights  $w_{r}, r \in [m]$  for each data point  $x_{i}, i \in [n]$  by constructing  $n$  weight-data correlation trees. In each
iteration, QUERY finds the set of activated neurons  $S_{i,\mathrm{fire}}$  (Definition 2.4) efficiently for each data point  $x_{i}$  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.1). Given  $n$  data points in  $\mathbb{R}^d$ , gradient descent using the DTree data structure (Algorithm 8) for the neural network  $2\mathrm{NN}(m,b = \sqrt{0.4\log m})$  (Definition 2.1) takes  $O(m^{4/5}n^2 d)$  time per iteration in expectation.
4.2 Data Preprocessing
Algorithm 4 Training Neural Network based on Correlation WTree   
1: procedure TRAININGWITHWTREE((  $\{ ( x _ { i } , y _ { i } ) \} _ { i \in [ n ] } , n , m , d$  ） ▷ Theorem 4.2   
2: Initialize  $w_{r},a_{r}$  for  $r\in [m]$  and  $b$  according to Section 2   
3: WTREE.INIT(  $\{x_i\}_{i\in [n]},n,d$  ） ▷ Algorithm 13   
4:  $\widetilde{S}_{r,\mathrm{fire}}\gets \mathrm{WT.QUERY}(w_r(0),b)$  for  $r\in [m]$  ▷ Data points fire set   
5:  $S_{i,\mathrm{fire}}\gets \{r\mid i\in \widetilde{S}_r,\mathrm{fire}\}$  ▷ Hidden neurons fire set   
6: for  $t = 1\to T$  do   
7: Forward pass for  $x_{i}$  only on neurons in  $S_{i,\mathrm{fire}}$  for  $i\in [n]$    
8: Calculate gradient for  $x_{i}$  only on neurons in  $S_{i,\mathrm{fire}}$  for  $i\in [n]$    
9: for  $r\in \cup_{i\in [n]}S_{i,\mathrm{fire}}$  do   
10:  $\widetilde{S}_{r,\mathrm{fire}}\gets \mathrm{WTREE.QUERY}(w_r(t + 1),b)$    
11: Update  $S_{i,\mathrm{fire}}$  for each  $i\in \widetilde{S}_r,\mathrm{fire}$    
12: end for   
13: end for   
14: return Trained weights  $w_{r}(T + 1)$  for  $r\in [m]$    
15: end procedure
Preprocessing weights based on data points is a common practice for neural networks. Here we consider its dual form: preprocessing input data  $x_{i}, i \in [n]$  based on neural network weights  $w_{r}, r \in [m]$ . This can be easily done due to the symmetric property of the inner product that we used in the correlation tree structure.
Given a weight vector  $w_{r}$ , we can quickly find  $\widetilde{S}_{i,\mathrm{fire}}$  (Definition 2.4) which contains the indices of data points that "fire" for weight  $w_{r}$ . By the dual relationship between  $\widetilde{S}_{i,\mathrm{fire}}$  and  $S_{i,\mathrm{fire}}$ , we can recover  $S_{i,\mathrm{fire}}$  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  $n$  data points in  $\mathbb{R}^d$ , gradient descent algorithm using the WTree data structure (Algorithm 9) for the neural network  $2\mathrm{NN}(m,b = \sqrt{0.4\log m})$  takes  $O(m^{4/5}n \cdot \log n)$ -time per iteration to initialize  $\tilde{S}_{r,\mathrm{fire}}$ ,  $S_{i,\mathrm{fire}}$  for  $r \in [m]$ ,  $i \in [n]$ , and the total running time per iteration is
$$
O (m ^ {4 / 5} n ^ {2} d)
$$
in expectation.
Deep neural networks are becoming larger every year to 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  $\Theta(nmd)$  time to evaluate the  $m$  neurons on  $n$  data points with dimension  $d$ . Recent work [SYZ21] reduced the per-iteration cost to  $o(nmd)$ , but required exponential time to preprocess either the data or the neural weights. We develop a new method that reduces the preprocessing cost to  $O(nmd)$  while keeping the
per-iteration running time at  $o(nmd)$ . One limitation of our algorithm is that it has an  $n^2$  dependence in the cost-per-iteration. However, for very wide neural networks (with  $m \gg 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 1 and 14) can be parallelized, where we can update all the correlation trees using distributed computing simultaneously. It will improve the running time from  $O(nd)$  to  $O(d)$ , resulting in a total running time  $O(m^{4/5}nd)$  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, SZZ21], to speed up neural network training with more than two layers? Giving a provable, theoretical backing for quickly training multi-layer networks remains an open, difficult challenge.
- Second, many empirical results (e.g.,  $\mathrm{[CMF^{+}20,CLP^{+}21]}$ ) indicate that only approximately identifying the activated neurons (i.e., neurons with top- $k$  inner products) in each iteration 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 our data structures will work for more general 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.
Acknowledgements The authors would like to thank Lichen Zhang for his helpful discussions. JA was partly supported by a 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. DMS-1925919).
[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 Dagstuhl-Leibniz-Zentrum für Informatik, 2022.
$\left[\mathrm{ABH}^{+}18\right]$  Amir Abboud, Arturs Backurs, Thomas Dueholm Hansen, Virginia Vassilevska 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 lcs 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 61st Annual Symposium on Foundations of Computer Science (FOCS), pages 541-552. IEEE, 2020.
[ADBMS98] Pankaj K Agarwal, Mark De Berg, Jiri Matousek, and Ottfried Schwarzkopf. Constructing 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.
$\left[\mathrm{ZPD}^{+}20\right]$  Yi Zhang, Orestis Plevrakis, Simon S Du, Xingguo Li, Zhao Song, and Sanjeev Arora. Over-parameterized adversarial training: An analysis overcoming the curse of dimensionality. In NeurIPS. arXiv preprint arXiv:2002.06668, 2020.
[  $\mathrm{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.  $\mathrm{H}_2\mathrm{O}$  : Heavy-hitter oracle for efficient generative inference of large language models. In NeurIPS. arXiv preprint arXiv:2306.14048, 2023.
Roadmap. We restate our 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 a  $2\mathrm{NN}(m,b)$  using DTree and WTree. Section F shows a formal version of lower bound for dynamic detection of firing neurons.
A Preliminary
A.1 Basic Notation
For any positive integer  $n$ , we use  $[n]$  to denote the set  $\{1, 2, \dots, n\}$ . We use  $\mathbb{E}[X]$  to denote the expected value of a random variable  $X$ , and  $\operatorname*{Pr}[Y]$  to denote the probability of a random event  $Y$ . For a matrix  $M$ , we write  $M^{\top}$  to denote the transpose of  $M$ . We use  $x^{\top}y$  to denote the inner product between vectors  $x$  and  $y$ . We use  $I_d$  to denote a  $d$ -dimensional identity matrix. We use  $\mathcal{N}(\mu; \sigma^2)$  to denote the Gaussian distribution with mean  $\mu$  and variance  $\sigma^2$ .
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 [DZPS19], Lemma 3.8 in [SY19]). Let  $\mathrm{err}(i)$  be defined as Definition A.2. If  $\forall i\in [t]$ ,  $\| \mathrm{err}(i)\| _2^2\leq (1 - \eta \lambda /2)^i\cdot \| \mathrm{err}(0)\| _2^2$ , then
$$
\left\| W (t + 1) - W _ {r} (0) \right\| _ {\infty , 2} \leq 4 \lambda^ {- 1} m ^ {- 1 / 2} \cdot \sqrt {n} \cdot \| \operatorname {e r r} (0) \| _ {2} := D.
$$
This claim shows a uniform bound on the movement of weights.
Next, we introduce the definition of error of prediction.
Definition A.2 (Error of prediction). For each  $t \in \{0, 1, \dots, T\}$ , we define  $\operatorname{err}(t) \in \mathbb{R}^n$  to be the error of prediction  $\operatorname{err}(t) = y - u(t)$ , where  $u(t) \coloneqq f(W(t), a, X) \in \mathbb{R}^n$
A.3 Probabilities
We introduce the classical Bernstein inequality here.
Lemma A.3 (Bernstein inequality [Ber24]). Assume  $Z_{1}, \dots, Z_{n}$  are  $n$  i.i.d. random variables.  $\forall i \in [n], \mathbb{E}[Z_i] = 0$  and  $|Z_i| \leq M$  almost surely. Let  $Z = \sum_{i=1}^{n} Z_i$ . Then,
$$
\Pr \left[ Z > t \right] \leq \exp \left(- \frac {t ^ {2} / 2}{\sum_ {j = 1} ^ {n} \mathbb {E} \left[ Z _ {j} ^ {2} \right] + M t / 3}\right), \forall t > 0.
$$
Next, we show an inequality on a shifted small ball with a Gaussian distribution.
Claim A.4 (Theorem 3.1 in [LS01]). Let  $b > 0$  and  $r > 0$ . Then,
$$
\exp (- b ^ {2} / 2) \underset {x \sim \mathcal {N} (0, 1)} {\operatorname * {P r}} [ | x | \leq r ] \leq \underset {x \sim \mathcal {N} (0, 1)} {\operatorname * {P r}} [ | x - b | \leq r ] \leq \underset {x \sim \mathcal {N} (0, 1)} {\operatorname * {P r}} [ | x | \leq r ].
$$
We state the anti-concentration inequality here.
Lemma A.5 (Anti-concentration for Gaussian distribution). Let  $Z \sim \mathcal{N}(0, \sigma^2)$ . Then, for  $t > 0$ ,
$$
\Pr [ | Z | \leq t ] \leq \frac {2 t}{\sqrt {2 \pi} \sigma}.
$$
B Sparsity
In this section, we start by restating the result about sparsity after initialization in Section 2. Then we show how to bound the number of fire neurons per iteration in Section B.2.
B.1 Sparsity after initialization
The goal of this section is 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 a query point  $x \in \mathbb{R}^d$ , let  $S_{\mathrm{fire}}(x)$  denote the set of neurons that are "fire", i.e.,
$$
\mathcal {S} _ {\text {f i r e}} (x) := \{i \in [ m ]: \langle x, w _ {i} \rangle > b \}
$$
Next, we introduce the fire set for each training iteration.
Definition B.2 (fire set per iteration, Definition 3.7 in [SYZ21]). For each data point  $x_{i} \in \mathbb{R}^{d}, i \in [n]$ , weight  $w_{r} \in \mathbb{R}^{d}, r \in [m]$  and each iteration  $t \in \{0,1,\dots,T\}$ , we define
$$
S _ {i, \text {f i r e}} (t) := r \in [ m ]: \langle x _ {i}, w _ {r} (t) \rangle
$$
$$
\widetilde {S} _ {r, \text {f i r e}} (t) := i \in [ n ]: \langle x _ {i}, w _ {r} (t) \rangle
$$
Also, we define  $k_{i,t} \coloneqq |S_{i,\mathrm{fire}}(t)|$  and  $\widetilde{k}_{r,t} \coloneqq |\widetilde{S}_{r,\mathrm{fire}}(t)|$
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 > 0$  be a tunable parameter. If we use the  $\Phi_b$  as the activation function, then after the initialization, with probability at least  $1 - \exp(-\Omega(m \cdot \exp(-b^2/2))$ ), it holds that for input data  $x$ , the number of activated neurons  $k_x$  is at most  $O(m \cdot \exp(-b^2/2))$ , where  $m$  is the total number of neurons.
B.2 Bounding the number of fired neuron per iteration per level
In this section, we will show that for  $t = 0,1,\dots ,T,k = 0,1,\dots ,\log m$ , the number of fire neurons  $k_{i,k,t} = |\mathcal{S}_{i,k,\mathrm{fire}}(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  $i \in [n]$ , for each time  $t \in [T]$  let  $S_{i,\mathrm{flip}}(t) \subset [m]$  denote the set of neurons that are never flipped during the entire training process,
$$
\mathcal {S} _ {i, \text {f l i p}} (t) := \{r \in [ m ]: \operatorname {s g n} \left(\left\langle w _ {r} (t), x _ {i} \right\rangle - b\right) \neq \operatorname {s g n} \left(\left\langle w _ {r} (t - 1), x _ {i} \right\rangle - b\right) \}.
$$
Over all the iterations of training algorithm, there are some neurons that never flip states. We provide a mathematical formulation of that set,
Definition B.5 (noflip set, Definition C.9 in [SYZ21]). For each  $i \in [n]$ , let  $S_i \subset [m]$  denote the set of neurons that are never flipped during the entire training process,
$$
S _ {i} := \{r \in [ m ]: \forall t \in [ T ] \operatorname {s g n} \left(\left\langle w _ {r} (t), x _ {i} \right\rangle - b\right) = \operatorname {s g n} \left(\left\langle w _ {r} (0), x _ {i} \right\rangle - b\right) \}. \tag {1}
$$
In Lemma 2.5, we already show that  $k_{i,0} = O(m \cdot \exp(-b^2 / 2))$  for all  $i \in [n]$  with high probability. We can show that it also holds for  $t > 0$ .
Lemma B.6 (Bounding the number of fired neuron per iteration, Lemma C.10 in [SYZ21]). Let  $b \geq 0$  be a parameter, and let  $\sigma_{b}(x) = \max \{x, b\}$  be the activation function. For each  $i \in [n]$ ,  $t \in [T]$ ,  $k_{i,t}$  is the number of activated neurons at the  $t$ -th iteration. For  $0 < t \leq T$ , with probability at least  $1 - n \cdot \exp \left(-\Omega(m) \cdot \min \{R, \exp(-b^{2}/2)\}\right)$ ,  $k_{i,t}$  is at most  $O(m \exp(-b^{2}/2))$  for all  $i \in [n]$ .
In this section, we explain the procedures for two correlation tree data structure. We use the same setting as Section 3
For the DTree data structure, it contains  $n$  binary trees indexed by  $n$  data points and supports the following operations: