Title: Clustering the Sketch: Dynamic Compression for Embedding Tables
Venue: NeurIPS 2023
Accepted: yes
Keywords: Embedding table compression

Clustering and sketching

Memory-efficient training
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

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 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.

1 Introduction

Neural networks can efficiently handle various data types, including continuous, sparse, and sequential features. However, categorical features present a unique challenge as they require embedding a typically vast vocabulary into a 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 are employed, such as 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 or 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 et 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 [Weinberger et al., 2009], accepting the possibility that unrelated IDs may share the same representation. Excessively aggressive hashing can impair the model's ability to distinguish its inputs, as it 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 or using multi-dimensional 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 a 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 1 for

Train

Start

Hash

(a)

(b)

(c)

(d)

Cluster

Retrain

10

0

10

1

10

2

10

3

10

4

300

400

500

600

700

800

900

MSE

Optimal, 1 sparse

Optimal, 2 sparse

CCE

(b) CCE for Least Squares : The least squares problem is to find the matrix T that minimizes ‖ XT -Y ‖ F . To save space we may use codebook quantization, that is factorize T ≈ HM , where H is a sparse Boolean matrix, and M 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 to store in memory, we can find the compressed form directly using CCE, which we prove finds the optimal H and M .

For the plot, we sampled X ∈ R 10 4 × 10 3 and Y ∈ R 10 4 × 10 . We ran CCE and compared it with the optimal solution of first T and then factorised it with either one or two 1s per row in H .

an example. Clustering also plays a crucial role in the theoretical literature on 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 3

