Original: M ONGOOSE : A LEARNABLE LSH FRAMEWORK FOR EFFICIENT NEURAL NETWORK TRAINING
Revised: M ONGOOSE : A LEARNABLE LSH FRAMEWORK FOR EFFICIENT NEURAL NETWORK TRAINING

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Recent advances by practitioners in the deep learning community have breathed new life into Locality Sensitive Hashing (LSH), using it to reduce memory and time bottlenecks in neural network (NN) training.
Revised: Recent advances by practitioners in the deep learning community have breathed new life into Locality Sensitive Hashing (LSH), using it to reduce memory and time bottlenecks in neural network (NN) training.

Original: However, while LSH has sublinear guarantees for approximate near-neighbor search in theory, it is known to have inefﬁcient query time in practice due to its use of random hash functions.
Revised: However, while LSH has sublinear guarantees for approximate near-neighbor search in theory, it is known to have inefﬁcient query time in practice due to its use of random hash functions.

Original: Moreover, when model parameters are changing, LSH suffers from update overhead.
Revised: Moreover, when model parameters are changing, LSH suffers from update overhead.

Original: This work is motivated by an observation that model parameters evolve slowly, such that the changes do not always require an LSH update to maintain performance.
Revised: This work is motivated by an observation that model parameters evolve slowly, such that the changes do not always require an LSH update to maintain performance.

Original: This phenomenon points to the potential for a reduction in update time and allows for a modiﬁed learnable version of data-dependent LSH to improve query time at a low cost.
Revised: This phenomenon points to the potential for a reduction in update time and allows for a modiﬁed learnable version of data-dependent LSH to improve query time at a low cost.

Original: We use the above insights to build MONGOOSE , an end-to-end LSH framework for efﬁcient NN training.
Revised: We use the above insights to build MONGOOSE , an end-to-end LSH framework for efﬁcient NN training.

Original: In particular, MONGOOSE is equipped with a scheduling algorithm to adaptively perform LSH updates with provable guarantees and learnable hash functions to improve query efﬁciency.
Revised: In particular, MONGOOSE is equipped with a scheduling algorithm to adaptively perform LSH updates with provable guarantees and learnable hash functions to improve query efﬁciency.

Original: Empirically, we validate MONGOOSE on large-scale deep learning models for recommendation systems and language modeling.
Revised: Empirically, we validate MONGOOSE on large-scale deep learning models for recommendation systems and language modeling.

Original: We ﬁnd that it achieves up to 8% better accuracy compared to previous LSH approaches, with 3 × speed-up and 6 × reduction in memory usage.
Revised: We ﬁnd that it achieves up to 8% better accuracy compared to previous LSH approaches, with 6 . 5 × speed-up and 6 × reduction in memory usage.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: 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;
Revised: 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;

Original: Rae et al., 2016; Kitaev et al., 2020), computer vision (Chen et al., 2015) and recommendation systems (Spring & Shrivastava, 2017; Chen et al., 2020).
Revised: Rae et al., 2016; Kitaev et al., 2020), computer vision (Chen et al., 2015) and recommendation systems (Spring & Shrivastava, 2017; Chen et al., 2020).

Original: Speciﬁcally, giant matrix multiplications in linear layers preceding a softmax can be approximated using nearest neighbor search (NNS) techniques, which often rely on LSH.
Revised: Speciﬁcally, giant matrix multiplications in linear layers preceding a softmax can be approximated using nearest neighbor search (NNS) techniques, which often rely on LSH.

Original: However, LSH methods used in NNs are inefﬁcient.
Revised: However, LSH methods used in NNs are inefﬁcient.

Original: 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).
Revised: 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).

Original: 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.
Revised: 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.

Original: 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).
Revised: 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).

Original: These methods reduce query time by incurring a one-time ofﬂine cost to learn useful data input patterns in a preprocessing step.
Revised: These methods reduce query time by incurring a one-time ofﬂine cost to learn useful data input patterns in a preprocessing step.

Original: The learning techniques are often computationally complex, but can lead to a net reduction in overall query time.
Revised: The learning techniques are often computationally complex, but can lead to a net reduction in overall query time.

Original: However, in NN training, the expensive preprocessing procedure has to be repeated each time the parameters are updated.
Revised: However, in NN training, the expensive preprocessing procedure has to be repeated each time the parameters are updated.

Original: Naïvely applying these techniques would increase the LSH update overhead rather than reduce it.
Revised: Naïvely applying these techniques would increase the LSH update overhead rather than reduce it.

Original: 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 parameter, and (iii) have better query time while accurately approximating matrix multiplication.
Revised: 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 parameter, and (iii) have better query time while accurately approximating matrix multiplication.

Original: We argue that it is unnecessary to view evolving model parameters as streaming data and not every parameter update requires an LSH update.
Revised: We argue that it is unnecessary to view evolving model parameters as streaming data and not every parameter update requires an LSH update.

Original: In fact, LSH updates are necessary only when the NN gradient steps are large enough to cause the model parameters’ LSH hash codes to change.
Revised: In fact, LSH updates are necessary only when the NN gradient steps are large enough to cause the model parameters’ LSH hash codes to change.

Original: Therefore,
Revised: Therefore,

Original: NN Training
Revised: NN Training

Original: Bucket   
 Bucket   

Revised: Bucket   
 Bucket

Original: W/O Update
Revised: W/O Update

Original: Bucket   

Revised: Bucket

Original: Time t Time t
Revised: Time t Time t

Original: Update
Revised: Update

Original: LSHHash Functions H
Revised: LSHHash Functions H

Original: H ' 5 145 

Revised: H ' 5 145

Original: SCHEDULER t
Revised: SCHEDULER t

Original: TimeforUpdate
Revised: TimeforUpdate

Original: TimeW/OUpdate
Revised: TimeW/OUpdate

Original: LEARNING
Revised: LEARNING

Original: UpdatedHash Functions H'
Revised: UpdatedHash Functions H'

Original: OR
Revised: OR

Original: W
Revised: W

Original: W 1  
Revised: W 1

Original: W W 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).
Revised: W W 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).

Original: We ﬁnd that only 1% of the hash codes change after each epoch on average for Transformer, and 5% for a fullyconnected NN, so most weights do not change on a scale that is enough to trigger an LSH update.
Revised: We ﬁnd that only 1% of the hash codes change after each epoch on average for Transformer, and 5% for a fullyconnected NN, so most weights do not change on a scale that is enough to trigger an LSH update.

Original: Furthermore, the rate of hash code change initially decays exponentially and eventually plateaus, as shown later in Figure 3.
Revised: Furthermore, the rate of hash code change initially decays exponentially and eventually plateaus, as shown later in Figure 3.

Original: We calculate that a 100 × speed up is possible with an update oracle.
Revised: We calculate that a 100 × speed up is possible with an update oracle.

Original: In contrast, current approaches with LSH have not fully exploited this observation.
Revised: In contrast, current approaches with LSH have not fully exploited this observation.

Original: 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.
Revised: 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.

Original: 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
Revised: 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

Original: Figure 1).
Revised: Figure 1).

Original: First, given weight changes, an algorithm could adaptively schedule LSH updates to realize
Revised: First, given weight changes, an algorithm could adaptively schedule LSH updates to realize

Original: low update
Revised: low update

Original: overhead.
Revised: overhead.

Original: Second, given current weights, we could learn data
Revised: Second, given current weights, we could learn data

Original: -dependent hash
Revised: -dependent hash

Original: functions
Revised: functions

Original: for better
Revised: for better

Original: LSH
Revised: LSH

Original: bucket distribution
Revised: bucket distribution

Original: to achieve fast query time.
Revised: to achieve fast query time.

Original: However, this comes with three challenges: (i) how to characterize the slowly-changing
Revised: However, this comes with three challenges: (i) how to characterize the slowly-changing

Original: phenomenon
Revised: phenomenon

Original: 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.
Revised: 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.

Original: In this paper, we ﬁrst provide a general formulation of the problem as dynamic NNS in Section 2.
Revised: In this paper, we ﬁrst provide a general formulation of the problem as dynamic NNS in Section 2.

Original: In
Revised: In

Original: Section 3, we propose MONGOOSE , a framework for fast and memory-efﬁcient NN training to address the above challenges.
Revised: Section 3, we propose MONGOOSE , a framework for fast and memory-efﬁcient NN training to address the above challenges.

Original: We show one major observation, slow change , and two key components based on it in MONGOOSE .
Revised: We show one major observation, slow change , and two key components based on it in MONGOOSE .

Original: Speciﬁcally,
Revised: Speciﬁcally,

Original: • In Section 3.1, we measure the (cid:96) 2 norm of the weight changes and their hash code changes during training on a fully-connected NN.
Revised: • In Section 3.1, we measure the (cid:96) 2 norm of the weight changes and their hash code changes during training on a fully-connected NN.

Original: (More models are studied in Appendix A).
Revised: (More models are studied in Appendix A).

Original: We ﬁnd 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.
Revised: We ﬁnd 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.

Original: Therefore, we formally deﬁne the slow change of (cid:96) 2 norm of weight changes in Assumption 3.1, which is the key to build MONGOOSE .
Revised: Therefore, we formally deﬁne the slow change of (cid:96) 2 norm of weight changes in Assumption 3.1, which is the key to build MONGOOSE .

Original: • In Section 3.2, we present an algorithm for scheduling efﬁcient LSH updates.
Revised: • In Section 3.2, we present an algorithm for scheduling efﬁcient LSH updates.

Original: We closely analyze our scheduler’s theoretical guarantees and provide an upper bound of its running time, which is proportional to the (cid:96) 2 norm of the weight changes.
Revised: We closely analyze our scheduler’s theoretical guarantees and provide an upper bound of its running time, which is proportional to the (cid:96) 2 norm of the weight changes.

Original: Therefore, under the slow change assumption, we show that our scheduler is provably faster than previous approaches.
Revised: Therefore, under the slow change assumption, we show that our scheduler is provably faster than previous approaches.

Original: • In Section 3.3, we propose a method of learning parameterized LSH hash functions (e
Revised: • In Section 3.3, we propose a method of learning parameterized LSH hash functions (e

Original: g., SimHash (Charikar, 2002)) during NN training.
Revised: g., SimHash (Charikar, 2002)) during NN training.

Original: Our method utilizes intermediate activations during the forward pass to generate training inputs for tuning LSH hash functions with little overhead.
Revised: Our method utilizes intermediate activations during the forward pass to generate training inputs for tuning LSH hash functions with little overhead.

Original: Combining this with our scheduler further decreases the overhead, while providing hash functions that better separate data.
Revised: Combining this with our scheduler further decreases the overhead, while providing hash functions that better separate data.

Original: Finally
Revised: Finally

Original: , in Section 4, we
Revised: , in Section 4, we

Original: demonstrate the efﬁcacy of MONGOOSE on two LSH-NN systems, SLIDE and Reformer.
Revised: demonstrate the efﬁcacy of MONGOOSE on two LSH-NN systems, SLIDE and Reformer.

Original: For SLIDE applications, we show up to 3 × speedup intraining time and 8% higher accuracy on three datasets.
Revised: For SLIDE applications, we show up to 6 . 5 × speedup in time and 8% higher accuracy on three datasets.

Original: For Reformer applications, we show improvement in perplexity when training language modeling benchmarks from scratch.
Revised: For Reformer applications , we show improvement in perplexity when training lan- guage modeling benchmarks from scratch.

Original: Moreover, we provide two additional sets of experiments to show how MONGOOSE addresses the aforementioned three challenges separately.
Revised: Moreover, We provide two additional sets of experiments to show how MONGOOSE addresses the aforementioned three challenges separately.

Original:  R ELATED WORK AND PROBLEM SETTING
Revised:  R ELATED WORK AND PROBLEM SETTING

Original: In this section, we ﬁrst discuss applications of LSH in the NN training setting and introduce datadependent LSH techniques.
Revised: In this section, we ﬁrst discuss applications of LSH in the NN training setting and introduce datadependent LSH techniques.

Original: Then we formally deﬁne the problem we solve as dynamic NNS .
Revised: Then we formally deﬁne the problem we solve as dynamic NNS .

Original: 2. L SH FOR EFFICIENT NN TRAINING
Revised: 2. L SH FOR EFFICIENT NN TRAINING

Original: [Figure]
Revised: [Figure]

Original: Recent works take advantage of LSH as an efﬁcient NNS algorithm to speed up matrix multiplication in neural networks.
Revised: Recent works take advantage of LSH as an efﬁcient NNS algorithm to speed up matrix multiplication in neural networks.

Original: 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.
Revised: 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.

Original: In this way, in the backward pass gradients are only computed for neurons with estimated large gradients.
Revised: In this way, in the backward pass gradients are only computed for neurons with estimated large gradients.

Original: Reformer (Kitaev et al., 2020), a variant of Transformer, similarly uses LSH to reduce the memory bottleneck of self-attention layers over long sequences.
Revised: Reformer (Kitaev et al., 2020), a variant of Transformer, similarly uses LSH to reduce the memory bottleneck of self-attention layers over long sequences.

Original: One pain point of LSH in this setting is that model weights change throughout training.
Revised: One pain point of LSH in this setting is that model weights change throughout training.

Original: This necessitates constantly updating the LSH data structure.
Revised: This necessitates constantly updating the LSH data structure.

Original: Figure 2 illustrates what happens when we don’t perform such updates.
Revised: Figure 2 illustrates what happens when we don’t perform such updates.

Original: In short, failing to update the LSH data structure as the search data changes degrades its NNS performance.
Revised: In short, failing to update the LSH data structure as the search data changes degrades its NNS performance.

Original: This in turn worsens the quality of the matrix product approximation.
Revised: This in turn worsens the quality of the matrix product approximation.

Original: 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.
Revised: 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.

Original: (More related work is presented in Appendix E.1)
Revised: (More related work is presented in Appendix E.1)

Original: 2. P ROBLEM FORMULATION
Revised: 2. P ROBLEM FORMULATION

Original: In this section, we formulate LSH for effcient training as a dynamic NNS problem.
Revised: In this section, we formulate LSH for effcient training as a dynamic NNS problem.

Original: This formulation is closely built on the static NNS problem and the well-known lower bound on LSH complexity.
Revised: This formulation is closely built on the static NNS problem and the well-known lower bound on LSH complexity.

Original: In the static NNS problem, we are given a set of weights w 1 , · · · , w n ∈ R d and want to construct a data structure (NNS-ds) that supports the QUERY operation, deﬁned as follows:
Revised: In the static NNS problem, we are given a set of weights w 1 , · · · , w n ∈ R d and want to construct a data structure (NNS-ds) that supports the QUERY operation, deﬁned as follows:

Original: given x ∈ R d ,
Revised: given x ∈ R d ,

Original: Q UERY ( x ) returns a set S of weights w i that are all “close” to x in a distance measure.
Revised: Q UERY ( x ) returns a set S of weights w i that are all “close” to x in a distance measure.

Original: More precisely, we require the Q UERY ( x ) operation to be ( c 1 , c 2 ) -accurate:
Revised: More precisely, we require the Q UERY ( x ) operation to be ( c 1 , c 2 ) -accurate:

Original: Deﬁnition 2.1 ( ( c 1 , c 2 ) -accurate) .
Revised: Deﬁnition 2.1 ( ( c 1 , c 2 ) -accurate) .

Original: Denote S to be the (random) set returned by Q UERY ( x ) .
Revised: Denote S to be the (random) set returned by Q UERY ( x ) .

Original: We say
Revised: We say

Original: Q UERY ( x ) is ( c 1 , c 2 ) accurate if
Revised: Q UERY ( x ) is ( c 1 , c 2 ) accurate if

Original: (i) for any i ∈ [ n ] such that (cid:104) w i , x (cid:105) ≥ c 1 (cid:107) x (cid:107) 2 (cid:107) w i (cid:107) 2 , Pr[ w i ∈ S ] ≥ 1 − 1 / poly( n ) , (ii) for any i ∈ [ n ] such that (cid:104) w i , x (cid:105) <
Revised: (i) for any i ∈ [ n ] such that (cid:104) w i , x (cid:105) ≥ c 1 (cid:107) x (cid:107) 2 (cid:107) w i (cid:107) 2 , Pr[ w i ∈ S ] ≥ 1 − 1 / poly( n ) , (ii) for any i ∈ [ n ] such that (cid:104) w i , x (cid:105) <

Original: c 2 (cid:107) x (cid:107) 2 (cid:107) w i (cid:107) 2 , Pr[ w i ∈ S ] < 1 / poly( n ) .
Revised: c 2 (cid:107) x (cid:107) 2 (cid:107) w i (cid:107) 2 , Pr[ w i ∈ S ] < 1 / poly( n ) .

Original: In our application, we assume 1 > c 1 > c 2 > 0 to be some constants and denote ρ
Revised: In our application, we assume 1 > c 1 > c 2 > 0 to be some constants and denote ρ

Original: = ( c 2 /c 1 ) 2 .
Revised: = ( c 2 /c 1 ) 2 .

Original: In the dynamic NNS problem, the weights w 1 , · · · , w n ∈ R d can evolve over time, so we need to update the data structure.
Revised: In the dynamic NNS problem, the weights w 1 , · · · , w n ∈ R d can evolve over time, so we need to update the data structure.

Original: Lemma 2.2 ((Andoni & Indyk, 2006)) .
Revised: Lemma 2.2 ((Andoni & Indyk, 2006)) .

Original: Using LSH, one can achieve ( c 1 , c 2 ) accuracy with query time O ( dn ρ ) , preprocessing time O ( dn 1+ ρ ) , updating time O ( dn ρ ) .
Revised: Using LSH, one can achieve ( c 1 , c 2 ) accuracy with query time O ( dn ρ ) , preprocessing time O ( dn 1+ ρ ) , updating time O ( dn ρ ) .

Original:  M ONGOOSE : A FRAMEWORK FOR LEARNABLE LSH
Revised:  M ONGOOSE : A FRAMEWORK FOR LEARNABLE LSH

