MONGOOSE: A Learnable LSH Framework for Efficient Neural Network Training

1 INTRODUCTION . Locality Sensitive Hashing ( LSH ) has been adapted to address the computational and memory bottlenecks of large-scale neural network ( NN ) training in natural language processing ( Chandar et al. , 2016 ; Rae et al. , 2016 ; Kitaev et al. , 2020 ) , computer vision ( Chen et al. , 2015 ) and recommendation systems ( Spring & Shrivastava , 2017 ; Chen et al. , 2020 ) . Specifically , giant matrix multiplications in linear layers preceding a softmax can be approximated using nearest neighbor search ( NNS ) techniques , which often rely on LSH . However , LSH methods used in NNs are inefficient . Although LSH achieves sub-linear query time in theory , it is known to suffer from high query and pre-processing ( update ) overhead in practice ( Erik et al. , 2018 ) . In the setting of NN training , where data points for LSH are model parameters , this overhead is exacerbated by a high number of updates due to constantly evolving model parameters . The most established solution for reducing LSH query overhead is data-dependent or learningbased hashing , which uses adaptive hash functions to optimize the LSH bucket distribution for each dataset ( Andoni & Razenshteyn , 2015 ; Dong et al. , 2019 ) . These methods reduce query time by incurring a one-time offline cost to learn useful data input patterns in a preprocessing step . The learning techniques are often computationally complex , but can lead to a net reduction in overall query time . However , in NN training , the expensive preprocessing procedure has to be repeated each time the parameters are updated . Naïvely applying these techniques would increase the LSH update overhead rather than reduce it . A more appropriate data-dependent LSH framework would ideally ( i ) have a deeper understanding of the training dynamics of model parameters , a setting in which LSH has not been well-studied , ( ii ) be able to perform low-cost updates to account for evolving parameters , and ( iii ) have better query time while accurately approximating matrix multiplication . We argue that it is unnecessary to view evolving model parameters as streaming data and not every parameter update requires an LSH update . In fact , LSH updates are necessary only when the NN gra- dient steps are large enough to cause the model parameters ’ LSH hash codes to change . Therefore , we count the number of hash code changes that occur for important models such as Transformers and fully-connected NNs from previous work ( Chen et al. , 2020 ; Kitaev et al. , 2020 ) . We find that only 1 % of the hash codes change after each epoch on average for Transformers , and 5 % for a fullyconnected NNs , so most weights do not change on a scale that is enough to trigger an LSH update . Furthermore , the rate of hash code change initially decays exponentially and eventually plateaus , as shown later in Figure 3 . We calculate that a 100× speed up is possible with an update oracle . In contrast , current approaches with LSH have not fully exploited this observation . We show in Section 4.2 that they either blindly skip updates for speed , resulting in a 10-point accuracy drop , or suffer from a 20× slow-down . We demonstrate that these slowly changing hash codes provide two opportunities to realize the ideal data-dependent LSH framework described above ( shown as the intuition for the scheduler in Figure 1 ) . First , given weight changes , an algorithm could adaptively schedule LSH updates to realize low update overhead . Second , given current weights , we could learn data-dependent hash functions for better LSH bucket distribution to achieve fast query time . However , this comes with three challenges : ( i ) how to characterize the slowly-changing phenomenon without computing hash codes , ( ii ) how to schedule the updates without the oracle , and ( iii ) how to learn data-dependent hash functions for shorter query time without compromising on the update time . In this paper , we first provide a general formulation of the problem as dynamic NNS in Section 2 . In Section 3 , we propose MONGOOSE , a framework for fast and memory-efficient NN training to address the above challenges . We show one major observation , slow change , and two key components based on it in MONGOOSE . Specifically , • In Section 3.1 , we measure the ` 2 norm of the weight changes and their hash code changes during training on a fully-connected NN . ( More models are studied in Appendix A ) . We find that ( i ) the hash codes are slowly changing ( ii ) both quantities share similar trajectories , but the former ’ s slow change is the necessary condition of the latter ’ s . Therefore , we formally define the slow change of ` 2 norm of weight changes in Assumption 3.1 , which is the key to build MONGOOSE . • In Section 3.2 , we present an algorithm for scheduling efficient LSH updates . We closely analyze our scheduler ’ s theoretical guarantees and provide an upper bound of its running time , which is proportional to the ` 2 norm of the weight changes . Therefore , under the slow change assumption , we show that our scheduler is provably faster than previous approaches . • In Section 3.3 , we propose a method of learning parameterized LSH hash functions ( e.g. , SimHash ( Charikar , 2002 ) ) during NN training . Our method utilizes intermediate activations during the forward pass to generate training inputs for tuning LSH hash functions with little overhead . Combining this with our scheduler further decreases the overhead , while providing hash functions that better separate data . Finally , in Section 4 , we demonstrate the efficacy of MONGOOSE on two LSH-NN systems , SLIDE and Reformer . For SLIDE applications , we show up to 6.5× speedup in time and 8 % higher accuracy on three datasets . For Reformer applications , we show improvement in perplexity when training lan- guage modeling benchmarks from scratch . Moreover , We provide two additional sets of experiments to show how MONGOOSE addresses the aforementioned three challenges separately . 2 RELATED WORK AND PROBLEM SETTING . In this section , we first discuss applications of LSH in the NN training setting and introduce datadependent LSH techniques . Then we formally define the problem we solve as dynamic NNS . 2.1 LSH FOR EFFICIENT NN TRAINING Recent works take advantage of LSH as an efficient NNS algorithm to speed up matrix multiplication in neural networks . SLIDE ( Chen et al. , 2020 ) is an algorithm that retrieves neurons with maximum inner product during the forward pass via an LSH-based data structure . In this way , in the backward pass gradients are only computed for neurons with estimated large gradients . Reformer ( Kitaev et al. , 2020 ) , a variant of Transformer , similarly uses LSH to reduce the memory bottleneck of self-attention layers over long sequences . One pain point of LSH in this setting is that model weights change throughout training . This necessitates constantly updating the LSH data structure . Figure 2 illustrates what happens when we don ’ t perform such updates . In short , failing to update the LSH data structure as the search data changes degrades its NNS performance . This in turn worsens the quality of the matrix product approximation . In our experiments , we found that failing to update the LSH data structure in SLIDE causes a 28 % decrease in top-1 accuracy for a fully-connected NN . ( More related work is presented in Appendix E.1 ) 2.2 PROBLEM FORMULATION . In this section , we formulate LSH for effcient training as a dynamic NNS problem . This formulation is closely built on the static NNS problem and the well-known lower bound on LSH complexity . In the static NNS problem , we are given a set of weights w1 , · · · , wn ∈ Rd and want to construct a data structure ( NNS-ds ) that supports the QUERY operation , defined as follows : given x ∈ Rd , QUERY ( x ) returns a set S of weights wi that are all “ close ” to x in a distance measure . More precisely , we require the QUERY ( x ) operation to be ( c1 , c2 ) -accurate : Definition 2.1 ( ( c1 , c2 ) -accurate ) . Denote S to be the ( random ) set returned by QUERY ( x ) . We say QUERY ( x ) is ( c1 , c2 ) accurate if ( i ) for any i ∈ [ n ] such that 〈wi , x〉 ≥ c1‖x‖2‖wi‖2 , Pr [ wi ∈ S ] ≥ 1− 1/ poly ( n ) , ( ii ) for any i ∈ [ n ] such that 〈wi , x〉 < c2‖x‖2‖wi‖2 , Pr [ wi ∈ S ] < 1/poly ( n ) . In our application , we assume 1 > c1 > c2 > 0 to be some constants and denote ρ = ( c2/c1 ) 2 . In the dynamic NNS problem , the weights w1 , · · · , wn ∈ Rd can evolve over time , so we need to update the data structure . Lemma 2.2 ( ( Andoni & Indyk , 2006 ) ) . Using LSH , one can achieve ( c1 , c2 ) accuracy with query time O ( dnρ ) , preprocessing time O ( dn1+ρ ) , updating time O ( dnρ ) . 3 MONGOOSE : A FRAMEWORK FOR LEARNABLE LSH . We present the workflow of our main framework in Figure 1 . In Section 3.1 , we first show the key observation that encourages the design of MONGOOSE and formally define the slow change assumption . In Section 3.2 , we introduce our first component , an algorithm for scheduling LSH updates with provable guarantees based on the assumption . Finally , we present the second component , lowcost learnable LSH hash functions , in Section 3.3 . More details about the efficient implementation of MONGOOSE are in Appendix F . 3.1 CHARACTERIZATION OF “ SLOW CHANGE ” PHENOMENON . We first present an observation on model weights and how their LSH hash codes change during training . Based on this observation , we define the slow change assumption . We denote ∆W as the difference between weights across gradient steps , and ∆H as the Hamming distance between hash codes during training . Observation . We plot the ` 2 norm of ∆W and ∆H of the weights during training a one layered fully-connected NN in Figure 3 . Our most surprising finding is that only 5 % of hash codes change after each epoch on average , which implies that for most neurons , ∆W is not large enough to trigger hash code changes . For both ∆W and ∆H , there is a sharp drop at the early stages of training before they plateau . ( Similar observations on Transformers and further discussions are presented in Appendix A . ) Insights . In the dynamic NNS problem , input data ( model weights ) change over time . Without any assumptions about weight updates , naïvely applying LSH will require updating the LSH hash functions at every time step . However , the above observation suggests that we can reasonably assume ∆W is ( roughly ) upper-bounded by ∆H . Denote w as the weight matrix , n as number of neurons and wi as the weight vector for neuron i . Formally : Assumption 3.1 ( Slow change ) . Assume NN weights change slowly over time . In particular , we assume there is an upper bound on the expected movement of the weights of the neural networks ( C1 ) and an upper bound on the variance ( C2 ) . Specifically , we denote the initial weight matrix as w ( 0 ) ∈ Rn×d . We assume the ( random ) update sequence w ( 1 ) , · · · , w ( T ) ∈ Rn×d satisfies n∑ i=1 ∥∥∥E [ w ( k+1 ) i ] − w ( k ) i ∥∥∥2 2 ≤ C21 and n∑ i=1 ‖Var [ w ( k+1 ) i ] ‖2 ≤ C22 , ( 1 ) where the expectation and variance is conditioned on w ( k ) i for all k = 0 , 1 , · · · , T − 1 .

Summary: Some neural network runs involve layers with a large number of neurons. These require large matrix-vector or matrix-multiplication which can slow their training/inference. However, if the output of mat-vec/mul is dominated by a few neurons with which the activation has large inner product (a matmul can be thought of as a weighted sum of inner products), then the computation can be sped up by approximating mat-vec/mul by a limited weighted sum with the dominant terms. This requires maintaining an ANNS data structure that is up to data with the back prop. These updates to ANNS have to be done carefully -- too frequent and the training will slow down, or too infrequent and the results of the mat-mul are way off.  This paper studies how to do this in a principled way using data-dependent LSH updates and backs it up with experimental data.