SlowGuess/NeurIPS_2023_OCR huggingface dataset:
- Initialize It takes data points  $\{x_1, \dots, x_n\} \subset \mathbb{R}^d$  and weights  $\{w_1, \dots, w_m\} \subset \mathbb{R}^d$  as input and compute inner products of all weight-data pairs  $(w_r, x_i)$ . It uses these inner products to create  $n$  different trees. For the  $i$ -th tree based on data point  $x_i$ , it is constructed from  $m$  leaf nodes  $\langle w_r, x_i \rangle$ ,  $r \in [m]$  and satisfies the property that the value of parent node is the maximum value of its child nodes.
- Update It takes a new weight  $z \in \mathbb{R}^d$  and an index  $r \in [m]$  as input. For the  $i$ -th tree, it calculates the new inner product  $\langle z, x_i \rangle$  and stores the value into the  $r$ -th leaf node. Then it compares the new value with its parent node. It replaces parent node with new value if it is larger and continues this comparing process. Otherwise it stops. Repeat the same operation for all  $n$  trees.
- **Query** It takes a threshold  $b \in \mathbb{R}_{\geq 0}$  and an index  $i \in [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 repeats this searching process until stop. Then it finds all indices  $r \in [m]$  that satisfy  $\{w_r : \operatorname{sgn}(\langle w_r, x_i \rangle - b) \geq 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  $\{x_1,\dots ,x_n\} \subset \mathbb{R}^d$  and weights  $\{w_{1},\dots ,w_{m}\} \subset \mathbb{R}^{d}$  as input and compute inner products of all weight-data pairs  $(w_{r},x_{i})$ . It uses these inner products to create  $nm$  different trees. For the  $r$ -th tree based on weight  $w_{i}$ , it is constructed from  $n$  leaf nodes  $\langle w_r,x_i\rangle$ ,  $i\in [n]$  and satisfies the property that the value of parent node is the maximum value of its child nodes.
- Update It takes a new weight  $z \in \mathbb{R}^d$  and an index  $r \in [m]$  as input. Then it re-constructs the  $r$ -th tree with weight  $z$ .
- Query It takes a threshold  $b \in \mathbb{R}_{\geq 0}$  and an index  $r \in [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, it treats the child nodes as the root of a new subtree and repeats this searching process until stop. Then it finds all indices  $i \in [n]$  that satisfy  $\{w_r : \operatorname{sgn}(\langle w_r, x_i \rangle - b) \geq 0\}$ .
Algorithm 5 Correlation DTree Data Structure
1: data structure CORRELATIONDTREE  
2: procedures:  
3:  $\mathrm{INIT}(S\subset \mathbb{R}^d,W\subset \mathbb{R}^d,n,m,d)$  ▷ Initialize the data structure via building  $n$  trees  
4: QUERY(i,b) ▷  $i\in [n],b\in \mathbb{R}$ . Output the set  $\{r\in [m]:\operatorname {sgn}(\langle w_r,x_i\rangle -b)\geq 0\}$   
5: UPDATE(x,i) ▷ Update the  $i$ th point in  $\mathbb{R}^d$  with  $x$   
6: end data structure
Algorithm 6 Correlation WTree Data Structure
1: data structure CORRELATIONWTREE  
2: procedures:  
3: INIT(S ⊂ Rd, W ⊂ Rd, n, m, d) ▷ Initialize the data structure via building m trees  
4: QUERY(r, b) ▷ r ∈ [m], b ∈ R. Output the set {i ∈ [n] : sgn(⟨wr, xi⟩ - b) ≥ 0}  
5: UPDATE(w, r) ▷ Update the r'th point in Rd with w  
6: end data structure
We present MAKEMAXTREE algorithm (Algorithm 7) which shows how to construct a 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.
Algorithm 7 Make MaxTree
1: procedure MAKEMAXTREEINNER  $(r_1,\dots ,r_n)$    
2: if  $n = 1$  then   
3: return  $r_1$    
4: else   
5: for  $i\in [n / 2]$  do   
6: Create node  $r_i^{\prime}$    
7: if  $r_{2i - 1}.$  value  $>$  r2i.value then   
8:  $r_i^\prime \gets r_{2i - 1}$    
9: else   
10:  $r_i^\prime \gets r_{2i}$    
11: end if   
12: Insert  $r_{2i - 1}$  as left child   
13: Insert  $r_{2i}$  as right child   
14: end for   
15: If  $n$  is odd, create a parent node for the last node.   
16: return MAKEMAXTREEINNER(\{r'1,,r'}))   
17: end if   
18: end procedure   
19: procedure MAKEMAXTREE(u1,..,un)   
20: for  $i\in [n]$  do   
21: Create nodes  $r_i$    
22:  $r_i$  .value  $\leftarrow u_{i}$    
23: end for   
24: return MAKEMAXTREEINNER(r1,..,rn)   
25: end procedure
D.1 Correlation DTree data structure
We 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_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 datas  $x_1, x_2, \dots, x_n$  in  $d$ -dimensional space, it preprocesses 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$  with  $z$  in time  $O(n \cdot (d + \log m))$
-  $\mathrm{QUERY}(i \in [n], \tau \in \mathbb{R})$ . Given an index  $i$  indicating data point  $x_i$  and a threshold  $\tau$ , it finds all index  $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\}$
D.2 Running time for CORRELATIONDTREE
The goal of this section is 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 a series of weights  $\{w_{1}, w_{2}, \dots, w_{m}\} \subset \mathbb{R}^{d}$  and datas  $\{x_{1}, x_{2}, \dots, x_{n}\} \subset \mathbb{R}^{d}$ , it preprocesses in time  $O(nmd)$
Proof. The INIT consists of two independent forloop and two recursive forloops. The first forloop (start from line 8) has  $n$  iterations, which takes  $O(n)$  time. The second forloop (start from line 11) has  $m$  iterations, which takes  $O(m)$  time. Now we consider the recursive forloop. The outer loop (line 14) has  $n$  iterations. In inner loop has  $m$  iterations. In each iteration of the inner loop, line 16 takes  $O(d)$  time. Line 18 takes  $O(m)$  time. Putting it all together, the running time of INIT is
$$
O (n + m + n (m d + m))
$$
Algorithm 8 Training Neural Network via building  $n$  trees, where each tree is the correlation between one data point and all the weights
1: procedure TRAININGWITHPREPROCESSWEIGHTS(\{xi\}_{i \in [n]}, \{yi\}_{i \in [n]}, n, m, d) ▷ Theorem 4.1  
2: /*Initialization step*/  
3: Sample W(0) and a according to Definition  
4: b ← √0.4 log m.  
5: /*A dynamic data-structure*/  
6: CORRELATIONDTREE CDT ▷ Theorem 3.1  
7: CDT.INIT(\{xi\}_{i \in [n]}, \{wr(0)\}_{r \in [m]}, n, m, d) ▷ It takes Tinit(n, m, d) time. Alg. 10  
8: /*Iterative step*/  
9: for t = 0 → T do  
10: /*Forward computation step*/  
11: for i = 1 → n do  
12: Si,fire ← CDT QUERY(i, b) ▷ It takes Tquery(m, ki, t) time. Alg. 12  
13: u(t)i ← 1/√m ∑r∈Si,fire ar · σb(wr(t)^\top xi) ▷ It takes O(d · ki, t) time  
14: end for  
15: /*Backward computation step*/  
16: P ← 0^{n \times m} ▷ P ∈ R^{n \times m}  
17: for i = 1 → n do  
18: for r ∈ Si,fire do  
19: Pi,r ← 1/√m ar · σb'(wr(t)^\top xi)  
20: end for  
21: end for  
22: M ← X diag(y - u(t)) ▷ M ∈ Rd  $\times$  n, it takes O(n · d) time  
23: ΔW ← M P ▷ ΔW ∈ Rd  $\times$  m, it takes O(d · nnz(P)) time, nnz(P) = O(nm^4/5)  
24: W(t + 1) ← W(t) - η · ΔW.  
25: /*Update data structure*/  
26: Let Q  $\subset$  [m] where for each r ∈ Q, the ΔW*,r is not all zeros ▷ |Q| ≤ O(nm^4/5)  
27: for r ∈ Q do  
28: CDT UPDATE(wr(t + 1), r) ▷ Alg. 11  
29: end for  
30: end for  
31: return W ▷ W ∈ R^{d \times m}
$$
= O (n m d)
$$
Next, we analyze the running time of UPDATE.
Lemma D.3 (Running time of UPDATE). Given a weight  $z \in \mathbb{R}^d$  and index  $j \in [m]$ , it updates weight  $w_j$  with  $z$  in time  $O(n \cdot (d + \log m))$
Proof. The running time of UPDATE mainly comes from the forloop (line 23), which consists of  $n$  iterations. In each iteration, line 24 takes  $O(\log m)$  time, line 25 takes  $O(d)$  time and the while loop takes  $O(\log m)$  time since it goes through a path bottom up. Putting it together, the running time of UPDATE is  $O(n(d + \log m))$ .
Finally, we state the running time for QUERY procedure.
Lemma D.4 (Running time of QUERY). Given a query  $q \in \mathbb{R}^d$  and a threshold  $\tau > 0$ , it finds all index  $i \in [n]$  such that  $\langle w_i, q \rangle > \tau$  in time  $O(|S(\tau)| \cdot \log m)$ , where  $S(\tau) \coloneqq \{i : \langle w_i, q \rangle > \tau\}$
Proof. The running time comes from FIND with input  $\tau$  and  $\operatorname{root}(T_i)$ . 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 r.value  $> \tau$ , the sub-tree with root  $r$  must contains a leaf whose value is greater than  $\tau$  If not satisfied, all the values of the nodes in the sub-tree with root  $r$  is less than  $\tau$ . This guarantees that all the paths it search does not have any branches that leads to the leaf we don't want and it will report all the
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 TRAININGWITHPROCESSDATA(\{xi\}_{i \in [n]}, \{yi\}_{i \in [n]}, n, m, d) ▷ Theorem 4.2  
2: /*Initialization step*/  
3: Sample W(0) and a according to Definition  
4: b ← √0.4 log m.  
5: /*A static data-structure*/  
6: CORRELATIONWTREE CWT ▷ Algorithm 10, Part 2 of Theorem 3.1  
7: CWT.INIT(\{xi\}_{i \in [n]}, \{wr(0)\}_{r \in [m]}, n, m, d) ▷ It takes Tinit(n, m, d) time  
8: /*Initialize S_r, fire and Si, fire */  
9: ▷ It takes ∑_{r=1}^{m} T_query(n, k_r, t) = O(m^4/5 n · log n) time  
10: S_r, fire ← ∅ for r ∈ [m]. ▷ S_r, fire is the set of samples, for which neuron r fires  
11: Si, fire ← ∅ for i ∈ [n]. ▷ Si, fire is the set of neurons, which fire for xi  
12: for r = 1 → m do  
13: S_r, fire ← CWT QUERY(r, b)  
14: for i ∈ S_r, fire do  
15: Si, fire. ADD(r)  
16: end for  
17: end for  
18: /*Iterative step*/  
19: for t = 1 → T do  
20: /*Forward computation step*/  
21: for i = 1 → n do  
22: u(t)_i ← 1/√m ∑_r∈Si, fire ar·σb(w_t) x_i) ▷ It takes O(d·ki, t) time  
23: end for  
24: /*Backward computation step*/  
25: P ← 0^{n×m} ▷ P ∈ R^{n×m}  
26: for i = 1 → n do  
27: for r ∈ Si, fire do  
28: Pi_r ← 1/√m ar·σb(w_t) x_i)  
end for  
end for  
M ← X diag(y - u(t)) ▷ M ∈ R^{d×n}, it takes O(n·d) time  
ΔW ← M P d×n n×m ▷ ΔW ∈ R^{d×m}, it takes O(d·nnz(P)) time, nnz(P) = O(nm^4/5)  
W(t + 1) ← W(t) - η · ΔW.  
/*Update S_r, fire and Si, fire step*/  
S[n], fire ← ∪_{i \in [n]} Si, fire  
for r ∈ S[n], fire do  
for i ∈ S_r, fire do ▷ Removing old fired neuron indices. It takes O(k_r, t) time  
Si, fire. DEL(r)  
end for  
CWT_UPDATE(w_t(t + 1), r) ▷ It takes T_update(n, d) time  
S_r, fire ← CWT QUERY(r, b) ▷ It takes T_query(n, d, k_r, t + 1) time  
for i ∈ S_r, fire do ▷ Adding new fired neuron indices. It takes O(k_r, t + 1) time  
Si, fire. ADD(r)  
end for  
end for  
end for  
return W ▷ W ∈ R^{d×m}
indices  $i$  satisfying  $\langle w_i, q \rangle > 0$ . Note that the depth of  $T$  is  $O(\log n)$ , the running time of QUERY is  $O(|S(\tau)| \cdot \log n)$
D.3 Correlation WTree data structure
In this section, we state the main theorem of correlation wtree data structure.
Algorithm 10 Correlation DTree data structure
1: data structure CORRELATIONDTREE  $\triangleright$  Theorem D.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:   
8: public:   
9: 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)$  Lemma D.2   
10: for  $i = 1\rightarrow n$  do   
11:  $x_{i}\gets x_{i}$    
12: end for   
13: for  $j = 1\to m$  do   
14:  $w_{j}\leftarrow w_{j}$    
15: end for   
16: for  $i = 1\to n$  do  $\triangleright$  for data point, we create a tree   
17: for  $j = 1\rightarrow m$  do   
18:  $u_{j}\gets \langle x_{i},w_{j}\rangle$    
19: end for   
20:  $T_{i}\gets \mathrm{MAKETREE}(u_{1},\dots ,u_{m})$ $\triangleright$  Each node stores the maximum value for his two children   
21: end for   
22: end procedure   
23: end data structure
Algorithm 11 Correlation DTrees
1: data structure CORRELATIONTREE  $\triangleright$  Theorem D.1   
2: public:   
3: procedure UPDATE  $(z\in \mathbb{R}^d,r\in [m])$ $\triangleright$  Lemma D.3   
4:  $w_{r}\gets z$    
5: for  $i = 1\rightarrow n$  do   
6:  $l\gets$  the  $l$  -th leaf of tree  $T_{i}$    
7:  $l.\mathrm{value} = \langle z,x_i\rangle$    
8: while  $l$  is not root do   
9:  $p\gets$  parent of  $l$    
10:  $p.\mathrm{value}\gets \max \{p.\mathrm{value},l.\mathrm{value}\}$    
11:  $l\gets p$    
12: end while   
13: end for   
14: end procedure   
15: end data structure
Theorem D.5 (Correlation WTree data structure). There exists a data structure with the following procedures:
- 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 datas  $x_1, x_2, \dots, x_n$  in  $d$ -dimensional space, it preprocesses 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$  with  $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 index  $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\}$
Algorithm 12 Correlation DTrees
1: data structure CORRELATIONDTREE  $\triangleright$  Theorem D.1   
2: public:   
3: procedure QUERY(i ∈ [n],τ ∈ R≥0)  $\triangleright$  Lemma D.4   
4: return FIND(τ, root(Ti))   
5: end procedure   
6:   
7: private:   
8: procedure FIND(τ ∈ R≥0, r ∈ T)   
9: if r is leaf then   
10: return r   
11: else   
12:  $r_1\gets$  left child of  $r$ ,  $r_2\gets$  right child of  $r$    
13: if  $r_1.$  value ≥ τ then   
14:  $S_{1}\leftarrow \mathrm{FIND}(\tau ,r_{1})$    
15: end if   
16: if  $r_2.$  value ≥ τ then   
17:  $S_{2}\leftarrow \mathrm{FIND}(\tau ,r_{2})$    
18: end if   
19: end if   
20: return  $S_{1}\cup S_{2}$    
21: end procedure   
22: end data structure
Algorithm 13 Correlation WTree data structure
1: data structure CORRELATIONWTREE ▷ Theorem D.5  
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_{M}$  ▷ m binary search trees  
6: end members  
7:  
8: public:  
9: procedure 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$ ) ▷ Lemma D.6  
10: for  $i = 1 \rightarrow n$  do  
11:  $x_{i} \gets x_{i}$   
12: end for  
13: for  $j = 1 \rightarrow m$  do  
14:  $w_{j} \gets w_{j}$   
15: end for  
16: for  $i = 1 \rightarrow m$  do ▷ for weight, we create a tree  
17: for  $j = 1 \rightarrow n$  do  
18:  $u_{j} \gets \langle x_{i}, w_{j} \rangle$   
19: end for  
20:  $T_{i} \gets \text{MAKETREE}(u_{1}, \dots, u_{n})$  ▷ Each node stores the maximum value for his two children  
21: end for  
22: end procedure  
23: end data structure
D.4 Running time for Correlation WTree
As in DTree, we first show the running time for INIT.
Lemma D.6 (Running time of INIT). Given a series of weights  $\{w_{1}, w_{2}, \dots, w_{m}\} \subset \mathbb{R}^{d}$  and datas  $\{x_{1}, x_{2}, \dots, x_{n}\} \subset \mathbb{R}^{d}$ , it preprocesses in time  $O(nmd)$
Algorithm 14 Correlation WTrees
1: data structure CORRELATIONWTREE ▷ Theorem D.5  
2: public:  
3: procedure UPDATE  $(z\in \mathbb{R}^d,r\in [m])$  ▷ Lemma D.7  
4:  $w_{r}\gets z$   
5: for  $j = 1\to n$  do  
6:  $u_{j}\gets \langle x_{j},w_{r}\rangle$   
7:  $T_{i}\gets \mathrm{MAKETREE}(u_{1},\dots ,u_{n})$  ▷ Each node stores the maximum value for his two children  
8: end for  
9: end procedure  
10: end data structure
Algorithm 15 Correlation WTree
1: data structure CORRELATIONWTREE  
2: public:  
3: procedure QUERY  $(r\in [m],\tau \in \mathbb{R}_{\geq 0})$   
4: return FIND(τ, root  $(T_r))$   
5: end procedure  
6:  
7: private:  
8: procedure FIND(τ ∈ R≥0, r ∈ T)  
9: if r is leaf then  
10: return r  
11: else  
12:  $r_1\gets$  left child of r,  $r_2\gets$  right child of r  
13: if  $r_1.$  value ≥ τ then  
14:  $S_{1}\gets \mathrm{FIND}(\tau ,r_{1})$   
15: end if  
16: if  $r_2.$  value ≥ τ then  
17:  $S_{2}\gets \mathrm{FIND}(\tau ,r_{2})$   
18: end if  
19: end if  
20: return  $S_{1}\cup S_{2}$   
21: end procedure  
22: end data structure
Proof. The INIT consists of two independent forloop and two recursive forloops. The first forloop (start from line 10) has  $n$  iterations, which takes  $O(n)$  time. The second forloop (start from line 13) has  $m$  iterations, which takes  $O(m)$  time. Now we consider the recursive forloop. The outer loop (line 16) has  $m$  iterations. In inner loop has  $n$  iterations. In each iteration of the inner loop, line 18 takes  $O(d)$  time. Line 20 takes  $O(n)$  time. Putting it all together, the running time of INIT is
$$
\begin{array}{l} O (n + m + m (n d + n)) \\ = O (n m d) \\ \end{array}
$$
Next, we turn to the running time for UPDATE.
Lemma D.7 (Running time of UPDATE). Given a weight  $z \in \mathbb{R}^d$  and index  $r \in [m]$ , it updates weight  $w_j$  with  $z$  in time  $O(nd)$
Proof. In this procedure, it generates a new tree for weight  $w_r$  with  $n$  leaves, which takes  $O(nd)$  time. Thus, we complete the proof.
Finally, we present the running time of QUERY.
Lemma D.8 (Running time of QUERY). Given a query  $q \in \mathbb{R}^d$  and a threshold  $\tau > 0$ , it finds all index  $i \in [n]$  such that  $\langle w_i, q \rangle > \tau$  in time  $O(|S(\tau)| \cdot \log m)$ , where  $S(\tau) \coloneqq \{i : \langle w_i, q \rangle > \tau\}$
Proof. The running time comes from FIND with input  $\tau$  and  $\operatorname{root}(T_i)$ . 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 r.value  $> \tau$ , the sub-tree with root  $r$  must contains a leaf whose value is greater than  $\tau$ . If not satisfied, all the values of the nodes in the sub-tree with root  $r$  is less than  $\tau$ . This guarantees that all the paths it search does not have any branches that leads to the leaf we don't want and it will report all the index  $i$  satisfying  $\langle w_i, q \rangle > 0$ . Note that the depth of  $T$  is  $O(\log n)$ , the running time of QUERY is  $O(|S(\tau)| \cdot \log n)$
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 preprocessing weights for each data point.
Theorem E.1 (Running time part, formal version of Theorem 4.1). Given  $n$  data points in  $\mathbb{R}^d$ , Running gradient descent algorithm (Algorithm 8) on  $2\mathrm{NN}(m,b = \sqrt{0.4\log m})$  (Definition 2.1) the expected cost per-iteration of the gradient descent algorithm is
Proof. The per-step time complexity is
$$
\begin{array}{l} \mathcal {T} = \mathcal {T} _ {1} + \mathcal {T} _ {2} + \mathcal {T} _ {3} \\ = \sum_ {i = 1} ^ {n} \mathcal {T} _ {\text {Q U E R Y}} (m, d, k _ {i, t}) + \mathcal {T} _ {\text {U P D A T E}} \cdot | \cup_ {i \in [ n ]} S _ {i, \text {f i r e}} (t) | + d \sum_ {i \in [ n ]} k _ {i, t} \\ \end{array}
$$
The first term  $\mathcal{T}_1 = \sum_{i=1}^n \mathcal{T}_{\mathrm{QUERY}}(m, d, k_{i,t})$  corresponds to the running time of querying the active neuron set  $S_{i,\mathrm{fire}}(t)$  for all training samples  $i \in [n]$ . With the first result in Theorem 3.1, the complexity is bounded by  $O(m^{4/5} n \log m)$ .
The second term  $\mathcal{T}_2 = \mathcal{T}_{\mathrm{UPDATE}}\cdot |\cup_{i\in [n]}S_{i,\mathrm{fire}}(t)|$  corresponds to updating  $w_{r}$  in the high-dimensional search data-structure (Line 28). Again with the first result in Theorem 3.1, we have  $\mathcal{T}_{\mathrm{UPDATE}} = O(n(d + \log m))$ . Combining with the fact that  $|\cup_{i\in [n]}S_{i,\mathrm{fire}}(t)|\leq |\cup_{i\in [n]}S_{i,\mathrm{fire}}(0)|\leq O(m^{4 / 5}n)$ , the second term is bounded by  $O(m^{4 / 5}n^2 d)$ .
The third term is the time complexity of gradient calculation restricted to the set  $S_{i,\mathrm{fire}}(t)$ . With the bound on  $\sum_{i\in [n]}k_{i,t}$  (Lemma B.6), we have  $d\sum_{i\in [n]}k_{i,t}\leq O(m^{4 / 5}nd)$
Putting them together, we have
$$
\begin{array}{l} \mathcal {T} \leq O \left(m ^ {4 / 5} n \log m\right) + O \left(m ^ {4 / 5} n ^ {2} d\right) + O \left(m ^ {4 / 5} n d\right) \\ = O \left(m ^ {4 / 5} n ^ {2} d\right) \\ \end{array}
$$
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  $n$  data points in  $\mathbb{R}^d$ , Running gradient descent algorithm (Algorithm 9) on  $2\mathrm{NN}(m,b = \sqrt{0.4\log m})$ , the expected per-iteration running time of initializing  $\widetilde{S}_{r,\mathrm{fire}}$ ,  $S_{i,\mathrm{fire}}$  for  $r \in [m]$ ,  $i \in [n]$  is  $O(m^{4/5}n \cdot \log n)$ . The cost per-iteration of the training algorithm is  $O(m^{4/5}n^2 d)$ .
Proof. We analyze the initialization and training parts separately.
Initialization From Line 10 to Line 17, the sets  $\widetilde{S}_{r,\mathrm{fire}}$ ,  $S_{i,\mathrm{fire}}$  for  $r \in [m]$ ,  $i \in [n]$  are initialized. For each  $r \in [m]$ , we need to query the data structure the set of data points  $x$ 's such that  $\sigma_b(w_r(0)^\top x) > 0$ . Hence the running time of this step is
$$
\begin{array}{l} \sum_ {r = 1} ^ {m} \mathcal {T} _ {\text {Q U E R Y}} (n, \widetilde {k} _ {r, 0}) = O \left(\sum_ {r = 1} ^ {m} \widetilde {k} _ {r, 0} \cdot \log n\right) \\ = O \left(\sum_ {i = 1} ^ {n} k _ {i, 0} \cdot \log n\right) \\ = O \left(m ^ {4 / 5} n \cdot \log n\right) \\ \end{array}
$$
where the second step follows from  $\sum_{r=1}^{m} \widetilde{k}_{r,0} = \sum_{i=1}^{n} k_{i,0}$ .
Training Consider training the neural network for  $T$  steps. For each step, first notice that the forward and backward computation parts (Line 21 - Line 33) are the same as previous algorithms. The time complexity is  $O(m^{4/5}n)$ .
We next show that maintaining  $\widetilde{S}_{r,\mathrm{fire}}, r \in [m]$  and  $S_{i,\mathrm{fire}}, i \in [n]$  (Line 36 - Line 45) takes  $O(m^{4/5}nd)$  time. For each fired neuron  $r \in [m]$ , we first remove the indices of data in the sets  $S_{i,\mathrm{fire}}$ , which takes time
$$
O (1) \cdot \sum_ {r \in \cup_ {i \in | n |} S _ {i, \text {f i r e}}} \widetilde {k} _ {r, t} = O (1) \cdot \sum_ {r = 1} ^ {m} \widetilde {k} _ {r, t} = O \left(m ^ {4 / 5} n\right)
$$
Then, we find the new set of  $x$ 's such that  $\sigma_b(\langle w_r(t + 1),x\rangle) > 0$  by querying the correlation tree data structure. The total ruunning time for all fired neurons is
$$
\begin{array}{l} \sum_ {r \in \cup_ {i \in [ n ]} S _ {i, \text {f i r e}}} \mathcal {T} _ {\text {U P D A T E}} (n, d) + \mathcal {T} _ {\text {Q U E R Y}} (n, \widetilde {k} _ {r, t + 1}) \lesssim m ^ {4 / 5} n ^ {2} (d + \log m) + \sum_ {r \in \cup_ {i \in [ n ]} S _ {i, \text {f i r e}}} \widetilde {k} _ {r, t + 1} \cdot \log n \\ = O (m ^ {4 / 5} n ^ {2} d) \\ \end{array}
$$