Original: We present the workﬂow of our main framework in Figure 1.
Revised: We present the workﬂow of our main framework in Figure 1.

Original: In Section 3.1, we ﬁrst show the key observation that encourages the design of MONGOOSE and formally deﬁne the slow change assumption.
Revised: In Section 3.1, we ﬁrst show the key observation that encourages the design of MONGOOSE and formally deﬁne the slow change assumption.

Original: In Section 3.2, we introduce our ﬁrst component, an algorithm for scheduling LSH updates with provable guarantees based on the assumption.
Revised: In Section 3.2, we introduce our ﬁrst component, an algorithm for scheduling LSH updates with provable guarantees based on the assumption.

Original: Finally, we present the second component, lowcost learnable LSH hash functions, in Section 3.3.
Revised: Finally, we present the second component, lowcost learnable LSH hash functions, in Section 3.3.

Original: More details about the efﬁcient implementation of MONGOOSE are in Appendix F.
Revised: More details about the efﬁcient implementation of MONGOOSE are in Appendix F.

Original: 3. C HARACTERIZATION OF “ SLOW CHANGE ” PHENOMENON
Revised: 3. C HARACTERIZATION OF “ SLOW CHANGE ” PHENOMENON

Original: We ﬁrst present an observation on model weights and how their LSH hash codes change during training.
Revised: We ﬁrst present an observation on model weights and how their LSH hash codes change during training.

Original: Based on this observation, we deﬁne the slow change assumption.
Revised: Based on this observation, we deﬁne the slow change assumption.

Original: We denote ∆ W as the difference between weights across gradient steps, and ∆ H as the Hamming distance between hash codes during training.
Revised: We denote ∆ W as the difference between weights across gradient steps, and ∆ H as the Hamming distance between hash codes during training.

Original: [Figure]
Revised: [Figure]

Original: Observation.
Revised: Observation.

Original: We plot the (cid:96) 2 norm of ∆ W and ∆ H of the weights during training a one layered fully-connected NN in Figure 3.
Revised: We plot the (cid:96) 2 norm of ∆ W and ∆ H of the weights during training a one layered fully-connected NN in Figure 3.

Original: Our most surprising ﬁnding 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.
Revised: Our most surprising ﬁnding 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.

Original: For both ∆ W and ∆ H , there is a sharp drop at the early stages of training before they plateau.
Revised: For both ∆ W and ∆ H , there is a sharp drop at the early stages of training before they plateau.

Original: (Similar observations on Transformers and further discussions are presented in Appendix A.)
Revised: (Similar observations on Transformers and further discussions are presented in Appendix A.)

Original: Insights.
Revised: Insights.

Original: In the dynamic NNS problem, input data (model weights) change over time.
Revised: In the dynamic NNS problem, input data (model weights) change over time.

Original: Without any assumptions about weight updates, naïvely applying LSH will require updating the LSH hash functions at every time step.
Revised: Without any assumptions about weight updates, naïvely applying LSH will require updating the LSH hash functions at every time step.

Original: However, the above observation suggests that we can reasonably assume ∆ W is (roughly) upper-bounded by ∆ H .
Revised: However, the above observation suggests that we can reasonably assume ∆ W is (roughly) upper-bounded by ∆ H .

Original: Denote w as the weight matrix, n as number of neurons and w i as the weight vector for neuron i .
Revised: Denote w as the weight matrix, n as number of neurons and w i as the weight vector for neuron i .

Original: Formally:
Revised: Formally:

Original: Assumption 3.1 (Slow change) .
Revised: Assumption 3.1 (Slow change) .

Original: Assume NN weights change slowly over time.
Revised: Assume NN weights change slowly over time.

Original: In particular, we assume there is an upper bound on the expected movement of the weights of the neural networks ( C 1 ) and an upper bound on the variance ( C 2 ).
Revised: In particular, we assume there is an upper bound on the expected movement of the weights of the neural networks ( C 1 ) and an upper bound on the variance ( C 2 ).

Original: Speciﬁcally, we denote the initial weight matrix as w (0)
Revised: Speciﬁcally, we denote the initial weight matrix as w (0)

Original: ∈
Revised: ∈

Original: R n × d .
Revised: R n × d .

Original: We assume the (random) update sequence w (1) , · · · , w ( T ) ∈ R n ×
Revised: We assume the (random) update sequence w (1) , · · · , w ( T ) ∈ R n ×

Original: d satisﬁes n (cid:88) i = (cid:13)(cid:13)(cid:13)
Revised: d satisﬁes n (cid:88) i = (cid:13)(cid:13)(cid:13)

Original: [Equation]
Revised: [Equation]

Original: Cand n (cid:88) i =1 (cid:107)
Revised: Cand n (cid:88) i =1 (cid:107)

Original: Var[ w ( k +1)
Revised: Var[ w ( k +1)

Original: i ] (cid:107) 2 ≤ C 22 , (1) where the expectation and variance is conditioned on w ( k ) i for all k = 0 , 1 , · · · , T − 1 .
Revised: i ] (cid:107) 2 ≤ C 22 , (1) where the expectation and variance is conditioned on w ( k ) i for all k = 0 , 1 , · · · , T − 1 .

Original: 3. A SMART SCHEDULER FOR LEARNABLE LSH UPDATES
Revised: 3. A SMART SCHEDULER FOR LEARNABLE LSH UPDATES

Original: 1: Input: w ∈ R n × d , v ∈ R n × d , (cid:15) mds ∈ (0 , 1 / log 2 n ) , a = min { ρ, α } , S ⊆
Revised: 1: Input: w ∈ R n × d , v ∈ R n × d , (cid:15) mds ∈ (0 , 1 / log 2 n ) , a = min { ρ, α } , S ⊆

Original: [ n ] 2: procedure I NITIALIZE ( w, (cid:15) mds )
Revised: [ n ] 2: procedure I NITIALIZE ( w, (cid:15) mds )

Original: (cid:46)
Revised: (cid:46)

Original: Lemma B.
Revised: Lemma B.

Original: 3: w ← w , v ← w , (cid:15) mds ← (cid:15) mds , S ← ∅ , 4:
Revised: 3: w ← w , v ← w , (cid:15) mds ← (cid:15) mds , S ← ∅ , 4:

Original: Build LSH with accuracy ( c 1 − (cid:15) mds , c 2 + (cid:15) mds ) 5: end procedure 6: procedure U PDATE ( w new ) (cid:46)
Revised: Build LSH with accuracy ( c 1 − (cid:15) mds , c 2 + (cid:15) mds ) 5: end procedure 6: procedure U PDATE ( w new ) (cid:46)

Original: Lemma B.
Revised: Lemma B.

Original: 7: y i ← w new i − v i , ∀ i ∈ [ n ] 8: r ← the number of indices i that (cid:107)
Revised: 7: y i ← w new i − v i , ∀ i ∈ [ n ] 8: r ← the number of indices i that (cid:107)

Original: y i (cid:107) 2 ≥ (cid:15) mds / 2 .
Revised: y i (cid:107) 2 ≥ (cid:15) mds / 2 .

Original: 9: if r < n a then 10: v new ← v 11: else 12:
Revised: 9: if r < n a then 10: v new ← v 11: else 12:

Original: Let π : [ n ] → [ n ] , that (cid:107)
Revised: Let π : [ n ] → [ n ] , that (cid:107)

Original: y π
Revised: y π

Original: ( i ) (cid:107) 2 ≥ (cid:107)
Revised: ( i ) (cid:107) 2 ≥ (cid:107)

Original: y π
Revised: y π

Original: ( i +1) (cid:107) 2 13:
Revised: ( i +1) (cid:107) 2 13:

Original: while 1 . 5 · r < n and (cid:107) y π ( (cid:100) 1 . 5 · r (cid:101) )
Revised: while 1 . 5 · r < n and (cid:107) y π ( (cid:100) 1 . 5 · r (cid:101) )

Original: (cid:107) 2 ≥ (1 − 1 / log n ) (cid:107)
Revised: (cid:107) 2 ≥ (1 − 1 / log n ) (cid:107)

Original: y π ( r ) (cid:107) 2 do 14: r ← min( (cid:100) 1 . 5 · r (cid:101) , n ) 15: end while (cid:46)
Revised: y π ( r ) (cid:107) 2 do 14: r ← min( (cid:100) 1 . 5 · r (cid:101) , n ) 15: end while (cid:46)

Original: Finding a smooth cut 16: v new π ( i ) ← (cid:40) w new π ( i ) i ∈ { 1 , 2 , · · · , r } v π ( i ) i ∈ { r + 1 , · · · , n } 17:
Revised: Finding a smooth cut 16: v new π ( i ) ← (cid:40) w new π ( i ) i ∈ { 1 , 2 , · · · , r } v π ( i ) i ∈ { r + 1 , · · · , n } 17:

Original: LSH.U PDATE π (1) , · · · , π ( r ) 18: end if 19: w ← w new , v ← v new , 20:
Revised: LSH.U PDATE π (1) , · · · , π ( r ) 18: end if 19: w ← w new , v ← v new , 20:

Original: S =
Revised: S =

Original: { i ∈ [ n ] : (cid:107) w i − v i (cid:107) 2 ≥ (cid:15) mds / 2 } 21: end procedure 22: procedure Q UERY ( h ) (cid:46)
Revised: { i ∈ [ n ] : (cid:107) w i − v i (cid:107) 2 ≥ (cid:15) mds / 2 } 21: end procedure 22: procedure Q UERY ( h ) (cid:46)

Original: Lemma B.
Revised: Lemma B.

Original: LSH.
Revised: LSH.

Original: QUERY and check S 24: end procedure
Revised: QUERY and check S 24: end procedure

Original: In this section, we introduce a dynamic data structure to schedule updates for LSH.
Revised: In this section, we introduce a dynamic data structure to schedule updates for LSH.

Original: Our scheduler provides signiﬁcant speedup under two realistic assumptions without compromising the theoretical guarantee of LSH: (1) slow change deﬁned in Assumption 3.1, (2) batch speed up deﬁned in Assumption 3.2.
Revised: Our scheduler provides signiﬁcant speedup under two realistic assumptions without compromising the theoretical guarantee of LSH: (1) slow change deﬁned in Assumption 3.1, (2) batch speed up deﬁned in Assumption 3.2.

Original: The scheduler can generalize to any near-neighbor search data structure (NNS-ds), but we show a design for LSH as an example since that is the NNS-ds in MONGOOSE .
Revised: The scheduler can generalize to any near-neighbor search data structure (NNS-ds), but we show a design for LSH as an example since that is the NNS-ds in MONGOOSE .

Original: 3.2. S MART SCHEDULER
Revised: 3.2. S MART SCHEDULER

Original: The pseudo-code for our smart scheduler is shown in Algorithm 1.
Revised: The pseudo-code for our smart scheduler is shown in Algorithm 1.

Original: Our scheduler is inspired by a dynamic data structure proposed in a recent line of work (Cohen et al., 2019), which was originally designed for inverse maintainance in Linear Programming solvers.
Revised: Our scheduler is inspired by a dynamic data structure proposed in a recent line of work (Cohen et al., 2019), which was originally designed for inverse maintainance in Linear Programming solvers.

Original: We provide a novel application here for the dynamic
Revised: We provide a novel application here for the dynamic

Original: LSH problem.
Revised: LSH problem.

Original: Notation: We denote [ n ] = { 1 , · · · , n } .
Revised: Notation: We denote [ n ] = { 1 , · · · , n } .

Original: The (cid:101) O ( · ) notation hides polylogarithmic dependence on n, d, (cid:15) mds .
Revised: The (cid:101) O ( · ) notation hides polylogarithmic dependence on n, d, (cid:15) mds .

Original: We use w ∈ R n × d to denote the weights of the neural network and each row w i ∈
Revised: We use w ∈ R n × d to denote the weights of the neural network and each row w i ∈

Original: R
Revised: R

Original: d ( i ∈ [ n ] ) corresponds to the input weight of the i -th neuron.
Revised: d ( i ∈ [ n ] ) corresponds to the input weight of the i -th neuron.

Original: We maintain a copy of the neural network weights, denoted as v ∈ R n × d , that approximately equals w except for a few coordinates.
Revised: We maintain a copy of the neural network weights, denoted as v ∈ R n × d , that approximately equals w except for a few coordinates.

Original: We take (cid:15) mds to be a sufﬁciently small precision parameter, i.e. (cid:15) mds (cid:28) c 1 − c 2 , and at the same time, satisﬁes (cid:15) mds ∈ Ω(1 / log 2 n ) .
Revised: We take (cid:15) mds to be a sufﬁciently small precision parameter, i.e. (cid:15) mds (cid:28) c 1 − c 2 , and at the same time, satisﬁes (cid:15) mds ∈ Ω(1 / log 2 n ) .

Original: We use S to denote the set of indices where w and v differ signiﬁcantly (see Line 20).
Revised: We use S to denote the set of indices where w and v differ signiﬁcantly (see Line 20).

Original: LSH.U PDATE (see line 17) takes indices of the neural weights and updates the LSH hashcode correspondingly.
Revised: LSH.U PDATE (see line 17) takes indices of the neural weights and updates the LSH hashcode correspondingly.

Original: Overview of algorithm:
Revised: Overview of algorithm:

Original: During the training process, we maintain two copies of the NN weights w, v ∈ R n × d , where w are the actual weights that are updated via gradient descent and v approximately equals w except for a few coordinates.
Revised: During the training process, we maintain two copies of the NN weights w, v ∈ R n × d , where w are the actual weights that are updated via gradient descent and v approximately equals w except for a few coordinates.

Original: We set an empirical threshold of (cid:15) mds .
Revised: We set an empirical threshold of (cid:15) mds .

Original: Due to the robustness of the NNS problem, we can tolerate this small error (cid:15) mds such that only those coordinates with a difference greater than (cid:15) mds matter.
Revised: Due to the robustness of the NNS problem, we can tolerate this small error (cid:15) mds such that only those coordinates with a difference greater than (cid:15) mds matter.

Original: To exploit the beneﬁt of batch updating, we perform lazy updates on the LSH hash table.
Revised: To exploit the beneﬁt of batch updating, we perform lazy updates on the LSH hash table.

Original: In particular, we update the LSH hash table if and only if the difference between w, v exceeds some threshold number n ρ (see line 9 to line 11).
Revised: In particular, we update the LSH hash table if and only if the difference between w, v exceeds some threshold number n ρ (see line 9 to line 11).

Original: One key part of our update procedure is that we also update some low-error coordinates when there are too many of them (see line 13 to line 16).
Revised: One key part of our update procedure is that we also update some low-error coordinates when there are too many of them (see line 13 to line 16).

Original: This procedure can be seen as forecasting future updates.
Revised: This procedure can be seen as forecasting future updates.

Original: The intuition is that when there are too many coordinates close to the decision boundary, they are likely to change within a short amount of time.
Revised: The intuition is that when there are too many coordinates close to the decision boundary, they are likely to change within a short amount of time.

Original: Updating these all together could further take advantage of the batch updates.
Revised: Updating these all together could further take advantage of the batch updates.

Original: To answer an LSH query, we keep track of a short list S (see line 20), which records coordinates that differ signiﬁcantly between w, v .
Revised: To answer an LSH query, we keep track of a short list S (see line 20), which records coordinates that differ signiﬁcantly between w, v .

Original: We can interpret S as the indices of corrupted datapoints to keep track of when performing an LSH query, since they are no longer the same.
Revised: We can interpret S as the indices of corrupted datapoints to keep track of when performing an LSH query, since they are no longer the same.

Original: In the query procedure (see line 22 to line 24), besides performing the normal LSH query, we also check S and pick up weights that are close to x .
Revised: In the query procedure (see line 22 to line 24), besides performing the normal LSH query, we also check S and pick up weights that are close to x .

Original: Hence, we balance the overhead of query time with update time, where update time is usually the bottleneck.
Revised: Hence, we balance the overhead of query time with update time, where update time is usually the bottleneck.

Original: 3.2. T HEORETICAL ANALYSIS
Revised: 3.2. T HEORETICAL ANALYSIS

Original: Assumption 3.2 (batch speedup) .
Revised: Assumption 3.2 (batch speedup) .

Original: Denote T r as the updating time for LSH when the weights of r neurons change and T r = r · t r .
Revised: Denote T r as the updating time for LSH when the weights of r neurons change and T r = r · t r .

Original: That is, T r is the updating time when we update r entries of the
Revised: That is, T r is the updating time when we update r entries of the

Original: LSH hash table and t r is the average updating time.
Revised: LSH hash table and t r is the average updating time.

Original: We assume { t i } i =1 ,...n is non increasing, i.e., t 1 ≥ t 2 ≥ . . .
Revised: We assume { t i } i =1 ,...n is non increasing, i.e., t 1 ≥ t 2 ≥ . . .

Original: ≥ t n .
Revised: ≥ t n .

Original: We further assume T r = n ρ when r ≤ n α for some constant α .
Revised: We further assume T r = n ρ when r ≤ n α for some constant α .

Original: The batch speedup assumption is realistic, especially due to advances in support for massive parallelism on modern computation architectures.
Revised: The batch speedup assumption is realistic, especially due to advances in support for massive parallelism on modern computation architectures.

Original: Theorem 3.3 (LSH maintenance) .
Revised: Theorem 3.3 (LSH maintenance) .

Original: For any constant c 1 , c 2 ( c 1 > c 2 ), there is a dynamic data structure (Algorithm 1) for LSH maintenance that achieves ( c 1 , c 2 ) -accuracy.
Revised: For any constant c 1 , c 2 ( c 1 > c 2 ), there is a dynamic data structure (Algorithm 1) for LSH maintenance that achieves ( c 1 , c 2 ) -accuracy.

Original: The data structure takes (cid:101)
Revised: The data structure takes (cid:101)

Original: O ( dn 1+ ρ ) time to initialize and each call of Q UERY ( h ) takes time (cid:101)
Revised: O ( dn 1+ ρ ) time to initialize and each call of Q UERY ( h ) takes time (cid:101)

Original: O ( n ρ d ) .
Revised: O ( n ρ d ) .

Original: By taking a = min { ρ, α } and
Revised: By taking a = min { ρ, α } and

Original: [Equation]
Revised: [Equation]

Original: The amortized expected time per call of U PDATE ( w ) is at most
Revised: The amortized expected time per call of U PDATE ( w ) is at most

Original: [Equation]
Revised: [Equation]

Original: Remark 3.4.
Revised: Remark 3.4.

Original: Depending on the exact form of batch updating time t r , the (amortized) running time and the beneﬁt of our scheduler could be different.
Revised: Depending on the exact form of batch updating time t r , the (amortized) running time and the beneﬁt of our scheduler could be different.

Original: We provide concrete examples in Appendix B.7, showing advantages of our smart scheduler under different scenarios.
Revised: We provide concrete examples in Appendix B.7, showing advantages of our smart scheduler under different scenarios.

Original: In general, we remark that (1) Our scheduler is always better than the naïve sequential updating strategy.
Revised: In general, we remark that (1) Our scheduler is always better than the naïve sequential updating strategy.

Original: (2) Our scheduler is oblivious to the exact form of the update time t r and the weight change C 1 , C 2 .
Revised: (2) Our scheduler is oblivious to the exact form of the update time t r and the weight change C 1 , C 2 .

Original: That is to say, we do not need to ﬁne tune these parameters.
Revised: That is to say, we do not need to ﬁne tune these parameters.

Original: In NN training, the only parameter that requires ﬁne tuning is the precision parameter (cid:15) mds .
Revised: In NN training, the only parameter that requires ﬁne tuning is the precision parameter (cid:15) mds .

Original: 3. H ASH FUNCTION LEARNING
Revised: 3. H ASH FUNCTION LEARNING

Original: In this section, we provide a simple streaming algorithm to learn parameterized hash functions.
Revised: In this section, we provide a simple streaming algorithm to learn parameterized hash functions.

Original: Here, the learning objective is to track the weight distribution as the weights change to better separate the input data to the NNS problem.
Revised: Here, the learning objective is to track the weight distribution as the weights change to better separate the input data to the NNS problem.

Original: Recall that in last section, we show an algorithm for scheduling the
Revised: Recall that in last section, we show an algorithm for scheduling the

Original: LSH updates, which makes it possible to involve learning for LSH and improve the query efﬁciency.
Revised: LSH updates, which makes it possible to involve learning for LSH and improve the query efﬁciency.

Original: Therefore, the overall update overhead can be reduced even with learning overhead if scheduled well, since the update time is closely related to query time.
Revised: Therefore, the overall update overhead can be reduced even with learning overhead if scheduled well, since the update time is closely related to query time.

Original: 3.3. L EARNABLE LSH
Revised: 3.3. L EARNABLE LSH

Original: [Table] [Table]
Revised: [Table] [Table]

Original: We view the sequence of neural network weights throughout training as a matrix stream.
Revised: We view the sequence of neural network weights throughout training as a matrix stream.

Original: In each training iteration, if our scheduler schedules an update, we compute a training signal with which we can train our LSH parameters.
Revised: In each training iteration, if our scheduler schedules an update, we compute a training signal with which we can train our LSH parameters.

Original: Speciﬁcally, we compute the classical triplet loss (Chechik et al., 2010) using positive and negative samples from the current model weights for each
Revised: Speciﬁcally, we compute the classical triplet loss (Chechik et al., 2010) using positive and negative samples from the current model weights for each

Original: LSH query.
Revised: LSH query.

Original: Here, positive samples (denoted by P + ) are neurons that have higher inner product within the input retrieved by LSH, while negative samples (denoted by P − ) are neurons with lower inner product that are selected.
Revised: Here, positive samples (denoted by P + ) are neurons that have higher inner product within the input retrieved by LSH, while negative samples (denoted by P − ) are neurons with lower inner product that are selected.

Original: Denote a set of neurons in a particular layer as C = { v r | 0 ≤ r < m } and the input embedding as q i .
Revised: Denote a set of neurons in a particular layer as C = { v r | 0 ≤ r < m } and the input embedding as q i .

Original: Denote the set of neurons retrieved from the LSH hash tables using q i as S i .
Revised: Denote the set of neurons retrieved from the LSH hash tables using q i as S i .

Original: We formally present the learning algorithm in Algorithm 2.
Revised: We formally present the learning algorithm in Algorithm 2.

Original: The operation of updating or training LSH hash functions is deﬁned as
Revised: The operation of updating or training LSH hash functions is deﬁned as

Original: REHASH and the operation of updating the hash tables based on the updated hash functions and weights is deﬁned as REBUILD.
Revised: REHASH and the operation of updating the hash tables based on the updated hash functions and weights is deﬁned as REBUILD.

Original: The triplet loss is deﬁned as follows:
Revised: The triplet loss is deﬁned as follows:

Original: [Equation]
Revised: [Equation]

Original: where H denotes the LSH hash function at the current iteration and α is a margin between positive and negative pairs.
Revised: where H denotes the LSH hash function at the current iteration and α is a margin between positive and negative pairs.

Original: cos( H ( · ) , H ( · )) represents the cosine similarity.
Revised: cos( H ( · ) , H ( · )) represents the cosine similarity.

Original: Note that in the application setting of unit norm like (Kitaev et al., 2020), using inner product is equivalent to cosine similarity.
Revised: Note that in the application setting of unit norm like (Kitaev et al., 2020), using inner product is equivalent to cosine similarity.

Original: We provide a detailed justiﬁcation of this learning approach in Appendix C.1.
Revised: We provide a detailed justiﬁcation of this learning approach in Appendix C.1.

Original: 3.3. F URTHER SPEEDUP
Revised: 3.3. F URTHER SPEEDUP

Original: We demonstrate how to optimize learnable LSH for further acceleration in
Revised: We demonstrate how to optimize learnable LSH for further acceleration in

Original: NN training.
Revised: NN training.

Original: Here, we provide the main insights; further details are in Appendix A.
Revised: Here, we provide the main insights; further details are in Appendix A.

Original: Layer Selection:
Revised: Layer Selection:

Original: In previous LSH-NN methods such as SLIDE or Reformer, the designation of whether a layer uses LSH is treated as a tunable hyperparameter.
Revised: In previous LSH-NN methods such as SLIDE or Reformer, the designation of whether a layer uses LSH is treated as a tunable hyperparameter.

Original: For example, Reformer applies LSH in three out of twelve attention layers in their ﬁnal model based on empirical performance.
Revised: For example, Reformer applies LSH in three out of twelve attention layers in their ﬁnal model based on empirical performance.

Original: In MONGOOSE we make this designation in a more principled manner.
Revised: In MONGOOSE we make this designation in a more principled manner.

Original: The intuition behind the design can be presented by visualizing the statistics of attention weights in different layers of the Transformer model (Figure 4).
Revised: The intuition behind the design can be presented by visualizing the statistics of attention weights in different layers of the Transformer model (Figure 4).

Original: The three plots represent the distribution of the average number of coordinates with “heavy” attention for each layer.
Revised: The three plots represent the distribution of the average number of coordinates with “heavy” attention for each layer.

Original: We also calculate the size of the minimal set of attention weights that captures 90% of the total attention ((Ramsauer et al., 2020) perform a similar experiment).
Revised: We also calculate the size of the minimal set of attention weights that captures 90% of the total attention ((Ramsauer et al., 2020) perform a similar experiment).

Original: The median for the left plot is 2, meaning that they are well-separated and LSH can perform well on such a distribution.
Revised: The median for the left plot is 2, meaning that they are well-separated and LSH can perform well on such a distribution.

Original: On the other hand, learnable LSH can be very helpful for the right plot, because classical LSH is inefﬁcient for this distribution.
Revised: On the other hand, learnable LSH can be very helpful for the right plot, because classical LSH is inefﬁcient for this distribution.

Original: Therefore, we can use this observation when deciding between LSH or learnable LSH in each layer for later training stages or ﬁne-tuning on downstream tasks.
Revised: Therefore, we can use this observation when deciding between LSH or learnable LSH in each layer for later training stages or ﬁne-tuning on downstream tasks.

Original: [Table]
Revised: [Table]

Original: From a system design perspective
Revised: From a system design perspective

Original: [Table]
Revised: [Table]

Original: MONGOOSE improves the ﬂexibility of LSH’s usage on NN training.
Revised: MONGOOSE improves the ﬂexibility of LSH’s usage on NN training.

Original: We provide a smart scheduler to decide when to perform LSH updates and visualization on when to use LSH or learnable LSH.
Revised: We provide a smart scheduler to decide when to perform LSH updates and visualization on when to use LSH or learnable LSH.

Original: Therefore [Table] the training process has the following tunable parameters for LSH:
Revised: Therefore, the training process has the following tunable parameters for LSH:

Original: (i) the layers which we apply random LSH to, (ii) the layers for which we can use learning to improve LSH’s random hash functions, (iii) the runtime that LSH or learnable
Revised: (i) the layers which we apply random LSH to, (ii) the layers for which we can use learning to improve LSH’s random hash functions, (iii) the runtime that LSH or learnable

Original: LSH should perform the update in.
Revised: LSH should perform the update in.

Original: Setting these parameters appropriately will enable much more efﬁcient NN training or ﬁne-tuning, as supported by our theory and observations.
Revised: Setting these parameters appropriately will enable much more efﬁcient NN training or ﬁne-tuning, as supported by our theory and observations.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: In Section 4.1 we present the results of experiments that demonstrate the higher performance of
Revised: In Section 4.1 we present the results of experiments that demonstrate the higher performance of

Original: MONGOOSE over other LSH-NN frameworks during training.
Revised: MONGOOSE over other LSH-NN frameworks during training.

Original: Then we study the effectiveness of the two components separately: (i) smart scheduler in Section 4.2, (ii) learnable LSH in Section 4.3.
Revised: Then we study the effectiveness of the two components separately: (i) smart scheduler in Section 4.2, (ii) learnable LSH in Section 4.3.

Original: 4. M AIN RESULTS : MONGOOSE IN NN TRAINING
Revised: 4. M AIN RESULTS : MONGOOSE IN NN TRAINING

Original: We present the results of main experiments that demonstrate MONGOOSE ’s advantage in end-to-end
Revised: We present the results of main experiments that demonstrate MONGOOSE ’s advantage in end-to-end

Original: NN training over two different LSH-NN baselines on recommendation and language modeling tasks.
Revised: NN training over two different LSH-NN baselines on recommendation and language modeling tasks.

Original: 4.1. T ASK 1: EXTREME MULTI - LABEL CLASSIFICATION
Revised: 4.1. T ASK 1: EXTREME MULTI - LABEL CLASSIFICATION

Original: Settings: We compare the performance of MONGOOSE against SLIDE and networks without LSH (Full-Computation) on 3 large-scale multi-label classiﬁcation datasets: Wiki10-31k, Delicious200k,  and Wiki-325k.
Revised: Settings: We compare the performance of MONGOOSE against SLIDE and networks without LSH (Full-Computation) on 3 large-scale multi-label classiﬁcation datasets: Wiki10-31k, Delicious200k, Wiki-325k and Amz-670k.

Original: The datasets are obtained from the Extreme Classiﬁcation Repository (Bhatia et al., 2016), which is a benchmark for various recommendation systems.
Revised: The datasets are obtained from the Extreme Classiﬁcation Repository (Bhatia et al., 2016), which is a benchmark for various recommendation systems.

Original: The model is a fullyconnected network with one hidden-layer.
Revised: The model is a fully-connected network with one hidden-layer.

Original: The statistics of each dataset are in Appendix D.
Revised: The statistics of each dataset are in Appendix D.

Original: We report the best results of each framework in P @1 / 5 following metrics in (Bhatia et al., 2016).
Revised: We report the best results of each framework in P @1 / 5 following metrics in (Bhatia et al., 2016).

Original: [Figure]
Revised: [Figure]

Original: Results: Table 1 shows the test accuracy, wall-clock time to converge and memory consumption for MONGOOSE and other baselines.
Revised: Results: Table 1 shows the test accuracy, wall-clock time to converge and memory consumption for MONGOOSE and other baselines.

Original: Left plot in Figurevisualizes P @1 for MONGOOSE and other baselines during the training process.
Revised: Left plot in Figure 5 visualizes P @1 for MONGOOSE and other baselines during the training process.

Original: MONGOOSE achieves 1.8 × faster convergence and 1.5 × less memory consumption compared to FullComputation on Wiki10-31K,and 3 × faster convergence, 1.5 × less memory on Delicious-200K and Wiki-325K. SLIDE has a large accuracy drop without a substantial speedup for Wiki-325k, but our method has a 3 × speedup with no accuracy loss.
Revised: MONGOOSE achieves 1.8 × faster convergence and 1.5 × less memory consumption compared to FullComputation on Wiki10-31K, 3 × faster convergence, 1.5 × less memory on Delicious-200K and Wiki-325K and 6.5 × faster convergence, 2 × less memory on Amz670K. SLIDE has a large accuracy drop without a substantial speedup for Wiki-325k, but our method has a 3 × speedup with no accuracy loss.

Original: [Table]
Revised: [Table]

Original: Analysis: The above mainly focuses on MONGOOSE ’s end-to-end performance during training.
Revised: Analysis: The above mainly focuses on MONGOOSE ’s end-to-end performance during training.

Original: We also present closer analysis of a single linear layer during training in
Revised: We also present closer analysis of a single linear layer during training in

Original: Table 2.
Revised: Table 2.

Original: In isolation, we ﬁnd that MONGOOSE achieves up to 20 × faster convergence and 4.5 × less memory con- sumption compared to Full-Computation.
Revised: In isolation, we ﬁnd that MONGOOSE achieves up to 20 × faster convergence and 4.5 × less memory con- sumption compared to Full-Computation.

Original: The greater time and memory savings here compared to Table 1 are because we ignore the earlier large embedding layer of the model (Details are in Appendix D.2).
Revised: The greater time and memory savings here compared to Table 1 are because we ignore the earlier large em- bedding layer of the model (Details are in Appendix D.2).

Original: We make three observations to help explain the mem- ory and time savings from MONGOOSE .
Revised: We make three observations to help explain the memory and time savings from MONGOOSE .

Original: First, matrix multiplication is only performed on a small subset of neurons.
Revised: First, matrix multiplication is only performed on a small subset of neurons.

Original: The average number of neurons sampled in the output layer for MONGOOSE on Delicious-200K is around 6k compared to the total number of 200k.
Revised: The average number of neurons sampled in the output layer for MONGOOSE on Delicious-200K is around 6k compared to the total number of 200k.

Original: Second, our scheduler further reduces updating costs.
Revised: Second, our scheduler further reduces updating costs.

Original: On Wiki-325K, we observe updating frequency drops signiﬁcantly at the start of the second epoch.
Revised: On Wiki-325K, we observe updating frequency drops signiﬁcantly at the start of the second epoch.

Original: This is compatible with our observation of slow weight change in Figure 3.
Revised: This is compatible with our observation of slow weight change in Figure 3.

Original: Finally, learnable LSH improves LSH query efﬁciency and quality.
Revised: Finally, learnable LSH improves LSH query efﬁciency and quality.

Original: In addition, we observe that MONGOOSE obtains a slightly higher accuracy compared to Full-Computation, which is a side beneﬁt from activation sparsity (Ba & Frey, 2013).
Revised: In addition, we observe that MONGOOSE obtains a slightly higher accuracy compared to Full-Computation, which is a side beneﬁt from activation sparsity (Ba & Frey, 2013).

Original: 4.1. T ASK 2: EFFICIENT TRANSFORMER
Revised: 4.1. T ASK 2: EFFICIENT TRANSFORMER

Original: Settings: We run experiments on two tasks from Reformer (Kitaev et al., 2020).
Revised: Settings: We run experiments on two tasks from Reformer (Kitaev et al., 2020).

Original: The ﬁrst task is a synthetic sequence duplication task where inputs are of the form 0 w 0 w and w ∈ { 0 , ..., N } ∗ .
Revised: The ﬁrst task is a synthetic sequence duplication task where inputs are of the form 0 w 0 w and w ∈ { 0 , ..., N } ∗ .

Original: This task is useful for demonstrating the effectiveness of LSH attention: it requires non-local attention lookups and therefore cannot be solved by any model relying on sparse attention with a limited range ( e.g. , local attention).
Revised: This task is useful for demonstrating the effectiveness of LSH attention: it requires non-local attention lookups and therefore cannot be solved by any model relying on sparse attention with a limited range ( e.g. , local attention).

Original: We compare MONGOOSE with a 1-layer Reformer by replacing the dataindependent LSH in the attention layer with learnable LSH, updating the LSH parameters according to our smart scheduler.
Revised: We compare MONGOOSE with a 1-layer Reformer by replacing the dataindependent LSH in the attention layer with learnable LSH, updating the LSH parameters according to our smart scheduler.

Original: The second task is a standard language modeling task on the enwik8 dataset.
Revised: The second task is a standard language modeling task on the enwik8 dataset.

Original: For this second task, we use sequence length of 8192 and 10 LSH-attention layers.
Revised: For this second task, we use sequence length of 8192 and 10 LSH-attention layers.

Original: Results: We compare the test loss of MONGOOSE and Reformer on each task.
Revised: Results: We compare the test loss of MONGOOSE and Reformer on each task.

Original: The loss curve for the synthetic sequence duplication task is shown in Figure 5 (right).
Revised: The loss curve for the synthetic sequence duplication task is shown in Figure 5 (right).

Original: On this task, MONGOOSE achieves a loss value of 0.75, which is 0.5 better than Reformer.
Revised: On this task, MONGOOSE achieves a loss value of 0.75, which is 0.5 better than Reformer.

Original: Moreover, MONGOOSE only needs 20 % time to reach the same loss as Reformer (5 × speed-up).
Revised: Moreover, MONGOOSE only needs 20 % time to reach the same loss as Reformer (5 × speed-up).

Original: Similarly, on enwik8, Reformer only reaches 2.65 in perplexity, while MONGOOSE achieves 2.59 in the same setting.
Revised: Similarly, on enwik8, Reformer only reaches 2.65 in perplexity, while MONGOOSE achieves 2.59 in the same setting.

Original: Both MONGOOSE and Reformer (with LSH) save 6 × memory over the original Transformer in above two tasks (more details are in Appendix D.3).
Revised: Both MONGOOSE and Reformer (with LSH) save 6 × memory over the original Transformer in above two tasks (more details are in Appendix D.3).

Original: 4. S TUDY ON SMART SCHEDULER
Revised: 4. S TUDY ON SMART SCHEDULER

Original: [Table]
Revised: [Table]

Original: In this section, we demonstrate the superiority of our smart scheduling algorithm over the two baselines we mentioned in Section 1 on Wiki-325K.
Revised: In this section, we demonstrate the superiority of our smart scheduling algorithm over the two baselines we mentioned in Section 1 on Wiki-325K.

Original: Effect of Smart Scheduler: The ﬁrst approach (“Infrequent Scheduler”) performs an update on a ﬁxed schedule, every 1000 iterations (6 times per epoch on Wiki-325K).
Revised: Effect of Smart Scheduler: The ﬁrst approach (“Infrequent Scheduler”) performs an update on a ﬁxed schedule, every 1000 iterations (6 times per epoch on Wiki-325K).

Original: The other baseline performs an update every iteration (“W/O Scheduler”).
Revised: The other baseline performs an update every iteration (“W/O Scheduler”).

Original: We compare speed (number of batches processed per second) and accuracy ( P @1 and P @5 ) of these schedulers in Table 3.
Revised: We compare speed (number of batches processed per second) and accuracy ( P @1 and P @5 ) of these schedulers in Table 3.

Original: On average, MONGOOSE Scheduler reduces more than 97 % of the update overhead compared to W/O Scheduler, leading an end-to-end 20 × speed-up.
Revised: On average, MONGOOSE Scheduler reduces more than 97 % of the update overhead compared to W/O Scheduler, leading an end-to-end 20 × speed-up.

Original: Meanwhile, although Infrequent Scheduler is faster in speed, it fails to reach the same task accuracy as the other approaches.
Revised: Meanwhile, although Infrequent Scheduler is faster in speed, it fails to reach the same task accuracy as the other approaches.

Original: In fact, MONGOOSE Scheduler achieves Infrequent Scheduler’s ﬁnal accuracy with 66 % less time.
Revised: In fact, MONGOOSE Scheduler achieves Infrequent Scheduler’s ﬁnal accuracy with 66 % less time.

Original: We can see that MONGOOSE substantially reduces LSH update overhead while maintaining performance.
Revised: We can see that MONGOOSE substantially reduces LSH update overhead while maintaining performance.

Original: Inﬂuence of the Parameters: Recall that in Algorithm 1, (cid:15) mds is the key parameter to our smart scheduler because the number of neurons exceeding this threshold determines if LSH updates are needed.
Revised: Inﬂuence of the Parameters: Recall that in Algorithm 1, (cid:15) mds is the key parameter to our smart scheduler because the number of neurons exceeding this threshold determines if LSH updates are needed.

Original: In practice, instead of choosing a single (cid:15) mds , we use some fraction of the (cid:96) 2 norm of the weights as the update threshold.
Revised: In practice, instead of choosing a single (cid:15) mds , we use some fraction of the (cid:96) 2 norm of the weights as the update threshold.

Original: We empirically observe from the previous experiments that
Revised: We empirically observe from the previous experiments that

Original: M ONGOOSE is robust to any choice of the fraction in the range 0.1-0.3.
Revised: M ONGOOSE is robust to any choice of the fraction in the range 0.1-0.3.

Original: Thus, it can consistently outperform other baselines without heavily tuning this parameter for different datasets.
Revised: Thus, it can consistently outperform other baselines without heavily tuning this parameter for different datasets.

Original: 4. S TUDY ON LEARNABLE LSH
Revised: 4. S TUDY ON LEARNABLE LSH

Original: Effect of Learnable LSH:
Revised: Effect of Learnable LSH:

Original: In this section, we compare the NNS performance of MONGOOSE ’s learnable LSH scheme to classical LSH on their NNS ability using the same number of hash functions.
Revised: In this section, we compare the NNS performance of MONGOOSE ’s learnable LSH scheme to classical LSH on their NNS ability using the same number of hash functions.

Original: In Figure 6 (left) we plot the difference between average inner products computed using neurons retrieved by learnable LSH ( MONGOOSE ) and random LSH ( SLIDE ) and inner products using the same number of randomly chosen neurons.
Revised: In Figure 6 (left) we plot the difference between average inner products computed using neurons retrieved by learnable LSH ( MONGOOSE ) and random LSH ( SLIDE ) and inner products using the same number of randomly chosen neurons.

Original: A negative value on these plots indicates that
Revised: A negative value on these plots indicates that

Original: LSH is performing worse than random sampling.
Revised: LSH is performing worse than random sampling.

Original: We can see that at the beginning of training, whenweights are initialized randomly, data-independent LSH is sufﬁcient to partition the data. However, [Figure] [Table]
Revised: We can see that at the beginning of training, when

Original: [Figure]
Revised: [Figure]

Original: [Table] [Figure] [Table] [Figure]
Revised: weights are initialized randomly, data-independent LSH is sufﬁcient to partition the data.

Original: [Table]

Original: [Figure] [Table] [Figure] [Table]
Revised: However, as training progresses, SimHash used in SLIDE fails to adapt to the evolving weight distribution

Original: [Figure]

Original: [Table]

Original: [Figure]

Original: [Table]

Original: [Figure]

Original: [Table]

Original: [Figure]

Original: [Table]

Original: [Figure]

Original: [Table]

Original: [Figure]

Original: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table] [Figure] [Table] as training progresses, SimHash used in SLIDE fails to adapt to the evolving weight distribution, while our selection better approximates the full softmax
Revised: while our selection better approximates the full softmax

Original: [Figure]
Revised: [Figure]

Original: We conduct the same experiments for enwik8 language modeling.
Revised: We conduct the same experiments for enwik8 language modeling.

Original: In Figure 6 (right), we can see that data-dependent hashing computes higher inner products within each bucket as training progresses.
Revised: In Figure 6 (right), we can see that data-dependent hashing computes higher inner products within each bucket as training progresses.

Original: For attention layers in the Transformer models, random LSH fails to approximate attention well in certain layers due to a skewed attention distribution (this phenomenon was also explained in Section 3.3.
Revised: For attention layers in the Transformer models, random LSH fails to approximate attention well in certain layers due to a skewed attention distribution (this phenomenon was also explained in Section 3.3.2 and

Original: and Figure 4).
Revised: Figure 4).

Original: In attention layers, our selection produces a better characterization of the full attention.
Revised: In attention layers, our selection produces a better characterization of the full attention.

Original: [Table]
Revised: [Table]

Original: Inﬂuence of the Parameters: For LSH-based NNS-ds, the number of hashes is one key parameter to balance NNS quality and efﬁciency.
Revised: Inﬂuence of the Parameters: For LSH-based NNS-ds, the number of hashes is one key parameter to balance NNS quality and efﬁciency.

Original: In Table 5 we study the effects of varying the number of hashes in MONGOOSE .
Revised: In Table 5 we study the effects of varying the number of hashes in MONGOOSE .

Original: We observe that compared to our best empirical choice ( 10 hashes), both a smaller (8) and larger number (12) of hashes increases convergence time.
Revised: We observe that compared to our best empirical choice ( 10 hashes), both a smaller (8) and larger number (12) of hashes increases convergence time.

Original: MONGOOSE with fewer hashes selects more unnecessary neurons and lowers training efﬁciency.
Revised: MONGOOSE with fewer hashes selects more unnecessary neurons and lowers training efﬁciency.

Original: A larger number of hashes might miss necessary neurons with higher probabilities and would similarly take longer to converge.
Revised: A larger number of hashes might miss necessary neurons with higher probabilities and would similarly take longer to converge.

Original: Another important parameter is the threshold for positive and negative examples to train hash functions.
Revised: Another important parameter is the threshold for positive and negative examples to train hash functions.

Original: In our experiments, top-1 and bottom-1 are the most effective choices (top-3 and top-5 do not make much difference).
Revised: In our experiments, top-1 and bottom-1 are the most effective choices (top-3 and top-5 do not make much difference).

Original: In general, this threshold in our design is robust and does not require heavy tuning.
Revised: In general, this threshold in our design is robust and does not require heavy tuning.

Original: Other data-dependent NNS data structures: In Section 1, we have brieﬂy discussed that most existing data-dependent NNS-ds have high update costs.
Revised: Other data-dependent NNS data structures: In Section 1, we have brieﬂy discussed that most existing data-dependent NNS-ds have high update costs.

Original: We choose an LSH-based data structure in MONGOOSE because update overhead for hash-table-based NNS-ds are less compared to other methods (e.g. graph-based) in practice (Wang et al., 2018;
Revised: We choose an LSH-based data structure in MONGOOSE because update overhead for hash-table-based NNS-ds are less compared to other methods (e.g. graph-based) in practice (Wang et al., 2018;

Original: Gong et al., 2020).
Revised: Gong et al., 2020).

Original: For example, in a recent paper, (Dong et al., 2019) proposed a k -nearest-neighbor graph-based algorithm for NNS that requires several hours to process 1 million data points.
Revised: For example, in a recent paper, (Dong et al., 2019) proposed a k -nearest-neighbor graph-based algorithm for NNS that requires several hours to process 1 million data points.

Original: Although a recent work, HNSW (Malkov & Yashunin, 2018), is relatively fast in preprocessing and updating step among graph-based datadependent NNS-ds empirically, it does not have guarantees (Fu et al., 2017).
Revised: Although a recent work, HNSW (Malkov & Yashunin, 2018), is relatively fast in preprocessing and updating step among graph-based datadependent NNS-ds empirically, it does not have guarantees (Fu et al., 2017).

Original: In Table 4 we compare the update times of the learnable LSH and HNSW , and their performance during training when both are equipped with the MONGOOSE smart scheduler.
Revised: In Table 4 we compare the update times of the learnable LSH and HNSW , and their performance during training when both are equipped with the MONGOOSE smart scheduler.

Original: We can see that HNSW ’s updating overhead grows exponentially and becomes over 20 × learnable LSH’s overhead in 350k neurons layer.
Revised: We can see that HNSW ’s updating overhead grows exponentially and becomes over 20 × learnable LSH’s overhead in 350k neurons layer.

Original: Moreover, HNSW achieves lower accuracy with much longer convergence time.
Revised: However, it is an ongoing work to reduce the updating overhead of graph-based NNS-ds, so it could be potentially included in MONGOOSE later.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: In this work, we make a novel observation on the slowly changing LSH hash codes during LSHNN training.
Revised: In this work, we make a novel observation on the slowly changing LSH hash codes during LSHNN training.

Original: Based on the observation, we present MONGOOSE , a learnable LSH framework that outperforms original LSH methods in efﬁcient neural network training.
Revised: Based on the observation, we present MONGOOSE , a learnable LSH framework that outperforms original LSH methods in efﬁcient neural network training.

Original: We demonstrate two key designs in the framework.
Revised: We demonstrate two key designs in the framework.

Original: One is a smart scheduling algorithm which reduces the LSH update overhead.
Revised: One is a smart scheduling algorithm which reduces the LSH update overhead.

Original: The other one is a low cost learnable LSH which can improve the query efﬁciency and hashing quality.
Revised: The other one is a low cost learnable LSH which can improve the query efﬁciency and hashing quality.

Original: We empirically verify the effectiveness of our framework on both recommendation and language modeling tasks.
Revised: We empirically verify the effectiveness of our framework on both recommendation and language modeling tasks.

Original: The slowly changing observation along with the smart scheduler in
Revised: The slowly changing observation along with the smart scheduler in

Original: MONGOOSE opens up the opportunity of designing new NNS-ds for dynamic similarity search for efﬁcient NN training.
Revised: MONGOOSE opens up the opportunity of designing new NNS-ds for dynamic similarity search for efﬁcient NN training.

Original: Appendix Table of Contents A
Revised: Appendix Table of Contents A

Original: Slow
Revised: Slow

Original: observations
Revised: observations

Original: B Dynamic maintenance data structure
Revised: B Dynamic maintenance data structure

Original: B.1 Illustration . .
Revised: B.1 Illustration . .

Original: B.2 Main result . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
Revised: B.2 Main result . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

Original: B.3 Proof of theorem B.1 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
Revised: B.3 Proof of theorem B.1 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

Original: B.4 Bounding w move . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
Revised: B.4 Bounding w move . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

Original: B.5 Bounding v move . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
Revised: B.5 Bounding v move . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

Original: B.6 Potential function ψ . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
Revised: B.6 Potential function ψ . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

Original: B.7 Examples . . .
Revised: B.7 Examples . . .

Original: C Learnable lsh
Revised: C Learnable lsh

Original: C.1 Learning hash functions . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
Revised: C.1 Learning hash functions . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

Original: C.2 Observations on attention distribution . . . . . . . . . . . . . . . . . . . . . . .
Revised: C.2 Observations on attention distribution . . . . . . . . . . . . . . . . . . . . . . .

Original: D Experiments details
Revised: D Experiments details

Original: D.1 Data statistics . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
Revised: D.1 Data statistics . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

Original: D.2 Additional results on the extreme-classiﬁcation task . . . . . . . . . . . . . . .
Revised: D.2 Additional results on the extreme-classiﬁcation task . . . . . . . . . . . . . . .

Original: D.3 Additional results on the language modeling task . . . . . . . . . . . . . . . . .
Revised: D.3 Additional results on the language modeling task . . . . . . . . . . . . . . . . .

Original: E Related work
Revised: E Related work

Original: E.
Revised: E.

Original: [Table]
Revised: [Table]

Original: E.
Revised: E.

Original: [Table] [Table]
Revised: [Table] [Table]

Original: [Table]
Revised: [Table]

Original: training . . .
Revised: training . . .

Original: F
Revised: F

Original: [Table]
Revised: [Table]

Original: GPU implementation
Revised: GPU implementation

Original: G Broader Impact
Revised: G Broader Impact

Original: A S LOW CHANGE OBSERVATIONS
Revised: A S LOW CHANGE OBSERVATIONS

Original: In this section, we analyze the relative distance between iterates of model parameters and connect it with LSH updates throughout training.
Revised: In this section, we analyze the relative distance between iterates of model parameters and connect it with LSH updates throughout training.

Original: We target at making full use of our slowly change characterization for MONGOOSE .
Revised: We target at making full use of our slowly change characterization for MONGOOSE .

Original: Settings.
Revised: Settings.

Original: We report ∆ W throughout training on several benchmark models.
Revised: We report ∆ W throughout training on several benchmark models.

Original: We follow the standard training procedures of fully-connected networks on Wiki325k (Partalas et al., 2015) and Transformer (Vaswani et al., 2017) on enwik8.
Revised: We follow the standard training procedures of fully-connected networks on Wiki325k (Partalas et al., 2015) and Transformer (Vaswani et al., 2017) on enwik8.

Original: We choose W in the last linear layer for the fully-connected model and W in the projection layer before attention for the transformer model to correspond experiments in section 4.
Revised: We choose W in the last linear layer for the fully-connected model and W in the projection layer before attention for the transformer model to correspond experiments in section 4.

Original: Results.
Revised: Results.

Original: We plot our results in Figure 3.
Revised: We plot our results in Figure 3.

Original: The left-two and right-most plots show that during the initial steps of the training, ∆ W is relatively high, but quickly drops and ﬂattens out afterwards.
Revised: The left-two and right-most plots show that during the initial steps of the training, ∆ W is relatively high, but quickly drops and ﬂattens out afterwards.

Original: The middle and right-most plots exhibit the hash code change of W in the hamming distance along with the training.
Revised: The middle and right-most plots exhibit the hash code change of W in the hamming distance along with the training.

Original: The pattern matches with ∆ W but has a direct connection with the LSH update overhead.
Revised: The pattern matches with ∆ W but has a direct connection with the LSH update overhead.

Original: This is also consistent with LSH theory that the hash code collision probability of two vectors equals to the vectors’ certain similarity, e.g. , angular.
Revised: This is also consistent with LSH theory that the hash code collision probability of two vectors equals to the vectors’ certain similarity, e.g. , angular.

Original: Note the above observations are made based on angular distance LSH.
Revised: Note the above observations are made based on angular distance LSH.

Original: The above phenomenon suggests that if there exists an optimal scheduler to update the data structures adaptively based on the actual demand, the overhead by LSH updates can be largely reduced.
Revised: The above phenomenon suggests that if there exists an optimal scheduler to update the data structures adaptively based on the actual demand, the overhead by LSH updates can be largely reduced.

Original: Furthermore, this opens the door to make LSH learnable in order to improve query efﬁciency.
Revised: Furthermore, this opens the door to make LSH learnable in order to improve query efﬁciency.

Original: Since the LSH update time is also closely related to the query time, the overall update overhead might still be reduced after considering learning costs, if the updates are well scheduled.
Revised: Since the LSH update time is also closely related to the query time, the overall update overhead might still be reduced after considering learning costs, if the updates are well scheduled.

Original: Several related works (Jastrzebski et al., 2018; Agarwal et al., 2020) have also observed the phenomenon of a sharp drop of weight changes at the early stages of training.
Revised: Several related works (Jastrzebski et al., 2018; Agarwal et al., 2020) have also observed the phenomenon of a sharp drop of weight changes at the early stages of training.

Original: One (Jastrzebski et al., 2018) conjectures that initially SGD visits increasingly sharp regions, reaching a maximum sharpness determined by both the learning rate and the batch-size of SGD and infer that it optimizes faster while ﬁnding a good sharp region.
Revised: One (Jastrzebski et al., 2018) conjectures that initially SGD visits increasingly sharp regions, reaching a maximum sharpness determined by both the learning rate and the batch-size of SGD and infer that it optimizes faster while ﬁnding a good sharp region.

Original: It is also discussed in a recent paper (Agarwal et al., 2020), which connects the phenomenon to critical learning periods deﬁned in (Achille et al., 2017).
Revised: It is also discussed in a recent paper (Agarwal et al., 2020), which connects the phenomenon to critical learning periods deﬁned in (Achille et al., 2017).

Original: This phenomenon can be observed on text (Transformers), recommendation (Embedding Models) and image (CNN) data and thereby it is relatively general
Revised: This phenomenon can be observed on text (Transformers), recommendation (Embedding Models) and image (CNN) data and thereby it is relatively general

Original: [Figure]
Revised: [Figure]

Original: B D YNAMIC MAINTENANCE DATA STRUCTURE
Revised: B D YNAMIC MAINTENANCE DATA STRUCTURE

Original: The idea of dynamic maintenance has been successfully applied to matrix inversion problem, which served as a central component in linear program solvers (Cohen et al., 2019; Lee et al., 2019;
Revised: The idea of dynamic maintenance has been successfully applied to matrix inversion problem, which served as a central component in linear program solvers (Cohen et al., 2019; Lee et al., 2019;

Original: Brand et al., 2020; Jiang et al., 2020c), cutting plane method (Jiang et al., 2020b), and semi-deﬁnite program solvers (Jiang et al., 2020a).
Revised: Brand et al., 2020; Jiang et al., 2020c), cutting plane method (Jiang et al., 2020b), and semi-deﬁnite program solvers (Jiang et al., 2020a).

Original: In this work, we adopt it for dynamic LSH maintenance problem.
Revised: In this work, we adopt it for dynamic LSH maintenance problem.

Original: Since the underlying problem is completely different, it requires several non-trivial generalizations and more ﬁne-grained analysis (see Lemma B.2, Lemma B.3, Lemma B.4).
Revised: Since the underlying problem is completely different, it requires several non-trivial generalizations and more ﬁne-grained analysis (see Lemma B.2, Lemma B.3, Lemma B.4).

Original: Moreover, in our proof, we care about the absolute distance of between vectors (i.e. (cid:107) w i − v i (cid:107) 2 ) rather than the relative difference between real numbers Cohen et al.
Revised: Moreover, in our proof, we care about the absolute distance of between vectors (i.e. (cid:107) w i − v i (cid:107) 2 ) rather than the relative difference between real numbers Cohen et al.

Original: (2019), and therefore, we need a new construction of the potential function (see Lemma B.8).
Revised: (2019), and therefore, we need a new construction of the potential function (see Lemma B.8).

Original: The rest of this section is organized as follows
Revised: The rest of this section is organized as follows

Original: • In Section B.1, we provide an illustration of the smart scheduler.
Revised: • In Section B.1, we provide an illustration of the smart scheduler.

Original: • In Section B.2, we state our main results.
Revised: • In Section B.2, we state our main results.

Original: • In Section B.3, we give a high-level overview on the correctness and the running time of the algorithm.
Revised: • In Section B.3, we give a high-level overview on the correctness and the running time of the algorithm.

Original: In Section B.4, we give an analyze on the movement of w vector.
Revised: In Section B.4, we give an analyze on the movement of w vector.

Original: • In Section B.5 we give analyze on the movement of v vector.
Revised: • In Section B.5 we give analyze on the movement of v vector.

Original: • In Section B.6, we discussed over the choice of potential function and analysis its property.
Revised: • In Section B.6, we discussed over the choice of potential function and analysis its property.

Original: • In Section B.7, we compare the performance our algorithm with sequential updating strategy/batch updating strategy.
Revised: • In Section B.7, we compare the performance our algorithm with sequential updating strategy/batch updating strategy.

Original: B. I LLUSTRATION
Revised: B. I LLUSTRATION

Original: We present a demonstration of our smart scheduler at Figure 8.
Revised: We present a demonstration of our smart scheduler at Figure 8.

Original: As shown in the top, if the change of data vectors (denoted as points with colors) stays within the partitions given by hash functions, the scheduler would not update current LSH tables.
Revised: As shown in the top, if the change of data vectors (denoted as points with colors) stays within the partitions given by hash functions, the scheduler would not update current LSH tables.

Original: However, if the data vectors change such that they cross partition boundaries, the scheduler updates the hash functions.
Revised: However, if the data vectors change such that they cross partition boundaries, the scheduler updates the hash functions.

Original: B. M AIN RESULT
Revised: B. M AIN RESULT

Original: Theorem B.1 (Locality sensitive hashing maintenance, Formal statement of Theorem 3.3) .
Revised: Theorem B.1 (Locality sensitive hashing maintenance, Formal statement of Theorem 3.3) .

Original: For any constant c 1 , c 2 ( c 1 > c 2 ) there is a dynamic data structure (Algorithm 1) that achieves ( c 1 , c 2 ) - accuracy.
Revised: For any constant c 1 , c 2 ( c 1 > c 2 ) there is a dynamic data structure (Algorithm 1) that achieves ( c 1 , c 2 ) - accuracy.

Original: The data structure takes (cid:101)
Revised: The data structure takes (cid:101)

Original: O ( dn 1+ ρ ) time to initialize and each call of Q UERY ( h ) takes time (cid:101)
Revised: O ( dn 1+ ρ ) time to initialize and each call of Q UERY ( h ) takes time (cid:101)

Original: O ( n ρ d ) .
Revised: O ( n ρ d ) .

Original: By taking a = min { ρ, α } and
Revised: By taking a = min { ρ, α } and

Original: [Equation]
Revised: [Equation]

Original: The amortized expected time per call of U PDATE ( w ) is at most
Revised: The amortized expected time per call of U PDATE ( w ) is at most

Original: [Equation]
Revised: [Equation]

Original: B. P ROOF OF THEOREM B. 
Revised: B. P ROOF OF THEOREM B.

Original: We ﬁrst give a rough explanation on the proof of Theorem 3.3.
Revised: We ﬁrst give a rough explanation on the proof of Theorem 3.3.

Original: For intuition, we consider the case
Revised: For intuition, we consider the case

Original: C 1 = Θ(1) , C 2 = Θ(1) , and (cid:15) mp = Θ(1) in this explanation.
Revised: C 1 = Θ(1) , C 2 = Θ(1) , and (cid:15) mp = Θ(1) in this explanation.

Original: The amortized time analysis is based on a potential function that measures the distance of the approximate vector v and the target vector w .
Revised: The amortized time analysis is based on a potential function that measures the distance of the approximate vector v and the target vector w .

Original: We will show that
Revised: We will show that

Original: • The cost to update the LSH data structure is proportional to the decrease of the potential.
Revised: • The cost to update the LSH data structure is proportional to the decrease of the potential.

Original: • Each call to query increase the potential by a ﬁxed amount.
Revised: • Each call to query increase the potential by a ﬁxed amount.

Original: Combining both together gives the amortized running time bound of our data structure.
Revised: Combining both together gives the amortized running time bound of our data structure.

Original: Now, we explain the deﬁnition of the potential.
Revised: Now, we explain the deﬁnition of the potential.

Original: Consider the k -th round of the algorithm.
Revised: Consider the k -th round of the algorithm.

Original: For all i ∈ [ n ] , we deﬁne x ( k ) i =
Revised: For all i ∈ [ n ] , we deﬁne x ( k ) i =

Original: w ( k ) i − v ( k ) i .
Revised: w ( k ) i − v ( k ) i .

Original: Note that (cid:107) x
Revised: Note that (cid:107) x

Original: ( k ) i (cid:107) 2 measures the distance between w ( k ) i and v ( k ) i .
Revised: ( k ) i (cid:107) 2 measures the distance between w ( k ) i and v ( k ) i .

Original: Our algorithm ﬁxes the indices with largest error (cid:107) x ( k ) i (cid:107) 2 .
Revised: Our algorithm ﬁxes the indices with largest error (cid:107) x ( k ) i (cid:107) 2 .

Original: To capture the fact that updating in a larger batch is more efﬁcient, we deﬁne the potential as a weighted combination of the error where we put more weight to higher x ( k ) i .
Revised: To capture the fact that updating in a larger batch is more efﬁcient, we deﬁne the potential as a weighted combination of the error where we put more weight to higher x ( k ) i .

Original: Formally, we sort the coordinates of x ( k ) such that (cid:107) x
Revised: Formally, we sort the coordinates of x ( k ) such that (cid:107) x

Original: ( k ) i (cid:107) 2 ≥ (cid:107) x
Revised: ( k ) i (cid:107) 2 ≥ (cid:107) x

Original: ( k ) i +1 (cid:107) 2 and deﬁne the potential by
Revised: ( k ) i +1 (cid:107) 2 and deﬁne the potential by

Original: [Equation]
Revised: [Equation]

Original: where g i are positive decreasing numbers to be chosen and ψ is a symmetric ( ψ ( x ) = ψ ( − x ) ) positive function that increases on both sides.
Revised: where g i are positive decreasing numbers to be chosen and ψ is a symmetric ( ψ ( x ) = ψ ( − x ) ) positive function that increases on both sides.

Original: For intuition, one can think ψ ( x ) behaves roughly like | x | .
Revised: For intuition, one can think ψ ( x ) behaves roughly like | x | .

Original: B.3. P ROOF OF CORRECTNESS
Revised: B.3. P ROOF OF CORRECTNESS

Original: We prove the correctness of Theorem 3.3, we will defer some simple calculations into later sections.
Revised: We prove the correctness of Theorem 3.3, we will defer some simple calculations into later sections.

Original: Deﬁnition of x and y .
Revised: Deﬁnition of x and y .

Original: Consider the k -th round of the algorithm.
Revised: Consider the k -th round of the algorithm.

Original: For all i ∈ [ n ] , we deﬁne x ( k ) i ∈
Revised: For all i ∈ [ n ] , we deﬁne x ( k ) i ∈

Original: R d , x ( k +1) i ∈ R d and y ( k ) i ∈ R d as follows:
Revised: R d , x ( k +1) i ∈ R d and y ( k ) i ∈ R d as follows:

Original: [Equation]
Revised: [Equation]

Original: Note that the difference between x ( k ) i and y ( k ) i is that w is changing.
Revised: Note that the difference between x ( k ) i and y ( k ) i is that w is changing.

Original: The difference between y ( k ) i and x ( k +1)
Revised: The difference between y ( k ) i and x ( k +1)

Original: i is that v is changing.
Revised: i is that v is changing.

Original: Assume sorting.
Revised: Assume sorting.

Original: Assume the coordinates of vector x ( k ) ∈ R n × d are sorted such that (cid:107) x ( k ) i (cid:107) 2 ≥ (cid:107) x
Revised: Assume the coordinates of vector x ( k ) ∈ R n × d are sorted such that (cid:107) x ( k ) i (cid:107) 2 ≥ (cid:107) x

Original: ( k ) i +1 (cid:107) 2 .
Revised: ( k ) i +1 (cid:107) 2 .

Original: Let τ and π are permutations such that (cid:107) x
Revised: Let τ and π are permutations such that (cid:107) x

Original: ( k +1) τ
Revised: ( k +1) τ

Original: ( i ) (cid:107) 2 ≥ (cid:107) x
Revised: ( i ) (cid:107) 2 ≥ (cid:107) x

Original: ( k +1) τ
Revised: ( k +1) τ

Original: ( i +1) (cid:107) 2 and (cid:107) y ( k ) π
Revised: ( i +1) (cid:107) 2 and (cid:107) y ( k ) π

Original: ( i ) (cid:107) 2 ≥ (cid:107)
Revised: ( i ) (cid:107) 2 ≥ (cid:107)

Original: [Equation]
Revised: [Equation]

Original: Deﬁnition of Potential function.
Revised: Deﬁnition of Potential function.

Original: Let ψ : R → R be deﬁned by
Revised: Let ψ : R → R be deﬁned by

Original: [Equation]
Revised: [Equation]

Original: We deﬁne the potential at the k -th round by
Revised: We deﬁne the potential at the k -th round by

Original: [Equation]
Revised: [Equation]

Original: where τ k ( i ) is the permutation
Revised: where τ k ( i ) is the permutation

Original: such that (cid:107) x
Revised: such that (cid:107) x

Original: ( k ) τ k
Revised: ( k ) τ k

Original: ( i ) (cid:107) 2 ≥ (cid:107)
Revised: ( i ) (cid:107) 2 ≥ (cid:107)

Original: [Equation]
Revised: [Equation]

Original: Bounding the potential.
Revised: Bounding the potential.

Original: We can express Ψ k +1 − Ψ k as follows:
Revised: We can express Ψ k +1 − Ψ k as follows:

Original: [Equation]
Revised: [Equation]

Original: Now, using Lemma B.5 and B.7, and the fact that Ψ 0 = 0 and Ψ T ≥ 0 , with Eq. 6, we get
Revised: Now, using Lemma B.5 and B.7, and the fact that Ψ 0 = 0 and Ψ T ≥ 0 , with Eq. 6, we get

Original: [Equation]
Revised: [Equation]

Original: where the third step follows by Lemma B.5 and Lemma B.7 and r k is the number of coordinates we update during that iteration.
Revised: where the third step follows by Lemma B.5 and Lemma B.7 and r k is the number of coordinates we update during that iteration.

Original: Therefore, we get,
Revised: Therefore, we get,

Original: [Equation]
Revised: [Equation]

Original: B.3. I NITIALIZATION TIME , UPDATE TIME , QUERY TIME
Revised: B.3. I NITIALIZATION TIME , UPDATE TIME , QUERY TIME

Original: To formalize the amortized runtime proof, we ﬁrst analyze the initialization time (Lemma B.2), update time (Lemma B.3), and query time (Lemma B.4) of our maintenance data-structure.
Revised: To formalize the amortized runtime proof, we ﬁrst analyze the initialization time (Lemma B.2), update time (Lemma B.3), and query time (Lemma B.4) of our maintenance data-structure.

Original: Lemma B.2 (Initialization time) .
Revised: Lemma B.2 (Initialization time) .

Original: The initialization time of data-structure M AINTAIN (Algorithm 1) is O
Revised: The initialization time of data-structure M AINTAIN (Algorithm 1) is O

Original: ( dn 1+ ρ ) .
Revised: ( dn 1+ ρ ) .

Original: Lemma B.3 (Update time) .
Revised: Lemma B.3 (Update time) .

Original: The update time of data-structure M AINTAIN (Algorithm 1) is O ( rg r )
Revised: The update time of data-structure M AINTAIN (Algorithm 1) is O ( rg r )

Original: +
Revised: +

Original: O ( Sd + S log n ) where r is the number of indices we updated in v , S is the number of weights changed.
Revised: O ( Sd + S log n ) where r is the number of indices we updated in v , S is the number of weights changed.

Original: The later two terms are dominated by the updating time of neural network.
Revised: The later two terms are dominated by the updating time of neural network.

Original: Proof.
Revised: Proof.

Original: We change r ( r ≥ n a ) terms each time, and the running time for updating the LSH is rg r .
Revised: We change r ( r ≥ n a ) terms each time, and the running time for updating the LSH is rg r .

Original: We need to update v and calculate its (cid:96) 2 norm every time, and the computational cost is Sd .
Revised: We need to update v and calculate its (cid:96) 2 norm every time, and the computational cost is Sd .

Original: We also need to maintain an order on the error y ( k ) i , using some standard data structure (like Fibonacci heap), this can be done in S log n .
Revised: We also need to maintain an order on the error y ( k ) i , using some standard data structure (like Fibonacci heap), this can be done in S log n .

Original: Thus the total running time per update is O ( rg r )
Revised: Thus the total running time per update is O ( rg r )

Original: + O ( Sd + S log n ) , the second term is bounded by the updating cost of neural network training, since calculate the gradient takes at least O ( Sd ) and we usually have d (cid:29) log n .
Revised: + O ( Sd + S log n ) , the second term is bounded by the updating cost of neural network training, since calculate the gradient takes at least O ( Sd ) and we usually have d (cid:29) log n .

Original: Lemma B.4 (Query time) .
Revised: Lemma B.4 (Query time) .

Original: The query time of data-structure M AINTAIN (Algorithm 1) is O ( n ρ d + n a d ) .
Revised: The query time of data-structure M AINTAIN (Algorithm 1) is O ( n ρ d + n a d ) .

Original: B. B OUNDING w MOVE
Revised: B. B OUNDING w MOVE

Original: Lemma B.5 ( w move) .
Revised: Lemma B.5 ( w move) .

Original: We have
Revised: We have

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: Observe that since the errors (cid:107) x ( k ) i (cid:107) 2 are sorted in descending order, and ψ ( x ) is symmetric and non-decreasing function, thus ψ ( x ( k ) i ) is also in decreasing order.
Revised: Observe that since the errors (cid:107) x ( k ) i (cid:107) 2 are sorted in descending order, and ψ ( x ) is symmetric and non-decreasing function, thus ψ ( x ( k ) i ) is also in decreasing order.

Original: In addition, note that g is decreasing, we have
Revised: In addition, note that g is decreasing, we have

Original: [Equation]
Revised: [Equation]

Original: Hence we have
Revised: Hence we have

Original: [Equation]
Revised: [Equation]

Original: Thus
Revised: Thus

Original: [Equation]
Revised: [Equation]

Original: we complete the proof of w move Lemma.
Revised: we complete the proof of w move Lemma.

Original: It remains to prove the following Lemma,
Revised: It remains to prove the following Lemma,

Original: Lemma B.6.
Revised: Lemma B.6.

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: We separate the term into two:
Revised: We separate the term into two:

Original: [Equation]
Revised: [Equation]

Original: For the ﬁrst term
Revised: For the ﬁrst term

Original: [Equation]
Revised: [Equation]

Original: we have
Revised: we have

Original: [Equation]
Revised: [Equation]

Original: where the ﬁrst step follows from the Mean value theorem, the second step follows from the deﬁnition of L 2 (see Part 4 of potential lemma B.8), the last step follows from the deﬁnition of y ( k ) i .
Revised: where the ﬁrst step follows from the Mean value theorem, the second step follows from the deﬁnition of L 2 (see Part 4 of potential lemma B.8), the last step follows from the deﬁnition of y ( k ) i .

Original: Next
Revised: Next

Original: [Equation]
Revised: [Equation]

Original: we denote γ i =
Revised: we denote γ i =

Original: Var[ w ( k +1) i ] .
Revised: Var[ w ( k +1) i ] .

Original: Summing over i and taking conditional expectation given w ( k ) on both sides, we get n (cid:88) i =
Revised: Summing over i and taking conditional expectation given w ( k ) on both sides, we get n (cid:88) i =

Original: [Equation]
Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: The ﬁrst step follows from Eq.
Revised: The ﬁrst step follows from Eq.

Original: 9, the second step follows from the linearity of expectation and the deﬁnition of γ i , the third step follows from Cauchy-Shwarz inequality, the fourth step follows from
Revised: 9, the second step follows from the linearity of expectation and the deﬁnition of γ i , the third step follows from Cauchy-Shwarz inequality, the fourth step follows from

Original: L 2 =
Revised: L 2 =

Original: O (1 /(cid:15) mds )
Revised: O (1 /(cid:15) mds )

Original: (see Lemma B.8) and Eq. 1.
Revised: (see Lemma B.8) and Eq. 1.

Original: For the second term, conditioning on w ( k ) i , we have
Revised: For the second term, conditioning on w ( k ) i , we have

Original: [Equation]
Revised: [Equation]

Original: The ﬁrst inequality follows from the L 1 -Lipschitz continuity of Ψ (see part 4 of Lemma B.8).
Revised: The ﬁrst inequality follows from the L 1 -Lipschitz continuity of Ψ (see part 4 of Lemma B.8).

Original: The second step follows from the deﬁnition of y ( k ) i and x ( k ) i .
Revised: The second step follows from the deﬁnition of y ( k ) i and x ( k ) i .

Original: Summing over i , we get
Revised: Summing over i , we get

Original: [Equation]
Revised: [Equation]

Original: The ﬁrst step follows from Eq.
Revised: The ﬁrst step follows from Eq.

Original: 11, the second step follows from Cauchy Shwarz inequality, the last step follows from L 1 ≤ 2 (see part 4 of Lemma B.8) and Eq. 1.
Revised: 11, the second step follows from Cauchy Shwarz inequality, the last step follows from L 1 ≤ 2 (see part 4 of Lemma B.8) and Eq. 1.

Original: Combining Eq.
Revised: Combining Eq.

Original: 81012, we have
Revised: 81012, we have

Original: [Equation]
Revised: [Equation]

Original: Thus completing the proof.
Revised: Thus completing the proof.

Original: B. B OUNDING v MOVE
Revised: B. B OUNDING v MOVE

Original: The goal of this section is to prove Lemma B.7.
Revised: The goal of this section is to prove Lemma B.7.

Original: Lemma B.7 ( v move) .
Revised: Lemma B.7 ( v move) .

Original: We have,
Revised: We have,

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: We split the proof into two cases.
Revised: We split the proof into two cases.

Original: We ﬁrst understand some simple facts which are useful in the later proof.
Revised: We ﬁrst understand some simple facts which are useful in the later proof.

Original: Note that from deﬁnition of x ( k +1)
Revised: Note that from deﬁnition of x ( k +1)

Original: i , we know that x ( k +1) has r k coordinates are (cid:126) 0 .
Revised: i , we know that x ( k +1) has r k coordinates are (cid:126) 0 .

Original: Basically, (cid:107) y ( k ) − x ( k +1) (cid:107) 

Revised: Basically, (cid:107) y ( k ) − x ( k +1) (cid:107)

Original: =
Revised: =

Original: r k .
Revised: r k .

Original: The difference between those vectors is, for the largest r k coordinates in y ( k ) , we erase them in x ( k +1) .
Revised: The difference between those vectors is, for the largest r k coordinates in y ( k ) , we erase them in x ( k +1) .

Original: Then for each i ∈ [ n − r k ] , x ( k +1)
Revised: Then for each i ∈ [ n − r k ] , x ( k +1)

Original: τ ( i ) =
Revised: τ ( i ) =

Original: y ( k ) π
Revised: y ( k ) π

Original: ( i + r k ) .
Revised: ( i + r k ) .

Original: For convenience, we deﬁne y ( k ) π ( n + i ) =
Revised: For convenience, we deﬁne y ( k ) π ( n + i ) =

Original: (cid:126) 0 , ∀ i ∈
Revised: (cid:126) 0 , ∀ i ∈

Original: [ r k ] .
Revised: [ r k ] .

Original: Case 1.
Revised: Case 1.

Original: We exit the while loop when 1 . 5 r k ≥ n .
Revised: We exit the while loop when 1 . 5 r k ≥ n .

Original: Let u ∗ denote the largest u satisfying (cid:107) y ( k ) π ( u ) (cid:107) 2 ≥ (cid:15) mds / 2 .
Revised: Let u ∗ denote the largest u satisfying (cid:107) y ( k ) π ( u ) (cid:107) 2 ≥ (cid:15) mds / 2 .

Original: If u ∗ ≥ r k , then we have that (cid:107) y
Revised: If u ∗ ≥ r k , then we have that (cid:107) y

Original: ( k ) π
Revised: ( k ) π

Original: ( r k ) (cid:107) 2 ≥ (cid:15) mds / 2 ≥ (cid:15) mds / 100 .
Revised: ( r k ) (cid:107) 2 ≥ (cid:15) mds / 2 ≥ (cid:15) mds / 100 .

Original: Otherwise, the condition of the loop shows that
Revised: Otherwise, the condition of the loop shows that

Original: [Equation]
Revised: [Equation]

Original: where we used that n ≥ 4 .
Revised: where we used that n ≥ 4 .

Original: According to the deﬁnition of x ( k +1)
Revised: According to the deﬁnition of x ( k +1)

Original: τ ( i ) , we have
Revised: τ ( i ) , we have

Original: [Equation]
Revised: [Equation]

Original: where the ﬁrst step follows from x ( k +1)
Revised: where the ﬁrst step follows from x ( k +1)

Original: τ ( i ) =
Revised: τ ( i ) =

Original: y ( k ) π
Revised: y ( k ) π

Original: ( i + r k ) , the second step follows from ψ ( x ) is nondecreasing (see part 2 of Lemma B.8) and (cid:107)
Revised: ( i + r k ) , the second step follows from ψ ( x ) is nondecreasing (see part 2 of Lemma B.8) and (cid:107)

Original: y ( k ) π
Revised: y ( k ) π

Original: ( i ) (cid:107) 2 is non-increasing, the third step follows from 1 .
Revised: ( i ) (cid:107) 2 is non-increasing, the third step follows from 1 .

Original: 5 r k > n and hence ψ ( y ( k ) π ( i + r k ) )
Revised: 5 r k > n and hence ψ ( y ( k ) π ( i + r k ) )

Original: = 

Revised: =

Original: for i ≥ n/ 3+1 , the fourth step follows from ψ ( x ) =
Revised: for i ≥ n/ 3+1 , the fourth step follows from ψ ( x ) =

Original: f ( (cid:107) x (cid:107) 22 ) is non-decreasing and (cid:107)
Revised: f ( (cid:107) x (cid:107) 22 ) is non-decreasing and (cid:107)

Original: y ( k ) π
Revised: y ( k ) π

Original: ( i ) (cid:107) 2 ≥ (cid:107)
Revised: ( i ) (cid:107) 2 ≥ (cid:107)

Original: y
Revised: y

Original: ( k ) π
Revised: ( k ) π

Original: ( r k ) (cid:107) 2 ≥ (cid:15) mds / 100 for all i < 2 n/ 3 , and the last step follows by g is non-increasing and the part 3 of Lemma B.8.
Revised: ( r k ) (cid:107) 2 ≥ (cid:15) mds / 100 for all i < 2 n/ 3 , and the last step follows by g is non-increasing and the part 3 of Lemma B.8.

Original: Case 2.
Revised: Case 2.

Original: We exit the while loop when 1 . 5 r k < n and (cid:107) y
Revised: We exit the while loop when 1 . 5 r k < n and (cid:107) y

Original: [Equation]
Revised: [Equation]

Original: By the same argument as Case 1, we have that (cid:107) y
Revised: By the same argument as Case 1, we have that (cid:107) y

Original: ( k ) π
Revised: ( k ) π

Original: ( r k ) (cid:107) 2 ≥ (cid:15) mds / 100 .
Revised: ( r k ) (cid:107) 2 ≥ (cid:15) mds / 100 .

Original: Part 3 of Lemma B.8 together with the fact (cid:107)
Revised: Part 3 of Lemma B.8 together with the fact (cid:107)

Original: [Equation]
Revised: [Equation]

Original: ) , indicates that
Revised: ) , indicates that

Original: [Equation]
Revised: [Equation]

Original: Putting things together, we have
Revised: Putting things together, we have

Original: [Equation]
Revised: [Equation]

Original: by
Revised: by

Original: [Equation]
Revised: [Equation]

Original: by
Revised: by

Original: [Equation]
Revised: [Equation]

Original: by
Revised: by

Original: [Equation]
Revised: [Equation]

Original: where the third step follows by (cid:107) y ( k ) π
Revised: where the third step follows by (cid:107) y ( k ) π

Original: ( i ) (cid:107) 2 is non-increasing and ψ is non-decreasing (see part 2 of
Revised: ( i ) (cid:107) 2 is non-increasing and ψ is non-decreasing (see part 2 of

Original: Lemma B.8).
Revised: Lemma B.8).

Original: B. P OTENTIAL FUNCTION ψ
Revised: B. P OTENTIAL FUNCTION ψ

Original: Lemma B.8 (Properties of potential function ψ ) .
Revised: Lemma B.8 (Properties of potential function ψ ) .

Original: Let function ψ : R d → R (deﬁned in Eq. 5).
Revised: Let function ψ : R d → R (deﬁned in Eq. 5).

Original: Then function ψ satisﬁes the following properties:
Revised: Then function ψ satisﬁes the following properties:

Original: Symmetric ( ψ ( − x ) = ψ ( x )) and ψ (0) = 0 ; 2.
Revised: Symmetric ( ψ ( − x ) = ψ ( x )) and ψ (0) = 0 ; 2.

Original: If (cid:107) x (cid:107) 2 ≥ (cid:107) y (cid:107) 2 , then ψ ( x ) ≥ ψ ( y ) ; 3.
Revised: If (cid:107) x (cid:107) 2 ≥ (cid:107) y (cid:107) 2 , then ψ ( x ) ≥ ψ ( y ) ; 3.

Original: | f (cid:48)
Revised: | f (cid:48)

Original: ( x ) | = Ω(1 /(cid:15) mds ) , ∀| x | ∈
Revised: ( x ) | = Ω(1 /(cid:15) mds ) , ∀| x | ∈

Original: [(0 . 01 (cid:15) mds ) 2 , (cid:15) 2mds ] ; 4.
Revised: [(0 . 01 (cid:15) mds ) 2 , (cid:15) 2mds ] ; 4.

Original: L 1 def = max x D x ψ [ h ] (cid:107) h (cid:107) 2 = 

Revised: L 1 def = max x D x ψ [ h ] (cid:107) h (cid:107) 2 =

Original: and L 2 def = max x D 2 x ψ [ h,h ] (cid:107) h (cid:107)= 10 /(cid:15) mds ; 5.
Revised: and L 2 def = max x D 2 x ψ [ h,h ] (cid:107) h (cid:107)= 10 /(cid:15) mds ; 5.

Original: ψ ( x ) is a constant for (cid:107) x (cid:107) 2 ≥ 2 (cid:15) mds
Revised: ψ ( x ) is a constant for (cid:107) x (cid:107) 2 ≥ 2 (cid:15) mds

Original: Proof.
Revised: Proof.

Original: Recall f : R + → R is deﬁned as
Revised: Recall f : R + → R is deﬁned as

Original: [Equation]
Revised: [Equation]

Original: We can see that
Revised: We can see that

Original: [Equation]
Revised: [Equation]

Original: and
Revised: and

Original: [Equation]
Revised: [Equation]

Original: It implies that max x | f ( x ) (cid:48) | ≤ 1 (cid:15) mds and max x | f ( x ) (cid:48)(cid:48) | ≤ 1 (cid:15) 3mds .
Revised: It implies that max x | f ( x ) (cid:48) | ≤ 1 (cid:15) mds and max x | f ( x ) (cid:48)(cid:48) | ≤ 1 (cid:15) 3mds .

Original: Let ψ ( x ) =
Revised: Let ψ ( x ) =

Original: f ( (cid:107) x (cid:107) 22 ) .
Revised: f ( (cid:107) x (cid:107) 22 ) .

Original: Proof of Part 1, 2 and 5.
Revised: Proof of Part 1, 2 and 5.

Original: These proofs are pretty standard from deﬁnition of ψ .
Revised: These proofs are pretty standard from deﬁnition of ψ .

Original: Proof of Part 3.
Revised: Proof of Part 3.

Original: This is trivially following from deﬁnition of scalar function f .
Revised: This is trivially following from deﬁnition of scalar function f .

Original: Proof of Part 4.
Revised: Proof of Part 4.

Original: By chain rule, we have
Revised: By chain rule, we have

Original: [Equation]
Revised: [Equation]

Original: We can upper bound
Revised: We can upper bound

Original: [Equation]
Revised: [Equation]

Original: Thus, we have
Revised: Thus, we have

Original: [Equation]
Revised: [Equation]

Original: =
Revised: =

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: (cid:15)
Revised: (cid:15)

Original: [Equation]
Revised: [Equation]

Original: + ∞ ) .
Revised: + ∞ ) .

Original: It implies that | D x ψ [ h ] | ≤ 2 (cid:107) h (cid:107) 2 , ∀ x .
Revised: It implies that | D x ψ [ h ] | ≤ 2 (cid:107) h (cid:107) 2 , ∀ x .

Original: By case analysis, we have
Revised: By case analysis, we have

Original: [Equation]
Revised: [Equation]

Original: ∞ ) .
Revised: ∞ ) .

Original: We can also upper bound
Revised: We can also upper bound

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: =
Revised: =

Original: [Equation]
Revised: [Equation]

Original: B. E XAMPLES Comparing with sequential updating algorithm
Revised: B. E XAMPLES Comparing with sequential updating algorithm

Original: We present another version of our main result, which is particularly useful when we compare with the sequential updating algorithm.
Revised: We present another version of our main result, which is particularly useful when we compare with the sequential updating algorithm.

Original: Theorem B.9 (Locality sensitive hashing maintenance, worst case bound) .
Revised: Theorem B.9 (Locality sensitive hashing maintenance, worst case bound) .

Original: For any constant c 1 , c( c 1 > c 2 ) there is a dynamic data structure (Algorithm 1) that achieves ( c 1 , c 2 ) -accuracy.
Revised: For any constant c 1 , c( c 1 > c 2 ) there is a dynamic data structure (Algorithm 1) that achieves ( c 1 , c 2 ) -accuracy.

Original: The data structure takes (cid:101)
Revised: The data structure takes (cid:101)

Original: O ( dn 1+ ρ ) time to initialize and each call of Q UERY ( h ) takes time (cid:101)
Revised: O ( dn 1+ ρ ) time to initialize and each call of Q UERY ( h ) takes time (cid:101)

Original: O ( n ρ d ) .
Revised: O ( n ρ d ) .

Original: By taking a = min { ρ, α } and g r = t r , after T iterations, the amortized expected time per call of U PDATE ( w ) is at most
Revised: By taking a = min { ρ, α } and g r = t r , after T iterations, the amortized expected time per call of U PDATE ( w ) is at most

Original: [Equation]
Revised: [Equation]

Original: The proof is almost indentical to Theorem B.1, we only need to replace Lemma B.5 by the following one.
Revised: The proof is almost indentical to Theorem B.1, we only need to replace Lemma B.5 by the following one.

Original: Lemma B.10 ( w move) .
Revised: Lemma B.10 ( w move) .

Original: We have
Revised: We have

Original: [Equation]
Revised: [Equation]

Original: We omit the dependence on (cid:15) mds since we assume it to be constant in the scheduler.
Revised: We omit the dependence on (cid:15) mds since we assume it to be constant in the scheduler.

Original: Proof.
Revised: Proof.

Original: Again, the proof is almost identical to Lemma B.5, the only difference is that we use (cid:96) ∞ /(cid:96)form of Cauchy-Shwarz instead of the (cid:96) 2 /(cid:96) 2 form in the previous proof.
Revised: Again, the proof is almost identical to Lemma B.5, the only difference is that we use (cid:96) ∞ /(cid:96)form of Cauchy-Shwarz instead of the (cid:96) 2 /(cid:96) 2 form in the previous proof.

Original: We only point out the difference here.
Revised: We only point out the difference here.

Original: We replace Eq.
Revised: We replace Eq.

Original: 10 with n (cid:88)
Revised: 10 with n (cid:88)

Original: [Equation]
Revised: [Equation]

Original: The ﬁrst step follows from Eq.
Revised: The ﬁrst step follows from Eq.

Original: 10, the second step follows from { g i } i =1 , ··· ,n are decreasing and the last step follows from the deﬁnition of γ i .
Revised: 10, the second step follows from { g i } i =1 , ··· ,n are decreasing and the last step follows from the deﬁnition of γ i .

Original: Similarly, we replace Eq.
Revised: Similarly, we replace Eq.

Original: 12 with n (cid:88) i =
Revised: 12 with n (cid:88) i =

Original: [Equation]
Revised: [Equation]

Original: The rest of the proof are the same and we omit it here.
Revised: The rest of the proof are the same and we omit it here.

Original: The sequential updating time for LSH is n ρ .
Revised: The sequential updating time for LSH is n ρ .

Original: The total number of updating is (roughly) at least
Revised: The total number of updating is (roughly) at least

Original: [Equation]
Revised: [Equation]

Original: It then easy to see that our scheduler always perform better than naive sequential update.
Revised: It then easy to see that our scheduler always perform better than naive sequential update.

Original: Comparing with naive batch updating algorithm
Revised: Comparing with naive batch updating algorithm

Original: Below we give a concrete example to show the effectiveness of our algorithm.
Revised: Below we give a concrete example to show the effectiveness of our algorithm.

Original: We take α = ρ for simplicity.
Revised: We take α = ρ for simplicity.

Original: Remember in Assumption 3.2, we already assume for any r ≤ n a = n ρ , T r = n ρ , i.e., the average LSH updating time t r = n ρ /r decays linearly in r when r ≤ n ρ .
Revised: Remember in Assumption 3.2, we already assume for any r ≤ n a = n ρ , T r = n ρ , i.e., the average LSH updating time t r = n ρ /r decays linearly in r when r ≤ n ρ .

Original: It remains to specify the rest t r ( r ≥ n ρ ).
Revised: It remains to specify the rest t r ( r ≥ n ρ ).

Original: Example B.11.
Revised: Example B.11.

Original: Assuming the average running time decays faster than 1 / √ r , i.e. t r = n − (1 − β ) ρ r − β for some β ∈ [ 12 , 1) when r ≥ n ρ .
Revised: Assuming the average running time decays faster than 1 / √ r , i.e. t r = n − (1 − β ) ρ r − β for some β ∈ [ 12 , 1) when r ≥ n ρ .

Original: We then have
Revised: We then have

Original: [Equation]
Revised: [Equation]

Original: The running time of our scheduler is then at most
Revised: The running time of our scheduler is then at most

Original: [Equation]
Revised: [Equation]

Original: As long as the weight changes slowly, say C 1 ≈ C 2 ≤ n ρ/ 2 , we know the amortized updating time of our scheduler is strictly better than the naive approach that updates LSH every time, which has running time n ρ each iteration.
Revised: As long as the weight changes slowly, say C 1 ≈ C 2 ≤ n ρ/ 2 , we know the amortized updating time of our scheduler is strictly better than the naive approach that updates LSH every time, which has running time n ρ each iteration.

Original: [Equation]
Revised: [Equation]

Original: C.
Revised: C.

Original: L EARNING HASH FUNCTIONS
Revised: L EARNING HASH FUNCTIONS

Original: Denote a set of neurons in a particular layer as C = { v r | 0 ≤ r < m } , where each neuron v r has weight and bias.
Revised: Denote a set of neurons in a particular layer as C = { v r | 0 ≤ r < m } , where each neuron v r has weight and bias.

Original: Given an input embedding q i from the previous hidden layer, the output of a forward pass through neuron r is σ ( (cid:104) q
Revised: Given an input embedding q i from the previous hidden layer, the output of a forward pass through neuron r is σ ( (cid:104) q

Original: i , v r (cid:105) ) , where σ is some activation functions.
Revised: i , v r (cid:105) ) , where σ is some activation functions.

Original: For each input embedding q i , both SLIDE and REFORMER select a set of neurons, denoted as S i , from the LSH hash tables for the forward pass.
Revised: For each input embedding q i , both SLIDE and REFORMER select a set of neurons, denoted as S i , from the LSH hash tables for the forward pass.

Original: We collect the training samples from q i and its S i to improve the performance of LSH.
Revised: We collect the training samples from q i and its S i to improve the performance of LSH.

Original: Formally, the pairwise training samples ( q, v ) are collected according to the following criterion:
Revised: Formally, the pairwise training samples ( q, v ) are collected according to the following criterion:

Original: • positive pairs P + = ( q, v ) if v ∈ S and (cid:104) q, v (cid:105) > t + • negative pairs P − =
Revised: • positive pairs P + = ( q, v ) if v ∈ S and (cid:104) q, v (cid:105) > t + • negative pairs P − =

Original: ( q, v ) if v ∈ S and (cid:104) q, v (cid:105) <
Revised: ( q, v ) if v ∈ S and (cid:104) q, v (cid:105) <

Original: t −
Revised: t −

Original: And the loss is deﬁned as:
Revised: And the loss is deﬁned as:

Original: [Equation]
Revised: [Equation]

Original: Here H is the hash functions of all L hash tables.
Revised: Here H is the hash functions of all L hash tables.

Original: H ( x ) generates a K · L vector containing the projection of x in each hash table.
Revised: H ( x ) generates a K · L vector containing the projection of x in each hash table.

Original: cos(
Revised: cos(

Original: H ( x ) , H ( y )) represents the cosine similarity of two projected vectors.
Revised: H ( x ) , H ( y )) represents the cosine similarity of two projected vectors.

Original: The major contribution of this learning approach is that it optimizes the hash functions and the hash table indices together.
Revised: The major contribution of this learning approach is that it optimizes the hash functions and the hash table indices together.

Original: Our method focuses on learning a useful indexing scheme for retrieval efﬁciency while other learnable hash methods (Wang et al., 2017) focus on binary sketches that have higher precision.
Revised: Our method focuses on learning a useful indexing scheme for retrieval efﬁciency while other learnable hash methods (Wang et al., 2017) focus on binary sketches that have higher precision.

Original: We present a demonstration of our loss in Figure 9.
Revised: We present a demonstration of our loss in Figure 9.

Original: Two random hash functions separatedata points into four parts.
Revised: Two random hash functions separatedata points into four parts.

Original: It is obvious that point 1 (pink) should be partitioned together with point 2 instead of point 5 and 6.
Revised: It is obvious that point 1 (pink) should be partitioned together with point 2 instead of point 5 and 6.

Original: On the other hand, point 3 and 4 should not be separated.
Revised: On the other hand, point 3 and 4 should not be separated.

Original: Therefore, during the hash function training, we choose point (5,6) as a positive pair.
Revised: Therefore, during the hash function training, we choose point (5,6) as a positive pair.

Original: Meanwhile, we chose point (5,1) and (5,4) as negative pairs.
Revised: Meanwhile, we chose point (5,1) and (5,4) as negative pairs.

Original: Then, the loss would update the hash functions to accommodate these new partitions.
Revised: Then, the loss would update the hash functions to accommodate these new partitions.

Original: C. O BSERVATIONS ON ATTENTION DISTRIBUTION
Revised: C. O BSERVATIONS ON ATTENTION DISTRIBUTION

Original: In this section, we present the visualization with analysis on the distribution of the minimal quantiﬁes of neurons that sum up to have 0.9 softmax values in attention.
Revised: In this section, we present the visualization with analysis on the distribution of the minimal quantiﬁes of neurons that sum up to have 0.9 softmax values in attention.

Original: On Figure 10, we present the distribution of each head of attention in each layer from a transformer model trained on Enwiki dataset.
Revised: On Figure 10, we present the distribution of each head of attention in each layer from a transformer model trained on Enwiki dataset.

Original: We classify the patterns(Ramsauer et al., 2020) of the distribution into 3 categories by their median values.
Revised: We classify the patterns(Ramsauer et al., 2020) of the distribution into 3 categories by their median values.

Original: With this observation, we are able to determine the layers that LSH or learnable LSH can apply in MONGOOSE framework.
Revised: With this observation, we are able to determine the layers that LSH or learnable LSH can apply in MONGOOSE framework.

Original: D E XPERIMENTS DETAILS D. D ATA STATISTICS
Revised: D E XPERIMENTS DETAILS D. D ATA STATISTICS

Original: [Table]
Revised: [Table]

Original: We present statistics on the 3 datasets we test on from the Extreme Classiﬁcation Repository (Bhatia et al., 2016).
Revised: We present statistics on the 3 datasets we test on from the Extreme Classiﬁcation Repository (Bhatia et al., 2016).

Original: While the number of datapoints in each dataset is not large (on the order of 200K at most), the key feature is the sheer size of the input and output dimensions.
Revised: While the number of datapoints in each dataset is not large (on the order of 200K at most), the key feature is the sheer size of the input and output dimensions.

Original: In particular, each dataset has over 10,000 output classes, which, using a conventional nueral network, requires a matrix multiplication involving over 10,000 neurons at the ﬁnal layer.
Revised: In particular, each dataset has over 10,000 output classes, which, using a conventional nueral network, requires a matrix multiplication involving over 10,000 neurons at the ﬁnal layer.

Original: D. A DDITIONAL RESULTS ON THE EXTREME - CLASSIFICATION TASK
Revised: D. A DDITIONAL RESULTS ON THE EXTREME - CLASSIFICATION TASK

Original: [Figure]
Revised: [Figure]

Original: We present the results of additional experiments comparing the classiﬁcation performance of MONGOOSE against SLIDE and FULL in Figure 11.
Revised: We present the results of additional experiments comparing the classiﬁcation performance of MONGOOSE against SLIDE and FULL in Figure 11.

Original: The two metrics, P@1 and P@5, are presented in (Bhatia et al., 2016) as follows:
Revised: The two metrics, P@1 and P@5, are presented in (Bhatia et al., 2016) as follows:

Original: [Equation]
Revised: [Equation]

Original: where rank k ( · ) extracts the top k indices from a vector, (cid:98) y refers to the predicted labels, and y refers to the ground truth label.
Revised: where rank k ( · ) extracts the top k indices from a vector, (cid:98) y refers to the predicted labels, and y refers to the ground truth label.

Original: We also present the results of time proﬁling to demonstrate the speedups of MONGOOSE in each layer.
Revised: We also present the results of time proﬁling to demonstrate the speedups of MONGOOSE in each layer.

Original: As shown in Table 7, Time 1 represents the total convergence time while Time 2 represents the execution time in the wide output layer (including the forward and backward pass).
Revised: As shown in Table 7, Time 1 represents the total convergence time while Time 2 represents the execution time in the wide output layer (including the forward and backward pass).

Original: Here we observe that MONGOOSE provides up to 20 × speedup in the second layer.
Revised: Here we observe that MONGOOSE provides up to 20 × speedup in the second layer.

Original: The major reason for the
Revised: The major reason for the

Original: [Table]
Revised: [Table]

Original: increase in convergence speedup ( 3 × to 20 × ) is that the ﬁrst embedding layer for high dimensional and sparse input is time and memory consuming Medini et al.
Revised: increase in convergence speedup ( 3 × to 20 × ) is that the ﬁrst embedding layer for high dimensional and sparse input is time and memory consuming Medini et al.

Original: Therefore, in the end-to-end result, the advantage of MONGOOSE has been diluted
Revised: Therefore, in the end-to-end result, the advantage of MONGOOSE has been diluted

Original: [Table]
Revised: [Table]

Original: D. A DDITIONAL RESULTS ON THE LANGUAGE MODELING TASK
Revised: D. A DDITIONAL RESULTS ON THE LANGUAGE MODELING TASK

Original: We present additional experiments comparing the training loss of MONGOOSE and Reformer on the synthetic copy task in Section 4.
Revised: We present additional experiments comparing the training loss of MONGOOSE and Reformer on the synthetic copy task in Section 4.

Original: In Figure 12 we present additional experiments comparing the training loss of MONGOOSE and Reformer on the synthetic copy task in Section 4.1.2.
Revised: In Figure 12 we present additional experiments comparing the training loss of MONGOOSE and Reformer on the synthetic copy task in Section 4.1.2.

Original: The titles of each graph denote the hyperparameters of the Transformer model being tested: h1_s2048_t16 refers to a model with 1 round of hashing, a sequence length of 2048, and a token size of 16.
Revised: The titles of each graph denote the hyperparameters of the Transformer model being tested: h1_s2048_t16 refers to a model with 1 round of hashing, a sequence length of 2048, and a token size of 16.

Original: Overall we can see that MONGOOSE makes a marked improvement in loss over Reformer in every case.
Revised: Overall we can see that MONGOOSE makes a marked improvement in loss over Reformer in every case.

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: In Figure
Revised: In Figure

Original: [Figure]
Revised: [Figure]

Original: we present further results comparing MONGOOSE to Reformer on the enwik8 characterlevel language modelling task.
Revised: we present further results comparing MONGOOSE to Reformer on the enwik8 characterlevel language modelling task.

Original: In particular, we train two sizes of Transformer (3 and 6 layers) with a maximum sequence length of 8192 and 4096 respectivelyh.
Revised: In particular, we train two sizes of Transformer (3 and 6 layers) with a maximum sequence length of 8192 and 4096 respectivelyh.

Original: In both settings mongoose achieves lower loss than Reformer.
Revised: In both settings mongoose achieves lower loss than Reformer.

Original: Note that the goal of this experiments is to prove the superiority of learnable LSH over classical
Revised: Note that the goal of this experiments is to prove the superiority of learnable LSH over classical

Original: LSH in NN training (or MONGOOSE vs. naive LSH-NN framework) rather than improve the stateof-the-art perplexity.
Revised: LSH in NN training (or MONGOOSE vs. naive LSH-NN framework) rather than improve the stateof-the-art perplexity.

Original: For 6 layer transformer experiment on enwik8, Reformer only reaches 2.71 in perplexity, while MONGOOSE achieves 2.65 in the same setting.
Revised: For 6 layer transformer experiment on enwik8, Reformer only reaches 2.71 in perplexity, while MONGOOSE achieves 2.65 in the same setting.

Original: Both MONGOOSE and Reformer (with LSH) save 6 × memory than the original transformers in above two tasks.
Revised: Both MONGOOSE and Reformer (with LSH) save 6 × memory than the original transformers in above two tasks.

Original: E R ELATED WORK E. D ATA STRUCTURES FOR DYNAMIC SIMILARITY SEARCH
Revised: E R ELATED WORK E. D ATA STRUCTURES FOR DYNAMIC SIMILARITY SEARCH

Original: Similarity search is a well-studied problem with wide applications in recommendation system (Xue et al., 2017; Severyn & Moschitti, 2015; Hall & Attenberg, 2015), question answering (Boytsov et al., 2016; Seo et al., 2019; Ahmad et al., 2019;
Revised: Similarity search is a well-studied problem with wide applications in recommendation system (Xue et al., 2017; Severyn & Moschitti, 2015; Hall & Attenberg, 2015), question answering (Boytsov et al., 2016; Seo et al., 2019; Ahmad et al., 2019;

Original: Chang et al., 2020), multi-label classiﬁcation (Covington et al., 2016; Jain et al., 2016; Tagami, 2017) and natural language processing (Bengio et al., 2003; Gao et al., 2014; Lee et al., 2016).
Revised: Chang et al., 2020), multi-label classiﬁcation (Covington et al., 2016; Jain et al., 2016; Tagami, 2017) and natural language processing (Bengio et al., 2003; Gao et al., 2014; Lee et al., 2016).

Original: There are two major categories of similarity measures: (1) metric similarity (cosine similarity, euclidean distance), (2) non-metric similarity (inner product,
Revised: There are two major categories of similarity measures: (1) metric similarity (cosine similarity, euclidean distance), (2) non-metric similarity (inner product,

Original: KL divergence, neural network).
Revised: KL divergence, neural network).

Original: The brute-force approach to solving similarity search is computationally expensive; in response, researchers have developed novel indexing structures to accelerate the search process, with trade-offs on search accuracy.
Revised: The brute-force approach to solving similarity search is computationally expensive; in response, researchers have developed novel indexing structures to accelerate the search process, with trade-offs on search accuracy.

Original: Based on these indexing structures, similarity search algorithms can be broadly categorized as (1) hashing (Shrivastava & Li, 2014; 2015), (2) quantization (Guo et al., 2016; Jegou et al., 2011), (3) tree-based (Ram & Gray, 2012), or (4) graph-based methods (Malkov et al., 2012; 2014; Malkov & Yashunin, 2018).
Revised: Based on these indexing structures, similarity search algorithms can be broadly categorized as (1) hashing (Shrivastava & Li, 2014; 2015), (2) quantization (Guo et al., 2016; Jegou et al., 2011), (3) tree-based (Ram & Gray, 2012), or (4) graph-based methods (Malkov et al., 2012; 2014; Malkov & Yashunin, 2018).

Original: Most similarity search scenarios studied by these papers are static (search data does not change).
Revised: Most similarity search scenarios studied by these papers are static (search data does not change).

Original: In the experiments section, the developed similarity search methods are compared on a ﬁxed dataset such as Sift (Jegou et al., 2011) or Glove (Pennington et al., 2014).
Revised: In the experiments section, the developed similarity search methods are compared on a ﬁxed dataset such as Sift (Jegou et al., 2011) or Glove (Pennington et al., 2014).

Original: However, in current similarity search applications such as the work of (Fan et al., 2019), the search distribution (e.g. product vectors) changes over time due to the activation of new products and the expiration of old products.
Revised: However, in current similarity search applications such as the work of (Fan et al., 2019), the search distribution (e.g. product vectors) changes over time due to the activation of new products and the expiration of old products.

Original: Therefore, some similarity search methods in metric space have been modiﬁed for online settings, such as hashing (Coleman et al., 2019) and quantization (Xu et al., 2018) methods; these methods sacriﬁce search speed for data-adaptiveness.
Revised: Therefore, some similarity search methods in metric space have been modiﬁed for online settings, such as hashing (Coleman et al., 2019) and quantization (Xu et al., 2018) methods; these methods sacriﬁce search speed for data-adaptiveness.

Original: The training phase of Deep Learning models provides a natural setting for Dynamic LSH.
Revised: The training phase of Deep Learning models provides a natural setting for Dynamic LSH.

Original: During training, the weight matrices are slowly modiﬁed via gradients derived from objective functions.
Revised: During training, the weight matrices are slowly modiﬁed via gradients derived from objective functions.

Original: If we consider the weights as the search data and the training sample as queries, we can view DNN training as a Dynamic Similarity Search problem.
Revised: If we consider the weights as the search data and the training sample as queries, we can view DNN training as a Dynamic Similarity Search problem.

