Clustering the Sketch: Dynamic Compression for Embedding Tables
Authors: Henry Tsang
Thomas Dybdahl Ahle
Venue: NeurIPS 2023 poster (2023)
Keywords: Embedding table compression
Clustering and sketching
Memory-efficient training
TL;DR: The CQR algorithm effectively combines clustering and sketching techniques to achieve memory-efficient training and compression rates close to post-training quantization methods in large-scale embedding tables.

Abstract: Embedding tables are used by machine learning systems  to work with categorical features.
In modern Recommendation Systems, these tables can be very large, necessitating the development of new methods for fitting them in memory, even during training.
We suggest Clustered Compositional Embeddings (CCE) which combines clustering-based compression like quantization to codebooks with dynamic methods like The Hashing Trick and Compositional Embeddings [Shi et al., 2020].
Experimentally CCE achieves the best of both worlds: The high compression rate of codebook-based quantization, but \emph{dynamically} like hashing-based methods, so it can be used during training.
Theoretically, we prove that CCE is guaranteed to converge to the optimal codebook and give a tight bound for the number of iterations required.

Clustering the Sketch: Dynamic Compression for Embedding Tables
Henry Ling-Hei Tsang Meta henrylatsang@meta.. com
Thomas Dybdahl Ahle' Metal
Normal Computing thomas@ahle dk
Abstract
Embedding tables are used by machine learning systems to work with categorical features. In modern Recommendation Systems, these tables can be very large, necessitating the development of new methods for fitting them in memory, even during training. We suggest Clustered Compositional Embeddings (CCE) which combines clustering-base based compression like quantization to codebooks with dy- namic methods like The Hashing Trick and Compositional Embeddings (Shi al. 2020]. Experimentally CCE achieves the best both worlds: The high com- pression rate of codebook-based quantization, but dynamically like hashing-bass methods, so can used during training. Theoretically, we prove that CCE is guaranteed to converge to the optimal codebook and give tight bound for the number of iterations required.
1 Introduction
Neural networks can efficiently handle various data types, including continuous, sparse, and sequential features. However, categorical features present unique challenge as they require embedding typically vast vocabulary into smaller vector space for further calculations. Examples of these features include user IDs, post IDs on social networks, video IDs, and IP addresses commonly encountered in Recommendation Systems.
In some domains where embeddings employed, such Natural Language Processing (Mikolov et al. 2013), the vocabulary can be significantly reduced by considering "subwords" or "byte pair encodings". In Recommendation Systems like Matrix Factorization DLRM (see Figure 2 it is typically not possible to factorize the vocabulary this way, leading to large embedding tables that demand hundreds of gigabytes of GPU memory (Naumov al., 2019). This necessitates splitting models across multiple GPUs, increasing cost and creating a communication bottleneck during both training and inference.
The traditional solution is to hash the IDs down to a manageable size using the Hashing Trick [Wein berger et al. 2009], accepting the possibility that unrelated IDs may share the representation. Excessively aggressive hashing can impair the model's ability to distinguish its inputs, as may mix up unrelated concepts, ultimately reducing model performance.
Another option for managing embedding tables is quantization. This typically involves reducing the precision to 4 or 8 bits using multi-dimension methods like Product Quantization and Residual Vector Quantization, which rely on clustering (e.g., K-means) to identify representative "code words" for each original ID. (See Gray and Neuhoff [1998] for survey of quantization methods.) For instance, vectors representing "red". "orange", and "blue" may be stored as simply "dark orange" and "blue", with the first two concepts pointing to the same average embedding. See Section for
'Equal contribution. Work done mostly at Probability at Meta.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

900
(b)
Optimal, 1 sparse
Start
Optimal, 2 sparse CCE
(a)
800
700
Train
MSE
600
Cluster
Retrain
500
400
(d)
Hash
(c)
300
10
10
10a
10
10
Chuster Learn iterations
(a) Single iteration of CCE: (a) Starting from a ram- dom embedding table, each ID is hashed to vector
(b) CCE for Least Squares: The least squares prob- lem is to find the matrix T that minimizes IXT Y| To save space we may use codebook quantization, that factorize T = M, where H is sparse Boolean matrix, and is a small dense matrix. If we already know T this is easy to do using the K-means algorithm, but if T is too large store in memory, we can find the compressed form directly using CCE, which prove finds the optimal and M.
in each of small tables. (b) During training, the em-
bedding of referenced code words. (c) After training for an epoch, the vectors for all (or sample of) the IDs are computed and clustered. This leaves new small table in which similar IDs are represented by the same vector. (d) We can choose to combine the cluster centers with new random table (and new hash function), after which the process can be repeated for an increasingly better understanding of which ID should be combined.
is taken to be the mean of the two
For the plot, we sampled 100*x10 x10,
1000"x10 and ran CCE and compared with the opti- mal solution of first T and then factorised it with either one or two Is per row in H.
Clustered Compositional Embeddings is a simple algorithm which you run interspersed with your normal model training, such every epoch of SGD. While the theoretical bound (and the least squares setting shown here) requires a lot of iterations for perfect convergence, in practice we get substantial gains from running 1-6 iterations.
Figure 1:
an example. Clustering also plays crucial role in the theoretical literature vector compression [Indyk and Wagner 2022). However, a significant drawback of these quantization methods is that the
model is only quantized after training, leaving memory utilization during training unaffected")
Recent authors have explored more advanced uses of hashing to address this challenge: Tito Svenstrup et al. [2077] (Shi et al. [2020), (Desai et al. [2022), Yin et al. [202] I. |Kng et al. [2021]. A common theme is to employ multiple hash functions, enabling features to have unique representations, while still mapping into small shared parameter table. Although these methods outperform the traditional Hashing Trick certain scenarios, they still enforce random sharing of parameters between unrelated concepts, introducing substantial noise into the subsequent machine learning model has to overcome.
Clearly, there is an essential difference between "post-training" compression methods like Product Quantization which can utilize similarities between concepts and "during training" techniques based on hashing, which are forced to randomly mix up concepts. This paper's key contribution is to bridge that gap: We present novel compression approach we call "Clustered Compositional Embeddings" (or CCE for short) that combines hashing and clustering while retaining the benefits of both methods. By continuously interleaving clustering with training, we train recommendation models with performance matching post-training quantization, while using a fixed parameter count and computational cost throughout training, matching hashing- based methods.
In spirit, our effort can be likened methods like RigL (Evei et al. 2020). which discovers the wiring of a sparse neural network during training rather than pruning a dense network post-training. Our work can also be seen as form of "Online Product Quantization" Xu et al. 22018) though prior work
RReducing the precision to 16-bit floats is often feasible during training, but this work aims for memory reductions much larger than that.

---

focused only on updating code words already assigned to the concept. Our goal is more ambitious: We want learn which concepts to group together without ever knowing the "true" embedding for the concepts.
Why is this hand? Imagine you are training your model and at some point decide vector for IDs i and j. For the remaining duration of the training, you can never distinguish the two IDs again, and thus any decision you make is permanent. The more you cluster, the smaller your table gets. But we are interested in keeping a constant number of parameters throughout training, while continuously improving the clustering.
the same
In summary, our main contributions are:
A new dynamic quantization algorithm (CCE) that combines clustering and sketching. Particularly well suited to optimizing the compression of embedding tables in Recommendation Systems.
We use CCE to train the Deep Learning Recommendation System (DLRM) to baseline performance with less than 50% of the table parameters required by the previous state of the art, and a wobbling 11,000 times fewer parameters than the haseline models without compression.
Using slightly more parameters, but still significantly less than the baseline, we can also improve the Binary Cross Entropy by 0.66%. Showing that CCE helps combat overfitting problems.
We prove theoretically that version of our method for the linear least-squares problem always suc- ceeds in finding the optimal embedding table in number of steps logarithmic in the approximation accuracy desired.
An implementation of our methods and related work is available at github.ccom/hhmmsshllllcce
2 Background and Related Work
We show how most previous work on table compression can be seen in the theoretical framework of linear dimensionality reduction. This allows us to generalize many techniques and guide our intuition how to choose the quality and number of hash functions in the system.
Gltk
M,2
We omit standard common preprocessing tricks, such as weight- ing entities by frequency, using separate tables and precision for common vs uncommon elements, or completely pruning rare entities. We also don't cover the background of "post- training" quantization, but refer to the survey by Gray and Neuhoff 19988.
Peiriis Intrrciiion
Entawiiiig g
Embestting
Theoretical work by et al.
20033 suggests "Learning to Count sketch",, but these methods require a very large number of full training runs of the model. We only consider methods that are practical scale to very large Recommendation Systems. See also Indyk and Wagner (2022 on metric compression.
Numerfol foature
Numerical feature N
Categorical feature
Cotegoriial foature
Figure 2: Typical Recommenda- tion System Architecture: The DLRM Moodel|Naumen et al. [2019] embeds each categorical feature separately and combines the result- ing vectors with pair-wise dot prod- ucts. Other architectures use dif- ferent interaction layers or a single embedding table for all categorical features, but the central role of the embedding table is universal. (Picture credit: Nvidia).
2.1 Embedding Tables as Linear Maps
An embedding table is typically expressed as tall skinny matrix T E R xd2, where each ID i E [d1] is mapped to the i-th row, T(d. Alternatively, can be expressed as one-hot row-vector e E {0, 1) in which case Tij = e, E Rd2
Most previous work in the area of table compression is based on the idea of sketching: We introduce (typically sparse) matrix H E {0, 1] xk and a dense matrix M E I &xxds where k << d2, and take T HM. In other words, to compute T[E] we compute M. Since H and e; are both sparse, this requires very little memory and takes only constant time. The vector e; E R called "the sketch of i" and M is the "compressed embedding table" that is trained with gradient descent.
3

---

In this framework, we can also express most other approaches to training- time table compression. Some previous work has focused on the "problem" of avoiding hash collisions, which intuitively makes sense as they make the model completely blind to differences in the colliding concepts. However, from our experiments, hashing does nearly as well as said proposed methods, suggesting that a different approach is needed. Sketching is a more general way to understand this.
The Hashing Trick [Weinberger et al. 2009] is normally described by a hash function h [di] [k]] such that i is given the vector [T]], where M is a table with just A << dj rows, Alternatively, we can think of this trick as multiplying e; with a random matrix H E {0, 1] x which has exactly one I in each row. Then the embedding of is [[(ii]] = e, M,H, where HM E gdixda
Hash Embeddings (Tito Svenstrup et al. 2017] map each ID i E V to the sum of a few table rows. For example, if mapped to two rows, then its embedding vector is M [h] + M [hg (i)]. Using the notation of H E {0, 1 1 one can check that this corresponds to each row having exactly two Is. In the paper, the authors also consider weighted combinations, which simply means that the non-zero entries of H can be some real numbers.
Compositional Embeddings (CE "Quotient Remainder", Shh al. 2020 define hn hi) Li/pl and (t) = mod p for integer p. and then combines 7 (i) and (1) in various ways. As mentioned by the authors, this choice is, however, not of great importance, and more general hash functions can also be used, which allows for more flexibility in the size and number of tables. Besides using sums, like Hash Embeddings, the authors propose element-wise multiplication and concatenation. Concatenation [T (1)), T" [i]]] can again be described with a matrix H E (0. 1 13 x where each row has exactly one 1 in the top half of H and one in the bottom half of H, as well as a block diagonal matrix M. While this restricts the variations in embedding matrices T that are allowed, we usually compensate by picking a larger m, so the difference in entropy is not much different from Hash Embeddings, and the practical results are very similar as well.
ROBE embeddings [esai et al. 2022] are essentially Compositional Embeddings with concatena- tion as described above, but add some more flexibility in the indexing from the ability of pieces to "wrap around" in the embedding table. In our experiments, ROBE was nearly indistinguishable from CE with concatenation for large models, though it did give some measurable improvements for very small tables.
Deep Hashing Embeddings (DHE, Kang et al. 2021 picks 1024 hash functions hi..... 111024 [di] + [-1,1] and feed the vector (i).... (1)2 (i) into a multi-layer perceptron. While the idea of using an MLP to save memory at the cost of larger compute is novel and departs from the sketching framework, the first hashing step of DHE is just sketching with a dense random matrix H E 11 x 1022 While this is less efficient than sparse matrix, it can still be applied efficiently to sparse inputs, e,, and stored in small amounts of memory. Indeed in our experiments, for a fixed parameter budget, the fewer layers of the MLP, the better DHE performed. This indicates to that the sketching part of DHE is still the most important part.
Tensor Train Yii et al. [021] doesn't use hashing, but like CE it splits the input in a deterministic way that can generalized to a random hash function if so inclined. Instead of adding or concatenat- ing chunks, Tensor Train multiplies them together as matrices, which makes it not strictly linear operation. However, like DHE, the first step in reducing the input size is sketching.
Learning to Collide In recent parallel work. Ghaemmaghami et al. [022 propose an alternate method for learning a clustering based on low dimensional embedding table. This is like horizontal sketch, rather than a vertical, which unfortunately means the potential parameter savings is substantially smaller.
Our method introduces a novel approach to dynamic compression by shifting from random sketching to learned sketching. This process can be represented as M,, where H is a sparse matrix and M a small dense matrix. The distinguishing factor is that we derive H from the data, instead of relying on a random or fixed matrix. This adaptation, both theoretically and empirically, allows learning the same model using less memory.
'While combining vectors with element- wise multiplication is not a linear operation. from personal commu- nication with the authors, it is unfortunately hard to train such embeddings in practice. Hence we focus on the two linear variants.

---

500
#hidden
1000
1000
1000
500
MLP
blue (a) The Hashing Trick: (b) Hash Embeddings: (c) ROBE:
Similar to CE (d) Deep Hash Em Also known as Count Each ID hashed to two with concatenation,
beddings: Computes
Sketch, each ID is hashed rows, one per table, and stead of using separate
weighted sum of all the
to one location in table
its embedding vector is as- columns, it uses contin-
table entries, where the
(here with 1000 rows) and
signed to be the sum of uous array from which the
weights come from hash- ing the input 1000 times
it is assigned the embed-
those
Here, (in this dim=4) pieces
ding vector stored the lo- cation. Many IDs are likely
we use two separate tables are picked. The pieces may
to [-1, The produced
unlike in Tito Svenstrup) even overlap.
(hidden) vector is then sent through an MLP for further refinement.
to share the same vector.
et al. 2017].
4
500
1000
500
(e) CE with concatenation: In the hashing version of compositional embeddings (CE), each ID hashed location each say, different tables. The four vectors stored there are concatenated into final embedding vector. Given the large number of possible combinations (here 1000*), ), it is unlikely that two IDs get assigned the exact embedding vector, even if they may share each part with some other IDs.
(f) Clustered CE: We can combine the sum hashing method of Tito Svenstrup al. 2017 with the concate- nation method of Shi al.] [202)] Each then gets assigned a vector that is the concatenation smaller sums. This is cnhanced with the clustering idca shown in Section [T] In each of the four blocks, we apply clustering every epoch, setting the results in the green tables, replacing the hash functions tables with new random values.
Figure 3: The evolution of hashing- based methods for embedding tables. The Hashing Trick and Hash Embeddings shown at the top, side by side with an equal amount of parameters. Next we introduce the idea of splitting the space into multiple concatenated subspaces. This a classic idea from product quantization and reminiscent of multi-head attention in transformers. Finally in CCE we combine both methods in a way that allows iterative improvement using clustering.

---

3 Sparse and Dense CCE
The goal this section is to give the background of Theorem 3.1 where we prove that CCE converges to the optimal assignments we would get from training the full embedding tables without hashing and clustering those.
We can't hope to prove this in a black box setting, where the recommendation model on top of the tables can be arbitrary, since there pathological functions where only full table will work. Instead, we pick a simple linear model, where the data is given by matrix E Rnxd and we want to find matrix T that minimizes the sum of squares, XT I (x)((X)) -y,)²²
We give two versions of the algorithm, a sparse method, which is what we build our experimental results on; and dense method which doesn't 't use clustering, and replaces the Count Sketch with a dense normal distributed random matrix for which we prove optimal convergence. The dense algorithm itself is interesting since it constitutes novel approximation algorithm for least squares regression with lower memory use.
Algorithm 1 Dense CCE for Least Squares
Algorithm Sparse CCE for Least Squares
Require: X E Rexd1 Y E Rnxd, k E E
Require: X E Raxd1,
Y E Rnxd, k E N Initialize assignments & Initialize codebook
I: Hn 0 E Rd
Hn E Rd x2k
2: My = 0) E IR. 222xx
2: Mo = =0 E 2Rxxxd2
3: for = 0, 1..... do
for 0, do
4:
T; H.M,
4:
T, = H,M, E Rd1xd2
5:
N ~ (0, 1)d1ix
5:
A kmeans (T:) E {0, 11
6:
Hi+1 = [ N]
C countsketch() E {-1, 1]dxxk
7: M5-1 arg min | H4+1 M YI?
H [A] CJ Rd x2k
8: end for
8: M,-1 = arg minM | H YI?
end for
In the Appendix we show the following theorem guaranteeing the convergence of Algorithm D Theorem 3.1. Assume d, > k > d2 and T E Rdh xdz be the matrix that minimizes
then T, = H;M, from Algorithm exponentially approaches the optimal loss in the sense
[ - Y s (1-p)* + +IXXT~>X
where p= | = 1/d1 is the smallest singular value of X squared divided by the sum of singular values squared.
We also show how to modify the algorithm to get an improved bound of (1-11114 (1-1/d1) by conditioning the random part H by the eigenspace of X. This means that after i = 0 log(1/=)) iterations we have a 1 + approximation to the optimal solution. Note that the standard least squares problem can be solved in (nd d2) time, but one iteration of our algorithm only takes ((nkdz) time. Repeating for di /k iterations is thus no slower than the default algorithm for the general least squares problem, but uses less memory.
Some notes about Algorithm 2 In line we compute HM e Edixdz all the main goal of CCE is to avoid storing this full matrix in memory.) Luckily K-means in practice works fine on just sample of the full dataset, which is what we do in the detailed implementation in Section 4.2
xd2 which may be very large. (After
In line note that K-means normally returns both set of cluster assignments and cluster centroids. For our algorithm we only need the assignments. We write those as the matrix which has Aıd.5 = 1 if id is assigned to cluster j and 0 otherwise. In the Appendix (Figure 5) show how A is sparse approximation to T's column space. Using the full, dense column space we' 'd recover the dense algorithm, Algorithm
In line 6. countsketch() is the distribution of {-1, 0, 1]dix matrices with one non-zero per row.
A matrix C' is sampled based on hash functions h; [dj] + [k] and s; [di] + {-1, 1] such that Cse = 3)) if h₂()) = and 0 otherwise. We see that with and C sparse, X M can be computed efficiently.
We can ignore the sign, #1 in most models if M has mean 0. See Charikar al. 2002

---

epocks
Full Emhedding Tahle
Product Quantization
Deep Hash Embeddings
The Hashing Trick
CE with Concatenation
0.456
CCE (This paper)
BCB
0.454
0,452
0.450
10a
Number of parameters for Largest Tables
(a) Best of 10 Epochs Test Loss on Criteo Kaggle. We trained DLRM on the Critco Kaggle datasct with different compression algorithms. Each 27 categorical features was given its own embedding table, where we limited the number of parameters in the largest table as shown the x-axis. Out of 10 epochs, we did early stopping at the minimum validation loss, plotting the test loss. We did independent repetitions with seeds each, plotting the min, max and mean test losses. Binary Cross-Entropy was used as both the training and test loss. We also evaluated the AUC in Appendix G but the results were indistinguishable. The Full Embedding Table used up to 16 10" parameters per table, as each categorical value is mapped to a unique embedding vector. This model over-fits fits immediately if trained for more than 1 epoch, which is why it does poorly compared to more parameter-ree restrrete methods. For Clustered Compositional Embeddings we ran clustering once every epoch for the first 6 epochs, unless stopping earlier. Product Quantization, heing post-training quantization method, never able do better than the baseline model is trained on. We tried fine-tuning the tables after PQ, but this approach immediately overfitted and gave terrible results. Interestingly this suggests CCE works as a regularization method for PQ.
Eod of 1 opoch
Bad 1 opoch
Pull Emhedking Table
Full Emhedding Table
Deep Hash Embeddings
0.1255
Deep Hash Emheddings
The Hashing Trick
Product Quamization
CE with Concatenation
The Hashing Trick
Product Quamizution
0.1255
CE with Concatenation
CCE paper)
CCE paper)
0.1280
0.12%
012258
10
Number of parameters for Largest Tables
Number of parameters for Largst Tables
(b) Kaggle dataset, 1 epoch: Previous work, like Shi
(c) Terabyte dataset, 1 epoch: We ran similar experi- ments Criteo Terabyte dataset, which expensive enough that we could only perform one repetition of cach algorithm. Like on the Kaggle dataset, running for just cpoch is not enough to show an improvement in BCE over the baseline. but are still able to save space. Curiously Product Quantization was basically no better than sketch-based methods for this dataset, which may be another reason CCE didn't perform great either.
al. (2000) only trained each method for one epoch. following DLRM, which did this to avoid over-fitting the baseline model. To compare with these results, we
ran CCE clustering after 1/4 and 1/2 of an epoch. In this setting, the hashing-base based methods were not able compete with the baseline and product quantization, but CCE was able to do slightly better. Allowing 300x parameter reduction.
Figure 4: CCE also improves models even when trained for just one epoch. This shows useful information is available from clustering even when only parts of the data have been seen.
7

---

Table 1: Memory Reduction Rates Across all Datasets. For each algorithm. dataset and epoch limit we measured the necessary number of parameters to reach baseline BCE. The compression ratios with ranges are estimated using degree and polynomial extrapolation. We can compare these results with reported compression rates from Desai et al. [2022) (ROBE) which gets x with multi epoch training; [Yin et. 202T] (Tensor Train) which reports 112x reduction epoch
on Kaggle; and Yin et al. [2021] which reports a 16x reduction at 1 epoch with "Mixed Dimension methods" on Kaggle.
Method
Dataset
Epochs
Embedding Compression
CCE (This Paper)
Criteo Kaggle
S 10
8,500
CE with Concatenation
Criteo Kaggle
S 10
3,800
The Hashing Trick
Criteo Kaggle
S 10
4,600
Deep Hash Embeddings
Criteo Kaggle
S 10
1,300
CCE (This Paper)
Criteo Kaggle
212x
CE with Concatenation
Criteo Kaggle
127 155x
The Hashing Trick
Criteo Kaggle
78 - 122x
Deep Hash Embeddings
Criteo Kaggle
7 -5x
CCE (This Paper)
Criteo TB
101x
CE with Concatenation
Criteo TB
25 48x
The Hashing Trick
Criteo TB
23 32x
Deep Hash Embeddings
Criteo TB
2 - 6x
Experiments and Implementation Details
Our primary experimental finding, illustrated in Table and Figure 4a indicates that CCE enables training model with Binary Cross Entropy matching full table baseline, using only half the parameters required by the next best compression method. Moreover, when allocated optimal parameters and trained to convergence, CCE can yield not insignificant 0.66% lower BCE.
4.1 Experimental Setup
In our experiments, we adhered to the setup from the open-source Deep Learning Recommendation Model (DLRM) by Naumov al. [2019], including the choice optimizer (SGD) and learning rate. The model uses both dense and sparse features with an embedding table for each sparse feature. We modified only the embedding table portion of the DLRM code. We used two public click log datasets from Criteo: the Kaggle and Terabyte datasets. These datasets comprise 13 dense and 26 categorical features, with the Kaggle dataset consisting of around 45 million samples over days, and the Terabyte dataset containing about billion samples over 24 days.
We the Kaggle dataset experiments we ran them on two A 100 GPUs using model parallelism. This was done mainly for memory reasons. With this setup, training for epoch the Kaggle dataset takes around 4 hours, while training for one epoch on the Terabyte dataset takes around 4 days. There was not big difference in training time between the algorithms we tested. total we used about 11,000 GPU hours for all the experiments across 5 algorithms, 3 seeds, 10 epochs and 9 different parameter counts.
single A100 GPU. For the Terabyte dataset experiments,
4.2 CCE Implementation Details
In the previous section we gave pseudo code for "Sparse CCE for Least Squares" (Algorithm) 2 In general model, an embedding table can be seen as data-structure with two procedures. Below give pseudo-code for an embedding table with vocabulary [d1] and output dimension d2, using 2kd parameters. The CCE algorithm is applied to each of columns, mentioned in Figure 3f
The value = was chosen match (Shi et al. 2020), but larger values are generally better, as long as the h; functions don't become too expensive to store. See Appendix E The random hash functions I are very cheap to store using universal hashing. See Appendix D The number of calls to Cluster was determined by grid search. See Appendix F for the effect of more less clustering.

---

Algorithm 3 Clustered Compositional Embeddings with e columns and 2k: rows
1: class CCE:
2: method Initialize:
3: for i = 1 to do:
4: h4, h44 +- -iid. random functions from [di] [k]
H in Algorithm
5: M4, M < iid. N 1)dx/
6:
7:
method GedEmbedding(id))
8:
return CONCAT(M (Ii + (i(h((id) for i= co
9:
10:
method Cluster(ilemss:
11: for: to c do:
12:
T [id]) + (i[]((dd] for id E [d1]]
See discussion below Find & clusters and assign to them
13:
centroids, assignments +- K-MEANS(T)
14:
h; f assignments
15:
M; < centroids
16: h < random function from [d1] to [k]]
17:
M Okxda/c
In line 12 we likely don't 't want to actually compute the embedding for every id in the vocabulary, but instead mini batch K-Means with oracle access the embedding table. In practice we follow the suggestion from FAISS K-means[Johnsse et al. 2019) and just sample 256k ids from [d1] and run K-means only for this subset. The assignments from ad to nearest cluster center are easy compute for the full vocabulary after rumning K-means. The exact number of samples per cluster didn't 't have a big impact on the final performance of CCE.
5 Conclusion
We have shown the feasibility of compressing embedding tables at training time using clustering. Our method, CCE, outperforms the state of the art on the largest available recommendation data sets. While there is still work to be done in expanding our theoretical understanding and testing the method in more situations, we believe this an exciting new paradigm for dynamic sparsity in neural networks and recommender systems in particular.
Previous studies have presented diverging views regarding the feasibility of compressing embedding tables. Our belief is that Figure 4a Figure 44b and Figure 44c shed light on these discrepancies. At standard learning rates and with one epoch of training, it's challenging to make significant improvements over the DLRM baseline, corroborating the findings of Na et al. (2019). However, upon training until convergence, it's possible to achieve parity with the baseline using thousand times fewer parameters than typically employed, even with the straightforward application of the hashing trick. Nevertheless, in the realm of practical recommendation systems, training to convergence isn 't a common practice. Our experiment, as illustrated in Figure 4a proposes a potential reason: an excessive size of embedding tables may lead overfitting for most methods. This revelation is startling, given the prevailing belief that "bigger is always better" and that ideal scenarios should allow each concept have its private embedding vectors. We contend that these experimental outcomes highlight the necessity for further research into overfitting within DLRM-style models.
In this paper we analyzed only the plain versions of each algorithm. There number of practical and theoretical improvements one may add. All methods are naturally compatible with float 6 and float8 reduced mixed precision. The averaging of multiple embeddings may even help smooth out some errors. It is also natural to consider pruning the vocabularies. In particular in an offline setting we may remove very rare values, or give them a smaller weight in the clustering an averaging. However, in an online setting this kind of pruning is harder to do, and it is easier to rely on hash collisions and SGD ignore the unimportant values. our experiments we used 27 embedding tables, one for each categorical feature. A natural compression idea is to map all features to the same embedding table (after making sure values don't 't collide between features.) We didn't experiment with this, but it potentially could reduce the need for tuning embedding table sizes separately. A later paper by Coleman et al. (0233 report good results with this method. For things we tried, but didn't 't work, Appendix

---

Acknowledgment
We would like to thank Ben Ghaemmaghami from the University of Texas Austin for fruitful discussions of learned hash functions. Thanks to Chunyin Siu for helpful discussions on the conver- gence of multi-step CCE. And finally to Michael Shi, Peter Tang, Summer Deng, Christina You, Erik Meijer, and the rest of the Probability group at Meta for helpful discussions of table embeddings and earlier versions of this manuscript.
10

---

References
Andrew R Barron, Albert Cohen, Wolfgang Dahmen, and Ronald Vore. Approximation and learning by greedy algorithms. The annals of statistics, 36(1):64-94, 2008.
J Lawrence Carter and Mark N Wegman. Universal classes of hash functions. In Proceedings of the ninth annual ACM symposium on Theory of computing, pages 106-112, 1977.
Moses Charikar, Kevin Chen, and Martin Farach-Colton Finding frequent items in data streams, In International Colloquium on Automata, Languages, and Programming, pages 693-703. Springer, 2002.
Michael Cohen, TS Jayram, and Jelani Nelson. Simple analyses of the sparse linnestrl transform. In 1st Symposiun on Simplicity in Algorithms (SOSA 2018). Schloss Dagstuhl-Leibnizz Zentrum fuer Informatik, 2018.
Benjamin Coleman, Wang-Cheng Kang, Matthew Fahrbach, Ruoxi Wang, Lichan Hong, Ed Chi, and Derek Zhiyuan Cheng. Unified embedding: Battle-tested feature representations for web-scale ml systems. arXiv preprint arXiv:2305. 102 2023.
Aditya Desai, Li Chou, and Anshumali Shrivastava. Random offset block embedding (robe) for compressed embedding tables in deep learning recommendation systems. Proceedings of Machine Learning and Systems, 4:762-778, 2022.
Martin Dietzfelbinger, Torben Hagerup, Jyrki Katajainen, and Martti Penttonen. reliable random- ized algorithm for the closest-pair problem. Journal of Algorithms, 25(1):19-51, 1997.
Utku Evci, Trevor Gale, Jacob Menick, Pablo Samuel Castro, and Erich Elsen. Rigging the lottery: Making all tickets winners. In International Conference on Machine Learning, pages 2943-2952.
PMLR, 2020.
Benjamin Ghaemmaghami, Mustafa Ozdal, Rakesh Komuravelli, Dmitriy Korchev, Dheevatsa Mudi- gere, Krishnakumar Nair, and Maxim Naumov. Learning collide: Recommendation system model compression with learned hash functions. arXiv preprint arXiv:2203. 15837, 2022.
Robert M. Gray and David L. Neuhoff. Quantization. IEEE transactions on information theory, 44 (6):2325-2383, 1998.
Piotr Indyk and Tal Wagner. Optimal (euclidean) metric compression. SIAM Journal Computing, 31(3):467-49) 2022.
Jeff Johnson, Matthijs Douze, and Hervé Jégou, Billion-scale similarity search with GPUs. IEEE Transactions Big Data, 7(3):335-544) 2019.
Wang-Cheng Kang, Derek Zhiyuan Cheng, Tiansheng Yao. Xinyang Yi, Ting Chen, Lichan Hong, and Ed H Chi. Learning to embed categorical features without embedding tables for recommendation. In Proceedings of the 27th ACM SIGKDD Conference Knowledge Discovery & Data Mining, pages 840-850, 2021.
Yi Li, Honghao Lin, Simin Liu, Ali Vakilian, and David Woodruff. Learning the positions in countsketch. In Proceedings of the International Conference on Learning Representations (ICLR), Feb 2023.
Tomas Mikolov, Kai Chen, Greg Corrado, and Jeffrey Dean. Efficient estimation word representa- tions in vector space. arXiv preprint arXiv: 2013.
Maxim Naumov, Dheevatsa Mudigere, Hao-Jun Michael Shi, Jianyu Huang, Narayanan Sundaraman, Jongsoo Park, Xiaodong Wang, Udit Gupta, Carole-Jean Wu, Alisson G Azzolini, al. Deep learning recommendation model for personalization and recommendation systems. arXiv preprint arXiv:1906.0000 00091, 2019.
Pedregosa, G. Varoquaux, A. Gramfort, V. Michel, B. Thirion, O. Grisel, M. Blondel, Pretten- hofer, R. Weiss, V. Dubourg, J. Vanderplas, A. Passos, D. Cournapeau, M. Brucher, M. Perrot, and E. Duchesnay. Scikit-learn: Machine learning Python. Journal of Machine Learning Research, 12:2825-288 2011.

---

Hao-Jun Michael Shi, Dheevatsa Mudigere, Maxim Naumov, and Jiyan Yang. Compositional embeddings using complementary partitions for memory--fficiese recommendation systems. In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, pages 165-175, 2020.
Dan Tito Svenstrup, Jonas Hansen, and Ole Winther. Hash embeddings for efficient word representa- tions. Advances in neural information processing systems, 30, 2017.
Kilian Weinberger, Anirban Dasgupta, John Langford, Alex Smola, and Josh Attenberg. Feature hashing for large scale multitask learning. In Proceedings of the 26th annual international conference on machine learning, pages 11 1113-1120, 2009.
David P. Woodruff. Sketching tool for numerical linear algebra. Found. Trends Theor. Comput. Sci., 10(1-2)11-177, 2014. doi: 10.1561/04000000600 URL https://doi. www.doi..g.ggggg org/o0....... 1561/ 0400000060
Donna Xu, Ivor W Tsang, and Ying Zhang. Online product quantization. IEEE Transactions on Knowledge and Data Engineering, 30(11):118-2219, 2018.
Chunxing Yin, Bilge Acun, Xing Liu, and Carole-Jcan Wu. Tt-rec: Tensor train compression for deep learning recommendation models. CoRR, abs/2101.11714, 2021. URL https: //rrxiaxvvvvv//// org/abs 2101 11714.
12

---

Reproducibility
The backbone recommendation model, DLRM by Naumov et al. 2019], has an open-source PyTorch implementation available on Github which includes an implementation of CE. For CCE you need a fast library for K-means. We recommend the open- open-sourced implementation by Dohnson et al.] [2019) for better performance, but you can also use the implementation Scikit-lear [Pedregosa al.) 2011]. The baseline result should be straightforward to reproduce as we closely follow the instructions provided by NNumov et al. [2119]. For the CE methods, we only need to change two functions in the code: create_ emb and apply- emb. We suggest using a class for each CE method; see Figure 3 For the random hash function, one could use a universal hash function or numpy. . random. randint.
Measuring the Embedding Compression f'actor The most important number from our experi- mental section is the "Embedding Compression" factor in Table We measure this by training the model with different caps on of parameters in the embedding tables (See e.g. the x-axis in Figure 4a E.g. if the Criteo Kaggle dataset has categorical features with vocabularies of sizes 10, 100 and 10", we try e.g. to cap this at 8000, using a full embedding table for the small features, and CCE table with 8000/16 /6 500 rows (since each row has 16 parameters). This corresponds rate of (10 + 100 + 1010 + 100 + 500) R 1639.5. Or if we measure only the compression of the largest table, 10 /500 2000. Unfortunately there's discrepancy in the article, which we only found after the main deadline, that uses the second measure in the introduction (hence the number 11,000x compression) where as Figure 4aluses, the first measure (and thus the lower number 8, 8,5000x).
compression
For the experiments where we only train for I epoch, some methods never reach baselinn BCE within the number of parameters we test. Hence the Compression Rates we report are based on extrapolations. For each algorithm we report a range, e.g. 127-155x for CE with Concatenation on Criteo Kaggle epoch. Since the loss graphs tend be convex, the upper bound (155x) is based on a linear interpolation (being optimistic about when the method will hit baseline BCE) and the lower bound (127x) is based on a quadratic interpolation, which only intersects the baseline higher parameter count.
K-means For the K-means from FAISS, we max seax__ponns__pe prooints_per_eeee _tt_ppppeeeiddd///ee1 centroid=256 and ni niter=50. The first parameter sub-samples the number of points to 256 times the number of cluster centroids (k), and is the recommended rate after which "no benefit is expected" according to the library maintainers. In practice we predict the right value will depend on the dimensionality of your data, so using the split into lower dimensional columns is beneficial. For niter we initially tried a larger value (300), but found it didn't 't improve the final test loss. We found on Kaggle for PQ, niter=50, BCE=0.445554 and niter=300, CCE=0.455531 For CCE (single epoch, single clustering half an epoch), niter=50 gave ACE=0.4592S and niter=300 gave BCE=0.4550 45905, so very slight improvement, but not enough to make up for the extra training time.
Datasets For our experiments, we sub-sa an eighth of the Terabyte dataset and pre-hashed them using simple modulus match the categorical feature limit of the Kaggle dataset. For both Kaggle and Terabyte dataset, we partitioned the data from the final day into validation and test sets. Using the benchmarking setting of the DLRM code, the Kaggle dataset has around 300,000 batches while the Terabyte dataset has around 2,000,000 batches.
Early stopping Early stopping is used when running the best of 10 epochs on the Kaggle dataset. We measure the performance of the model in BCE every 50,000 batches (around one-sixth of one epoch) using the validation set. If the minimum BCE of the previous epoch is less than the minimum BCE of the current epoch, we early stop.
Deep Hash Embeddings We follow Kang et al. [2021] in using fixed-width MLP with Mish activation. However, DHE is only described one version in the paper: 5 layers of 1024 nodes per layer. For our experiments, we need to initialize DHE with different parameter budgets. We found that in general, DHE performs better with fewer layers when the number of parameters is fixed. However, we cannot use just a single layer, since that would be linear embedding table, not an MLP. As compromise, we fix the number of hidden layers to 2 and set the number of hashes to be the same as the dimension of the hidden layers.
13

