SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Josh Alman* Jiehao Liang† Zhao Song‡ Ruizhe Zhang§ Danyang Zhuo¶
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.
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 to train these models quickly has become an important research challenge.
Training a neural network is an iterative algorithm, and in each iteration, we need to process each of the  $m$  neurons on each of the  $n$  data points. Assuming each data point has a length of  $d$  (e.g.,  $d$  could be the size of an input image), this means the per-iteration training time of the straightforward algorithm is at least  $\Omega(nmd)$  just to compute the activations. As we train larger neural networks on more training data, this running time can become a significant obstacle.
Recent work by Song, Yang, and Zhang [SYZ21] gave the first training algorithm that reduces this per iteration training time to  $o(nmd)$ . The high-level idea of their algorithm is to use a 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
$^{*}$  josh@cs.columbia.edu. Columbia University.
† jiehao.liang@berkeley.edu. University of California, Berkeley.
$\ddagger$  zsong@adobe.com. Adobe Research.
$^{\S}$ ruizhe@utexas.edu. Simons Institute for the Theory of Computing.
danyang@cs.duke.edu. Duke University.
weight of the iteration. However, their algorithm's initial preprocessing time to set up the data structure is exponential in the dimension  $d$ , making it too slow in most applications. This raises a 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(nmd)$  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 to 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 over-parameterization regime (which has  $m \gg n$ ):
Theorem 1.1 (Main result). There is a data structure which preprocesses  $n$  data points in  $d$ -dimensional space, and  $m$  initialization weights points in  $d$ -dimensional space, in  $O(mnd)$  preprocessing time and  $O(mn + md + nd)$  space, which can be used to speed up neural network training: Running the gradient descent algorithm on a two-layer,  $m$ -width, over-parameterized ReLU neural network, which will minimize the training loss to zero, can be performed with an expected running time (of the gradient descent algorithm per iteration) of
$$
\widetilde {O} (m ^ {4 / 5} n ^ {2} d).
$$
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.1 and Part 1 of Corollary B.6 in [SYZ21]) has  $O(2^d)$  preprocessing time and uses  $O(m^{1 - 1 / d}nd)$  cost per iteration. Their second result (see Theorem 6.2 and Part 2 of Corollary B.6 of [SYZ21]) has  $O(n^{d})$  preprocessing time and uses  $O(m^{4 / 5}nd)$  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], we do not use any complicated geometric data structure in previous work, and our algorithms are much easier to implement (see Algorithms 1 and 2). Moreover, as we discussed in Section 5, they can be parallelized to further reduce the cost-per-iteration to  $\widetilde{O}(m^{4 / 5}nd)$ .
Our key observation is that in each iteration of the training process, the weight updates are mostly sparse, and only a small fraction of neurons are activated for each training data point. Given this observation, we construct a 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 = \{x_{1},\ldots ,x_{n}\} \subset \mathbb{Z}^{d}$ ,  $Y = \{y_{1},\dots,y_{m}\} \subset \mathbb{Z}^{d}$  and a threshold  $b\in \mathbb{R}$ , design a data structure to support the following operations:
- UPDATE  $(j \in [m], z \in \mathbb{Z}^d)$ , set  $y_j$  to  $z$
- QUERY(), either output the set
$$
Q = \{(i, j) \in [ n ] \times [ m ] \mid \langle x _ {i}, y _ {j} \rangle \geq b \},
$$
or report that  $|Q| > m^{4/5} n$ .
We give a data structure for DDFN which takes  $O(mnd)$ -time for preprocessing,  $\widetilde{O}(nd)$ -time per update, and  $O(\min\{|Q|, m^{4/5}n\})$ -time per query. At a high level, our data structure works as follows.
Preprocessing We build  $n$  binary search trees to maintain the  $(x_{i},y_{j})$  pairs for  $i\in [n]$  and  $j\in [m]$ . More specifically, the  $i$ -th tree has  $m$  leaf nodes, storing the inner-products between  $x_{i}$  and  $\{y_j\}_{j\in [m]}$ . Each 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  $O(nmd)$  time and  $O(mn)$  space.
Update Suppose we will update  $y_{j}$ . Then, for each  $i \in [n]$ , we need to update a path from the leaf node corresponding to  $(x_{i},y_{i})$  to the root, which contains  $O(\log m)$  nodes. Hence, the running time of each update is  $O(nd\log 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 at most  $b$ , we will not move further. Otherwise, we 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 a leaf if and only if its value is greater than  $b$ . Hence, the total query cost is  $O(\min\{|Q|, m^{4/5}n\})$ .
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 our running time for Dynamic Detection of Firing Neurons (DDFN). Prior work [SYZ21] got around this by using exponential preprocessing time to 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 = 2^{O(\log^* n)}$ , and assume the OVC or SETH conjecture from complexity theory. For every constant  $\varepsilon > 0$ , there is no data structure for DDFN with  $O(m^{4/5}n^{1/5 - \varepsilon})$  update time and  $O(m^{4/5}n^{6/5 - \varepsilon})$  query time.
Here,  $\log^* n$  denotes the iterated logarithm function, which grows incredibly slowly, such that the dimension  $2^{O(\log^* n)}$  is barely larger than a constant, and one would typically pick a much larger  $d$ . The complexity-theoretic 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 a 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.
Orthogonal Vector Conjecture The orthogonal vector problem (OV) is a fundamental problem in fine-grained complexity which asks, given  $X, Y \subset \{0,1\}^d$  of size  $|X| = |Y| = n$ , whether there are  $x \in X$  and  $y \in Y$  with  $\langle x,y \rangle = 0$ . The state-of-the-art algorithm [AWY14, CW16] runs in time  $n^{2 - 1 / O(\log c)}$  in dimension  $d = c\log n$  for all  $c \geq 1$ ; as the dimension  $d$  increases, its running time approaches the trivial bound  $n^2$ . The orthogonal vector conjecture (OVC) conjectures an  $n^{2 - o(1)}$  lower bound for OV when  $d = \omega (\log n)$ . It is also known that the popular Strong Exponential Time Hypothesis (SETH) regarding the hardness of  $k$ -SAT implies OVC. This conjecture has been used to obtain conditional lower bounds for other important problems with polynomial-time algorithms in a wide variety of areas, including pattern matching [AWW14, Bri14, BI15, BI16, BM16, BGL17, BK18, CW19], kernel sparsification [ACSS20, AA22], softmax attention computation [AS23a, AS23b, DSZ23], graph algorithms [RVW13, ABH+18, GIKW18, KT18, DLW22, CWX22], and computational geometry [BBK+16, Rub18, Wil18a, Che20, KM20]; see also the survey [Wil18b].
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 (e.g., small  $\ell_2$  distance [DIIM04, AR15, AIL $^{+}$ 15, ARN17, Raz17, AIR18, BIW19, DIRW20] or large inner product [SL14, SL15b, SL15a]) to a query  $q \in \mathbb{R}^d$  in a given set of points  $S \subset \mathbb{R}^d$ . LSH-based algorithms typically run quickly in practice, but only support approximate nearest-neighbor queries. The other approach is based on space partitioning data structures, such as partition trees [Mat92a, Mat92b, AEM92, AC09, Cha12],  $k$ - $d$  trees/range trees [CT17, TOG17, Cha19], 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  $\left[\mathrm{CMF}^{+}20\right]$  uses LSH-based methods to efficiently find neurons with maximum inner product in the forward pass; Reformer
[KKL20] also applies LSH to save the space complexity such that the neural networks are able to process very long sequences; MONGOOSE  $\left[\mathrm{CLP}^{+}21\right]$  combines the learnable LSH-based data structure [Cha02] with the lazy update framework [CLS19] to speedup the neural network training. Theoretically, [SYZ21] gives the first provable sublinear-time training algorithm for 2-layer overparameterized neural networks using the HSR data structures in [AEM92].
The goal of our paper is to 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-parameterization Over-parameterization, where the trainable parameters are much larger than the number of training data points (i.e.,  $m \gg n$ ), is a very natural and common regime in deep learning. It plays a 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 networks in the over-parameterization regime, e.g., [LL18, DZPS19, AZLS19a, AZLS19b, ADH+19a, ADH+19b, SY19, CGH+19, ZMG19, CG19, ZG19, OS20, JT20, LSS+20, HLSY21, ZPD+20, BPSW21, SZZ21, Zha22, HSWZ22, MOSW22, YJZ+23, GMS23, LSY23, QSS23, QSY23]. A key observation is that when the width ( $m$ ) of the neural network tends to infinity, the neural network is equivalent to a neural tangent kernel (NTK) [JGH18], 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 ((S)GD) can train a sufficiently wide neural network with random initialization, converging to a small training error in a polynomial number of steps.
Roadmap This paper is organized as follows: In 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.
Before describing our new data structure, we first present the notation we will use, and formulate the problem setting.
Basic Notation. For  $n \in \mathbb{N}_+$ , we use  $[n]$  to denote the set  $\{1, 2, \dots, n\}$ . We write  $\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 the  $d \times d$  identity matrix. We use  $\mathcal{N}(\mu; \sigma^2)$  to denote the Gaussian distribution with mean  $\mu$  and variance  $\sigma^2$ .
Problem Formulation In this section, we introduce the neural network model we study in this paper. We consider a two-layer ReLU-activated neural network  $f$  that has width  $m$  and uses an  $\ell_2$  loss function.
Definition 2.1 (Prediction function and loss function). For a threshold parameter  $b \in \mathbb{R}$ , data point  $x \in \mathbb{R}^d$ , weight matrix  $W \in \mathbb{R}^{d \times m}$ , and weights  $a \in \mathbb{R}^m$ , the prediction function  $f(W, x, a)$  and loss function  $L(W)$  are given by
$$
f (W, x, a) := \frac {1}{\sqrt {m}} \sum_ {r = 1} ^ {m} a _ {r} \sigma_ {b} (\langle w _ {r}, x \rangle), \quad L (W) := \frac {1}{2} \sum_ {i = 1} ^ {n} (f (W, x _ {i}, a) - y _ {i}) ^ {2},
$$
where  $\sigma_b(x) \coloneqq \max \{x - b, 0\}$  is the ReLU function with threshold parameter  $b$ . Following the prior work [SYZ21], we write  $2\mathrm{NN}(m,b)$  to denote this function  $f$  for simplicity.
Remark 2.2. In the neural network,  $w_{1}, \dots, w_{m} \in \mathbb{R}^{d}$  are the weight vectors of the edges between input nodes and hidden neuron nodes, and  $a_{1}, \dots, a_{m} \in \mathbb{R}$  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 \in \mathbb{R}^{d \times m}$  to minimize the loss  $L(W)$ , and will leave  $a \in \mathbb{R}^m$  unchanged after the initialization.
In this work, we study the following training process:
- Initialization: For each hidden neuron, we sample  $w_{r}(0) \sim \mathcal{N}(0, I_{d})$ , and sample  $a_{r}$  from  $\{-1, +1\}$  uniformly at random.
Gradient computation: For each neuron, we have
$$
\begin{array}{l} \frac {\partial f (W , x , a)}{\partial w _ {r}} = \frac {a _ {r}}{\sqrt {m}} x \mathbf {1} _ {w _ {r} ^ {\top} x \geq b}, \text {a n d} \\ \frac {\partial L (W)}{\partial w _ {r}} = \frac {a _ {r}}{\sqrt {m}} \sum_ {i = 1} ^ {n} \left(f \left(W, x _ {i}, a\right) - y _ {i}\right) x _ {i} \mathbf {1} _ {\left\langle w _ {r}, x _ {i} \right\rangle \geq b}. \\ \end{array}
$$
- Weight update: We follow the standard update rule of the GD algorithm from iteration  $k$  to iteration  $k + 1$ :
$$
W (k + 1) = W (k) - \eta \cdot \delta W (k),
$$
where  $W(k)$  denotes the weights at iteration  $k$ , and
$$
\delta W (k) = \frac {\partial L (W (k))}{\partial W (k)}.
$$
Sparsity phenomenon in the training process As observed in many experimental and theoretical works [CLP+21, SYZ21, SZZ21, GS22, GSZ23, DCL+22, DMS23, GSYZ23, LWD+23, GSY23, ZSZ+23, SWYZ23, BSZ23], for a randomly initialized over-parameterized neural network (Definition 2.1), given input data  $x$ , only a small fraction  $(o(m))$  of the neurons will be activated when evaluating  $2\mathrm{NN}(m,b)$  (Definition 2.1) in each training iteration. We refer to this phenomenon as "Sparsity". We exploit this property to design an efficient data structure to identify the sparse activated neurons, achieving sublinear training time in terms of  $m$ , the number of neurons in the hidden layer.
To be specific, the sparsity of activated neurons during the training process is bounded by choosing a proper threshold  $b$  for the ReLU function. Because of a concentration phenomenon of the randomized initialization, we can upper-bound the number of activated neurons just after initialization, which we refer to as "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 a small quantity.
For simplicity, we define the "fire set"  $S_{\mathrm{fire}}(x)$  first, which is the set of neurons that is activated when the neural network's input is  $x$ .
Definition 2.3 (Fire set). Let the neural network be defined as in Definition 2.1. For a data point  $x \in \mathbb{R}^d$ , let  $S_{\mathrm{fire}}(x)$  denote the set of neurons that are activated on input  $x$ , i.e.,
$$
\mathcal {S} _ {\text {f i r e}} (x) := \{i \in [ m ]: \sigma_ {b} (\langle x, w _ {i} \rangle) > 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 each data point  $x_{i} \in \mathbb{R}^{d}$  with  $i \in [n]$  and each iteration  $t \in \{0,1,\dots,T\}$ , let  $w_{r}(t) \in \mathbb{R}^{d}$  be the weight vector of the  $r$ -th neuron at the  $t$ -th iteration for  $r \in [m]$ . Define
$$
S _ {i, \text {f i r e}} (t) := \{r \in [ m ]: \sigma_ {b} (\langle x _ {i}, w _ {r} (t) \rangle) > 0 \},
$$
$$
\widetilde {S} _ {r, \text {f i r e}} (t) := \{i \in [ n ]: \sigma_ {b} (\langle x _ {i}, w _ {r} (t) \rangle) > 0 \}.
$$
We further denote the sizes of these sets by  $k_{i,t} \coloneqq |S_{i,\mathrm{fire}}(t)|$  and  $\widetilde{k}_{r,t} \coloneqq |\widetilde{S}_{r,\mathrm{fire}}(t)|$ .
The following lemma upper bounds the sparsity after initialization.
Lemma 2.5 (Sparsity after initialization, informal version of Lemma B.3, [SYZ21]). Let  $b > 0$  be a tunable parameter. If we setup the neural network as in Definition 2.1, then after the randomized initialization, with probability at least  $1 - \exp(-\Omega(m \cdot \exp(-b^2/2))$ ), it holds that for any input data  $x$ , the number of activated neurons is at most  $O(m \cdot \exp(-b^2/2))$ , where  $m$  is the total number of neurons.
Remark 2.6. This suggests that if we take  $b = \sqrt{0.4\log m}$ , we achieve a sublinear number,  $O(m^{4/5})$  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 a neural network  $2\mathrm{NN}(m,b)$  (Definition 2.1) with  $n$  data points. We let  $\{w_1,\dots ,w_m\} \subset \mathbb{R}^d$  be the weights,  $\{x_{1},\dots ,x_{n}\} \subset \mathbb{R}^{d}$  be the data points, and  $\{(w_r,x_i)\}_{r\in [m],i\in [n]}\subset \mathbb{R}^{m + n}$  be the weight-data pairs.
We propose two data structures: Correlation DTree and Correlation WTree. The DTree data structure has  $n$  trees, and its  $i$ -th tree has  $m$  leaf nodes corresponding to the set of inner-products between  $x_{i}$  and all hidden neurons, i.e.,  $\{\langle w_{r},x_{i}\rangle \}_{r\in [m]}$ . Similarly, the WTree data structure consists of  $m$  trees, and its  $r$ -th tree has  $n$  leaf nodes corresponding to the set of inner-products between the  $r$ -th neuron and all data points, i.e.,  $\{\langle w_{r},x_{i}\rangle \}_{i\in [n]}$ .
The Correlation Tree is a simple binary tree data structure. At a high level, it works as follows: