MONGOOSE: A Learnable LSH Framework for Efficient Neural Network Training
Authors: Beidi Chen
Zichang Liu
Binghui Peng
Zhaozhuo Xu
Jonathan Lingjie Li
Tri Dao
Zhao Song
Anshumali Shrivastava
Christopher Re
Venue: ICLR 2021 Oral (2021)
Keywords: Large-scale Deep Learning
Large-scale Machine Learning
Efficient Training
Randomized Algorithms

Abstract: 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. However, while LSH has sub-linear guarantees for approximate near-neighbor search in theory, it is known to have inefficient query time in practice due to its use of random hash functions. Moreover, when model parameters are changing, LSH suffers from update overhead. 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. This phenomenon points to the potential for a reduction in update time and allows for a modified learnable version of data-dependent LSH to improve query time at a low cost. We use the above insights to build MONGOOSE, an end-to-end LSH framework for efficient NN training. In particular, MONGOOSE is equipped with a scheduling algorithm to adaptively perform LSH updates with provable guarantees and learnable hash functions to improve query efficiency. Empirically, we validate MONGOOSE on large-scale deep learning models for recommendation systems and language modeling. We find that it achieves up to 8% better accuracy compared to previous LSH approaches, with $6.5 \times$ speed-up and $6\times$ reduction in memory usage.

Published as a conference paper at ICLR 2021
MONGOOSE: A LEARNABLE LSH FRAMEWORK FOR EFFICIENT NEURAL NETWORK TRAINING
Beidi Chen', Zichang Liu", Binghui Peng?, Zhaozhuo Xu², Jonathan Lingjie Li', Tri Dao', Zhao Song', Anshumali Shrivastava', Christopher Re
Stanford University, 2 Rice University, Columbia University, Princeton University (beidic, c, gili, trid, chri smre @stanford. edu, 2p 60126010 col cmlib cumbia.edi edu {2171, zx22, ans anshumall @rice. edu, zhaos@princeton.ccmm
ABSTRACT
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 neural network (NN) training. However, while LSH has sub- linear guarantees for approximate near-ncighbor search in theory, it is known have inefficient query time in practice due to its use of random hash functions. changing, LSH suffers from update over-
Moreover, when model parameters head. This work is motivated by an observation that model parameters evolve slowly, such that the changes do not always require LSH update to maintain performance. This phenomenon points to the potential for a reduction in update time and allows for modified learnable version of data-dependent LSH to im- prove query time at a low cost. We use the above insights to build MONGOOSE, an end-to-end LSH framework for efficient NN training. particular, MONGOOSE is equipped with a scheduling algorithm to adaptively perform L.SH updates with provable guarantees and learnable hash functions improve query efficiency. Empirically, we validate MONGOOSE on large-scale deep learning models for rec- ommendation systems and language modeling. We find that achieves up to 8% better accuracy compared to previous LSH approaches, with 6.5x speed-up and 6x reduction in memory usage.
INTRODUCTION
Locality Sensitive Hashing (LSH) has been adapted address the computational and memory bot- tlenecks of large-scale neural network (NN) training in natural language processing (Chandar et al.) 2016 Rae et al.] | 2016 Kitaev et al. 2020), computer vision (Chen et al. 2015) and recommendation systems (Spring & Shrivastava 2017 Chen et al. 2020). Specifically, giant matrix multiplications in linear layers preceding a softmax can be approximated using nearest neighbor search (NNS) tech- niques, which often rely on LSH. However, LSH methods used in NNs are inefficient. Although LSH achieves sub-linear query time in theory, it is known to suffer from high query and pre-processing (update) overhead in practice (Erik al. (2018). In the setting of NN training, where data points for LSH are model parameters, this overhead is exacerbated by a high number of updates due to constantly evolving model parameters.
The most established solution for reducing LSH query overhead is data-dependent or learning- based hashing, which uses adaptive hash functions optimize the LSH bucket distribution for each dataset (Andoni & Razenshteyn, 2013 Dong et al. 2019). These methods reduce query time by incurring a one-time offline cost learn useful data input patterns learning techniques are often computationally complex, but can lead to a net reduction in overall query time. However, in NN training, the expensive preprocessing procedure has to be repeated each time the parameters overhead rather than reduce it. A more appropriate data-dependent LSH framework would ideally (i) have deeper understanding of the training dynamics of model parameters, setting in which LSH has not been well-studied, (ii) be able to perform low-cost updates to account for evolving parameters, and (iii) have better query time while accurately approximating matrix multiplication.
preprocessing step. The
updated. Naïvely applying these techniques would increase the LSH update
We argue that it is unnecessary to view evolving model parameters as streaming data and not every parameter update requires LSH update. In fact, LSS updates necessary only when the NN gra-

---

Published as a conference paper at ICLR 2021
Time
Lww SCHEDULER
Bucket
Updatc
LSH Hash Fanction
Update
W;
LEARNING
NN Training
WVO Updare
Update
Updnal lbab Functices
Time t
Bucket
Bucket
2
OR
2
3
Figure 1: MONGOOSE workflow For each layer in the NN. besides the usual setup for LSH, we add two components: (i) scheduler deciding if LSH update should be performed in this iteration (ii) if updating, learning (data-adaptive) mechanism will be triggered for tuning the LSH hash functions based the current weights, and then LSH updates will be performed.
dient steps are large enough cause the model parameters' LSH hash codes change. Therefore, we count the number of hash code changes that occur for important models such as Transformers and fully-comnected NNs from previous work (Chen al.) 2020 Kitaev et al.) 2020). We find that only 1 % of the hash codes change after each epoch on average for Transformers, and 5% for fully- connected NNs, so most weights do not change on scale that is enough to trigger an LSS update. Furthermore, the rate of hash code change initially decays exponentially and eventually platcaus, as shown later in Figure 3 We calculate that a 100x speed up is possible with an update oracle. In contrast, current approaches with LSH have not fully exploited this observation. We show in Section 4.2 that they either blindly skip updates for speed, resulting in a 10-point accuracy drop. or suffer from 20x slow-down.
We demonstrate that these slowly changing hash codes provide two opportunities to realize the ideal data-dependent LSH framework described above (shown the intuition for the scheduler in Figure D First, given weight changes, an algorithm could adaptively schedule LSH updates to realize low update overhead, Second, given current weights, we could learn data-dependent hash functions for better LSH bucket distribution to achieve fast query time. However, this comes with three challenges: (i) how characterize the slowly-changing phenomenon without computing hash codes, (ii) how to schedule the updates without the oracle, and (iii) how to learn data-dependen hash functions for shorter query time without compromising on the update time.
In this paper, we first provide general formulation of the problem as dynamic NNS in Section 2 In Section g] we propose MONGOOSE, framework for fast and memory-efficient NN training to ad- dress the above challenges. We show one major observation, slow change, and two key components based on MONGOOSE. Specifically,
we measure the Ee norm of the weight changes and their hash code changes during training on fully-connected NN. (More models are studied in Appendix [A). We find that (i) the hash codes are slowly changing (ii) both quantities share similar trajectories. but the former's 's slow change is the necessary condition of the latter's. Therefore, we formally define the slow change of &g norm of weight changes in Assumption 3.1 which the key to build MONGOOSE.
In Section 3.1
In Section 3.2 we present an algorithm for scheduling efficient LSH updates. We closely analyze our scheduler's theoretical guarantees and provide an upper bound of its running time, which is proportional to the & norm of the weight changes. Therefore, under the slow change assumption, we show that our scheduler is provably faster than previous approaches.
In Section 3.3 we propose method of learning parameterized LSH hash functions (e.g.. SimHash (Charikar) 2002)) during NN training. Our method utilizes intermediate activations during the forward pass to generate training inputs for tuning LSS hash functions with little overhead. Com- bining this with our scheduler further decreases the overhead, while providing hash functions that better separate data.
demonstrate the efficacy of MONGOOSE two LSH-NN systems, SLIDE and Reformer. For SLIDE applications, we show up to 6.5x speedup in time and 8% higher accuracy on three datasets. For Reformer applications, show improvement perplexity when training lan-
Finally, in Section 4
2

---

Published as a conference paper at ICLR 2021
guage modeling benchmarks from scratch. Moreover, We provide two additional sets of experiments to show how MONGOOSE addresses the aforementioned three challenges separately.
RELATED WORK AND PROBLEM SETTING
In this section, we first discuss applications of LSH in the NN training setting and introduce data- dependent LSH techniques. Then we formally define the problem we solve as dynamic NNS.
2.1 LSH FOR EFFICIENT NN TRAINING
Recent works take advantage of LSH as an efficient NNS algorithm to speed up matrix multiplication in neural networks. SLIDE (Chen ett al., 20200) is algorithm that retrieves neurons with maximum inner product during the forward pass via an LSH-based data struc- ture. In this way, in the backward pass gradients only com- puted for neurons with estimated large gradients. Reformer (Kitaev al. 2020), variant of Transformer, similarly uses LSH to reduce the memory bottleneck of self-attention layers over long sequences. One pain point of LSH in this setting that model weights change throughout training. This necessitates constantly updating the LSH data structure. Figure illustrates what happens when we don't 't per- form such updates. In short, failing to update the LSH data structure as the search data changes degrades its NNS performance. This in turn worsens the quality of the matrix product approximation. In our experiments, we found that failing to update the LSH data structure in SLIDE causes 28% decrease in top-1 accuracy for fonneccon NN. (More related work is presented in Appendix E.I]
Dla
11s..
h(xy/=0
h(x,) =1
Visualization of why LSH updates are essential when data changes. h a hyperplane separating two hash buckets. 21 is updated in training, its hash value may change.
Figure 2:
2
2.2 PROBLEM FORMULATION
In this section, we formulate LSH for effcient training as dynamic NNS problem. This formulation closely built on the static NNS problem and the well-known lower bound LSH complexity. In the static NNS problem, we are given a set of weights .1..... Wy E R and want to construct data structure (NNS-ds) that supports the QUERY operation, defined follows: given x Rd (x) returns a set S of weights w; that are all "close" to x in a distance measure.
More precisely, we require the (UERY(+) operation to be (c1, C2) aacuurttt
Definition 2.1 ((c1, e2) aa)cuueut),) Denote S to be the (random) set returned by QUERY(x). We say QUERY (z) ) (c1, 02) accurate if
for any i E [n] such that (wi, æ)) 2 ei @||l Pr/w, E S] 2 1 1/ (ii) for any i [n] such that (wi, .r) C2 | Pr|w; S] < 1/ 1/po)y(n).
In our application, we assume >> C1 > C > O to be some constants and denote (c2/c))2. In the dynamic NNS problem, the weights W1, ... Wn I2d can evolve over time, so we need to update the data structure.
Lemma 2.2 ([Andoni & Indyk, 2006)). Using ISH, one achieve (cI, accuracy with query time (dd"), preprocessing time D(dntt)) updating time O(dn)).
3 MONGOOSE: A FRAMEWORK FOR LEARNABLE LSH
We present the workflow of our main framework in Figure ! In Section 3.1 we first show the key observation that encourages the design of MONGOOSE and formally define the slow change assump- tion. In Section 3.2 we introduce our first component, an algorithm for scheduling LSH updates with provable guarantees based on the assumption. Finally, present the second component, low- cost learnable LSH hash functions, in Scction 3.3 More details about the efficient implementation of MONGOOSE in Appendix
3.1 CHARACTERIZATION OF "SLOW CHANGE" PHENOMENON
We first present an observation on model weights and how their LSH hash codes change during training. Based this observation, define the slow change assumption. We denote AW as the
3

---

Published as a conference paper at ICLR 2021
difference between weights across gradient steps, and AH the Hamming distance between hash codes during training.
Observation. We plot the &g norm of AW and AH of the weights during training a one lay- ered fully-connected NN in Figure 3 Our most surprising finding is that only 5% of hash codes change after each epoch on average, which im- plies that for most neurons, AW is not large enough to trigger hash code changes. For both AW and AH, there is a sharp drop at the early stages of training before they plateau. (Similar observations on Transformers and further dis- cussions are presented in Appendix A
Ang Wiight Change
Ave Halli Clear Charge
Epoch
Figure 3: AW and AH both start out relatively high at the beginning of training, but quickly drop off and flatten out.
Insights. In the dynamic NNS problem, input data (model weights) change over time. Without any assumptions about weight updates, naïvely applying LSH will require updating the LSH hash func- tions at every time step. However, the above observation suggests that we can reasonably assume bounded by AH. Denote w as the weight matrix, n as number of neurons
A AW is (roughly) and wy the weight vector for neuron i. Formally:
Assumption 3.1 (Slow change). Assume NN weights change slowly over time. In particular, assume there an upper bound on the expected movement of the weights of the neural networks (C₂) ) and upper bound on the variance (C2). Specifically, we denote the initial weight matrix as (())) E paxd
We assume the (random) update sequence w(1), ...),....... (T( E Rxdd satisfies
(8)) W
and Σ ||aar ( 1112 C22
where the expectation and variance is conditioned on
(k) for all 0,1, -1
3.2 SMART SCHEDULER FOR LEARNABLE LSH UPDATES
In this section, introduce dy- namic data structure to schedule updates for LSH. Our scheduler provides signifi- cant speedup under two realistic assump- tions without compromising the theoreti- cal guarantee of LSH: (1) slow change de- fined in Assumption 3.1, (2) batch speed
Algorithm 1 General Maintenance Data Structure
Input: w E R Raxd V E Rxaxd Emds E(0, (0, 1/lll"oo"@0 n),
a = min(p, a}, S C [n]]
2: procedure INITIALIZE(W; Emda)
> Lemma B.2
3:
w
W. Emda f- Cadda, S 0.
4: Build LSH with accuracy (c1 Ende, C2 fmds)
5:
end procedure
up defined in Assumption 3.2 The sched- uler can generalize to any neir-neigh search data structure (NNS-ds), but we show a design for LSH as an example since that is the NNS-ds in MONGOOSE.
6: procedure UPDATE(W)
b Lemma B.3
7:
Vi w.""
Vi [n]
8:
the number of indices that ||w/||2 emdaa
if < n then o new +v
10:
11:
else
12:
Let [n] [m]]. that Ilu-collz lla 2 11
3.2.1 SMART SCHEDULER
13:
while 1.5-
and ||ue(pp.a>r)>l> 2 (1
1/ log n) (r) (rille lla do
The pseudo-code for our smart scheduler
14:
r & min( [[1.5-11]]] rl.
shown Algorithm|| [ Our scheduler
15:
end while
> Finding smooth cut
is inspired by dynamic data structure proposed in a recent line of work (Co-) hen et al. 2019), which was originally de- signed for inverse maintainance Programming solvers. We provide
W ((())
i {1, 2, ,r]
16:
(T))
(T(()
iE (r .......,1)
17:
LSH.UPDATE &T (1), (()r
Lin-
18:
end if
19:
w 4- w new V new
a novel application here for the dynamic LSH problem.
20:
S fi e [n]
||w
vzlla 2 Emds/22
21:
end procedure
22: procedure QUERY(A)
> Lemma B.4
Notation: We denote [n] = (1, ,..
23: LSH.QUERY and check S
The O(-) notation hides polylogarithmic
24: end procedure
4

---

Published as a conference paper at ICLR 2021
dependence d, Emds. We each row wy E IR' (i E [n]) corresponds to the input weight of the i-th neuron. We maintain a copy of the neural network weights, denoted as v E Rxxd that approximately equals except for a few coordinates. We take émds, to be a sufficiently small precision parameter, i.e. mmds & 61 C2. which satisfies Cuds E /2(1/ log n2) We use S to denote the set indices where w and v differ signifi- cantly (see Line 20) LSH.UPDATE (see line 17 takes indices of the neural weights and updates the LSH hashcode correspondingly.
Rexd to denote the weights of the neural network and
Overview of algorithm: During the training process, we maintain two copies of the NN weights w, where w are the actual weights that are updated via gradient descent and v approximately equals w except for a few coordinates. We set an empirical threshold of émds. Due to the robustness of the NNS problem, we can tolerate this small error émds such that only those coordinates with a difference greater than mnds matter. To exploit the benefit of batch updating, we perform lazy updates on the LSH hash table. In particular, difference between w, v exceeds some threshold number n (see line to line [T). One key part of update procedure is that we also update some low-error coordinates when there too many of them (see line to line line[T6.. This procedure can be seen as forecasting future updates. The intuition that when there are many coordinates close the decision boundary, they are likely to change within a short amount of time. Updating these all together could further take advantage of the batch updates. To answer an LSH query, we keep track of a short list S (see line 20). which records coordinates that differ significantly between w, V. 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. the query procedure (see line to line 24). besides performing the normal LSH query, we also check S and pick up weights that are close to æ. Hence, we balance the overhead of query time with update time, where update time is usually the bottleneck.
update the LSH hash table if and only if the
3.2.2 THEORETICAL ANALYSIS
Assumption 3.2 (batch speedup). Denote T, as the updating time for LSH when the weights of neurons change and T, = r..rte ty- That is, Tr is the updating time when we update entries of the LSH hash table and t, is the average updating time. We assume (ti)i=i,.... is non increasing, i.e., t 2 2...2...2 2 tn- We further assume T, = n' when r S n for some constant a.
The batch speedup assumption is realistic, especially due advances in support for massive paral- lelism on modern computation architectures.
Theorem 3.3 (LSH maintenance). For any constant C1, (c1 > 02). there is dynamic data structure (Algorithm !} for LSH maintenance that achieves (2), aaccuracyy The data structure takes Ödn (P) time to initialize and each call of QUERY(h) takes time Ó(n By taking a = min(p, p. ) and
[np_, r r n
9r
tr
>> n²
The amortized expected time per call of UPDATE(W) is most
Ö((C) + (19-119112)
(2)
Remark 3.4. Depending on the exact form of batch updating time t,, the (amortized) running time and the benefit of our scheduler could be different. We provide concrete examples in Appendix BB. showing advantages of our smart scheduler under different scenarios. In general, we remark that (1) Our scheduler always better than the naïve sequential updating strategy. (2) Our scheduler is oblivious to the exact form of the update time ip and the weight change C1, Cz. That is to say, we do not need to fine tune these parameters. In NN training. the only parameter that requires fine tuning is the precision parameter Emds-
3.3 HASH FUNCTION LEARNING
In this section, we provide a simple streaming algorithm to learn parameterized hash functions. Here, the learning objective is to track the weight distribution the weights change to better separate the
5

---

Published as a conference paper at ICLR 2021
input data the NNS problem. Recall that in last section, show algorithm for scheduling the LSH updates, which makes it possible to involve learning for LSH and improve the query efficiency. 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.
3.3.1 LEARNABLE LSH
We view the sequence of neural network weights throughout training as a matrix stream. In each training iteration, if our scheduler schedules an update, we compute a training signal with which we can train our LSH parameters. Specifically, we com- pute the classical triplet loss (Chechik et al.) 2010) using positive and negative samples from the current model weights for each LSH query. 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 se- lected. Denote a set of neurons in partic- ular layer as C {ur 0 S < m) and the input embedding as qs- Denote the set of neurons retrieved from the LSH hash tables using qi as S, We formally present the learning al- gorithm in Algorithm 2 The operation of updating or training LSH hash functions is defined as REHASH and the operation of updating the hash tables based on the updated hash functions and weights is defined as REBUILD. The triplet loss is defined as follows:
Algorithm 2 Learnable Hash Function
1: Input: query set {q. ke|N) query size N
2: hash table G = { hell
3: hash function 4Q E [74, Hi = {/I, H.
Uie(z Hi
4: Retrieved Neurons Sets (S.)iEN), S. (u>]re[MM]
5: parameters 4,
6: P+ f 0, P_
A Pt,, E R22+2
7: for == 1- -
8:
P+ t P+ U fig vx) Ur E vr) t+
9:
P_ + P- (( Ur) er Si, (9:, Uy) < xt-
end for
11: size new &- m
12: RESIZE(P,,, P size
13: Huew +- REEAASH(P,,) P_ L) L loss function Gnow + REBUILD(C, Heew)
15: return gnow, men
C(H, P+, P-) = max(0,
Σ
- -oo (L())) H(w)) + Σ
cos(HL(g), + aa) (3)
(q.")EP;
(g.N)P
where H denotes the LSH hash function the current iteration and is margin between positive and negative pairs. (o))), H(()) represents the cosine similarity. Note that in the application setting of unit norm like (Kitaev et al. 2020), using inner product is equivalent cosine similarity. We provide a detailed justification of this learning approach in Appendix C.T
3.3.2 FURTHER SPEEDUP
We demonstrate how to optimize learn- able LSH for further acceleration in NN training. Here, we provide the
Layer 1
Layer 3
Layer
No. Attentons
201
main insights; further details are in Ap- pendix A
Layer Selection: In previous LSH-NN methods such as SLIDE or Reformer, the designation of whether layer uses LSH is treated as tunable hyperparameter. For example, Reformer applies LSH in three out of twelve attention layers in their final model based on empirical performance. In MON- GOOSE we make this designation be presented by visualizing the statistics of attention weights in different layers of the Transformer model (Figure 4.. The three plots represent the distribution of the average number of coordinates with "heavy" attention for each layer. 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 experi- ment). The median for the left plot is 2, meaning that they are well-separated and LSH can perform well on such a distribution. On the other hand, learnable LSH can be very helpful for the right plot, because classical LSH is inefficient for this distribution. Therefore, when deciding between LSH or learnable LSH in each layer for later training stages or fine-tuning on downstream tasks.
Figure 4: Attention distribution examples different layers: Y-axis the least number of attentions summing to 90% of full attentions. The annotations represent the median. The total number of attentions (or sequence length) is 512.
more principled manner. The intuition behind the design can
this observation

---

Published as a conference paper at ICLR 2021
Table I: This table summarizes the performance of MONGOOSE, SLIDE and Full-Computation implemented with PyTorch (Paszkc et al. 2019). P@1/5 is top-1/5 precision. Time represents convergence time and Mem represents memory consumption.
Full-Compatation
SLIDE
MONGUOSE
Datusets
PS1l
PM5
Time (s)
Mem (GB)
7116
Time
Mem
PS1
5785
Time
Mem
WIk
0.824
0.578
63
0.3
0.824
0.556
(1.3
(1.3
0.824
0,618
(1.8%)
(12
Delicious
0.356
483
(1.5x)
1.7 (L3xx)
0.469
162 (3x)
(1.5x)
Wiki-32KK
0.501
0,235
5702
3.9
0.438
0.205
4680
(1.2x)
(1.xx
0.519
0,255
1906 (3x)
(1.5x)
Amz-670K
0.332
0.273
8310
6.7
0,335
0.276 3224 (2.6x)
(L.6x)
0.336 0.281
1279 (6.5%)
3,3 (2x)
From a system design perspective, MONGOOSE improves the flexibility of LSH's usage in NN train- ing. We provide a smart scheduler to decide when to perform LSH updates and visualization on when to use LSH or learnable LSH. Therefore, the training process has the following tunable pa- rameters for L.H:: (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 LSH should perform the update in. Setting these parameters appropriately will enable much more efficient NN training or fine-tuning, supported by our theory and observations.
4 EXPERIMENTS
4.1 we present the results of experiments that demonstrate the higher performance of MONGOOSE over other LSH-NN frameworks during training. Then we study the effectiveness of the two components: (i) smart scheduler in Section 4.2 (ii) learnable LSH in Section 4.3
In Section
4.1 MAIN RESULTS: MONGOOSE IN NN TRAINING
We present the results of main experiments that demonstrate MONGOOSE's ' advantage in and-to-end NN training over two different LSH-NN baselines on recommendation and language modeling tasks.
4.1.1 TASK1: EXTREME MULTI-LABEL CLASSIFICATION
Settings: We compare the performance of MONGOOSE, against SLIDE and networks without LSH ) on large-scale multi-label classification datasets: Wiki10-31k, Delicious- 200k, Wiki-325k and Amz-670k. The datasets are obtained from the Extreme Classification Repos- itory (Bhatia et al. 2016), which is a benchmark for various recommendation systems. The model is fully-connected network with one hidden-layer. The statisties of each dataset are in Appendix D We report the best results of each framework in P@1/5 following metrics Bhhtii et al. 2016).
Results: Table [ shows the test accuracy, wall-clock time to converge and memory consump- tion for MONGOOSE and other baselines. Left plot in Fig- ure [ visualizes P@l for MON- GOOSE and other baselines dur- ing the training process. MON- GOOSE achieves 1.8x faster con- vergence and 1.5x less memory consumption compared Full- Computation Wikil0-31KK KK, 3x faster convergence, 1.5x less
0.5
Refoemer
MONGDOSE
0.43
1.75
0.35
1.50
0.33
Fill-Comgutatin
0.5
SLIDE MONGOUSE
0.33
1000
2000
2000
40301
50003
2000
50031
T5003
200031
125000
lirations
Figure 5: Left plot: P01 versus training iteration on Wiki-325k. Right plot: Loss versus training iteration on synthetic copying task.
memory Delicious-200k and Wiki-325K and 6.5x faster convergence, 2x less memory Amz- 670K. SLIDE has a large accuracy drop without a substantial speedup for Wiki-325k, but our method has 3x speedup with no accuracy loss.
Table 2: Time and memory advantage com-
The above mainly focuses on MONGOOSE's end-to-end performance during training. We also present closer analysis of a single linear layer during training Table In isolation, we find that MONGOOSE achieves up to 20x faster convergence and 4.5x less memory con- sumption compared Full-Computation. The greater
Analysis:
pared to Coll-Computatin on one layer.
SLIDE
MONGOOSE
Datasets
Time
Mem
Time
Mem
Wiki1033IK
1.8x
2x
6.6x
4x
Delicious -00K
1.8x
2.2x
8.6x
4.5x
Wiki-325K
14x
1.7x
20
4x
MMNGooSE code is available http://////ttttt.../////eeeeeeee

---

Published as a conference paper at ICLR 2021
time and memory savings here compared to Table I are because we ignore the earlier large em- bedding layer of the model (Details are in Appendix (D33) We make three observations to help explain the memory and time savings fromMONGOOSE. First, matrix multiplication is only per- formed on a small subset of neurons. The average number of neurons sampled in the output layer for MONGOOSE on Delicious-2000 200K is around 6k compared the total number of 200k. Second, our scheduler further reduces updating costs. On Wiki-325K, we observe updating frequency drops significantly at the start of the second epoch. This is compatible with our observation of slow weight change in Figure B] Finally, learnable LSH improves LSH query efficiency and quality. In addi- tion, we observe that MONGOOSE obtains a slightly higher accuracy compared Full- Colputatttt,,nn,, which is a side benefit from activation sparsity (Ba & Frey 2013
4.1.2 TASK2: EFFICIENT TRANSFORMER
Settings: We run experiments on two tasks from Reformer (Kitaev et al. 2020). The first task is a synthetic sequence duplication task where inputs of the form Ow:w and E {0....,N] N)",,""" 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 limited range (e.g., local attention). We compare MONGOOSE with a I-layer Reformer by replacing the data- independent LSH in the attention layer with learnable LSH, updating the LSH parameters according to our smart scheduler. The second task is standard language modeling task on the enwik8 dataset. For this second task, we use sequence length of 8192 and 10 L.SH-attention layers.
We compare the test loss of MONGOOSE and Reformer on each task. The loss curve for the synthetic sequence duplication task is shown in Figure 5 (right). On this task, MONGOOSE achieves a loss value of 0.75, which is 0.5 better than Reformer. Moreover, MONGOOSE only needs 20% time to reach the same loss as Reformer (5x speed-up). Similarly, on enwik8, Reformer only reaches 2.65 in perplexity, while MONGOOSE achieves 2.59 in the same setting. Both MONGOOSE and Reformer (with LSH) save 6x memory over the original Transformer in above two tasks (more details are in Appendix D.4).
Results:
4.2 STUDY ON SMART SCHEEDLER
Table 3: Speed with respect to accuracy
In this section, we demonstrate the superiority of our smart scheduling algorithm over the two base- lines we mentioned in Section on Wiki-325K.
P8l
Pa5
Speed (hatchls)
Infrequent Scheduler
0.39
0.193
31.2
w/O Scheduler
0.521
0.254
0.4
MONGOOSE Scheduler
0.519
0.255
7.7
Effect of Smart Scheduler:
The first approach
("Infrequent Scheduler") performs an update on a fixed schedule, every 1000 iterations (6 times per epoch on Wiki-325K). The other baseline performs an update every iteration ("W/O Sched- uler"). We compare speed (number of batches processed per second) and accuracy (P@l and P@5) of these schedulers in Table 3] On average, MONGOOSE Scheduler reduces more than 97% of the update overhead compared to W/O Scheduler, leading end-to-end 20x speed-up. Meanwhile, although Infrequent Scheduler faster in speed, it fails to reach the same task accuracy as the other approaches. In fact, MONGOOSE Scheduler achieves Infrequent Scheduler's final accuracy in 66% less time. We can see that MONGOOSE substantially reduces LSH update overhead while maintain- ing performance.
4.3 STUDY ON LEARNABLE LSH
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 func- tions. 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 us- ing the same number of randomly chosen neurons. A negative value these plots indicates that LSH is performing worse than random sampling. We can see that at the beginning of training, when weights are initialized randomly, data-independdet LSH is sufficient partition the data. However, as training progresses, SimHash used in SLIDE fails to adapt to the evolving weight distribution, while our selection better approximates the full softmax.
Effect of Learnable LSH:
We conduct the same experiments for enwik8 language modeling, In Figure (right), we can see that data-dependent hashing computes higher inner products within each bucket as training progresses.
8

---

Published as a conference paper at ICLR 2021
Table 4: This table compares MONGOOSE and HNSW. P5 is the top- 1/5 precision.
Rebuild Time(s) vs Number of Neurons
Wiki 10k Training
Wiki-35kk Training
10k
100k
300k
P81
P'55
Time (s)
Pg1
P
Time (s)
HNSW
0.4983
5.1 168
20.5933
0.829
0.602
60
0.4625
0,2052
10106
MONGOOSE
0.4625 (LIX) 0.4780 (IIx)
0.8610 (23x)
0.824
0.618
35 (1.73)
0.519
0.255
1906 (5.3))
For attention layers in the Transformer models, random LSH fails to approximate attention well in certain layers due to skewed attention distribution (this phenomenon was also explained in Section 3.3.2 and Figure 4 In attention layers, our selection produces better characterization of the full attention.
Other data-dependent NNS data structures: In Section D we have briefly discussed that most exist- ing data-dependent NNS-ds have
Arg. more product
SLDDF MONGGOSE
1.0
high update costs. We choose an LSH-based data structure in MONGOOSE because update over- head for hash-table-based NNS-ds less compared to other meth- ods (e.g. graph-based) in prac- tice (Wang et al. 2018 Gong eet al., 2020). For example. in a
1.0
1
Befanner
MONGOO6E
50000
300301
176000
200EC
2000
53301
10C30
100030
Ieerotiozs
Eterations
Figure 6: The normalized inner product computed by NNS-ds neu- rons. In hoth Wiki-325k (left) and enwik8 (right), MONGOOSE has much higher average inner product than baselines.
recent paper, (Dong et al. 2019)
proposed a n-nearest-neighbor graph-based algorithm for NNS that requires several hours to pro- cess million data points. Although recent work, HNSW (Malkov & Yashunin) 2018), relatively fast in preprocessing and updating step among graph-based data-dependent NNS-ds empirically, it does not have guarantees (Fu et al., 2017). In Table we compare the update times of the learnable LSH and HNSW, and their performance during training when both equipped with the MONGOOSE smart scheduler. We can see that HNSW's ss updating overhead grows exponentially and becomes over x0 learnable LSH's overhead in 350k neurons layer. However, it is an ongoing work to reduce the updating overhead of graph-based NNS-ds, so could be potentially included in MONGOOSE later.
CONCLUSION
In this work, we make a novel observation on the slowly changing LSH hash codes during LSH- NN training. Based on the observation, we present MONGOOSE, a learnable LSH framework that outperforms original LSH methods in efficient neural network training. We demonstrate two key designs the framework. One is a smart scheduling algorithm which reduces the L.SH update overhead. The other one is low cost learnable LSH which can improve the query efficiency and hashing quality. We empirically verify the effectiveness of our framework on both recommendation and language modeling tasks. The slowly changing observation along with the smart scheduler in MONGOOSE opens up the opportunity of designing new NNS-ds for dynamic similarity search for efficient NN training.
ACKNOWLEDGMENTS
We would like to thank Xun Huang, Fred Sala, Avner May, Ben Coleman and the anonymous re- viewers for helpful discussions and feedback. We gratefully acknowledge the support of NIH under No. U54EB020405 (Mobilize), NSF under Nos. CCF1763315 (Beyond Sparsity), CCF1563078 (Volume Velocity), and 1937301 (RTML); ONR under No. N000141712266 (Unifying Weak Su- pervision); the Moore Foundation, NXP, Xilinx, LETI-CEA, Intel, IBM, Microsoft, NEC, Toshiba, TSMC, ARM, Hitachi, BASF, Accenture, Ericsson, Qualcomm, Analog Devices. the Okawa Foun- dation, American Family Insurance, Google Cloud, Swiss Re, Total, the HAI-AWS Cloud Credits for Research program, the Stanford Data Science Initiative (SDSI), and members of the Stanford DAWN project: Facebook, Google, and VMWare. The Mobilize Center is Biomedical Technol- ogy Resource Center, funded by the NIH National Institute Biomedical Imaging and Bioengineer- ing through Grant P4IE00200... The U.S. Government is authorized to reproduce and distribute reprints for Governmental purposes notwithstanding any copyright notation thereon. Any opinions, findings, and conclusions or recommendations expressed in this material are those of the authors and do not necessarily reflect the views, policies, or endorsements, either expressed or implied, of NIH, ONR, the U.S. Government.

---

Published as a conference paper at ICLR 2021
REFERENCES
Alessandro Achille, Matteo Rovere, and Stefano Soatto. Critical learning periods in deep neural networks. arXiv preprint arXiv: 2017.
Saurabh Agarwal, Hongyi Wang, Kangwook Lee, Shivaram Venkataraman, and Dimitris Papail- iopoulos. Accordion: Adaptive gradient communication via critical learning regime identifica- tion. arXiv preprint arXiv:2010. 16248, 2020.
Amin Ahmad, Noah Constant, Yinfei Yang. and Daniel Cer. Reqa: An evaluation for end-to-end answer retrieval models. In MRQA @ EMNLP, pp. 137-146, 2019.
Alexandr Andoni and Piotr Indyk. Near-optimal hashing algorithms for approximate nearest neigh- bor in high dimensions. In 47th annual IEEE symposium on foundations of computer science (FOCS), pp. 459-468, 2006.
Alexandr Andoni and Iya Razenshteyn. Optimal data-dependent hashing for approximate near neighbors. In Proceedings of the forty-seventh annual ACM symposium on Theory of computing (STOC), pp. 793-801, 2015.
Alexandr Andoni, Thijs Laarhoven, Ilya Razenshteyn, and Erik Waingarten. Optimal hashing-based time-space trade-offs for approximate near neighbors. In Proceedings of the Twenty-Eightt An- nual ACM-SIAM Symposium on Discrete Algorithms,, pp. 47-66. SIAM, 2017.
Alexandr Andoni, Assaf Naor, Aleksandar Nikolov, Ilya Razenshteyn, and Erik Waingarten. Data- dependent hashing via nonlinear spectral gaps. In Proceedings of the 50th Annual ACM SIGACT Symposium on Theory of Computing (STOC), pp. 787-800, 2018
Jimmy Ba and Brendan Frey. Adaptive dropout for training deep neural networks. In Advances in neural information procexsing systems (NeurIPS), pp. 3084-3092, 2013.
Yoshua Bengio, Réjean Ducharme, Pascal Vincent, and Christian Jauvin. A neural probabilistic language model. Journal of machine learning research (JMLR), 3(Fcb): 2003.
K. Bhatia, K. Dahiya, H. Jain, A. Mittal, Y. Prabhu, and M. Varma. The extreme classificatton repos- itory: Multi-labe datasets and code, 2016. URL http::ttpppmmtet htt////maaaaaaaa.aamaaaaa wtwp:/awwwaa...ood/doooo..aaaa.
CC//MMRepos. wtw.ostttttooooooooooooooo html
Lconid Boytsov, David Novak, Yury Malkov, and Eric Nyberg. Off the beaten path: Let's replace term-based retrieval with k-nn search. In Proceedings of the 25th ACM international confer- ence on information and knowledge management (CIKM), pp. 1099-1108, 2016
Jan den Brand, Yin-Tat Lee, Danupon Nanongkai, Richard Peng, Thatchaphol Saranurak, Aaron Sidford, Zhao Song, and Di Wang. Bipartite matching in nearly-linear time on moderately dense graphs. In 2020 IEEE 61st Annual Symposium on Foundations of Computer Science (FOCS), pp. IEEE, 2020a.
Jan van den Brand, Yin Tat Lee, Aaron Sidford, and Zhao Song. Solving tall dense linear programs in nearly linear time. In Proceedings of the 52nd Annual ACM SIGACT Symposium on Theory of Computing (STOC). pp. 778,788, 2020b.
Jan van den Brand, Yin Tat Lee, Yang P Liu, Thatchaphol Saranurak, Aaron Sidford, Zhao Song, and Di Wang. Minimum cost flows, mdps, and -regresssoo in nearly linear time for dense instances. In Proceedings of the 53rd Annual ACM SIGACT Symposium on Theory of Computing (STOC). Xrviv preprint (ixi:21111005711 2021.
Tom B Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. arXiv preprint arXiv:2005.14165. 2020.
M. Burtscher, R. Nasre, and K. Pingali. A quantitative study of irregular programs on gpus. In 2012 IEEE International Symposium Workload Characterization (ISWC), pp. 141-151, 2012.
10

---

Published as a conference paper at ICLR 2021
Sarath Chandar, Sungjin Ahn, Hugo Larochelle, Pascal Vincent, Gerald Tesauro, and Yoshua Ben- gio. Hierarchical memory networks. arXiv preprint arXiv:160S.07423 2016.
Wei-Cheng Chang, Felix X Yu, Yin- Wen Chang, Yiming Yang, and Sanjiv Kumar. Pre-training tasks for embedding- based large-scale retrieval. arXiv preprint arXiv:2002.03933 2020.
Moses S Charikar. Similarity estimation techniques from rounding algorithms. In Proceedings of the thiry-fourth annual ACM symposium on Theory of computing (STOC), pp. 380-388, 2002.
Gal Chechik, Varun Sharma, Uri Shalit, and Samy Bengio. Large scale online learning of image similarity through ranking. Journal of Machine Learning Research (JMLR), 11(3), 2010.
Beidi Chen, Tharun Medini, James Farwell, Sameh Gobriel, Charlie Tai, and Anshumali Shrivas- tava. Slide: In defense of smart algorithms over hardware acceleration for large-scale deep learn- ing systems. In MLSys. https: http///axivv.vv. org http://..rgppp///////// 03129 2020.
Wenlin Chen, James Wilson, Stephen Tyree, Kilian Weinberger, and Yixin Chen. Compressing neural networks with the hashing trick. In International conference on machine learning (ICML), pp. 2285-2294, 2015.
Michael Cohen, Yin Tat Lee, and Zhao Song. Solving linear programs in the current matrix multiplication time. In Proceedings of the 51st annual ACM SIGACT symposium on theory of computing (STOC), pp. 938 942, 2019.
Benjamin Coleman, Anshumali Shrivastava, and Richard G Baraniuk. Race: Sub-linear memory sketches for approximate near-neighbor search on streaming data. arXiv preprint arXiv:1902.06687, 2019.
Paul Covington, Jay Adams, and Emre Sargin. Deep neural networks for youtube recommendations. In Proceedings of the 10th ACM Conference on Recommender Systems, pp. 191-198, 2016.
Mayur Datar, Nicole Immorlica, Piotr Indyk, and Vahab Mirrokni. Locality-sentt sensstvvee hashing scheme based p-stable distributions. In Proceedings of the twentieth annual symposium Computational geometry, pp. 253-262. ACM, 2004.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In IEEE conference on computer vision and pattern recognition (CVPR), pp. 248-255, 2009.
Yihe Dong, Piotr Indyk, Ilya Razenshteyn, and Tal Wagner. Learning space partitions for nearest neighbor search. In International Conference on Learning Representations (ICLR), 2019.
Bernhardsson Erik, Aumüller Martin, and Faithfull Alexander. ANN Benchmarks. https://
github. wwtp://ww/e/eee/eeeeeeaeemae aar.reeeeeaaaaaaeeenaaaeeee 2018.
Miao Fan, Jiacheng Guo, Shuai Zhu, Shuo Miao, Mingming Sun, and Ping Li. Mobius: Towards the next generation of query-ad matching baidu's 's sponsored search. In Proceedings of the 25th ACM SIGKDD International Conference on Knowledge Discovery and Data Mining (KDD), pp- 2509-2517, 2019.
Cong Fu, Chao Xiang, Changxu Wang, and Deng Cai. Fast approximate nearest neighbor search with the navigating spreading-out graph. arXiv preprint arXiv: 1707.00177 2017.
Jianfeng Gao, Xiaodong He, Wen-tau Yih, and Li Deng. Learning continuous phrase representations for translation modeling. Proceedings of the 52nd Annual Meeting of the Association for Computational Linguistics (ACL), pp. 699-709, 2014.
Long Gong, Huayi Wang, Mitsunori Ogihara, and Jun Xu. idec: Indexable distance estimating codes for approximate nearest neighbor search. Proc. VLDBB Endow, 1389):14831149, 2020.
Ruiqi Guo, Sanjiv Kumar, Krzysztof Choromanski, and David Simcha:. Quantization based fast inner product search. In Artificial Intelligence and Statisfics (AISTATS), pp. 482-490, 2016.
11

---

Published as a conference paper at ICLR 2021
Rob Hall and Josh Attenberg. Fast and accurate maximum inner product recommendations map- reduce. In Proceedings of the 24th International Conference on World Wide Web (WWW). pp- 1263-1268, 2015.
Xiangnan He and Tat-Seny Seng Chua. Neural factorization machines for sparse predictive analytics. In SIGIR. pp. 355-364, 2017.
Baihe Huang, Shunhua Jiang, Zhao Song, and Runzhou Tao. Solving tall dense sdps in the current matrix multiplication time. arXiv preprint arXiv:2 2021.
Himanshu Jain, Yashoteja Prabhu, and Manik Varma. Extreme multi-label loss functions for rec- ommendation, tagging, ranking & other missing label applications. In Proceedings of the 22nd ACM SIGKDD International Conference on Knowledge Discovery and Data Mining (KDD), pp- 935-944, 2016.
Stanislaw Jastrzebski, 7achary Kenton, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amos Storkey. On the relation between the sharpest directions of dnn loss and the sgd step length. arXiv preprint arXiv: 0,5031, 2018
Herve Jegou. Matthijs Douze, and Cordelia Schmid. Product quantization for nearest neighbor search. IEEE Transactions on Pattern Analysis and Machine Intelligence (PAMI), 33(1):117- 128. 2011.
Haotian Jiang, Tarun Kathuria, Yin Tat Lee, Swati Padmanabhan, and Zhao Song. A faster interior point method for semidefinite programming. FOCS, 2020a.
Haotian Jiang, Yin Tat Lee, Zhao Song, and Sam Chiu-wai Wong, An improved cutting plane method for convex optimization, convex-concave games and its applications. In STOC, 2020b.
Shunhua Jiang, Zhao Song, Omri Weinstein, and Hengjic Zhang. Faster dynamic matrix inverse for faster Ips. Proceedings of the 53rd Annual ACM SIGACT Symposium on Theory of Computing (STOC). arXiv preprint arXiv:0004.07400, 2021.
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. ICLR, 2020.
Tim Kraska, Alex Beutel, Ed H Chi, Jeffrey Dean, and Neoklis Polyzotis. The case for learned index structures. In SIGMOD, pp. 489-504, 2018.
Moontae Lee, Xiaodong He, Wen-tau Yih, Jianfeng Gao, Li Deng. and Paul Smolensky. Reasoning in vector space: An exploratory study of question answering. In ICLR, 2016.
Yin Tat Lee, Zhao Song, and Qiuyi Zhang. Solving empirical risk minimization the current matrix multiplication time. In COLT, 2019.
Yury Malkov, Alexander Ponomarenko, Andrey Logvinov, and Vladimir Krylov. Scalable dis- tributed algorithm for approximate nearest neighbor search problem in high dimensional general metric spaces. In International Conference on Similarity Search and Applications, pp. 132-147. Springer, 2012.
Yury Malkov, Alexander Ponomarenko, Andrey Logvinov, and Vladimir Krylov. Approximate near- est neighbor algorithm based on navigable small world graphs. Information Systems, 45:61-68, 2014.
Yury A Malkov and Dmitry A Yashunin. Efficient and robust approximate nearest neighbor search using hierarchical navigable small world graphs. IEEE Transactions on Pattern Analysis and Machine Intelligence (PAMI), 2018
Tharun Kumar Reddy Medini, Qixuan Huang, Yiqiu Wang, Vijai Mohan, and Anshumali Shrivas- tava. Extreme classification in log memory using count-min sketch: A case study of amazon search with 50m products. In Advances in Neural Information Processing Systems (NeurIPS), pp- 13244-13333 13254, 2019.
12

---

Published as a conference paper at ICLR 2021
Ioannis Partalas, Aris Kosmopoulos, Nicolas Baskiotis, Thierry Artieres, George Paliouras, Eric Gaussier, lon Androutsopoulos, Maasii--Reza Amini, and Patrick Galinari. Lshtc: A benchmark for large-scale text classification. arXiv preprint arXiv:1503.0858. 2015.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang. Zachary Vetito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library. In Advances in Neural Information Processing Systems (NeurIPS), pp. 8024- 8035. 2019.
Jeffrey Pennington, Richard Socher, and Christopher D Manning. Glove: Global vectors for word representation. In Proceedings of the 2014 conference on empirical methods in natural language processing (EMNLP), pp. 1532-1543, 2014.
Jack Rae, Jonathan Hunt, Ivo Danihelka, Timothy Harley, Andrew W Senior, Gregory Wayne, Alex Graves, and Timothy Lillicrap. Scaling memory-augmented neural networks with sparse reads and writes. In Advances in Neural Information Processing Systems (NeurlPS), pp. 3621- 3629, 2016.
Parikshit Ram and Alexander G Gray. Maximum inner-product search using cone trees. In Pro- ceedings of the 18th ACM SIGKDD international conference on Knowledge discovery and data mining (KDD), pp. 19119331 ACM, 2012.
Hubert Ramsauer, Bernhard Schäfl, Johannes Lehner, Philipp Seidl, Michael Widrich, Lukas Gru- ber, Markus Holzleitnrr, Milena Pavlovic, Geir Kjetil Sandve, Victor Greiff, David P. Kreil, Michael Kopp, Günter Klambauer, Johannes Brandstetter, and Sepp Hochreiter. Hopfield net- works is all you need. arXiv preprint arXiv:2008.02211 2020.
Alexandre Sablayrolles, Matthijs Douze, Cordelia Schmid, and Hervé Jégou. Spreading vectors for similarity search. In ICLR, 2019.
Minjoon Seo, Jinhyuk Lee, Tom Kwiatkowski, Ankur P Parikh, Ali Farhadi, and Hannaneh Ha- jishirzi. Real-time open-domain question answering with dense-sparse phrase index. In ACL, pp- 4430-4441, 2019.
Aliaksei Severyn and Alessandro Moschitti. Learning to rank short text pairs with convolutional deep neural networks. In Proceedings of the 38th International ACM SIGIR Conference Re- search and Development in Information Retrieval (SIGIR), pp. 373-382, 2015.
Anshumali Shrivastava and Ping Li, Asymmetric Ish (alsh) for sublinear time maximum inner prod- uct search (mips). In Advances in Neural Information Processing Systems (NeurIPS), pp. 2321- 2329, 2014.
Anshumali Shrivastava and Ping Li. Improved asymmetric locality sensitive hashing (ALSH) for maximum inner product search (MIPS). In UAI, pp. 812-821, 2015.
Zhao Song and Zheng Yu. Oblivious sketching- based central path method for solving linear pro- gramming problems. In manuscript, 2021.
Ryan Spring and Anshumali Shrivastava. Scalable and sustainable deep learning via randomized hashing. In Proceedings of the 23rd ACM SIGKDD International Conference on Knowledge Discovery and Data Mining (KDD), pp. 445-454, 2017.
Yukihiro Tagami. Annexml: Approximate nearest neighbor search for extreme multi-label clas- sification. In Proceedings of the 23rd ACM SIGKDD International Conference Knowledge Discovery and Data Mining (KDD), pp. 455 464, 264,2017.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems (NeurIPS), pp. 1998-6008, 2017.
Jingdong Wang, Ting Zhang, Nicu Sebe, and Heng Tao Shen. A survey on learning to hash. JEEE Transactions Pattern Analysis and Machine Intelligence (PAMI), 7064),669-79,, 2017.
13

---

Published as a conference paper at ICLR 2021
Yiqiu Wang, Anshumali Shrivastava, Jonathan Wang, and Junghee Ryu. Randomized algorithms accelerated over cpu gpu for ultra-high dimensional similarity search. In Proceedings of the 2018 International Conference Management of Data, pp. 889-903, 2018.
Donna Xu, Ivor W Tsang, and Ying Zhang. Online product quantization. IEEE Transactions on Knowledge and Data Engineering (TKDE), 30(1 30(11)2218--1198, 2018.
Hong-Jian Xue, Xinyu Dai, Jianbing Zhang, Shujian Huang, and Jiajun Chen. Deep matrix factor- ization models for recommender systems. In IJCAI, pp. 3203-3209, 2017.
14

---

Published as a conference paper at ICLR 2021
Appendix
Table of Contents
Slow change observations
16
B
Dynamic maintemance data structure
17
Illustration
17
Main result
17
Proof theorem
B.4
Hounding move
19
B.5
Bounding move
21
B.6
Polential function
23
B.7 7
Examples
24
Learnable Ish
26
Learning hash functions)
26
Observations attention distribution
................................
27
D
Experiments details
28
D. Data statisties
28
D.2 Influence of the Parameters
28
D.3
Additional results on the excrrme--lassificateee
.a..............................
28
Additional results on the language modeling task
E
Related work
31
E.1
Data structures for dynamic similarity search
................................
31
E.2 Data dependent indexing
................................
31
E.3
Elficient neural network training
32
F
Efficient GPU implementation
33
Broader Impact
33
15

---

Published as a conference paper at ICLR 2021
A SLOW CHANGE OBSERVATIONS
In this section, we analyze the relative distance between iterates of model parameters and connect it with LSH updates throughout training. We target at making full use of our slowly change character- ization for MONGOOSE.
Settings. We report AW throughout training on several benchmark models. We follow the stan- dard training procedures of fulny-connected networks on Wiki325k (Partalas et al. 2013) and Trans- former (Vaswani al. (017) on enwik8. 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 ex- periments in section [
Results. We plot our results in Figure 3] The left-two and right-most plots show that during the initial steps of the training, AW relatively high, but quickly drops and flattens out afterwards. The middle and right-most plots exhibit the hash code change of W in the hamming distance along with the training. The pattern matches with AW but has a direct connection with the LSH update overhead. 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. Note the above observations are made based angular distance LSH. The above phenomenon suggests that if there exists optimal scheduler to update the data structures adaptively based on the actual demand, the overhead by LSH updates can be largely reduced. Furthermore, this opens the door to make LSH learnable in order to improve query efficiency. 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. Several related works (Jastrzebski et al. 20188] Agarwal et al. 2020) have also observed the phenomenon of a sharp drop of weight changes at the early stages of training One (Jastrzebski et al. (2018) conjectures that initially SGD visits increasingly sharp regions, reaching maximum sharpness determined by both the learning rate and the batch-size of SGD and infer that it optimizes faster while finding good sharp region. It is also discussed in recent paper (Agarwal al. 2020), which connects the phenomenon to critical learning periods defined in (Achille et al.) (017). This phenomenon can be observed on text (Transformers), recommendation (Embedding Models) and image (CNN) data and thereby it is relatively general.
Aug Weight Change
Avg Weight Charge
Avg Hash Code Change
0.30
0.43
4.49
Cooine Distance
0:30
0.30
0.30
0.53
0.50
0.82
033
0.31
Epoch
Emocl
Epoch
Avg Weight Clange
Avg Weight Change
Aog Heal Code Clange
113
Codaae Dodqumen
0.0
10
AIF
11.4
011
0.0
Egucc
Epoch
Epacl
Figure 7: We show the average change of weight (left), cosine similarity (middle), weight's hash code (right). Top row is reformer and the bottom row is for Wiki-325k.
16

---

Published as a conference paper at ICLR 2021
B DYNAMIC MAINTENANCE DATA STRUCTURE
The idea of dynamic maintenance has been successfully applied to matrix inversion problem. which served as central component in linear program solvers (Cohen al. 2019 Lee et al., 2019 Brand et al. 2020b Song & Yu 2021 Jiang et al. 2021), cutting plane method (Jiang et al., 2020b),
maximum matching Brand et 2020a max- max-flow Brand et al. (202I), and semi-defnite program
solvers (Jiang et al. 2020a Huang et al. 2021). In this work, we adopt it for dynamic LSH mainte-
nance problem. Since the underlying problem completely different, it requires several non-trivial
generalizations and more fine-grained analysis (see Lemma B.2 Lemma B.3] Lemma [.4]. More- over, in our proof, we care about the absolute distance between vectors (i.e. ||wr % rather than the relative difference between real numbers Cohen et al. (2019), and therefore, need new construction the potential function (see Lemma B.8).
The rest of this section is organized as follows
In Section
B.
we provide an illustration of the smart scheduler.
In Section
B.2
we state our main results.
In Section
B33
we give high-level overview on the correctness and the running time of
the algorithm.
In Section
B.4
we give an analyze on the movement of w vector.
In Section
B.5
we give analyze on the movement of u vector.
Section
B.6
we discussed over the choice of potential function and analysis its property.
compare the performance algorithm with sequential updating strat- egy/batch updating strategy.
In Section
B.7
B.1 ILLUSTRATION
SCHEDULER
We present a demonstration of smart scheduler Fig- ure 8 As shown in the top,
Fai
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. However. if the data vec- tors change such that they cross partition boundaries, the scheduler updates the hash functions.
Uptler
Updore
Figure 8: Illustration of the smart scheduler
B.2 MAIN RESULT
Theorem B.1 (Locality sensitive hashing maintenance, Formal statement of Theorem 3.3). For any constant 01,02 02 (c) (2) there is dynamic data structure (Algorithm 7 that achieves (c1, 02)- accuracy. The data structure takes Ö(dn +P) time initialize and each call of QUERY takes time dn. By taking a = min(p, a) and
nnPa S n
9r
n br
The amortized expected time per call of UPDATE(W) at most
Ō((C + 119l12).
(4)
B.3 PROOF OF THEOREM B.
We first give rough explanation on the proof of Theorem 3.3 For intuition, we consider the case
G @(1), C2 ((1), and 6mp (11) in this explanation. The amortized time analysis is based
17

---

Published as a conference paper at ICLR 2021
on potential function that measures the distance of the approximate vector and the target vector w. We will show that
The cost to update the LSH data structure is proportional to the decrease of the potential.
Each call to query increase the potential by a fixed amount.
Combining both together gives the amortized running time bound of our data structure.
Now, explain the definition of the potential. Consider the k-thh round of the algorithm. For all i E [n]] we define =1-h w/(h - u; (*( Note that ||=;">] 22 measures the distance between (x()) and of Our algorithm fixes the indices with largest error ||=]{* ||2- To capture the fact that updating in a larger batch is more efficient, we define the potential as weighted combination of the error
where put more weight to higher x²!². Formally, sort the coordinates of x:(k) such that
|| (k) 122 2 || and define the potential by
Σ Di {x{x) (1)
Vx
where 9i are positive decreasing numbers to be chosen and vb is symmetric (+b(x) v(-x)) positive function that increases on both sides. For intuition, one can think (x)x) behaves roughly like [x].
B.3.1 PROOF OF CORRECTNESS
We prove the correctness of Theorem 3.3 we will defer some simple calculations into later sections.
Definition of and yy- Consider the k-th round of the algorithm. For all i E [rx], define (k)
Rd. (k+1) IRd and y; (k)
E IRd as follows:
(k) (k) (N) (k)
(k+1) (R) (k+1)
w. (.++1) Uy
(k+1)
F w
12 U/s wi Us
Note that the difference between x; and (k+1) is that is changing.
(k) and yi
(k) is that is changing. The difference between y/
(k) Assume sorting, Assume the coordinates of vector (*)) E gaxd are sorted such that || 1|2 2 | Let T and æ are permutations such that -| 1ot (k) [1utt+t))ll-
112 2 || 7(++(ll llz and (1v) llz
Definition of Potential function. Let v R I. be defined by
||=112 2emda
||ml| E [0, fmds)
(xx))=
=
Cmdss
(4c2 mdl 11
18c
| (Emds mcnds
(5)
Cuudss
|| |alla (26mds, +00)
We define the potential at the k-th round by
Σ gi (b (b)
where (i) is the permutation such that || 111299122 2
|| (k)
Bounding the potential.
We can express M
D as follows:
Σ 9:
(b(2t(t1))
(1( -b(())))))
k++1
1=1
(5(2 (1))
Σ g:
(2 - (a (b+1))
(6)
1=1
18

---

Published as a conference paper at ICLR 2021
Now, using Lemma B.5 and B.7
and the fact that Vo 0 and Mr 0, with Eq. 6 get
0347 v0 Σ (MR+ - Mx)
(0(C, + C2/Emds) 11912 - (00smdeek..../ogl
7 O + (mds)) ||9||2
k log
where the third step follows by Lemma B.5 and Lemma B.7 and F3 is the number of coordinates we update during that iteration.
Therefore, we get,
T
Tk9re 0 (T (C1/ + C2/cnds) log l1gl12).
B.3.2 INITIALIZATION TIME, UPDATE TIME, QUERY TIME
To formalize the amortized runtime proof, we first analyze the initialization time (Lemma B.2 update time (Lemma B.3). and query time (Lemma B.4) of our maintenance data-structure.
Lemma B.2 (Initialization time). The initialization time of data-structure MAINTAIN (Algorithm Odd)+t)).
7
Lemma B.3 (Update time). The update time of data-strun strrccureuuure MAINTAIN (Algorithm) 7 ((rgr)++ O(Sd + S n) where is the number of indices updated in S is the number of weights changed. The later two terms are dominated by the updating time of neural network.
Proof. We change r (r 2 n²)) terms each time, and the running time for updating the LSH is rgr- We need to update and calculate its & norm every time, and the computational cost is Sd. We also need to maintain an order on the error y! using some standard data structure (like Fibonacci heap), this can done in S log n. Thus the total running time per update O(rg,) + +O(d log n), the second term is bounded by the updating cost of neural network training, since calculate the gradient
takes at least ((Sd) and we usually have d >> log n.
Lemma B.4 (Query time). The query time of data-structure MAINTAIN (Algorithm 7 dd.).
is O(ndd +
B.4 BOUNDING MOVE
Lemma B.5 (w move). We have
Ie E (1
+x,,((a)] < <0 + C2/Emds) ||
Proof. Observe that since the errors || (8) [|2 are sorted in descending order, and vl(x) is symmetric and non-decreasing function, thus (h (8)) is also in decreasing order. In addition, note that g is decreasing, we have
Σ
x
x(1) < Σ (+)))
(k)
(7)
1=1
Hence have
19

---

Published as a conference paper at ICLR 2021
Σ
9i
(N(V (k) - v(()1))
- EE Σ 9i - -(x
by Eq. 7
1=1
Σ gi Blu -
O(1 C2/Cmd -191122)
by Lemma B.6
Thus, we complete the proof of w move Lemma.
remains prove the following Lemma, Lemma B.6.
9: Elp(U -000--(220000) 0 C2/Emds) 119122.
Proof. We separate the term into two:
9h E(P(U ((2
(k)
Σ
E El(b( (k)) (B(v (1
9n 9m-1(())
i=1
IN 10)) (B(( (xx10)) )).
(8)
For the first term, have
v(u) (1)) (l((
= (w/ (Elv 1) /1 - Ely (1) +
2 (v.
Ely. 1) " E
S (w/ (E(u) (1)) Vi (k) Ely (K) }>>
T.2
(k)
y (k)
11
(E(v
(2) 1), w/ (f++)) E Elw! (1++ 11) + Ly || (1-1)))))) El
w
(k + 1) 111
(9)
where the first step follows from the Mean value theorem, the second step follows from the definition
of L2 (see Part 4 of potential lemma B.8) the last step follows from the definition of y; (k)
Next, we denote " = Var(w((**1) Summing over and taking conditional expectation given (e( on both sides, we get
In 9 E(v(y(*)) (=(Elv(('1)
([ (BLy, (1)))))))0)))) wi B[w
(k-1) (k+1)
(r)
1=1
L2 (k+1) - Blwo (k 1)
9x 1())
i=1
n
+ L2
Jn (0)) 7
1=1
1
L 1111112
7
2
$ O(1/(((11tt)
|19l|2 - C2
2
= O((2//Emds) 119112.
(10)
The first step follows from Eq. 9 the second step follows from the linearity of expectation and the definition of % the third step follows from Mtuchy-Shward inequality, the fourth step follows from Ta2 O1/ (see Lemma B.8 and Eq.
20

---

Published as a conference paper at ICLR 2021
For the second term, conditioning wi (k) have
()xx) + (B -p - L1
|E [1 2 112 L1
||E
(
w ll2 def
L1 Bz (11)
The first inequality follows from the -ipischhitz continuity of v (see part of Lemma B.8 The second step follows from the definition of y; (8) and x; Summing over i, we get
n
9m (1) (bb( v(x! s Σ -x-1(1) L13, L -19111-222222 ② 39) S O(1₁)
119112-
(12)
The first step follows from Eq. the second step follows from Cauchy Shwarz inequality, the last step follows from L1 (see part 4 of Lemma B.8) and Eq. Combining Eq.
G 2
we have
(1) Σ gi ((W())))) ((22(001 -v())) S SCC C2/Emas)]lgl1
Thus completing the proof.
B.5 BOUNDING v MOVE
The goal of this section is to prove Lemma B.7
Lemma B.7 move). We have,
gi (N(O (K) ((((tt))))) M(emdeTkSrrr/log
Proof. We split the proof into two cases.
We first understand some simple facts which are useful in the later proof. Note that from definition
of x (l++), we know that (l++) has r4k coordinates are Ö. Basically, [|y(k) (x++1) [|o rk- The
difference between those vectors is, for the largest r' coordinates in (*), we erase them in (*++1)
Then for each i E [n rx], =7) Vo(itre) For convenience, we define Vz(n+z) = Ó. Vi E [rx]-
(k+1) (k)
Case 1. We exit the while loop when 1.5rk 2 n.
(k) Let u denote the largest satisfying [[, ||2 emds/ 2.. If u² 2 Ykk then have that (Iy 112
(k)
Emda/ 2 2 emds
100, 100. Otherwise, the condition of the loop shows that
(1 2l2 2 (1 1/ log n los.. ra roosss. I1 (
(k)
(1 1/ log n) 1081.0 " Emaaa/
Edss 100.
where used that 2 4.
21

---

Published as a conference paper at ICLR 2021
According to the definition of o (k+1) have 7())
Σ ( (K) - 15(0 (1+1)) =
(((())))) (8) - +(
((((( (b)
imn/3+
geb(
im/33+
2n/3
Σ Ii
2 Qrrk-rrr (mnds)
in/3+1
where the first step follows from (k+1)
)x(++r))' the second step follows from (5)) is non-
(k)
(()
decreasing (see part 2 of Lemma B.8 and I| |1 ll2 is non-increasing, the third step follows from 1.5rg > n and hence (b ((+
0 for i 2 n/3+11 1, the fourth step follows from (x()) f(1]==
non-decreasing and |1 (1)) ll2 2 ( 2 Emds/ /100 for all i < 2n/3, and the last step follows
by g is non-increasing and the part of Lemma B.8
Case 2. We exit the while loop when 1.5r and (1.) ll2 < (1 - log n
m)llu ( 112.
By the same argument as Case 1, we have that (1// l|l 2 mmds/ 100. Part 3 of Lemma B.8 together with the fact
(K) ( Il2 < min(Cmas/ 11 1l (1 n))
indicates that
(k) (( (( 52(mmds/ l)..).
(13)
Putting things together, we have
Si (V-WH())) (2((tt)))
(K)
(vb(v (k) - ((
by (k+1)
x(1)) = (k)
gi
=Vz(itra)
>
Σ
. (1(11))) (k)
((w
by (( -vb( 2 0
imrt/2
1kk
Σ
gi (W(w))))) ((
(8)
i=r//2
Σ
( Fidds
by 13
log
i=rr/2
Fik
D Emds
gra.
by gi decreasing
log
i=rr/2
= Ω (EmasTkgrr log n).
where the third step follows by 1/22 is non-increasing and ta is non-decreasing (see part 2 of Lemma B.8
22

---

Published as a conference paper at ICLR 2021
B.6 POTENTIAL FUNCTION
Lemma B.8 (Properties of potential function v.). Let function vb Rdd - I& (defined in Eq. 5 Then
function vb satisfies the following properties:
1. Symmetric (+((-x) )- = 1f(x)) and (0) = 0;
If |||| 2 ||.l|2 then (xx (yy):
3. If (x) = (md/) V/zl E [00.0lmmde) ), 2mddll:
4. J max, D. and L2 max,
10/ mdssi
(x)) is a constant for || mnnds
Proof. Recall f :R + R is defined as
2. E [0, mmaali
2rande
f(x) :=
(4c 2
E (< 4e mddll
mmds
18c
x E (4e udda t00).
Emds:
We can see that
z
1 E [0, Endsli
f(x)
(eider 4e mddel; (4emdss +00).
9c
mda
0,
and
x E cmdal:
((x)"
x E (e 4e 4e
|x|| E (4e mds: +00)
It implies that max [xx)]
<
and max [1x) - - Let =x)= 1(1==11))
Cadda
-mdo
Proof of Part 1, 2 and 5. These proofs are pretty standard from definition of sb.
Proof of Part 3. This is trivially following from definition of scalar function f.
Proof of Part 4. By chain rule, we have
Dz 2 (11== (x, Di (b. M| // (11==13) ((x, (1==))= (h. h)
We can upper bound
[a I 2) 12)1 (a, M = 2) (1==11)
||ا|| 11/1122
Thus, we have
| / I 1, (4e || 11x111)0 (9 <
||x||| E [0, Emdel: | lull2 E (Emder 2c ); | 2| E (26mde +o0).
[f((1=1111 ||
0.
It implies that |D2b[h]] < 2|1h|12, Vx.
By case analysis, we have
||0||2 S 4/emds:
|| E [0, 4e2 2c ||= E (4emds: +00).
|-( - 1| <
0,
We can also upper bound
DD h]]l
< (1x=(( (x, 122 + 211=(( ||A112
$ 2 (||=1(²) (||= ||>|2=) + (1(=))) ||A1/2
< 2
|/ 12
Emels |111/
6mdls
10
|/
<mds
23

---

Published as a conference paper at ICLR 2021
B.7 EXAMPLES
Comparing with sequential updating algorithm We present another version of our main result, which is particularly useful when we compare with the sequential updating algorithm.
Theorem B.9 (Locality sensitive hashing maintenance, worst case bound). For any constant C1. C2 (ci > 02) there is dynamic data structure (Algorithm [ that achieves (ci, 62) -accurccy The data structure takes O(dn'+p) +P) time to initialize and each call of QUERY (h) takes time ((n d). By taking min(p, af and gr tr, after T iterations, the amortized expected time per call of UPDATE(W) is at most
n T-
1
(b+1)) (llthttt
Varlw,
(14)
T
k=
1=1
The proof is almost indentical Theorem one. Lemma B.10 (w move). We have
B.I
we only need to replace Lemma B.5 by the following
(K) Σ 9i E (V
v(a!1))
so
91
Blw
(k+1)
(+ + Varlw
i=1
n2
(++1)
w (1)) Varlw,
We omit the dependence on Emds since we assume it to be constant in the scheduler.
only difference that we use / form of Cauchy-Shwarz Shwarz instead of the l form in the previous proof. We only point out the difference here.
Proof. Again, the proof is almost identical Lemma B.5
We replace Eq.
10 with
n In 1(6)
In 1()
(1)) (1EN)(()1)
=
7
L2g1 Σ
7/
2
i=1
n
0 91
Σ
Var[w[+1]
The first step follows from Eq. 10 the second step follows from {gi)i=l.... nn are decreasing and the last step follows from the defimitonn %. Similarly, we replace Eq. 12 with
Σ 9x-1() (F()) 1=1
((((1)
Σ JN (m-1() T. 0
91 Σ
Blw
(
m:
The first step follows from Eq. 12 the second step follows from the definition of Bi, Li 2 and {gi hg.)=,,.... are decreasing.
The rest of the proof are the same and we omit it here.
The sequential updating time for LSH is nz" The total number of updating is (roughly) at least
Σ
B[we
(++1)
wi (6)
+ Var[w[+1)
24

---

Published as a conference paper at ICLR 2021
It then easy that our scheduler always perform better than naive sequential update.
Comparing with naive batch updating algorithm Below we give a concrete example to show the effectiveness of our algorithm. We take we already assume for any r S n = n², T, n², i... the average LSH updating time t, 12" n²"r decays linearly in when I nf. It remains to specify the rest & (r n²).
p for simplicity. Remember in Assumption 3.2
Example B.11. Assuming the average running time decays faster than 1/vr, ie. 4 = -0-811er-- for some B E 1) when r 2 nP. We then have
11
Σ g R n P/2
The running time of our scheduler is then at most
(C1 + C2)lllla = (C1 + C2)m
As long the weight changes slowly, say C1 C2 S np12, 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.
25

---

Published as a conference paper at ICLR 2021
C LEARNABLE LSH
C.I LEARNING HASH FUNCTIONS
m), where each neuron er has weight and bias. Given an input embedding & from the previous hidden layer, the output of forward pass through neuron is ((g, vy)), where some activation functions. For each input embedding di, both SLIDE and REFORMER select set of neurons, denoted S,, from the LSH hash tables for the forward pass. We collect training samples from qi and its S; to improve performance of LSH. Formally, the pairwise training samples (g, v) are collected according to the following criterion:
Denote set neurons in particular layer C {or 0
positive pairs P+ (g, 2) if E S and (q. v) t+ negative pairs P- = (q, v) i E and v) < L-
And the loss is defined
L(H, (4, max {o, Σ coos(H(1), H(v)) + (g.))P;
Σ cos(H(()) H(u)) aa. (15) (q.")E
Here H. the hash functions of all L hash tables. H(x) generates a K vector containing the projection of in each hash table. cos(H(x), H(()) represents the cosine similarity of two projected vectors. The major contribution of this learning approach that optimizes the hash functions and the hash table indices together. Our method focuses learning useful indexing scheme for retrieval efficiency while other learnable hash methods (Wang et al. 2017) focus on binary sketches that have higher precision.
We present a demonstration of our loss in Figure ? Two ran- dom hash functions separate 6 data points into four parts. It is obvious that point (pink) should be partitioned together with point 2 instead of point and 6. On the other hand, point 3 and 4 should not be separated. Therefore, dur- ing the hash function training, we choose point (5,6) as positive pair. Meanwhile, we chose point (5,1) and (5,4) as negative pairs.
Figure Illustration of learning hash functions
Then, the loss would update the hash functions to accommodate these new partitions.
26

---

Published as a conference paper at ICLR 2021
C.2 OBSERVATIONS ON ATTENTION DISTRIBUTION
In this section, we present the visualization with analysis on the distribution of the minimal quantifies of neurons that up have 0.9 softmax values in attention. On Figure TO present the distribution of each head attention in each layer from a transformer model trained on Enwiki8 dataset. We classify the patterns (amssuurr al. (020) of the distribution into 3 categories by their median values. With this observation, we are able to determine the layers that LSH or learnable LSH can apply in MONGOOSE framework.
N
IN
WHAL
X
NM
19
30
38
32
2N
2H
TH
IM
T
Figure 10: The distrihution of the minimal quantifies of neurons that sum to have 0.9 softmax values in attention in each head of attention in each layer of transformer for Enwiki8
27

---

Published as a conference paper at ICLR 2021
D EXPERIMENTS DETAILS
D.1 DATA STATISTICS
Table 5: Statistics for benchmark dataset
Dataset
Wikil0-31k
Delicious-200K
Wiki-325K
Amz-670K
Output Dimension
30938
205443
325056
670091
Input Dimension
101938
782585
1617899
135909
Training Samples
14146
6616
1778351
490449
Testing Samples
196606
100095
587084
153025
We present statistics on the datasets we test on from the Extreme Classification Repository (Bhatia et al. 2016). 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. In particular, each dataset has over 10,000 output classes, which, using conventional nueral network, requires matrix multiplication involving over 10,000 neurons at the final layer.
D.2 INFLUENCE OF THE PARAMETERS
Recall that Algorithm Gmds is the key parameter to our smart scheduler because the number of neurons exceeding this threshold determines if LSH updates are needed. In practice, instead of choosing a single fmdss we use some fraction of the &2 norm of the weights as the update threshold. We empirically observe from the previous experiments that MONGOOSE is robust to any choice the fraction in the range 0.1-0.3. Thus, can consistently outperform other baselines without heavily tuning this parameter for different datasets.
Smart Scheduler Parameters:
Learnable LSH Parameters: For LSH-based NNS-ds, the number of hashes is one key parameter to balance NNS qual- ity and efficiency. In Table 6 we study the effects varying the number of hashes in MONGOOSE. We observe that com- pared to our best empirical choice (10 hashes), both smaller (8) and larger number (12) of hashes increases convergence time. MONGOOSE with fewer hashes selects more unnecessary neurons and lowers training effi- ciency. A larger number of hashes might miss necessary neurons with higher probabilities and would similarly take longer to converge. Another important parameter is the threshold for positive and negative examples to train hash functions. our experiments, top-1 and bottom-1 are the most effective choices (top-3 and top-5 do not make much difference). In general, this threshold in our design robust and does not require heavy tuning.
Table 6: MONGOOSE performance as function of number of hashes.
Number Humee
Pla1
P05
(buccbs)
0.492
0.211
3.79
0.519
12
0.469
0.24
9.92
D.3 ADDITIONAL RESULTS ON THE EXTREME- EXTREME-CLASSIFEEEE CLASSIFICATION TASK
We present the results of additional experiments comparing the classification performance of MON- GOOSE against SLIDE and FULL in Figure
The two metrics, P@ and P@S, are presented in
(Bhatia et al. 2016) follows:
P@k
Σ y:
tErank (5)
where rank;()) extracts the top l: indices from vector, y refers to the predicted labels, and y refers to the ground truth label.
We also present the results of time profiling to demonstrate the speedups of MONGOOSE in each layer. As shown Table 7 Time; represents the total convergence time while Time represents the execution time in the wide output layer (including the forward and backward pass). Here we observe that MONGOOSE provides up to 20x speedup in the second layer. The major reason for the increase in convergence speedup (3x 20x) is that the first embedding layer for high dimensional and sparse input is time and memory consuming Medini et al. (2019). Therefore, the end-to-end result, the advantage of MONGOOSE has been diluted.
28

---

Published as a conference paper at ICLR 2021
Table 7: This table summarizes the performance of MONGOOSE, SLIDE and Full-Computation implemented with PyTorch [Paszke al. 2019). P@ I the lop-1 accuracy and P@5 the top-5 accuracy. Time represents convergence time.
Full Colpppattoooonnn
Untasets
Port
PM6
Time;
Time;
Mem;
Mem;
P10
Time)
Ment:
WIR00 31K
0.824
0,578
63
33
0.3
0.12 0.524
0.556
(1.3 X)
15.33
0.4(1.3)) 0.0
0,524
0,618 35 (1.8) (666)) 0.2 (155x (403
De0coms-2000
11:06
0.356
685
161
233
09
0,403
0,572
(1.50
197 (1.8)
(20330 (1.3) 8:40
0.371
(Xx) 42 (86x1 (1.5%. (.264
WIL-325K
1.501
5700
5000
1.5
(6.0(1.2x))
(14x) (3())s) (0(0.)))
1906
(1.5) 04(4x))
Wiln-1sKK
Delicton 201
Wikito-11
0..0
0.0
0.11
0.13
01:27
0.43
11:30
FullCoompiaaion
0.39
FullCCompuaaii
Full-Comppatiiii. SLIDE SIONGOOSE
SLIDE
SLIDE MONGOOOS
MONGOOSE
0.34
01:00
0
50301
30000
39031
5030
5000
1000
1300
5000
350
500
750
1000
1950
1540
Iterations
Itecattioss
Iterations
Wib-325K
Wikito-11
0.55
0.35
04:00
0.0
PNG
0:17
0.35
Computation SLIRE MONGOOSE
Pull-Compttaton
Full Coll-CCotttoootcc
11:30
0.30
SUDE
SUDD
MONGOOSE
MONGOOSE
20000
19031
50303
5000
251
500
Thb
Ioeeations
Tterations
Figure 11: Comparison of MONGOOSE against SLIDE and FULL during the training. The two metrics (P@1 the top row and P@5 in the bottom row) are the same as (Bhatia et al. 2016).
D.4 ADDITIONAL RESULTS ON THE LANGUAGE MODELING TASK
We present additional experiments comparing the training loss of MONGOOSE and Reformer on the synthetic copy task in Section 4 In Figure 12 present additional experiments comparing the training loss of MONGOOSE and Reformer on the synthetic copy task in Section 4.1.2 The titles of each graph denote the hyperparameters of the Transformer model being tested: h1 $2048 t16 refers to model with round of hashing, sequence length of 2048, and a token size of 16. Overall we can see that MONGOOSE makes a marked improvement in loss over Reformer in every case.
bl 110
h2 2218.166
hla1024.61
2.50
Reforize
MONGOOSE
MONGOOSE
Reforaee
4.7
MONGOOSR
Belienne
0:2
1.78
2.5
1.50
1,25
3.0
Bantiion
Herations
Figure 12: Comparison MONGOOSE against Reformer during training on synthetic copy task.
In Figure B we present further results comparing MONGOOSE to Reformer on the enwik8 character- level language modelling task. In particular, we train sizes of Transformer (3 and layers) with
29

---

Published as a conference paper at ICLR 2021
wwwikS-llyye--000 (176
MONDOOSE
4.4
MONGOOSE
6.7
0.
3.3
1090 12200 1908)
11200 20031
20000
2903
1000
22300
10400
27502
2000
21303
2000
Iiarations
Teerations
Figure 13: Comparison of MONGOOSE against Reformer during training enwik8.
maximum sequence length of 8192 and 4096 respectivelyh. In both settings mongoose achieves lower loss than Reformer.
Note that the goal of this experiments is to prove the superiority of learnable LSH over classical LSH in NN training (or MONGOOSE vs. naive LSH-NN framework) rather than improve the state- of-the-art perplexity. For 6 layer transformer experiment on enwik8, Reformer only reaches 2.71 in perplexity, while MONGOOSE achieves 2.65 in the same setting. Both MONGOOSE and Reformer (with LSH) save 6x memory than the original transformers in above two tasks.
30

---

Published as a conference paper at ICLR 2021
E RELATED WORK
E. DATA STRUCTURES FOR DYNAMIC SIMILARITY SEARCII
Similarity search is a well-studied problem with wide applications in recommendation system (Xue jet al. 2017 Severyn & Moschitti 2015 Hall Attenberg 2015), question answering (Boytsov et al 2016 Seo et al. 2019; Ahmad et al. 2019 2 et al. 2020), multi-label classification (Cov-
ington et al. [0016] J Jain et al. 2016 Tagami 2017) and natural language processing (Bengio et al,, 2003 Gao al.] 2014 Lee et al. 2016). There are two major categories of similarity measures: (1) metric similarity (cosine similarity, euclidean distance), (2) non-metric similarity (inner product, KL divergence, neural network). The brute-force approach solving similarity search is computa- tionally expensive; in response, researchers have developed novel indexing structures to accelerate the search process, with trade-off's on search accuracy. Based on these indexing structures, simi-
larity search algorithms can be broadly categorized as (1) hashing (Shrivastava & Li 2014 2015),
(2) quantization (Guo et al. 2016, Jeeou et al. 201 I). (3) tree-based (Ram & Gray 2012). or (4)
graph-based methods (Malkov et al. 2012 2014
Malkov & Yashumin 2018).
Most similarity search scenarios studied by these papers are static (search data does not change). In the experiments section, the developed similarity search methods are compared on fixed dataset such as Sift (Jegou et al., [201T) or Glove (Pennington et al.) 2014). 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. Therefore, some similarity search methods in metric space have been modified for online settings, such as hashing (Coleman et al. 2019) and quantization (Xu et al., 2018) methods; these methods sacrifice search speed for data- aata-aadativveeess.
The training phase of Deep Learning models provides a natural setting for Dynamic LSH. During training, the weight matrices are slowly modified via gradients derived from objective functions. 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. Recent works take advantage of this view of DNN training by introducing LSH data structures to the NN training process. (Chen et al.) 2020) propose an algorithm (SLIDE) that retrieves neurons with maximum inner product in each step via an LSH based data structure. In this way. the backward pass of NN training is concentrated on the neurons with estimated large gradients. Their CPU implementation able to outperform a traditional GPU implementation. Similar hashing based algorithms have also been used in Transformer models: (Kitaev et al., (020) (Reformer) propose an LSH structure to reduce the memory bottleneck of self- attention modules especially over long sequences in Transformer.
Since DNN weights are the search data, the distribution the weights among different hash buckets changes throughout the training. This necessitates constantly updating the LSH data structure: fail- ing to update the LSH data structure as the search data changes degrades its sourest-neighbo search performance, which in turn worsens the quality of the DNN approximation. In experiments, we found that failing to update the LSH data structure in SLIDE caused a 28% decrease in t accuracy.
E.2 DATA DEPENDENT INDEXING
Data dependent hashing methods (DDH) focus on adapting hashing schemes to the data distribution. They often relate indexing via an objective function. Most literature concentrates on hashing and quantization methods. Although previous works have achieved promising results in learning B-
Trees (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, Complex index designs introduce computation overhead. In the theory of data dependent indexing, (Andoni et al.) 2018) present fruitful insights on hashing-based Approximate Nearest Neighbor search on metric space. However, there is not much literature that puts these insights into practice. (Dong et al., 2019) propose k-NN graph based algorithm for efficiently learning data dependent LSH in Euclidean space based on the insights given by (Andoni et al. 20188). However, their method requires pre- computing a k-NN graph, which is computationally expensive scale. In the experiments, (Dong jet al. 2019) only compare their method with k- k-means; their performance compared to major ANN benchmarks (Erik al. 2018) remains unknown.
31

---

Published as a conference paper at ICLR 2021
Many practical applications of LSH eschew DDH methods for vanilla LSH due to the aforemen- tioned computational bottlenecks. However, despite having sub-linear query time in theory, query complexity for even vanilla LSH is still high in practice Datar et al. 2004);[Andom et al. (2017).
When designing the dynamic data structure, we want to minimize the total running time and balance the overhead and benefit brought by the data structure. In general, the matrix multiplication and backpropagation procedures comes from neural network training, which we can not control, hence, want balance the accuracy-efficinncy trade-offs by selecting neurons.
This is especially pronounced in the Dynamic LSH regime, when weights are evolving, besides query time, LSH updates incur more significant overhead that harms the overall efficiency. (Chen et al. 2020) introduce a method to control update frequency, but their method requires heavy hyper- parameter tuning and not guaranteed work for each benchmark. In contrast, our findings in the following section show that the update overhead can be significantly reduced while maintaining nearest- neirest-neighbor search quality during NN training. To our knowledge, this is the first time DDH techniques have been successfully applied to the Deep Learning setting.
E.3 EFFICIENT NEURAL NETWORK TRAINING
The weights of a neural network dynamically change during the training process, which brings great challenges the LSH implementation. Since the data update the hash table, which incurs extra computation cost that could harm the overall performance. We present a dynamic data structure handle this issue, and our data structure achieves signifi- cant speedup over naive implementation under mild assumptions, without compromising the worst case guarantee. The dynamic data structure (shown in Algorithm [T] borrows ideas from the work of (Cohen et al. 2019 Jiang et al., 2021) (which are originally designed for linear programming), and we adapt it the L.H setting. Our data structure generalizes several practical insights and turns practical heuristics into rigorous theory, which guarantees significant speed up under natural assumptions for training neural networks and ensure the worst guarantee at the same time. In partic- ular, we generalize and provide the following three practical heuristics (i) only update significantly changed coordinates, (ii) batch updating LSH instead of sequential updating, (iii) using "prediction" on those "marginal" coordinates.
longer static, need to constantly
In this work, we primarily study two LSH based efficient NN training methods: SLIDE and RE- FORMER. SLIDE introduces LSH to select neurons in the forward pass and then only do gradient descent on the chosen neurons. The major trade-off of SLIDE is the rebuild overhead versus the ac- curate neuron selection. To accurately retrieve the neurons with high inner products, the hash tables are required to be updated. This rehashing and rebuilding will be the major overhead caused by SLIDE. Therefore, the SLIDE'S speed over full NN training determined by the relative magnitude of rebuild overhead compared to the saved prok-propngttton time. The major goal of REFORMER to reduce the memory consumption of the transformer so that GPU based hardware can support sequence to sequence tasks with longer sequence length. REFORMER also shares the trade-off when rebuilding the hash tables for better retrieval of attention weights. Besides, learning hash functions of REFORMER is more challenging than SLIDE. From an information retrieval's I's perspective, in at- tention settings, each data vector is also a query. This unique setting increases the hardness for learning better hash functions.
32

---

Published as a conference paper at ICLR 2021
F EFFICIENT GPU IMPLEMENTATION
As noted by (Chen et al. 2020), each input in one batch sam- ples a different set of weights(neurons). These irregular oper- ations significantly reduce the benefits of using GPUs: they correspond to different sizes of rows for each input, which makes the forward pass much more difficult to parallelize and thus, more challenging to map to GPUs than regular programs (Burtscher et al., 2012). For this reason, (Chen) al. 2020j build their SL system in C++ from scratch on CPU. Even though their implementation achieved remark- able speed up, their impact is limited as they implemented their system from scratch in C++, making it difficult for the community to adopt SLIDE in practice. To verify that such randomized algorithm not GPU friendly, we implement the exact same algorithm for GPUs. We show in Section 4.1.1 that it indeed fails gain any speed up.
1 1
2 a 1000
2,00
Figure 14: Visualization of how union selection within the batch instead of in- dependent sets of neurons with variant length can avoid irregular memory ac- cess.
We design a variant of the algorithm to exploit fast matrix multiplications on GPUs. In the original SLIDE, each train- ing example in a batch retrieves its own subset of weights. Here, we take union of the retrieved subsets in each batch to avoid irregular and unbalanced mem- ory access, shown in figure [4] Our implementation of this proposal is written in python under Pytorch framework with Cython compiled LSH and CUDA kernels for hashcode efficient computa- tion. We believe this implementation would be more beneficial for the community as it can be easily plugged into any deep learning models.
Practical Implementation: In our design, we borrow insights from the theoretical guarantee from above and aim to find the sweet spot between theory limitations and practical challenges. First, instead of detecting weight changes the cost of an extra copy, we reduce the problem to detect low quality of retrieved weights. In NN training setting, we measure quality as the inner product and low quality indicates a low inner product. We argue these two approaches are similar that they both signal the necessity of updating the data structure. Besides, rather than a soft margin for detecting data changes, which can be ambiguous under dynamic setting, inner product is better measurement as it clearly indicates the performance of current data structure. More importantly from an efficiency perspective, quality detection comes almost for free because the inner product between the query embedding and retrieved neurons are necessary for the forward pass.
G BROADER IMPACT
With the exponentially growing data volume, the scale of neural network (NN) models keeps in- creasing 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.) [203)] Vaswani et al. 2017) and computer vision (Deng et al., 2009) Vaswani al., 2017). However, train- ing those large-scale models imposes challenges on the computational and memory resources even with advanced modern hardware. For example, the recent GPT-3 model (Brown et al.] (0200) has capacity of 175 billion parameters. Therefore, there is an emergent need to reduce the cost of training those giant models.
We now quantify the potential impacts of our work. MONGOOSE sheds lights on efficient NNS-ds for efficient training of deep neural networks. 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. We show applications of MONGOOSE on language modeling and recommendation tasks. But since our key or fundamental observations are general, MONGOOSE could potentially generalize to more deep learning models and tasks. On the other hand, more NNS-ds equipped with our observations and scheduler could be involved in the deep learning community to tackle the efficiency issue.
33


Reviews:
[{"review_id": "pdxkWtrShCE", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary:\nThe authors make a good insight into the slowly changing of Locality-Sensitive Hashing (LSH) hash codes for the weights (or model parameters) during the Neural Network (NN) training. With this new insight, they introduce a framework Mongoose with a newly designed schedule mechanism to reduce the LSH update overhead. The authors also analyse their model and show some bounds for batch speedup and LSH maintenance. Experimental results validate the efficiency and effectiveness of Mongoose over original LSH methods in NN training.\n\nPros:\n1. First of all, I like the problem the authors focus on. Efficient NN training is very necessary. The idea of using LSH for acceleration is a good direction. \n2. The observation of the slow change of LSH hash codes is good and impressive. The authors also conduct many experiments to validate this observation. \n3. The idea of using a scheduler for lazy updating is interesting, and they add theoretical analysis about it. And based on the experiments, it seems that this scheduler can capture the changing of model parameters. \n4. The proposed framework Mongoose seems to be general for different deep learning model.\n\nCons:\n1. Compared to the good insight and the promising framework, the practical improvement is fair. Although the authors provide many experiments to demonstrate the effectiveness of Mongoose, I still suggest the authors conduct more experiments about parameter studies and memory usage which may be good to improve the paper quality. More details can be found in minor comments (1 & 2) later. \n2. Since the idea of the scheduler is inspired by a former work (Cohen et al., 2019), I suggest the authors add a discussion about their connection and difference. \n3. The presentation is not very clear. Many notations are used without pre-defined. More details can also be found in minor comments (3 ~ 6) later.\n4. For the learnable LSH (Section 3.3.1), when selecting positive/negative samples, the authors use the inner product, while for the loss function, they consider cosine similarity. I suggest the authors make an illustration about the use of these two different measures. \n\nMinor Comments or Typos:\n1. The framework Mongoose consists of many parameters. A fresh user may do not know how to set up them. A discussion of parameter settings or some experiments about parameters studies seems to be necessary.\n2. The authors claim Mongoose is memory-efficient for NN training. However, I cannot find analysis about space overhead or experiments about memory usage. It should be done to correspond to this claim. \n3. For Assumption 3.1, what are the definitions of C_1 and C_2? Is user-defined or auto determined? Since Theorem 3.3 and some Lemma in the appendix also use these two notations, are they the same? For a rigorous expression, I think some declarations are necessary. Also, in Theorem 3.3, the condition of r for g_r is not clear.  \n4. More problems can be found in Algorithm 1. It seems the core idea of this paper, but there are many typos and errors. For example, \nThe parameter ‘A’ is not defined and may be unnecessary in Initialize function (line 2). \nIn line 4, what is the definition of \\epsilon_{mp}? Does it have any connection with \\epsilon_{mds}? In order to have guarantee for LSH, c_1 - \\epsilon_{mp} should be larger than c_2 + \\epsilon_{mp}. How to ensure this?\nWhy setting 1.5 for a smooth cut (lines 13-15)? Is it an empirical value or has any benefit?\nIn line 17, LSH may update with w^{new}_{\\pi([r])} instead of \\pi([r]).  \n5. The y-axis of Figure 3 should be \\Delta H (left-subfigure) and \\Delta W (right-subfigure) instead of Hamming and L2 distance. \n6. The caption of Figure 5 is not clear. And for the sentence “Figure 5 shows P@1 and P@5 for Mongoose and other baselines…” in the paragraph of results in Section 4.1.1, “and P@5” should be removed. \n\nIn summary, I like the motivation and the observation in this paper, and the method Mongoose looks good, but I still have many concerns about the idea. So at this stage, I first give a borderline for this paper. I hope the authors can address my concerns. \n\n====================================================================================================\n\nUpdate: Thank you for your new experiments and detailed feedback. Most of the concerns have been addressed and the experimental results look better. I believe this paper will provide new insight into efficient neural network training. Thus, I raise my rating and recommend this paper to be accepted.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "F-O03j6eqq", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary of the paper:\n\nThis paper introduces a framework that uses an LSH sketches to improve time an memory bottlenecks neural network training. Specifically, an LSH sketch is used to approximate the matrix multiplications involved in training. It is observed that networks' weights get stable after small number of epochs therefore frequent updates to LSH sketches (which is expensive) are not required. This paper uses data dependent/learnable LSH methods that better adapt to data in order to improve the performance (query and update) of the sketches. Ample experiments are provided to validate the results.\n\nQuality:\n\nNeeds improvements in notations. For example, in assumption 3.1, sum is over what? If the indexing is over $i$, then does that means over rows of $w$s (since $w \\in \\mathbb{R}^{n \\times d}$)? Then should the quantities be L2 norms since they are vectors?\n\nClarity and the presentation can be improved significantly. For example, \"Rehash\", \"Rebuild\" functions in algorithm 2 needs to be defined or explained. In general it is better to explain the intuition behind both algorithms 1 and 2.  I believe the figure 1 should clarify the LSH update scheduling, but it is not very clear.\n\nOriginality and significance:\n\nI believe that the ideas presented in this paper adds nice contributions to the ICLR community. The idea of using learnable LSH that adapts to the data together with the observation that the weights stabilize after a few epochs is a clever approach to improve the bottlenecks associated with using vanilla LSH sketches.\n\nOther comments:\n\nIn regards to the observation with figure 3, I am curious what properties of a dataset leads to this kind of behavior? are there any quantifiable properties? Is this true for any dataset? Are there previous works that explains why this is the case?\n\nFinal feedback:\nI understand the idea of the proposed framework and the theoretical claims look natural and believable, but the presentation needs to be improved. I am willing to increase my score if the concerns mentioned above are properly addressed.\n\n\n=====================================================================================================\n\nAdded after author response\n\n----------------------------------------\n\nI believe authors have clarified many things I asked and addressed the issues I and other reviewers raised. Therefore, I increase my score. I believe the idea of using LSH for efficient training has a lot of promise and this paper brings a possible way to do this into light.", "extra_scores": {}, "extra_text": {}}, {"review_id": "Ar-iE15aBR1", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "+++Pros.  \n-----The observation that model parameters evolve slowly is quite inspiring for more efficient neural network training.  \n-----The paper proposes MONGOOSE, which is equipped with a scheduler to adaptively perform LSH updates and learnable hash functions to improve query efficiency.  \n-----Experiments demonstrates the effectiveness of the proposed method, and ablation studies give the readers further insights.  \n\n+++Cons.  \n-----The paper is overall good, but with some minors, such as “Figure 5 shows P@1 and P@5 for MONGOOSE and other baselines during the training process.” in Section 4.1.1.  \n-----Besides, there are several mathematical symbols should be explained clearly when they first appeared, such as “w” in definition 2.1, “C1, C2” in assumption 3.1, “t_r” in assumption 3.2.  \n\n+++Conclusion.  \n-----Based on the above analysis, I would prefer to make an “ACCEPT” recommendation.  \n-----By the way, I’m curious about why you named your method “MONGOOSE”? Could you give some reasons?  \n\n+++Suggestions.  \n-----Better make the mathematical symbols more clearly for readers.  \n\n\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "1R222uVdrJG", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Some neural network runs involve layers with a large number of neurons. These require large matrix-vector or matrix-multiplication which can slow their training/inference. However, if the output of mat-vec/mul is dominated by a few neurons with which the activation has large inner product (a matmul can be thought of as a weighted sum of inner products), then the computation can be sped up by approximating mat-vec/mul by a limited weighted sum with the dominant terms. This requires maintaining an ANNS data structure that is up to data with the back prop. These updates to ANNS have to be done carefully -- too frequent and the training will slow down, or too infrequent and the results of the mat-mul are way off.  This paper studies how to do this in a principled way using data-dependent LSH updates and backs it up with experimental data.\n\nThe ideas and algorithms explored in the papers are as follows:\n1. The weights changed in a limited manner over time, so it should be possible to take advantage of this.\n2. Concentrated changes to a subset of weights can be tracked and patched upon.\n3. An LSH update rule to make these changes effectively\n4. An earlier algorithm that is reused to decide when the LSH scheme is updated.\n\nThe paper also talks about how to identify the layers that benefit the most from this scheme. then it goes on to show the training time benefits of the smart scheduler and the update scheme on extreme classifications tasks as well as transfomers.\n\n\nFew questions:\n1. Why no serious consideration of graph based ANNS? They are data dependent and SOTA for search efficiency and it is possible to update them. Why is LSH the better choice for ANNS here? This needs a rigorous argument. \n2. Is this really a general and serious enough problem amongst practitioners that a solution merits publication at a popular conference? It might be, in which case a better quantification of potential impact can help. \n3.   Is wiki-325K really the largest dataset for XC? What about larger language models -- sec 4.1.2 seems to study more medium sized networks. Larger scale experiments could make this paper more compelling.\n\nI would more strongly recommend this paper if these questions  can be addressed.", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: Thanks for your submission to ICLR.

When the initial reviews were written, three of the four reviewers were positive about the paper.  Everyone felt it was overall a solid contribution, but there were some concerns about the clarity and presentation, as well as some suggestions for additional experiments.  During the rebuttal/response period, the authors did a very nice job in responding to the concerns of the reviewers.  Ultimately, all of the reviews were in agreement after discussion that the paper is strong and ready for publication.   I also like this paper a lot, and find it to be a nice way to combine LSH with NN training.  I am happy to recommend this paper for publication.
Decision: Accept (Oral)