---

For example, if we were allowed to use 64, 000 parameters with an embedding dimension of 64, then by solving a quadratic equation we get that the number of hashes and the dimension of the hidden layers are both 136. This gives us
hashes hidden dim + 2 hidden dim hidden dim embedding dim 64192
parameters.
14

---

A
What didn't 't work
Here are the ideas we tried but didn't work the end.
Using multiple helper tables It is natural idea use more than one helper table. However, in our experiments, the effect of having more helper tables is not apparent.
Circular clustering Based on the CE concat method, the circular clustering method would use information from other columns to do clustering. However, the resulting index pointer functions are too similar each other, meaning that this method is essentially the hashing trick. We further discuss this issue in Appendix H
Continuous clustering We originally envisioned our methods in tight loop between training and (re)clustering. It turned out that reducing the number of clusterings didn't impact perfor- mance, so we eventually reduced it all the way down to just one. In practical applications, with distribution shift over time, doing more clusterings may still be useful, Section 3
discuss in
Changing the number of columns In general, increasing the number of columns leads to better results. However the marginal benefits quickly decrease, and as the number of hash functions grow, so does the training and inference time. We found that columns hash-functions good spot.
Residual vector guantization The CCE method combines Product Quantization (PQ) with the CE concat method. We tried combining Residual vector quantization (RVQ) with the Hash Embeddings method from Tito Svenstrup et al. (2017). This method does not perform significantly better than the Hash Embeddings method.
Seeding with PQ We first train a full embedding table for one epoch, and then do Product Quantiza- tion (PQ) on the table to obtain the index pointer functions.
We then use the index pointer functions instead of random hash functions in the CE concat method. This method turned out performing badly: The training loss quickly diverges from the test loss after training on just few batches of data.
Here are some variations of the CCE method:
Earlier clustering We currently have two versions of the CCE method: CCE half, where clustering happens at the middle of the first epoch, and CCE, where clustering happens at the end of the first epoch. We observe that when we cluster earlier. the result is slightly worse. Though in our case the CCE half method still outperforms the CE concat method.
More parameters before clustering The CCE method allows using two number of parameters, one in Step where follow the CE hybrid method get sketch, and one in Step 3 where follow the CE concat method. We thought that by using more parameters at the beginning, would be able to get better set of index pointer tables. However, the experiment suggested that the training is faster but the terminal performance is not significantly better.
Smarter initialization after clustering In Algorithm 3 we initialize M, with the cluster centroids from K-means and the "helper table" M/ + O. We could instead try to optimize Mí to match the residuals of T as well as possible. This could reduce the discontinuity during training more than initializing zeros. However, didn't large effect in either training loss smoothness or the ultimate test score.
15

---

B Proof of the main theorem
T0.417
0.720
0
0
00
07
0.000
0.302
0
1
0
TO.411
0.648
0.147
0.092
0
0 1
0 0 0 0
0.093
0.324
0.186
0.346
0
0.204
0.878
0.397
0.539
0
[0.147
0.092
0.419
0.685
0
0
10.204
0.878.
0
0_
Figure 5: K-means matrix factorization. A central part of the analysis of CCE the simple observation that K-means factors matrix into tall sparse matrix and a small dense one. In other words, finds a sparse approximation the column space of the matrix.
Let's remind ourselves of the "Dense CCE algorithm" from Section Given X E Ruxd and E RRnxd2, pick such that n > dj > > d2. We want to solve find matrix T of size d1 x d2 such that ||XT* Y| is minimized the classical Least Squares problem. However, we want to use memory less than the typical ndf. We thus use this algorithm:
Dense CCE Algorithm: Let To
E Rd xd2 For i m:
Sample
G; (0, 1)d1x(k=xd)
Compute
H, = [Ti-1 Gi] E RR xk
M. arg inf || H.H, Y E Rxxd, M
T; = H,M,
We will now argue that Tm a good approximation to in the sense that X YI? is not much bigger than |XT - V
Tag
Let's consider non-optimal choice of Mj first. Suppose we set M, such that || |,H,M, IF is minimized. By direct multiplication, we have H,M, = T--1 + G,M, Hence in this case minimizing |H, - T is equivalent to finding Mf at each step such that ||G.M|- - (T - TI: is minimized.
where M/ is chosen
M²
In other words, we are trying to estimate 7 with E, G,M,, where each G; is random and each M' greedily chosen at each step. This is similar to, for example, the approaches in Barron et al. [2008], though they use a concrete list of G!'s. In their case, by the time we have d1 /k such G s, we are just multiplying X with d, x d, random Gaussian matrix, which of course will have full rank, and so the concatenated M matrix can basically ignore it. However, in our case we do local, not global optimization over the Mj.
Recall the theorem:
Theorem B.O. Given E BRxdd and Y Raxds, Let T arg migmineeerxx0 |IT YI be an optimal solution to the least squares problem. Then
E [IIT: - [ s (1 p))(p))k---- (k do) XTT
-1I
where p |X||2/|
Here we use the notation that || the smallest singular value of X. Corollary B.1. In the setting of the theorem, all singular values of X are equal, then
E [IIT: - Y12 Se- - |XT* + |XT - YI
Proof of Theorem B.7 Note that | is the sum of the d, singular values squared: || =
Since all singular values are equal, say to E IR, then | = dio?. Similarly in this setting, || = o² so p 1/d,. Using the inequality 1 1/d, S e 1/d1 gives the corollary.
Proof of Theorem B.O
First split Y into the part that's in the column space of X and the part that's XT Z, where 1" arg min XT Y [| is the solution to the least
not, Z. We have
16

---

squares problem. By Pythagoras theorem we then have
E. [ -
-YI:] = |TT; - (XT* + 2)1 = (T) T + +
so it suffices show
E [| S (1
d2) XT |
We will prove the theorem by induction over i. In the case i 0 have 7; 0, [
T*) | = -||| trivially. For i 2 we insert T: = H,M, and optimize over Mf:
[[]
| - = (HI - 1
S [II M| - 1
+ 1
EIIX - (T (T-
= EIEIII((,,M (T=)))- TT-1]]
- (1 - 0) de -(T- 1
-
- (1 (( d2) |T
where the last step followed by induction. The critical step here was bounding
Rcrlinf] - 1) < () d2 |XT|
for fixed T. We will do this in series of lemmas below.
We show the lemma first in the "vector case", corresponding to k = 2, d2 = 1. The general matrix case follow below, and mostly case of induction on the vector case.
Lemma B.2. Let X Raxd be matrix with singular values 01 2 od 2 0. Define p == 0= E, o? then for any t E Rd.
(g
inf || t)11 S (1- (1 mER
Proof. Setting
( we get
[| t)112 m x + |x
2m Xt)
(1 | X gg
(Xt, x)
||xt
We use the singular value decomposition of X = U.VT. Since ~ N(O, 1)d and VT unitary, we
have VTg ~ NO, 1)dd and hence we can assume V = I. Then
(Xt, Xg)²
UEg'
(3)
| X961
|
())
(4)
|E EglB
( t
(5)
E O, ofg
where Equation follows from UTU = in the SVD. We expand the upper sum to get
Σn, 4tzofoj999
Fyg
= Frg
(6)
(E,
Σx Got
E
17

---

Here we use the fact that the 's are symmetric, so the cross terms of the sum have mean 0. By scaling, we can assume E, oft and define p; o.tt?. Then the sum just convex combination:
of
Σ Ps Egg
(8)
Σ
of
Since 04 2 aa and ge are IID, by direct comparison we have
of
a
Bg
Eg
Σ, o7g
o9
Hence
2 Eg
0 a
Σ Pi = Eg
o
E, 0
2
E
remains to bound
0 Σ; o
Fgg
Σ,
P:
but this follows from cute, but rather technical lemma, which we will postpone to the end of this section. (Theorem B.4
is interesting to notice how the improvement we make each step (that is - p) could be increased 1/d by picking G from distribution other than IID normal.
If X UEVE, can also take -V--1 g', where o' N(0, 1)th x(k x(k- da), In that case we get
XXt,Xg)
2 E
(VVEV" g {
VVV
1 1
E
[Xllal|
X
|
di 1 di
| T S 11/d. Of course comes with the negative side of having to compute the SVD of X. But since this is just theoretical algorithm, it's still interesting and shows how we would ideally update T.. See Figure for the effect of this change experimentally.
So this way we recreate the ideal bound from Theorem B.1 Note that
It's interesting problem how it might inspire better CCE algorithm. Somehow we'd have get information about the the SVD of X into our sparse super-space approximations.
We now show how to extend the vector case to general matrices.
Lemma B.3, Let X E Raxd1 be matrix with singular values 01
2 ... 2 Od, 2 Define
p = / Eg o7 then for any E [gaxd,
Ea~N(0.14dxx
inf MERRXd
IX(GM-T)I - (11-)) 11--))
Proof. The case k = 1, d2 is proven above in Theorem B.2
Case l: = 1: We first consider the case where k = 1, but d2 can be any positive integer (at most k). Let [4 [///2//.. [ag be the columns of 7 and M [mi [ol]...|ma,) be the columns of M. Then the ith column of X (GM T) is X (Gm, tz), and since the squared Frobenius norm of
18

---

n= 11=22, d1-125, d2-5, k-25
0.0175
noise
smart noise
half smart noise
half noise
0.0125
0.0100
MSE
0.0075
0.0025
0.0000
10
Cluster Learn Iterations
Figure SVD aligned noise converges faster. In the discussion we mention that picking the random noise in H, g VV-1g' g', where o' N(0, 1)d ((R-da), can improve the convergence rate from p) to (1 1/d)"h, which is always better. In this graph experimentally compare this approach (labeled "smart noise") against the IID gaussian noise (labeled "noise"), and find that the smart noise indeed converges faster at least once get close repetitions where X is a random rank-10 10 matrix plus some low magnitude noise.
noise. The graph is over 40
We also investigate how much we lose in the theorem by only considering M the form M7] rather than a general M that could take advantage of last rounds T:- The plots labeled "half noise" and "half smart noise" are limited in this way, while the two others not. We observe that the effect of this is much larger in the "non-smart" case, which indicates that the optimal noise distribution we found might accidentally be tailored analysis.
matrix is simply the of the squared column 12 norms, have
d2 X [(Gm; - -1))13
E| X - 7)
E
dg
E [|]x
-4)
d2
E[[[X
(10)
de
(1 p) E
Xt, 113
(1 p) E XT
where in (10) we applied the single vector case.
Case ls > 1: This time, let 91, 92..... 9s the columns G and let mt. mm....., of M.
I be the
We prove the lemma by induction over We already proved the base-case k = so all we need is the induction step. We use the expansion of the matrix product G. as sum of outer products
19

---

2.0
1.8
1.8
Exponential Distribution
1.4
ChiSquare Distribution
1.2
1.0
0.2
0.4
0.6
0.8
1.0
when x. y IID with Exponential (blue) Chi Square distribution (Orange). In both cases the expectation is 2 1 when p S 1/2, just predicts.
Figure 7: Expectation, E
pa++ ply
Theorem B.4
GM Ei=1 g:m
(IIX(GMMT) 1CCM-TT)1E1 E
X Σ gim!
-T
b Σ gemo
E
X
giml T +
T
Σ gim! T |1
k
S1 1) X
(1I)
S (1 f E []
where 11] used the k = 1 case shown above, and 12 used the inductive hypothesis. This completes the proof for general and d2 that we needed for full theorem.
B.1 Technical lemmas
It remains show interesting lemma used for proving the vector case in Theorem B.2
Lemma B.4. Let ai On 2 0 be IID random variables and assume some values p: 2 0 st.
Σ, P: and pr. < 1/n. Then
E
2 1.
Piai
o² This completes the original proof with p; Σ, oz
and a; 92
20

---

Proof. Since the are IID, it doesn't 't matter if we permute them. In particular, if π is random permutation of ,I....,77 1}.
an
an Pnlnn
E
Ea
Ex
(12)
E, Piaii
Ei PiaT.
On
2 Ea
(13)
Pir Punnn Eien Pil
Ea
an
(14)
Pnan + Eien Jen (n)
On
En
(15)
Pnan + (1 (1-Pul Pa)
E Licn
n-1
where Equation 13
uses Jensen's ss inequality on the convex function 1/x.
Now define
ag- By permuting on with the other variables, get:
ag
En
En
(16)
an + (1 - Pn
x-
Pn
1- n-1
(a an
ai
En
(17)
im Pn
+
1-pn (a ai)
n-1
a./a
Ea
(18)
I-p Ps
1-Pa
Pr 1naa
i=
1-1
n-1
Fa
bla;
(19)
1=1
where
4i
6(%)
Pr
(1-Pa
Pm )}
n-
n-
is convex whenever 1-pe / 1-pp means
pn)
1-p
which is true when 0 S Pr. 1/n. That
n-1
1-np
Jensen's 's again:
n
04 -
1
(
pla,/a) 20
=6
1,
71
a
which is what we wanted to show.
C Correspondence between theory and practice
Theorem 3.1 quite tightly matches the empirical behavior of Algorithm (Dense CCE for Least Squares). In Figure g we show the convergence when solving the least squares problem argming XT Y|? using the method.
21

---

Least squares 4 1000, 100, d2 - P 0.000468
Algonithes CCE Squarss
Theaeen 3.1
(1 1X40" |XT YE
96.00
LX22 VI
90.00
90000
10P
Iterations
Figure 8: Convergence of the least squares problem. We show and compare the convergence of the least square methods using different methods. We found that the convergence of Theorem 3.1 is close to that of Algorithm
As discussed in Appendix B we can often strengthen the result from using P (smallest singular value / the squared Frobenius norm) to just using 1/d1. This is because the smallest singular value issue only becomes relevant when chosen to exactly align with the worst possible subspace of X. In practice this is rarely the case. In the particular case of Figure g we chose X and Y as iid standard normals.
D Hashing
If h : [m]] - [m] and [n]] [-1,1} random functions, Count Sketch is a matrix H E {0, 10,-1,1]mx where Hi,- = =((i) if h(i) = and otherwise. Charikar et al. [2002] showed that if m large enough, the matrix is dimensionality reduction in the sense that the norm ||>||2 of any vector in R. is approximately preserved, ||H|l = ||>||| [
This gives simple theoretical way to think about the algorithms above: The learned matrix T H"T is simply a lower dimensional approximation to the real table that we wanted to learn. While the theoretical result requires the random "sign function" for the approximation to be unbiased, in practice this appears to not be necessary when directly learning 7". Maybe because the vectors can simply be slightly shifted to debias the result.
There are many strong theoretical results on the properties of Count Sketches. For example, Woodruff (2014) showed that they are called "subspace embeddings" which means the dimensionality reduction is "robust" and docsn't have blind spots that SGD may accidentally walk into. However, the most practical result is that one only needs h to be "universal hash function" ala Carter and Wegman 1977). which can be as simple and fast as the "multiply shift" hash function by Dietzfelbinger et al. 1997].
If Count Sketch shows that hashing each i [n]] to a single row [m]] we may wonder why methods like Hash Embeddings use multiple hash functions (or DHE uses more than a thousand.) The answer can be seen in the theoretical analysis of the "Johnson Lindenstrauss" transformation and in particular the "Sparse Johnson Lindenstrauss" as analyzed by Cohen et al. (2018). The analysis shows that the data being hashed is not very uniform, it is indeed better to use more than one hash function (more than non-zero per column in H.) The exact amount depends distribution, but one can always get away with a sparsity of e when looking for a ++ dimensionality reduction. Hence speculate that DHE could in general replace the 1024 hash functions with something more like Hash Embeddings with an MLP on top. Another interesting part of the Cohen et al. (2018) analysis is that should ideally split [m] in segments, and have one hash function into each segment. This matches the implementations we based our work on below.
characteristics in the data
'This also implies that inner products approximately preserved by the dimensionality reduction.
22

---

E How to store the hash functions
We note that unlike the random hash functions used Step the index pointer functions obtained from clustering takes space linear in the amount of training data or at least in the ID universe size. At first this may seem like a major downside of our method, and while isn't different from the index tables needed after Product Quantization, it definitely is something extra not needed by purely sketching based methods.
We give three reasons why storing this table is not an issue in practice:
1.
The index pointer functions can be stored the CPU rather than the GPU, since they used as the first step of the model before the training/inference data has been moved from the CPU to the GPU. Furthermore the index lookup is typically done faster on CPUs, since it doesn't 't involve any dense matrix operations.
2. The index pointers can replace the original IDs. Unless we are working in a purely streaming setting, the training data has to be stored somewhere. If IDs are 64 bit integers, replacing them with four 16-bit index pointers net neutral.
3. Some hashing and pruning can be used as prepossessing step, reducing the universe size of the IDs and thus the size of the index table needed.
F Different strategies for CCE
We include other graphs about CCE in Figure 9 They are all on the Kaggle dataset and were run three times. These graphs helped develop insights CCE and choose the correct versions for Figure and Figure 4b
G AUC Graphs
We also evaluate the models using AUC, which is another commonly used metric for gauging the effectiveness of recommendation model. For example, it was used in [Kang et al. 2021] AUC provides the probability of getting a correct prediction when evaluating test sample from balanced dataset. Therefore, a better model is implied by a larger AUC. In this section, we plot the graphs again using AUC; see Figure I0 and Figure [I]
H Table Collapse
Table collapsing was problem encountered for the circular clustering method as described in Appendix A We describe the problem and the metric we used to detect here, since we think they may be of interest to the community.
Suppose we are doing k-means clustering on a table of 3 partitions order obtain 3 index pointer functions hg. These functions can be thought as table, where the (i. -entryy is given by (i).
There multiple failure modes have to be aware of. The first one is column-wise collapse:
0
2
1
0
3
In this table the first column has collapsed to just cluster. Because of the way k-means clustering works, this exact case of complete collapse isn't actually possible, but we might get arbitrarily low entropy measured by Hi, which define follows: For each column j. its column entropy is defined to be the entropy the probability distribution py: h5([m)) > [o, 1] defined by
#(i: h)(()) x}
(=(x)=
n
23

---

CCE best 10 epochs
Strategy for epoch
0.467
Full Embedding
Table
0.404
Full Emteddding
Tabl
CCE cto
CCE cto
CCE ct1 c1300000
0.482
CCE ct1 c0100000
0.455
CCE ct2 c1300000
CCE ct2 cf75000
CCE ct3 cf300000
CCE c13 c150000
CCE ct4 cf300000
CCE ct4 cf37500
0.488
CCE ct5 cf300000
0.488
CCE ct5 c130000
CCE ct8 c1300000
CCE ct8 c125000
0.86
0.451
0.450
0.448
Number of parameters for Largest Tables
Number of parameters for Lergest Tables
(a) Kaggle dataset, CCE, best of 10 epoch: We ran different versions of CCE for 10 epochs. Here ct means the number of clustering done, and ef refers to the her of batches between clusterings. Since each epoch has around 300, 000 batches, we essentially clustered once every epoch. The performance increases with more clustering. Another observation is that as in- creases, few lines were merged due early stopping. We found that CCE ctó c1300000 performs the best, which becomes the CCE model in Figure 4a
(b) Kaggle dataset, CCE, Strategy I: We ran differ- ent versions of CCE for epoch under the constraint that all clusterings must finish before half of an epoch. It turns out that there is balance between the num- her of clusterings and the "quality" of the clustering, represented by the number of batches between cluster- ings. We found that CCE c12 c175000 performs the best, which becomes the CCE model in Figure 4b
Strategy for epoch
Strategy for epoch
0.484
Full Embedding Table
0.400
Full Embedding Table
CCE ct1 cf150000
CCE ct1 c1300000
CCE ct2 cf100000
0.464
CCE ct1 c0100000
CCE at4 cf60000
CCE et2 c5100000
0.482
0.400
CCE cté cf42500
CCE ct cf75000
CCE ct2 c175000
0.438
0458
0.44
0.45
Number of parameters for Largest Tables
Number of parameters for Largest Tables
(c) Kaggle dataset, CCE, Strategy 2: Strategy 2 here tries to have all the clustering finish by 2/3 of an epoch. These runs did not perform well. It turned out that need to let the model have time to converge after clusterings.
(d) Kaggle dataset, CCE, Strategy 3: This strategy perfectly summarizes all the previous findings. Increas- ing the number of clusterings in general gives better performance; Letting the model 'rest' after clustering increases the performance; Increasing the interval be- tween clusterings give better result.
Figure 9: Strategies for CCE that gave us insight.
Then we define H1 to be the minimum entropy of the (here 3) column- -ettopiessroppes.. The second failure mode is pairwise collapse:
1
2
2
3
1
3
3
0
2
24

---

Best of 10 epochs
0.500
0.04
0.502
AUC
Embedding Table
0.788
CCE [This paper)
with Concatenation
Hashing Trick
0.796
Deep Hash Embeddings
Product Cuantization
Number of parameters Largest Tables
(a) Kaggle dataset, Best of 10 epoch, AUC.
End of 1 epoch
End of epoch
0.76
0.790
Full Embedding Table
Full Embedding Table
OCE (This paper)
OCE (This paper)
Product Quantization
0.785
CE with Concatenation
CE with Concatenation
The Hashing Trick
The Hashing Trick
Product Quantization
Deep Hash Embeddings
0.700
Deep Hash Embeddings
112
10a
Number of parameters for Largest Tables
Number of parameters for Largest Tables
(b) Kaggle dataset, epoch, AUC.
(c) Terabyte dataset, 1 epoch, AUC.
Figure 10: The AUC version of Figure
In this case the second column is just a permutation of the first column. This means the expanded set of possible vectors is much smaller than would expect. We measure pairwise collapse by computing the entropy of the histogram of pairs, where the entropy of the column pair (2,,2) is (ef,()) (-). Then we define H2 to be the minimum defined by the column entropy of hy, (+) + max(hş, )hf, of such pair-entropies for all (>) pairs of columns.
Pairwise entropy can be trivially generalized to triple-wise wise and If we have columns may compute each H,.... He. In practice Hj and H2 may contain all the information we need.
H.1 What entropies are expected??
The maximum value for H, is log kt in the case of a uniform distribution over clusters. The maximum value for Ha is log (5) 2 log k. (Note log also an upper bound, where the number of points in the dataset rows in the table.)
With the CE method we expect all the entropies to be near their maximum. However, for the Circular Clustering method this is not the case! That would mean we haven't been able to extract any useful cluster information from the data.
Instead we expect entropies close to what one gets from performing Product Quantization (PQ) on complete datasct. In short:
25

---

CCE best 10 epochs
Strategy for epoch
0.968
0.002
0.800
0.834
Full Embedding Table
Full Embediing Table
CCE cto
CCE cto
0.822
CCE cf300000
0.790
CCE ct1 cf150000
CCE c1300000
CCE ct2 c175000
0.800
c1300000
0.794
CCE ct3 c150000
c1300000
0.792
CCE c14 ct37500
CCE ct5 c1300000
CCE ct5 c180000
CCE ct6 c1300000
0.70
CCE cts c125000
Number of parameters Largest Tables
Number of parameters for Largest Tables
(a) Kaggle dataset, CCE, best of 10 epoch, AUC.
(b) Kaggle dataset, CCE, Strategy 1, AUC.
Strategy epoch
Strategy for epach
0.822
0.800
0.800
0.788
AUG
0.790
Embedding Table
0.794
Embedding Table
CCE at c1300000
CCE ct1 c1150000
0.792
CCE et1 50000
0.792
CCE ct2 c1100000
0.750
CCE ct2 50000
CCE ct4 cf60000
CCE ct1 cf75000
0.700
CCE ct6 cf42500
CCE ct2 cf75000
Number of parameters for Largest Tables
Number of parameters for Largest Tables
(c) Kaggle dataset, CCE, Strategy 2, AUC.
(d) Kaggle dataset, CCE, Strategy 3, AUC.
Figure 11: The AUC version of Figure
1. Too high entropy: We are just doing CE more slowly.
2. Too low entropy: We have table collapse.
3. Golden midpoint: Whatever entropy normal PQ gets.
26


Reviews:
[{"review_id": "JvRN4eDo87", "reviewer": "Reviewer_Xht4", "summary": "It is obvious that this paper should be rejected due to its poor writing. There are many typos and grammatical errors in the paper. Furthermore, the content is poorly organized.  ", "questions": "It is obvious that this paper should be rejected due to its poor writing. There are many typos and grammatical errors in the paper. Furthermore, the content is poorly organized.  ", "limitations": "yes", "rating": 1, "confidence": 5, "soundness": 1, "presentation": 1, "contribution": 1, "strengths": "no strengths", "weaknesses": "It is obvious that this paper should be rejected due to its poor writing. There are many typos and grammatical errors in the paper. Furthermore, the content is poorly organized.  ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": ""}, {"review_id": "hkZwH77nua", "reviewer": "Reviewer_5kPE", "summary": "The embedding of categorical sparse features is a very important component of modern recommendation systems (RecSys). These models replace categorical feature values with a d-dimensional vector embedding that is fed into the model. This problem is not unique to recommendation (it also exists in NLP and other areas), but in RecSys the vocabularies are massive and require advanced compression methods for the embedding table learning to be tractable.\n\nThis paper proposes the Clustered Compositional Embeddings (CCE) method, which co-learns a hash projection matrix and a set of embeddings. The hash projections are learned by occasionally performing kmeans clustering on the embeddings (to figure out which embeddings to merge together), and the embeddings are learned by standard methods. The result is an algorithm that learns a smart hash while also learning the embedding values. The paper justifies the algorithm with theory and a lot of experiments, showing better performance for CCE on major recommendation benchmarks.", "questions": "\n- How might this work in dynamic settings, where the vocabulary drifts during training / updating? This kind of setting is common in recommendation systems and is one of the reasons why the hashing trick is so widely used (it can easily accommodate more vocabulary entries on the fly).\n\nOther notes:\n- Thank you for including your negative results in the appendix. Too few papers do this.\n- Figure 5 is referenced in the main text but appears in the appendix. This is likely because a longer form of the paper was cut down to page limit, but it might be good to fix.\n- L185: \"In line 5 note we only assignments\" -> \"In line 5, note that we only have assignments\"\n- L256: Regarding the use of the same embedding table for all features, this idea was recently explored in [[arXiv 2023]](https://arxiv.org/abs/2305.12102) with positive results.", "limitations": "Yes, no issues.", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "**S1. Good results on a significant problem.** Embeddings are the core driver of performance for some of the largest (up to 12T parameters), most complex, most economically significant (nearly all search, advertisement, and recommendation) applications of machine learning today. Embedding learning has been an active research area for about 20 years and it is very hard to make a significant improvement (as evidenced by all the \"approaches that didn't work\" in the appendix).\n\n**S2. Great experiments.** This paper implements and compares with (nearly) all of the major baselines over the last 8-10 years.\nThe ablations and wealth of empirical results are highly informative, not only in understanding the CCE method but also prior methods. \n\n**S3. Interesting ideas.** Clustering normally only works after training the embedding representations - even just a few training iterations after clustering will cause performance to degrade. The common approach is \"smart hashing\" where we know (apriori) which IDs to merge, but this is hard to do in dynamic production setting. This method seems to get around these challenges.\n\nThe theory is also a strong point, as it goes beyond the \"dimension reduction\" arguments that are common in embedding research and analyzes the embeddings in the context of a supervised learning problem. Though, there is one clarification that would improve this (see W2 below).", "weaknesses": "The paper is quite strong. The weaknesses are minor and easily addressed.\n\n**W1. Clarify the implementation.** For example, Algorithm 3 shows a few methods of the CCE table class, but does not show a full learning loop. How often / when should the clustering occur? According to Algorithm 2, it happens many times in a loop, but according to the experiment section we can get away with doing it just once per epoch - or even once per several epochs. I understand that these results are in the appendices but it would be great to summarize the takeaway in the main text. Many of these questions could be also addressed by providing a public implementation alongside the paper.\n\n\n**W2. Clarify the theory.** The theory applies only to the dense CCE method, which is not actually used in practice. It seems that the dense CCE method might have been developed solely because it is analyzable - this is fine, but it would be nice to see an explicit discussion of why we expect Algorithm 2 to approximate Algorithm 1. If my understanding is correct, this is because we suppose that $T_i$ is low-rank, and use kmeans to approximate the column space of $T_i$ so that the argmin in Line 8 is cheaper (but this point should be explicit). The theory also applies to a regression problem that is not the same as the classification problems used in recommendation, so this limitation should be clarified.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the encouraging remarks. They are truly appreciated.\n\nRe Weakness 1: We have submitted a public implementation of the algorithm, which hopefully can clarify some technical details.\n\nBut to clarify the clustering issue: If we think of the leastsq line in Algorithm 2 as completing a full training run to convergence, clustering once per epoch in Algorithm 3 is actually _more often_ than the theory suggests, rather than less. We do this because large scale ML systems are often not trained to convergence even once.\nWe have to appeal to an intuition that most of a full training run only finetunes the weights, but we can learn something about the relevant “similarity structure” even from a partially trained model.\nSimilarly, most information of the previous epoch intuitively doesn’t need to be thrown away. By seeding the tables with the cluster centroids after clustering, we speed up training in practice.\n\nPerhaps even more surprising is the fact that we can start clustering after _less_ than a full epoch. This is of course necessary for CCE to be relevant in the regime of “single epoch” trained models, like the terabyte Criteo.\nPresumably the reason this works is that even though much of the training data still hasn’t been seen, most of the vocabulary still has. As big as the vocabulary is, the number of data points is even larger.\n\nRe Weakness 2: Our original intuition for why T_i can be approximated by k-means is the success of product-quantization, PQ, and codebook learning methods. If we assume it’s possible to compress a model trained with a full embedding table using kmeans, some of this structure should also be available in the approximate table, T_i, learned using hashing.\n\nYou are right that this is akin to T_i being low rank, or even stronger, having a sparse, low-rank basis. However, since we assume k > d2, the sparse basis we find using k-means actually has more columns than T_i. (See Figure 5 in the appendix.) We suppose it makes sense that requiring this level of sparsity may require some more basis vectors.\nIt’s curious to notice that this “compression” of T_i into a sparse basis is apparently successful in pertaining the relevant parts of the column space, and keep improving until the optimal sparse column space is found. One might expect this process to converge to some local optimum, but in our experience (some of which Figure 1b shows) this doesn’t happen, and we really find the globally optimal sparse basis.\n\nThere is definitely room for more theoretical investigation in Algorithm 2.\nWe hope the analysis of Algorithm 1 and the comments above sheds light on why we designed it this way. We will include more of these intuitions into the final paper. \n\nRe Question 1: This is a good question. We have two main solutions to the “dynamic vocabulary problem”.\n1. You can simply assign the new id randomly in both tables. This matches how every id is assigned at the beginning of training. If we store h_i and h’_i as hashmaps, it’s easy to initialize them specifically for the new vocabulary as it arrives. We can also prune vocab that hasn’t been seen/used for a long time.\n2. You can use a “pre-hash” function, which maps any incoming id to a manageable space. In fact this is already the case for our DLRM experiments, as it is part of the standard preprocessing which we didn’t try to reverse.\nThis hash function will typically map to a space which is in between the size of your table and the size of your vocabulary. E.g., if you have a 10^12 vocabulary, and want a 10^6 table size, you may use an outer hash to reduce the vocabulary to 10^9.\nIf you do this, there is no such thing as a “new” id.\n\nIn both cases it is important to keep clustering regularly, such as every epoch, so the correct clustering of the new id will eventually be learned.\n\nRe Other notes: Thank you! We will make sure to incorporate these in the final manuscript. We will also cite Coleman et al. for the experiments on Unified Embeddings.\n"}, {"review_id": "lfWFADnzpS", "reviewer": "Reviewer_mVdt", "summary": "This paper discusses the challenges in practical ML when working with large categorical feature tables in modern Recommendation Systems. It introduces a novel approach called Clustered Compositional Embeddings (CCE) that addresses these challenges by combining clustering-based compression techniques with dynamic methods like The Hashing Trick and Compositional Embeddings.", "questions": "Does the empirical evaluation matches the expectation of theoretical results?", "limitations": "The authors adequately addressed the limitations", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "1. The proposed CCE offers the benefits of high compression rates from codebook-based quantization while still maintaining the dynamic nature of hashing-based methods. This combination allows CCE to be utilized during training, which is a significant advantage.\n\n2. Theoretical results showing CCE are guaranteed to converge to the optimal codebook. Additionally, it claims to provide a tight bound for the number of iterations required. These theoretical contributions add depth to the proposed method and suggest that CCE is not only practically effective but also supported by solid mathematical foundations.\n\n3. The paper is well-written with organized presentation of CCE.", "weaknesses": "NA", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the encouraging feedback! \n\nRegarding the question: Theorem 3.1 quite tightly matches the empirical behavior of Algorithm 1 (Dense CCE). We have submitted a plot (see Figure 3 in the pdf in the global response) showing the convergence when solving the least squares problem argmin_T||XT-Y||_F^2 using the method.\n\nAs discussed in the appendix, we can often strengthen the result from using rho (smallest singular value / Frobenius norm)^2 to just 1/d_1. This is because the smallest singular value issue only becomes relevant when Y is chosen to exactly align with the worst possible subspace of X. In practice this is rarely the case.\n\nIn the particular case of our plot we chose X and Y as iid standard normals.\n"}, {"review_id": "FaDfzqG3ZW", "reviewer": "Reviewer_jNoe", "summary": "In this paper, the authors propose a novel embedding table compression method for recommender models. The novelty of their method, CCE (Clustered Compositional Embeddings) is that it they learn the sparse matrix at training time, instead of using random matrices, as used by other embedding compression methods. It combines clustering and hashing techniques, so as to leverage the high compression rate of the former, and the latter's ability to use the method at training time. CCE achieves baseline accuracy of the full model, with 50% fewer parameters than the SOTA compression models, and 11000 times fewer parameters than the full model. They show experimental results on the click log datasets from Criteo: Kaggle and TB, and show that when trained for multiple epochs, CCE performs at par with the fully learned embedding table, and is always better than other embedding compression methods. The authors also provide theoretical analysis of the method, which shows that for the least squares regression problem, the method is guaranteed to converge to the optimal codebook, and the tight bound on the number of iterations required. The authors also discuss possible extensions of the method, and some interesting insights into how traditional deep recommender models are trained. ", "questions": "- Does this technique work well with Language Models? LLMs have huge embedding tables as well, and it seems that this technique could be tried for these models too, leading to good amount of compression. \n- What happens when the total number of parameters are the same as the full embedding table? In Fig. 4b, the CCE BCE loss goes below the baseline full embedding table loss. Why is this the case? shouldn't the loss converge at the loss for the full embedding table?\n- In Algorithm 3, should it be $M_i'(h_i'(id)$ instead of $M_i'(h_i(id)$? The $h_i'$ is not being used anywhere otherwise.", "limitations": "- Experiments on some more datasets would have been useful.\n- Discussion on the trade-offs of the method, in terms of any overhead of additional training time or computational resources would be good to see. ", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Good overview of the problem and related work. The connection of the technique with existing methods, like identifying sparse connections during training, and online product quantization was useful in developing a better intuition of the method.\n- The use of good illustrations and pseudo-code improved the readability of the paper.\n- Theoretical analysis on the least squares regression problem was sound and strengthened the case for the proposed method. ", "weaknesses": "- The experimental results only show the Binary Cross Entropy (BCE) loss, however at some places in the paper, the authors say that the model performs baseline accuracy. This is a little confusing. (Upon checking the dataset and its evaluation method, it seems that logarithmic loss is the metric used, so my comment is only about the use of \"accuracy\" when talking about the results in the paper)\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the encouraging feedback! \n\nRe Weakness 1: Thank you for noticing this issue. We do indeed focus on BCE in this paper (and occasionally AUC), but never accuracy. Accuracy is not a very useful measure on our datasets, as they are very biased towards negative samples. We will make sure to fix this confusion in the final paper.\n\nRe Question 1: This is a great question, and there are indeed people who have considered using embedding tables of LLMs, such as the arXiv preprint \"TensorGPT: Efficient Compression of the Embedding Layer in LLMs based on the Tensor-Train Decomposition\" by Xu et al. However if we consider GPT-2 has 50257 token embeddings of dimension 768, that makes 38.6M parameters. This is only 2.5% of the total 1.5B params. In the paper by Xu et al, they consider GPT-2 “small”, which has only 124M, and so the embeddings become more important, but in general it doesn’t seem to be the bottleneck.\n\nOn the flip side, if we have a more efficient way to handle large embedding tables, it might change the tokenization calculus and make it more interesting to use larger tokens / more embeddings than is currently done.\nIt is also possible that the regularizing effect of parameter sharing would be useful in avoiding \"anomalous tokens\" like the famous SolidGoldMagikarp in GPT-3.\n\nFinally, it could be that clustering based methods can help LLMs in other ways, such as giving them cheap access to large tables of information that thus don’t have to be stored in the feed forward layers. However, this is pure speculation and should be investigated in future work.\n\nRe Question 2: In Figure 4(b) the CCE BCE loss goes below the baseline full embedding table loss, even though we only train for 1 epoch. This does indeed look strange, since we don’t expect to overfit one a single epoch training. However, this is probably explained simply by the randomness in independent seeds. We should have added uncertainties to the baseline in addition to the other runs, which we do now in the global response, see Figure 1 of the pdf. Doing this we see that we don’t statistically significantly beat the baseline. \n\nRe Question 3: This is correct. Thank you for catching this.\n\nRe Limitation 1: We ran the open source implementation we’re providing with the paper on Movielens. In particular we implemented a simple DLRM model by embedding the user and the movie, taking the element-wise product and feeding it into an MLP. We then tried to learn whether a user would rate a movie >= 3. We summarize the results in Figure 2 in the globel response pdf, which show the BCE loss for different numbers of table sizes for the 100k and the 1M Movielens datasets.\n\nRe Limitation 2: We discuss some possible limitations of CCE here:\n* Overhead of training time: From our experiments, we did not observe a significant increase in training time: Compared to the hashing trick, CCE saw a 7.5% increase in training time; Compared to CE, CCE saw a 5.5% increase in time. \n* Extra step in the middle of training: CCE requires clusterings, which would pause training until all the clusterings are finished.\n* Need to store the hash functions: Unlike CE, CCE requires storing the index pointer functions. However, we addressed in detail in Appendix D how this is not a problem.\n"}]

Author rebuttal:
Adding a pdf of figures for rebuttal.

Meta-review: The paper present Clustered Compositional Embeddings  to compress embedding tables.  Reviewers were very excited about the idea, experiments and novelty of the papes.  The rebuttal improved the paper.  Reviewers support the acceptance of the paper strongly. One reviewer gave a very strong opinion solely based on writing and the review did not provide any further details, as a result the review was ignored.
Decision: Accept (poster)