Recent authors have explored more advanced uses of hashing to address this challenge: Tito Svenstrup et al. [2017], Shi et al. [2020], Desai et al. [2022], Yin et al. [2021], Kang et al. [2021]. A common theme is to employ multiple hash functions, enabling features to have unique representations, while still mapping into a small shared parameter table. Although these methods outperform the traditional Hashing Trick in 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 a 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 to methods like RigL [Evci 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 a form of 'Online Product Quantization' Xu et al. [2018], though prior work

focused only on updating code words already assigned to the concept. Our goal is more ambitious: We want to learn which concepts to group together without ever knowing the 'true' embedding for the concepts.

Why is this hard? Imagine you are training your model and at some point decide to use the same 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.

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 baseline 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 a version of our method for the linear least-squares problem always succeeds in finding the optimal embedding table in a number of steps logarithmic in the approximation accuracy desired.

An implementation of our methods and related work is available at github.com/thomasahle/cce.

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 on how to choose the quality and number of hash functions in the system.

We omit standard common preprocessing tricks, such as weighting 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 'posttraining' quantization, but refer to the survey by Gray and Neuhoff [1998].

Theoretical work by Li et al. [2023] 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 to scale to very large Recommendation Systems. See also Indyk and Wagner [2022] on metric compression.

2.1 Embedding Tables as Linear Maps

An embedding table is typically expressed as a tall skinny matrix T ∈ R d 1 × d 2 , where each ID i ∈ [ d 1 ] is mapped to the i -th row, T [ i ] . Alternatively, i can be expressed as a one-hot row-vector e i ∈ { 0 , 1 } d 1 in which case T [ i ] = e i T ∈ R d 2 .

Most previous work in the area of table compression is based on the idea of sketching: We introduce a (typically sparse) matrix H ∈ { 0 , 1 } d 1 × k and a dense matrix M ∈ R k × d 2 , where k << d 1 , and take T = HM . In other words, to compute T [ i ]

Click

probability

Top

MLP

Concat

Pairwise interaction

Bottom

MLP

Embedding

table 1

Embedding

tableM

Numerical

Numerical

Categorical

Categorical

feature1

feature N

feature 1

featureM

we compute ( e i H ) M . Since H and e i are both sparse, this requires very little memory and takes only constant time. The vector e i H ∈ R k is called 'the sketch of i ' and M is the 'compressed embedding table' that is trained with gradient descent.

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 : [ d 1 ] → [ k ] , such that i is given the vector M [ h ( i )] , where M is a table with just k < < d 1 rows. Alternatively, we can think of this trick as multiplying e i with a random matrix H ∈ { 0 , 1 } d 1 × k which has exactly one 1 in each row. Then the embedding of i is M [ h ( i )] = e i HM , where HM ∈ R d 1 × d 2 .

Hash Embeddings [Tito Svenstrup et al., 2017] map each ID i ∈ V to the sum of a few table rows. For example, if i is mapped to two rows, then its embedding vector is v = M [ h 1 ( i )] + M [ h 2 ( i )] . Using the notation of H ∈ { 0 , 1 } m × n , one can check that this corresponds to each row having exactly two 1s. 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 or 'Quotient Remainder', Shi et al., 2020), define h 1 ( i ) = glyph[floorleft] i/p glyph[floorright] and h 2 ( i ) = i mod p for integer p , and then combines T [ h 1 ( i )] and T [ h 2 ( i )] 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 4 and concatenation . Concatenation [ T [ h 1 ( i )] , T [ h 2 ( i )]] can again be described with a matrix H ∈ { 0 , 1 } d 1 × k 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 [Desai et al., 2022] are essentially Compositional Embeddings with concatenation 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 h 1 , . . . , h 1024 : [ d 1 ] → [ -1 , 1] and feed the vector ( h 1 ( i ) , . . . , h 1024 ( 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 ∈ [ -1 , 1] d 1 × 1024 . While this is less efficient than a sparse matrix, it can still be applied efficiently to sparse inputs, e i , 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 us that the sketching part of DHE is still the most important part.

Tensor Train [Yin et al., 2021] doesn't use hashing, but like CE it splits the input in a deterministic way that can be generalized to a random hash function if so inclined. Instead of adding or concatenating chunks, Tensor Train multiplies them together as matrices, which makes it not strictly a 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. [2022] propose an alternate method for learning a clustering based on a low dimensional embedding table. This is like a 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 e i HM , where H is a sparse matrix and M is 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.

16

1000

(a)

The

Hashing

Trick:

Also

known

as

Count

Sketch, each ID is hashed

to one location in a table

(here with 1000 rows) and

it is assigned the embed-

ding vector stored at the lo-

cation. Many IDs are likely

to share the same vector.

+

16

500

500

(b)

Hash

Embeddings:

Each ID is hashed to two

rows,

one per table,

and

its embedding vector is as-

signed

to

be the sum of

those two vectors.

Here,

we use two separate tables

unlike

in

Svenstrup

et al. [2017].

16

1000

(c)

Similar to CE

with concatenation, but in-

stead

of

using

separate

columns, it uses a contin-

uous array from which the

(in this case dim=4) pieces

are picked. The pieces may

even overlap.

Σ

1000

MLP

(d)

Deep

Hash

Em-

beddings:

Computes

a

weighted

sum of

all

the

table

entries,

where

the

weights come from hash-

ing the input 1000 times

to

[

-

1

,

1]

.

The produced

(hidden)

vector

is

then

sent through an MLP for

further refinement.

4

4

4

4

1000

+

4

+

4

+

4

+

4

500

500

(e) CE with concatenation: In the hashing version of compositional embeddings (CE), each ID is hashed to a location in each of, say, 4 different tables. The four vectors stored there are concatenated into the final embedding vector. Given the large number of possible combinations (here 1000 4 ), it is unlikely that two IDs get assigned the exact same 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 et al. [2017] with the concatenation method of Shi et al. [2020]. Each ID then gets assigned a vector that is the concatenation of smaller sums. This is enhanced with the clustering idea shown in Section 1: In each of the four blocks, we apply clustering every epoch, setting the results in the green tables, and replacing the hash functions in the blue tables with new random values.

3 Sparse and Dense CCE

The goal of 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 are pathological functions where only a full table will work. Instead, we pick a simple linear model, where the data is given by a matrix X ∈ R n × d 1 and we want to find a matrix T that minimizes the sum of squares, ‖ XT -Y ‖ 2 F = ∑ i,j (( XT ) i,j -Y i,j ) 2 .

We give two versions of the algorithm, a sparse method, which is what we build our experimental results on; and a dense method which doesn'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 a novel approximation algorithm for least squares regression with lower memory use.

9:

end for

In the Appendix we show the following theorem guaranteeing the convergence of Algorithm 1:

Theorem 3.1. Assume d 1 > k > d 2 and let T ∗ ∈ R d 1 × d 2 be the matrix that minimizes ‖ XT ∗ -Y ‖ 2 F , then T i = H i M i from Algorithm 1 exponentially approaches the optimal loss in the sense

where ρ = ‖ X ‖ 2 -2 / ‖ X ‖ 2 F ≈ 1 /d 1 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 -1 /d 1 ) ik by conditioning the random part H by the eigenspace of X . This means that after i = O ( d 1 k log(1 /ε )) iterations we have a 1 + ε approximation to the optimal solution. Note that the standard least squares problem can be solved in O ( nd 1 d 2 ) time, but one iteration of our algorithm only takes O ( nkd 2 ) time. Repeating it for d 1 /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 4 we compute HM ∈ R d 1 × d 2 which may be very large. (After all the main goal of CCE is to avoid storing this full matrix in memory.) Luckily K-means in practice works fine on just a sample of the full dataset, which is what we do in the detailed implementation in Section 4.2.

In line 5, note that K-means normally returns both a set of cluster assignments and cluster centroids. For our algorithm we only need the assignments. We write those as the matrix A which has A id,j = 1 if id is assigned to cluster j and 0 otherwise. In the Appendix (Figure 5) we show how A is a sparse approximation to T 's column space. Using the full, dense column space we'd recover the dense algorithm, Algorithm 1.

In line 6, countsketch() is the distribution of {-1 , 0 , 1 } d 1 × k matrices with one non-zero per row. A matrix C is sampled based on hash functions h i : [ d 1 ] → [ k ] and s i : [ d 1 ] →{-1 , 1 } such that C j,glyph[lscript] = s i ( j ) if h i ( j ) = glyph[lscript] and 0 otherwise. 5 We see that with A and C sparse, XHM can be computed efficiently.

Best of 10 epochs

10

3

10

4

10

5

10

6

Number of parameters for Largest Tables

0.450

0.452

0.454

0.456

0.458

BCE

Full Embedding Table

Product Quantization

Deep Hash Embeddings

The Hashing Trick

CE with Concatenation

CCE (This paper)

(a) Best of 10 Epochs Test Loss on Criteo Kaggle . We trained DLRM on the Criteo Kaggle dataset with different compression algorithms. Each of 27 categorical features was given its own embedding table, where we limited the number of parameters in the largest table as shown in 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 3 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 7 parameters per table, as each categorical value is mapped to a unique embedding vector. This model over-fits immediately if trained for more than 1 epoch, which is why it does poorly compared to more parameter-restricted methods. For Clustered Compositional Embeddings we ran clustering once every epoch for the first 6 epochs, unless stopping earlier. Product Quantization , being a post-training quantization method, is never able to do better than the baseline model it 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.

10

3

10

4

10

5

10

6

Number of parameters for Largest Tables

0.455

0.460

0.465

0.470

0.475

BCE

End of 1 epoch

Full Embedding Table

Deep Hash Embeddings

The Hashing Trick

CE with Concatenation

Product Quantization

CCE (This paper)

10

4

10

5

10

6

Number of parameters for Largest Tables

0.1265

0.1270

0.1275

0.1280

0.1285

0.1290

0.1295

BCE

End of 1 epoch

Full Embedding Table

Deep Hash Embeddings

Product Quantization

The Hashing Trick

CE with Concatenation

CCE (This paper)

4 Experiments and Implementation Details

Our primary experimental finding, illustrated in Table 1 and Figure 4a, indicates that CCE enables training a model with Binary Cross Entropy matching a full table baseline, using only a half the parameters required by the next best compression method. Moreover, when allocated optimal parameters and trained to convergence, CCE can yield a 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 et al. [2019], including the choice of 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 7 days, and the Terabyte dataset containing about 4 billion samples over 24 days.

We ran the Kaggle dataset experiments on a single A100 GPU. For the Terabyte dataset experiments, we ran them on two A100 GPUs using model parallelism. This was done mainly for memory reasons. With this setup, training for one epoch on the Kaggle dataset takes around 4 hours, while training for one epoch on the Terabyte dataset takes around 4 days. There was not a big difference in training time between the algorithms we tested. In total we used about 11,000 GPU hours for all the experiments across 5 algorithms, 3 seeds, 10 epochs and 9 different parameter counts.

4.2 CCE Implementation Details

In the previous section we gave pseudo code for 'Sparse CCE for Least Squares' (Algorithm 2). In a general model, an embedding table can be seen as a data-structure with two procedures. Below we give pseudo-code for an embedding table with vocabulary [ d 1 ] and output dimension d 2 , using 2 kd 2 parameters. The CCE algorithm is applied to each of c columns, as mentioned in Figure 3f.

The value c = 4 was chosen to match Shi et al. [2020], but larger values are generally better, as long as the h i functions don't become too expensive to store. See Appendix E. The random hash functions h ′ 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 or less clustering.

In line 12 we likely don't want to actually compute the embedding for every id in the vocabulary, but instead use mini batch K-Means with oracle access to the embedding table. In practice we follow the suggestion from FAISS K-means[Johnson et al., 2019] and just sample 256 k ids from [ d 1 ] and run K-means only for this subset. The assignments from id to nearest cluster center are easy to compute for the full vocabulary after running K-means. The exact number of samples per cluster didn'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 is 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 4b, and Figure 4c 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 Naumov et al. [2019]. However, upon training until convergence, it's possible to achieve parity with the baseline using a 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 to 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 to 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 are a number of practical and theoretical improvements one may add. All methods are naturally compatible with float16 and float8 reduced or 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 to ignore the unimportant values. In 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 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. [2023] report good results with this method. For things we tried, but didn't work, see Appendix A.

Acknowledgment

We would like to thank Ben Ghaemmaghami from the University of Texas at Austin for fruitful discussions of learned hash functions. Thanks to Chunyin Siu for helpful discussions on the convergence 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.

References

Andrew R Barron, Albert Cohen, Wolfgang Dahmen, and Ronald A DeVore. 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 B Cohen, TS Jayram, and Jelani Nelson. Simple analyses of the sparse johnson-lindenstrauss transform. In 1st Symposium on Simplicity in Algorithms (SOSA 2018) . Schloss Dagstuhl-LeibnizZentrum fuer Informatik, 2018.

Benjamin Coleman, Wang-Cheng Kang, Matthew Fahrbach, Ruoxi Wang, Lichan Hong, Ed H Chi, and Derek Zhiyuan Cheng. Unified embedding: Battle-tested feature representations for web-scale ml systems. arXiv preprint arXiv:2305.12102 , 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. A reliable randomized 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 Mudigere, Krishnakumar Nair, and Maxim Naumov. Learning to 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 on Computing , 51(3):467-491, 2022.

Jeff Johnson, Matthijs Douze, and Hervé Jégou. Billion-scale similarity search with GPUs. IEEE Transactions on Big Data , 7(3):535-547, 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 on 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 of word representations in vector space. arXiv preprint arXiv:1301.3781 , 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, et al. Deep learning recommendation model for personalization and recommendation systems. arXiv preprint arXiv:1906.00091 , 2019.

Pedregosa, G. Varoquaux, A. Gramfort, V. Michel, B. Thirion, O. Grisel, M. Blondel, P. Prettenhofer, R. Weiss, V. Dubourg, J. Vanderplas, A. Passos, D. Cournapeau, M. Brucher, M. Perrot, and E. Duchesnay. Scikit-learn: Machine learning in Python. Journal of Machine Learning Research , 12:2825-2830, 2011.

Hao-Jun Michael Shi, Dheevatsa Mudigere, Maxim Naumov, and Jiyan Yang. Compositional embeddings using complementary partitions for memory-efficient 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 representations. 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 1113-1120, 2009.

David P. Woodruff. Sketching as a tool for numerical linear algebra. Found. Trends Theor. Comput. Sci. , 10(1-2):1-157, 2014. doi: 10.1561/0400000060. URL https://doi.org/10.1561/ 0400000060 .

Donna Xu, Ivor W Tsang, and Ying Zhang. Online product quantization. IEEE Transactions on Knowledge and Data Engineering , 30(11):2185-2198, 2018.

Chunxing Yin, Bilge Acun, Xing Liu, and Carole-Jean Wu. Tt-rec: Tensor train compression for deep learning recommendation models. CoRR , abs/2101.11714, 2021. URL https://arxiv. org/abs/2101.11714 .

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-sourced implementation by Johnson et al. [2019] for better performance, but you can also use the implementation in Scikit-learn [Pedregosa et al., 2011]. The baseline result should be straightforward to reproduce as we closely follow the instructions provided by Naumov et al. [2019]. 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 factor The most important number from our experimental section is the 'Embedding Compression' factor in Table 1. 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 6 , we try e.g. to cap this at 8000 , using a full embedding table for the small features, and a CCE table with 8000 / 16 = 500 rows (since each row has 16 parameters). This corresponds to a compression rate of (10 + 100 + 10 6 ) / (10 + 100 + 500) ≈ 1639 . 5 . Or if we measure only the compression of the largest table, 10 6 / 500 = 2000 . Unfortunately there's a 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 4a uses the first measure (and thus the lower number 8,5000x).

For the experiments where we only train for 1 epoch, some methods never reach baseline 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 1 epoch. Since the loss graphs tend to 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 at a higher parameter count.

K-means For the K-means from FAISS, we use max_points_per_centroid=256 and 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 improve the final test loss. We found on Kaggle for PQ, niter=50, BCE= 0 . 455540 and niter=300, BCE= 0 . 455537 . For CCE (single epoch, single clustering at half an epoch), niter=50 gave BCE= 0 . 45928 and niter=300 gave BCE= 0 . 45905 , so a very slight improvement, but not enough to make up for the extra training time.

Datasets For our experiments, we sub-sampled an eighth of the Terabyte dataset and pre-hashed them using a simple modulus to 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 a fixed-width MLP with Mish activation. However, DHE is only described in 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 a linear embedding table, not an MLP. As a 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.

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 n _ hashes · hidden _ dim + 2 ∗ hidden _ dim 2 +hidden _ dim · embedding _ dim = 64192 parameters.

A What didn't work

Here are the ideas we tried but didn't work at the end.

Using multiple helper tables It is a 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 to 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 a tight loop between training and (re)clustering. It turned out that reducing the number of clusterings didn't impact performance, 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, as we discuss in Section 3.

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 4 columns / hash-functions was a good spot.

Residual vector quantization 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 Quantization (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 a 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 1 where we follow the CE hybrid method to get a sketch, and one in Step 3 where we follow the CE concat method. We thought that by using more parameters at the beginning, we would be able to get a 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 i with the cluster centroids from K-means and the 'helper table' M ′ i ← 0 . We could instead try to optimize M ′ i to match the residuals of T as well as possible. This could reduce the discontinuity during training more than initializing to zeros. However, we didn't see a large effect in either training loss smoothness or the ultimate test score.

B Proof of the main theorem

Let's remind ourselves of the 'Dense CCE algorithm' from Section 3: Given X ∈ R n × d 1 and Y ∈ R n × d 2 , pick k such that n > d 1 > k > d 2 . We want to solve find a matrix T ∗ of size d 1 × d 2 such that ‖ XT ∗ -Y ‖ F is minimized - the classical Least Squares problem. However, we want to use memory less than the typical nd 2 1 . We thus use this algorithm:

Dense CCE Algorithm: Let T = 0 ∈ R d 1 × d 2 . For i = 1 to m

We will now argue that T m is a good approximation to T ∗ in the sense that ‖ XT m -Y ‖ 2 F is not much bigger than ‖ XT ∗ -Y ‖ 2 F .

Let's consider a non-optimal choice of M i first. Suppose we set M i = [ I d 2 M ′ i ] where M ′ i is chosen such that ‖ H i M i -T ∗ ‖ F is minimized. By direct multiplication, we have H i M i = T i -1 + G i M ′ i . Hence in this case minimizing ‖ H i M i -T ∗ ‖ F is equivalent to finding M ′ i at each step such that ‖ G i M ′ i -( T ∗ -T i -1 ) ‖ F is minimized.

In other words, we are trying to estimate T ∗ with ∑ i G i M ′ i , where each G i is random and each M ′ i is 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 i 's. In their case, by the time we have d 1 /k such G i 's, we are just multiplying X with a d 1 × d 1 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 a local, not global optimization over the M i .

Recall the theorem:

Theorem B.0. Given X ∈ R n × d 1 and Y ∈ R n × d 2 . Let T ∗ = arg min T ∈ R d 1 × d 2 ‖ XT -Y ‖ 2 F be an optimal solution to the least squares problem. Then

where ρ = ‖ X ‖ 2 -2 / ‖ X ‖ 2 F .

Here we use the notation that ‖ X ‖ -2 is the smallest singular value of X .

Corollary B.1. In the setting of the theorem, if all singular values of X are equal, then

Proof of Theorem B.1. Note that ‖ X ‖ 2 F is the sum of the d 1 singular values squared: ‖ X ‖ 2 F = ∑ i σ 2 i . Since all singular values are equal, say to σ ∈ R , then ‖ X ‖ 2 F = d 1 σ 2 . Similarly in this setting, ‖ X ‖ 2 -2 = σ 2 so ρ = 1 /d 1 . Using the inequality 1 -1 /d 1 ≤ e -1 /d 1 gives the corollary.

Proof of Theorem B.0. First split Y into the part that's in the column space of X and the part that's not, Z . We have Y = XT ∗ + Z , where T ∗ = arg min T ‖ XT -Y ‖ F is the solution to the least

squares problem. By Pythagoras theorem we then have

so it suffices to show

We will prove the theorem by induction over i . In the case i = 0 we have T i = 0 , so E[ ‖ X ( T 0 -T ∗ ) ‖ 2 F ] = E[ ‖ XT ∗ ‖ 2 F ] trivially. For i ≥ 1 we insert T i = H i M i and optimize over M ′ i :

where the last step followed by induction. The critical step here was bounding

for a fixed T . We will do this in a series of lemmas below.

We show the lemma first in the 'vector case', corresponding to k = 2 , d 2 = 1 . The general matrix case follow below, and is mostly a case of induction on the vector case.

Lemma B.2. Let X ∈ R n × d be a matrix with singular values σ 1 ≥ ·· · ≥ σ d ≥ 0 . Define ρ = σ 2 d / ∑ i σ 2 i , then for any t ∈ R d ,

Proof. Setting m = 〈 Xt,Xg 〉 / ‖ Xg ‖ 2 2 we get

We use the singular value decomposition of X = U Σ V T . Since g ∼ N (0 , 1) d and V T is unitary, we have V T g ∼ N (0 , 1) d and hence we can assume V = I . Then

where Equation (4) follows from U T U = I in the SVD. We expand the upper sum to get

Here we use the fact that the g i 's are symmetric, so the cross terms of the sum have mean 0. By scaling, we can assume ∑ i σ 2 i t 2 i = 1 and define p i = σ 2 i t 2 i . Then the sum is just a convex combination:

Since σ i ≥ σ d and g i 's are IID, by direct comparison we have

but this follows from a cute, but rather technical lemma, which we will postpone to the end of this section. (Theorem B.4.)

It is interesting to notice how the improvement we make each step (that is 1 -ρ ) could be increased to 1 -1 /d by picking G from a distribution other than IID normal.

If X = U Σ V T , we can also take g = V Σ -1 g ′ , where g ′ ∼ N (0 , 1) d 1 × ( k -d 2 ) . In that case we get

So this way we recreate the ideal bound from Theorem B.1. Note that ‖ X ‖ 2 -2 ‖ X ‖ 2 F ≤ 1 /d 1 . Of course it comes with the negative side of having to compute the SVD of X . But since this is just a theoretical algorithm, it's still interesting and shows how we would ideally update T i . See Figure 6 for the effect of this change experimentally.

It's an interesting problem how it might inspire a better CCE algorithm. Somehow we'd have to 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 ∈ R n × d 1 be a matrix with singular values σ 1 ≥ · · · ≥ σ d 1 ≥ 0 . Define ρ = σ 2 d 1 / ∑ i σ 2 i , then for any T ∈ R n × d 2 ,

Proof. The case k = 1 , d 2 = 1 is proven above in Theorem B.2.

Case k = 1 : We first consider the case where k = 1 , but d 2 can be any positive integer (at most k ). Let T = [ t 1 | t 2 | . . . | t d 2 ] be the columns of T and M = [ m 1 | m 2 | . . . | m d 2 ] be the columns of M . Then the i th column of X ( GM -T ) is X ( Gm i -t i ) , and since the squared Frobenius norm of a

Hence

It remains to bound

0

5

10

15

20

25

30

0.0000

0.0025

0.0050

0.0075

0.0100

0.0125

0.0150

0.0175

MSE

n=625, d1=125, d2=5, k=25

noise

smart noise

half smart noise

half noise

matrix is simply the sum of the squared column l2 norms, we have

where in (10) we applied the single vector case.

Case k > 1 : This time, let g 1 , g 2 , . . . , g k be the columns of G and let m T 1 , m T 2 , . . . , m T k be the rows of M .

We prove the lemma by induction over k . We already proved the base-case k = 1 , so all we need is the induction step. We use the expansion of the matrix product GM as a sum of outer products



����

��

0.2

0.4

0.6

0.8

1.0

p

1.0

1.2

1.4

1.6

1.8

2.0

Exponential Distribution

ChiSquare Distribution

1

where (11) used the k = 1 case shown above, and (12) used the inductive hypothesis. This completes the proof for general k and d 2 that we needed for the full theorem.

B.1 Technical lemmas

It remains to show an interesting lemma used for proving the vector case in Theorem B.2.

Lemma B.4. Let a 1 . . . , a n ≥ 0 be IID random variables and assume some values p i ≥ 0 st. ∑ i p i = 1 and p n ≤ 1 /n . Then

This completes the original proof with p i = σ 2 i ∑ j σ 2 j and a i = g 2 i .

Proof. Since the a i are IID, it doesn't matter if we permute them. In particular, if π is a random permutation of { 1 , . . . , n -1 } ,

where Equation (13) uses Jensen's inequality on the convex function 1 /x .

Now define a = ∑ n i =1 a i . By permuting a n with the other variables, we get:

is convex whenever 1 -p n n -1 / ( 1 -p n n -1 -p n ) = 1 -p 1 -np > 1 , which is true when 0 ≤ p n < 1 /n . That means we can use Jensen's again:

which is what we wanted to show.

C Correspondence between theory and practice

Theorem 3.1 quite tightly matches the empirical behavior of Algorithm 1 (Dense CCE for Least Squares). In Figure 8, we show the convergence when solving the least squares problem argmin T || XT -Y || 2 F using the method.

where

Least squares with d1 =1000, k =100, d2 =10,p= 0.000468

F00000T

Algorithm 1:Dense CCE for Least Squares

Theorem3.1

(1 - 1/d)i*XT*I12 +IIXT* - YI12

98000

96000

1

LXIl

94000

92000

00006

100

101

102

Iterations

As discussed in Appendix B, we can often strengthen the result from using ρ (smallest singular value / the squared Frobenius norm) to just using 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. In the particular case of Figure 8, we chose X and Y as iid standard normals.

D Hashing

If h : [ n ] → [ m ] and s : [ n ] → {-1 , 1 } are random functions, a Count Sketch is a matrix H ∈ { 0 , -1 , 1 } m × n where H i,j = s ( i ) if h ( i ) = j and 0 otherwise. Charikar et al. [2002] showed that if m is large enough, the matrix H is a dimensionality reduction in the sense that the norm ‖ x ‖ 2 of any vector in R n is approximately preserved, ‖ Hx ‖ 2 ≈ ‖ x ‖ 2 . 6

This gives a simple theoretical way to think about the algorithms above: The learned matrix T ′ = H T 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' s for the approximation to be unbiased, in practice this appears to not be necessary when directly learning T ′ . 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 so called 'subspace embeddings' which means the dimensionality reduction is 'robust' and doesn't have blind spots that SGD may accidentally walk into. However, the most practical result is that one only needs h to be a '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 in [ 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 if the data being hashed is not very uniform, it is indeed better to use more than one hash function (more than 1 non-zero per column in H .) The exact amount depends on characteristics in the data distribution, but one can always get away with a sparsity of glyph[epsilon1] when looking for a 1 + glyph[epsilon1] dimensionality reduction. Hence we 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 one should ideally split [ m ] in segments, and have one hash function into each segment. This matches the implementations we based our work on below.

E How to store the hash functions

We note that unlike the random hash functions used in Step 1, 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 it 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:

The index pointer functions can be stored on the CPU rather than the GPU, since they are 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 involve any dense matrix operations.

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 is net neutral.

Some hashing and pruning can be used as a 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 us develop insights on CCE and choose the correct versions for Figure 4a and Figure 4b.

G AUC Graphs

We also evaluate the models using AUC, which is another commonly used metric for gauging the effectiveness of a recommendation model. For example, it was used in [Kang et al., 2021]. AUC provides the probability of getting a correct prediction when evaluating a test sample from a balanced dataset. Therefore, a better model is implied by a larger AUC. In this section, we plot the graphs again using AUC; see Figure 10 and Figure 11.

H Table Collapse

Table collapsing was a problem we encountered for the circular clustering method as described in Appendix A. We describe the problem and the metric we used to detect it 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 in order to obtain 3 index pointer functions h c j . These functions can be thought as a table, where the ( i, j ) -entry is given by h c j ( i ) .

There are multiple failure modes we have to be aware of. The first one is column-wise collapse:

In this table the first column has collapsed to just one 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 as measured by H 1 , which we define as follows: For each column j , its column entropy is defined to be the entropy of the probability distribution p j : h c j ([ n ]) → [0 , 1] defined by

10

3

10

4

10

5

10

6

Number of parameters for Largest Tables

0.449

0.450

0.451

0.452

0.453

0.454

0.455

0.456

0.457

BCE

CCE best of 10 epochs

Full Embedding Table

CCE ct0

CCE ct1 cf300000

CCE ct2 cf300000

CCE ct3 cf300000

CCE ct4 cf300000

CCE ct5 cf300000

CCE ct6 cf300000

(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 cf refers to the number 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 m increases, a few lines were merged due to early stopping. We found that CCE ct6 cf300000 performs the best, which becomes the CCE model in Figure 4a.

10

3

10

4

10

5

10

6

Number of parameters for Largest Tables

0.454

0.456

0.458

0.460

0.462

0.464

BCE

Strategy 2 for 1 epoch

Full Embedding Table

CCE ct1 cf150000

CCE ct2 cf100000

CCE ct4 cf60000

CCE ct6 cf42500

10

3

10

4

10

5

10

6

Number of parameters for Largest Tables

0.454

0.456

0.458

0.460

0.462

0.464

BCE

Strategy 1 for 1 epoch

Full Embedding Table

CCE ct0

CCE ct1 cf150000

CCE ct2 cf75000

CCE ct3 cf50000

CCE ct4 cf37500

CCE ct5 cf30000

CCE ct6 cf25000

10

3

10

4

10

5

10

6

Number of parameters for Largest Tables

0.454

0.456

0.458

0.460

0.462

0.464

0.466

BCE

Strategy 3 for 1 epoch

Full Embedding Table

CCE ct1 cf300000

CCE ct1 cf150000

CCE ct2 cf150000

CCE ct1 cf75000

CCE ct2 cf75000

Then we define H to be the minimum entropy of the (here 3 ) column-entropies.

1 The second failure mode is pairwise collapse:

10

3

10

4

10

5

10

6

Number of parameters for Largest Tables

0.796

0.798

0.800

0.802

0.804

0.806

AUC

Best of 10 epochs

Full Embedding Table

CCE (This paper)

CE with Concatenation

The Hashing Trick

Deep Hash Embeddings

Product Quantization

(a) Kaggle dataset, Best of 10 epoch, AUC .

10

3

10

4

10

5

10

6

Number of parameters for Largest Tables

0.780

0.785

0.790

0.795

0.800

0.805

AUC

End of 1 epoch

Full Embedding Table

CCE (This paper)

Product Quantization

CE with Concatenation

The Hashing Trick

Deep Hash Embeddings

(b)

Kaggle dataset, 1 epoch, AUC

.

10

4

10

5

10

6

Number of parameters for Largest Tables

0.780

0.785

0.790

0.795

0.800

AUC

End of 1 epoch

Full Embedding Table

CCE (This paper)

CE with Concatenation

The Hashing Trick

Product Quantization

Deep Hash Embeddings

(c)

Terabyte dataset, 1 epoch, AUC

.

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 we would expect. We can measure pairwise collapse by computing the entropy of the histogram of pairs, where the entropy of the column pair ( j 1 , j 2 ) is defined by the column entropy of h c j 1 ( · ) + max( h c j 1 ) h c j 2 ( · ) . Then we define H 2 to be the minimum of such pair-entropies for all ( 3 2 ) pairs of columns.

Pairwise entropy can be trivially generalized to triple-wise and so on. If we have c columns we may compute each of H 1 , . . . , H c . In practice H 1 and H 2 may contain all the information we need.

H.1 What entropies are expected?

The maximum value for H 1 is log k , in the case of a uniform distribution over clusters. The maximum value for H 2 is log ( k 2 ) ≈ 2 log k . (Note log n is also an upper bound, where n is 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 a complete dataset. In short:

10

3

10

4

10

5

10

6

Number of parameters for Largest Tables

0.798

0.800

0.802

0.804

0.806

AUC

CCE best of 10 epochs

Full Embedding Table

CCE ct0

CCE ct1 cf300000

CCE ct2 cf300000

CCE ct3 cf300000

CCE ct4 cf300000

CCE ct5 cf300000

CCE ct6 cf300000

(a) Kaggle dataset, CCE, best of 10 epoch, AUC .

10

3

10

4

10

5

10

6

Number of parameters for Largest Tables

0.790

0.792

0.794

0.796

0.798

0.800

0.802

AUC

Strategy 2 for 1 epoch

Full Embedding Table

CCE ct1 cf150000

CCE ct2 cf100000

CCE ct4 cf60000

CCE ct6 cf42500

10

3

10

4

10

5

10

6

Number of parameters for Largest Tables

0.790

0.792

0.794

0.796

0.798

0.800

0.802

AUC

Strategy 1 for 1 epoch

Full Embedding Table

CCE ct0

CCE ct1 cf150000

CCE ct2 cf75000

CCE ct3 cf50000

CCE ct4 cf37500

CCE ct5 cf30000

CCE ct6 cf25000

(b)

Kaggle dataset, CCE, Strategy 1, AUC

.

10

3

10

4

10

5

10

6

Number of parameters for Largest Tables

0.788

0.790

0.792

0.794

0.796

0.798

0.800

0.802

AUC

Strategy 3 for 1 epoch

Full Embedding Table

CCE ct1 cf300000

CCE ct1 cf150000

CCE ct2 cf150000

CCE ct1 cf75000

CCE ct2 cf75000

(c) Kaggle dataset, CCE, Strategy 2, AUC .

(d)

Kaggle dataset, CCE, Strategy 3, AUC

1.

Too high entropy: We are just doing CE more slowly.

Too low entropy: We have a table collapse.

Golden midpoint: Whatever entropy normal PQ gets.

.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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\nQUESTIONS:\n- 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.\n\nLIMITATIONS:\n- 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

comments: 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.

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.

2. 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.

3. The paper is well-written with organized presentation of CCE.

WEAKNESSES:
NA

QUESTIONS:
Does the empirical evaluation matches the expectation of theoretical results?

LIMITATIONS:
The authors adequately addressed the limitations
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2023-07-06

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n**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).\n\nWEAKNESSES:\nThe 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.\n\nQUESTIONS:\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.\n\nLIMITATIONS:\nYes, no issues.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-07"}

comments: 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.

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.

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
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 1
DATE: 2023-07-10