Original: Recent works take advantage of this view of DNN training by introducing LSH data structures to the NN training process.
Revised: Recent works take advantage of this view of DNN training by introducing LSH data structures to the NN training process.

Original: (Chen et al., 2020) propose an algorithm (SLIDE) that retrieves neurons with maximum inner product in each step via an LSH based data structure.
Revised: (Chen et al., 2020) propose an algorithm (SLIDE) that retrieves neurons with maximum inner product in each step via an LSH based data structure.

Original: In this way, the backward pass of NN training is concentrated on the neurons with estimated large gradients.
Revised: In this way, the backward pass of NN training is concentrated on the neurons with estimated large gradients.

Original: Their CPU implementation is able to outperform a traditional GPU implementation.
Revised: Their CPU implementation is able to outperform a traditional GPU implementation.

Original: Similar hashing based algorithms have also been used in Transformer models: (Kitaev et al., 2020) (Reformer) propose an LSH structure to reduce the memory bottleneck of selfattention modules especially over long sequences in Transformer.
Revised: Similar hashing based algorithms have also been used in Transformer models: (Kitaev et al., 2020) (Reformer) propose an LSH structure to reduce the memory bottleneck of selfattention modules especially over long sequences in Transformer.

Original: Since DNN weights are the search data, the distribution of the weights among different hash buckets changes throughout the training.
Revised: Since DNN weights are the search data, the distribution of the weights among different hash buckets changes throughout the training.

Original: This necessitates constantly updating the LSH data structure: failing to update the LSH data structure as the search data changes degrades its nearest-neighbor search performance, which in turn worsens the quality of the DNN approximation.
Revised: This necessitates constantly updating the LSH data structure: failing to update the LSH data structure as the search data changes degrades its nearest-neighbor search performance, which in turn worsens the quality of the DNN approximation.

Original: In our experiments, we found that failing to update the LSH data structure in SLIDE caused a 28% decrease in topaccuracy.
Revised: In our experiments, we found that failing to update the LSH data structure in SLIDE caused a 28% decrease in topaccuracy.

Original: E. D ATA DEPENDENT INDEXING
Revised: E. D ATA DEPENDENT INDEXING

Original: Data dependent hashing methods (DDH) focus on adapting hashing schemes to the data distribution.
Revised: Data dependent hashing methods (DDH) focus on adapting hashing schemes to the data distribution.

Original: They often relate to indexing via an objective function.
Revised: They often relate to indexing via an objective function.

Original: Most literature concentrates on hashing and quantization methods.
Revised: Most literature concentrates on hashing and quantization methods.

Original: Although previous works have achieved promising results in learning BTrees (Kraska et al., 2018) or Lattice quantization (Sablayrolles et al., 2019), there are two major bottlenecks for DDH: (1) Theoretical insights are few applied in practice, (2) Complex index designs introduce computation overhead.
Revised: Although previous works have achieved promising results in learning BTrees (Kraska et al., 2018) or Lattice quantization (Sablayrolles et al., 2019), there are two major bottlenecks for DDH: (1) Theoretical insights are few applied in practice, (2) Complex index designs introduce computation overhead.

Original: In the theory of data dependent indexing, (Andoni et al., 2018) present fruitful insights on hashing-based Approximate Nearest Neighbor search on metric space.
Revised: In the theory of data dependent indexing, (Andoni et al., 2018) present fruitful insights on hashing-based Approximate Nearest Neighbor search on metric space.

Original: However, there is not much literature that puts these insights into practice.
Revised: However, there is not much literature that puts these insights into practice.

Original: (Dong et al., 2019) propose a k -NN graph based algorithm for efﬁciently learning data dependent LSH in Euclidean space based on the insights given by (Andoni et al., 2018).
Revised: (Dong et al., 2019) propose a k -NN graph based algorithm for efﬁciently learning data dependent LSH in Euclidean space based on the insights given by (Andoni et al., 2018).

Original: However, their method requires precomputing a k -NN graph, which is computationally expensive at scale.
Revised: However, their method requires precomputing a k -NN graph, which is computationally expensive at scale.

Original: In the experiments, (Dong et al., 2019) only compare their method with k -means; their performance compared to major ANN benchmarks (Erik et al., 2018) remains unknown.
Revised: In the experiments, (Dong et al., 2019) only compare their method with k -means; their performance compared to major ANN benchmarks (Erik et al., 2018) remains unknown.

Original: Many practical applications of LSH eschew DDH methods for vanilla LSH due to the aforementioned computational bottlenecks.
Revised: Many practical applications of LSH eschew DDH methods for vanilla LSH due to the aforementioned computational bottlenecks.

Original: However, despite having sub-linear query time in theory, query complexity for even vanilla LSH is still high in practice Datar et al.
Revised: However, despite having sub-linear query time in theory, query complexity for even vanilla LSH is still high in practice Datar et al.

Original: ; Andoni et al. (2017).
Revised: ; Andoni et al. (2017).

Original: When designing the dynamic data structure, we want to minimize the total running time and balance the overhead and beneﬁt brought by the data structure.
Revised: When designing the dynamic data structure, we want to minimize the total running time and balance the overhead and beneﬁt brought by the data structure.

Original: In general, the matrix multiplication and backpropagation procedures comes from neural network training, which we can not control, hence, we want to balance the accuracy-efﬁciency trade-offs by selecting neurons.
Revised: In general, the matrix multiplication and backpropagation procedures comes from neural network training, which we can not control, hence, we want to balance the accuracy-efﬁciency trade-offs by selecting neurons.

Original: This is especially pronounced in the Dynamic LSH regime, when weights are evolving, besides query time, LSH updates incur a more signiﬁcant overhead that harms the overall efﬁciency.
Revised: This is especially pronounced in the Dynamic LSH regime, when weights are evolving, besides query time, LSH updates incur a more signiﬁcant overhead that harms the overall efﬁciency.

Original: (Chen et al., 2020) introduce a method to control update frequency, but their method requires heavy hyperparameter tuning and is not guaranteed to work for each benchmark.
Revised: (Chen et al., 2020) introduce a method to control update frequency, but their method requires heavy hyperparameter tuning and is not guaranteed to work for each benchmark.

Original: In contrast, our ﬁndings in the following section show that the update overhead can be signiﬁcantly reduced while maintaining nearest-neighbor search quality during NN training.
Revised: In contrast, our ﬁndings in the following section show that the update overhead can be signiﬁcantly reduced while maintaining nearest-neighbor search quality during NN training.

Original: To our knowledge, this is the ﬁrst time DDH techniques have been successfully applied to the Deep Learning setting.
Revised: To our knowledge, this is the ﬁrst time DDH techniques have been successfully applied to the Deep Learning setting.

Original: E. E FFICIENT NEURAL NETWORK TRAINING
Revised: E. E FFICIENT NEURAL NETWORK TRAINING

Original: The weights of a neural network dynamically change during the training process, which brings great challenges to the LSH implementation.
Revised: The weights of a neural network dynamically change during the training process, which brings great challenges to the LSH implementation.

Original: Since the data are no longer static, we need to constantly update the hash table, which incurs extra computation cost that could harm the overall performance.
Revised: Since the data are no longer static, we need to constantly update the hash table, which incurs extra computation cost that could harm the overall performance.

Original: We present a dynamic data structure to handle this issue, and our data structure achieves signiﬁcant speedup over naive implementation under mild assumptions, without compromising the worst case guarantee.
Revised: We present a dynamic data structure to handle this issue, and our data structure achieves signiﬁcant speedup over naive implementation under mild assumptions, without compromising the worst case guarantee.

Original: The dynamic data structure (shown in Algorithm 1) borrows ideas from the work of (Cohen et al., 2019; Jiang et al., 2020c) (which are originally designed for linear programming), and we adapt it to the LSH setting.
Revised: The dynamic data structure (shown in Algorithm 1) borrows ideas from the work of (Cohen et al., 2019; Jiang et al., 2020c) (which are originally designed for linear programming), and we adapt it to the LSH setting.

Original: Our data structure generalizes several practical insights and turns practical heuristics into rigorous theory, which guarantees signiﬁcant speed up under natural assumptions for training neural networks and ensure the worst guarantee at the same time.
Revised: Our data structure generalizes several practical insights and turns practical heuristics into rigorous theory, which guarantees signiﬁcant speed up under natural assumptions for training neural networks and ensure the worst guarantee at the same time.

Original: In particular, we generalize and provide the following three practical heuristics (i) only update signiﬁcantly changed coordinates, (ii) batch updating LSH instead of sequential updating, (iii) using “prediction” on those “marginal” coordinates.
Revised: In particular, we generalize and provide the following three practical heuristics (i) only update signiﬁcantly changed coordinates, (ii) batch updating LSH instead of sequential updating, (iii) using “prediction” on those “marginal” coordinates.

Original: In this work, we primarily study two LSH based efﬁcient NN training methods: SLIDE and REFORMER .
Revised: In this work, we primarily study two LSH based efﬁcient NN training methods: SLIDE and REFORMER .

Original: SLIDE introduces LSH to select neurons in the forward pass and then only do gradient descent on the chosen neurons.
Revised: SLIDE introduces LSH to select neurons in the forward pass and then only do gradient descent on the chosen neurons.

Original: The major trade-off of SLIDE is the rebuild overhead versus the accurate neuron selection.
Revised: The major trade-off of SLIDE is the rebuild overhead versus the accurate neuron selection.

Original: To accurately retrieve the neurons with high inner products, the hash tables are required to be updated.
Revised: To accurately retrieve the neurons with high inner products, the hash tables are required to be updated.

Original: This rehashing and rebuilding will be the major overhead caused by
Revised: This rehashing and rebuilding will be the major overhead caused by

Original: SLIDE .
Revised: SLIDE .

Original: Therefore, the SLIDE ’s speed over full NN training is determined by the relative magnitude of rebuild overhead compared to the saved back-propagation time.
Revised: Therefore, the SLIDE ’s speed over full NN training is determined by the relative magnitude of rebuild overhead compared to the saved back-propagation time.

Original: The major goal of REFORMER is to reduce the memory consumption of the transformer so that GPU based hardware can support sequence to sequence tasks with longer sequence length.
Revised: The major goal of REFORMER is to reduce the memory consumption of the transformer so that GPU based hardware can support sequence to sequence tasks with longer sequence length.

Original: REFORMER also shares the trade-off when rebuilding the hash tables for better retrieval of attention weights.
Revised: REFORMER also shares the trade-off when rebuilding the hash tables for better retrieval of attention weights.

Original: Besides, learning hash functions of REFORMER is more challenging than SLIDE .
Revised: Besides, learning hash functions of REFORMER is more challenging than SLIDE .

Original: From an information retrieval’s perspective, in attention settings, each data vector is also a query.
Revised: From an information retrieval’s perspective, in attention settings, each data vector is also a query.

Original: This unique setting increases the hardness for learning better hash functions.
Revised: This unique setting increases the hardness for learning better hash functions.

Original: F E FFICIENT GPU IMPLEMENTATION
Revised: F E FFICIENT GPU IMPLEMENTATION

Original: [Figure]
Revised: [Figure]

Original: ×
Revised: ×

Original: As noted by (Chen et al., 2020), each input in one batch samples a different set of weights(neurons).
Revised: As noted by (Chen et al., 2020), each input in one batch samples a different set of weights(neurons).

Original: These irregular operations signiﬁcantly reduce the beneﬁts of using GPUs: they correspond to different sizes of rows for each input, which makes the forward pass much more difﬁcult to parallelize and thus, more challenging to map to GPUs than regular programs (Burtscher et al., 2012).
Revised: These irregular operations signiﬁcantly reduce the beneﬁts of using GPUs: they correspond to different sizes of rows for each input, which makes the forward pass much more difﬁcult to parallelize and thus, more challenging to map to GPUs than regular programs (Burtscher et al., 2012).

Original: For this reason, (Chen et al., 2020) build their SLIDE system in C++ from scratch on CPU.
Revised: For this reason, (Chen et al., 2020) build their SLIDE system in C++ from scratch on CPU.

Original: Even though their implementation achieved remarkable speed up, their impact is limited as they implemented their system from scratch in C++, making it difﬁcult for the community to adopt SLIDE in practice.
Revised: Even though their implementation achieved remarkable speed up, their impact is limited as they implemented their system from scratch in C++, making it difﬁcult for the community to adopt SLIDE in practice.

Original: To verify that such randomized algorithm is not GPU friendly, we implement the exact same algorithm for GPUs.
Revised: To verify that such randomized algorithm is not GPU friendly, we implement the exact same algorithm for GPUs.

Original: We show in Section 4.1.
Revised: We show in Section 4.1.

Original: that it indeed fails to gain any speed up.
Revised: that it indeed fails to gain any speed up.

Original: We design a variant of the algorithm to exploit fast matrix multiplications on GPUs.
Revised: We design a variant of the algorithm to exploit fast matrix multiplications on GPUs.

Original: In the original SLIDE, each training example in a batch retrieves its own subset of weights.
Revised: In the original SLIDE, each training example in a batch retrieves its own subset of weights.

Original: Here, we take a union of the retrieved subsets in each batch to avoid irregular and unbalanced memory access, shown in ﬁgure 14.
Revised: Here, we take a union of the retrieved subsets in each batch to avoid irregular and unbalanced memory access, shown in ﬁgure 14.

Original: Our implementation of this proposal is written in python under
Revised: Our implementation of this proposal is written in python under

Original: Pytorch framework with Cython compiled LSH and CUDA kernels for hashcode efﬁcient computation.
Revised: Pytorch framework with Cython compiled LSH and CUDA kernels for hashcode efﬁcient computation.

Original: We believe this implementation would be more beneﬁcial for the community as it can be easily plugged into any deep learning models.
Revised: We believe this implementation would be more beneﬁcial for the community as it can be easily plugged into any deep learning models.

Original: Practical Implementation:
Revised: Practical Implementation:

Original: In our design, we borrow insights from the theoretical guarantee from above and aim to ﬁnd the sweet spot between theory limitations and practical challenges.
Revised: In our design, we borrow insights from the theoretical guarantee from above and aim to ﬁnd the sweet spot between theory limitations and practical challenges.

Original: First, instead of detecting weight changes at the cost of an extra copy, we reduce the problem to detect low quality of retrieved weights.
Revised: First, instead of detecting weight changes at the cost of an extra copy, we reduce the problem to detect low quality of retrieved weights.

Original: In NN training setting, we measure quality as the inner product and low quality indicates a low inner product.
Revised: In NN training setting, we measure quality as the inner product and low quality indicates a low inner product.

Original: We argue these two approaches are similar that they both signal the necessity of updating the data structure.
Revised: We argue these two approaches are similar that they both signal the necessity of updating the data structure.

Original: Besides, rather than a soft margin for detecting data changes, which can be ambiguous under dynamic setting, inner product is a better measurement as it clearly indicates the performance of current data structure.
Revised: Besides, rather than a soft margin for detecting data changes, which can be ambiguous under dynamic setting, inner product is a better measurement as it clearly indicates the performance of current data structure.

Original: More importantly from an efﬁciency perspective, quality detection comes almost for free because the inner product between the query embedding and retrieved neurons are necessary for the forward pass.
Revised: More importantly from an efﬁciency perspective, quality detection comes almost for free because the inner product between the query embedding and retrieved neurons are necessary for the forward pass.

Original: G B ROADER I MPACT
Revised: G B ROADER I MPACT

Original: With the exponentially growing data volume, the scale of neural network (NN) models keeps increasing to achieve unprecedented performance on tasks in many domains, such as recommendation systems (He & Chua, 2017; Medini et al., 2019), natural language processing (Bengio et al., 2003;
Revised: With the exponentially growing data volume, the scale of neural network (NN) models keeps increasing to achieve unprecedented performance on tasks in many domains, such as recommendation systems (He & Chua, 2017; Medini et al., 2019), natural language processing (Bengio et al., 2003;

Original: Vaswani et al., 2017) and computer vision (Deng et al., 2009; Vaswani et al., 2017).
Revised: Vaswani et al., 2017) and computer vision (Deng et al., 2009; Vaswani et al., 2017).

Original: However, training those large-scale models imposes challenges on the computational and memory resources even with advanced modern hardware.
Revised: However, training those large-scale models imposes challenges on the computational and memory resources even with advanced modern hardware.

Original: For example, the recent GPT-3 model (Brown et al., 2020) has a capacity of 175 billion parameters.
Revised: For example, the recent GPT-3 model (Brown et al., 2020) has a capacity of 175 billion parameters.

Original: Therefore, there is an emergent need to reduce the cost of training those giant models.
Revised: Therefore, there is an emergent need to reduce the cost of training those giant models.

Original: We now quantify the potential impacts of our work.
Revised: We now quantify the potential impacts of our work.

Original: MONGOOSE sheds lights on efﬁcient NNS-ds for efﬁcient training of deep neural networks.
Revised: MONGOOSE sheds lights on efﬁcient NNS-ds for efﬁcient training of deep neural networks.

Original: Especially, the slow change observation along with the smart scheduler demonstrate the possibility of applying NNS-ds with larger indexing overhead for dynamic similarity search where the distribution of data changes slowly.
Revised: Especially, the slow change observation along with the smart scheduler demonstrate the possibility of applying NNS-ds with larger indexing overhead for dynamic similarity search where the distribution of data changes slowly.

Original: We show applications of MONGOOSE on language modeling and recommendation tasks.
Revised: We show applications of MONGOOSE on language modeling and recommendation tasks.

Original: But since our key or fundamental observations are general, MONGOOSE could potentially generalize to more deep learning models and tasks.
Revised: But since our key or fundamental observations are general, MONGOOSE could potentially generalize to more deep learning models and tasks.

Original: On the other hand, more NNS-ds equipped with our observations and scheduler could be involved in the deep learning community to tackle the efﬁciency issue.
Revised: On the other hand, more NNS-ds equipped with our observations and scheduler could be involved in the deep learning community to tackle the efﬁciency issue.