C
LUSTERING E MBEDDING T ABLES ,
W
ITHOUT F IRST
L EARNING T HEM .
Anonymous authors
Paper under double-blind review
A BSTRACT
To work with categorical features, machine learning systems employ embedding tables.
These tables can become exceedingly large in modern recommendation systems, necessitating the development of new methods for ﬁtting them in memory, even during training.
Some of the most successful methods for table compression are Product- and Residual Vector Quantization (Gray & Neuhoff, 1998).
These methods replace table rows with references to k-means clustered “codewords.”
Unfortunately, this means they must ﬁrst know the table before compressing it, so they can only save memory during inference, not training.
Recent work has used hashingbased approaches to minimize memory usage during training, but the compression obtained is inferior to that obtained by “post-training” quantization.
We show that the best of both worlds may be obtained by combining techniques based on hashing and clustering.
By ﬁrst training a hashing-based “sketch”, then clustering it, and then training the clustered quantization, our method achieves compression ratios close to those of post-training quantization with the training time memory reductions of hashing-based methods.
We show experimentally that our method provides better compression and/or accuracy that previous methods, and we prove that our method always converges to the optimal embedding table for least-squares training.
I NTRODUCTION
Machine learning can model a variety of data types, including continuous, sparse, and sequential features.
Categorical features are especially noteworthy since they necessitate an “embedding” of a (typically vast) vocabulary into a smaller vector space in order to facilitate further calculations.
IDs of different types, such as user IDs, post IDs on social networks, video IDs, or IP addresses in recommendation systems, are examples of such features.
Natural Language Processing is another prominent use for embeddings (usually word embeddings such as Mikolov et al., 2013), however in NLP the vocabulary can be signiﬁcantly reduced by considering “subwords” or “byte pair encodings”.
In recommendation systems like Matrix Factorization or
DLRM (see ﬁg. 2) it is typically not possible to factorize the vocabulary this way, and embedding tables end up very big, requiring hundreds of gigabytes of GPU memory (Naumov et al., 2019).
This in effect forces models to be split across may GPUs which is expensive and creates a communication bottleneck during training and inference.
The traditional solution has been to hash the IDs down to a manageable universe size using the Hashing Trick (Weinberger et al., 2009), and accepting that unrelated IDs may wind up with the same representation.
Too aggressive hashing naturally hurts the ability of the model to distinguish its inputs by mixing up unrelated concepts and reducing model accuracy.
Another option is to quantize the embedding tables.
Typically, this entails rounding each individual parameter to 4 or 8 bits.
Other quantization methods work in many dimensions at the same time, such as Product Quantization and Residual Vector Quantization.
(See Gray & Neuhoff (1998) for a survey of quantization methods.)
These multi-dimensional methods typically rely on clustering (like k -means) to ﬁnd a set of representative “code words” to which each original ID is assigned.
For example, vectors representing “red”, “orange” and “blue” may be stored as simple “dark orange” Two [Figure] hash functions
Single pointer to Nearest cluster center
Initial random Expand rows using hashing. Cluster table to K rows Concatenatewithnewrandomtable
Repeatfromstep ...
← (cid:45) and “blue” with the two ﬁrst concepts pointing to the same average embedding.
See ﬁg. 1 for an example. Even in the theoretical literature on optimal vector compression, such clustering plays a crucial role (Indyk & Wagner, 2022). All these quantization methods share one obvious drawback compared to hashing: the model is only quantized after training, thus memory utilization during training is unaffected. (Note: While it is common to do some “ﬁnetuning” of the model after, say, product quantization, the method remains primarily a “post-training” approach.)
Recent authors have considered more advanced ways to use hashing to overcome this problem:
Tito
[Equation]
The common theme has been using multiple hash functions which allow features to take different representations with high probability, while still mapping into a small shared table of parameters.
While these methods can work better than the hashing trick in some cases, they still fundamentally mix up completely unrelated concepts in a way that introduces large amounts of noise into the remaining machine learning model.
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 CQR for short) that combines hashing and clustering while retaining the beneﬁts of both methods.
By continuously interleaving clustering with training, we train recommendation models with accuracy matching post-training quantization, while using a ﬁxed 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”, though prior work like Xu et al. (2018) focused only on updating code words already assigned to 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
Numericalfeatures
[Figure]
Click probability table gets. But we are interested in keeping a constant number of parameters throughout training, while continuously improving the
In summary, our main contributions are:
• We show that it is possible to use clustering at training time, even though we don’t have memory enough to learn or even store the concepts be clustered.
• Our method bridges the state of the art for quantizing “post-training” (clustering based methods) and during training (hashing methods).
• We show experimentally that our technique allows training the Deep Learning Recommendation System (DLRM) to baseline accuracy with less than 50% of the table parameters of the best previously suggested methods.
• We prove theoretically that a version of our method always succeeds in ﬁnding the optimal embedding table in the context of Least Squares learning.
B ACKGROUND AND R ELATED W ORK
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 “post training” quantization, but refer to the survey, Gray &
Neuhoff (1998).
Finally, we keep things reasonably heuristic, but for a deep theoretical understanding of metric compression, we recommend Indyk & Wagner (2022).
2. E MBEDDING T ABLES AS L INEAR M APS
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 ] 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. [Figure]
[Figure]
In this framework we can also express most other approaches to training-time table compression:
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 (or QR embeddings, Shi et al., 2020), deﬁne h 1 ( i )
= (cid:98) i/p (cid:99) 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 ﬂexibility in the size and number of tables.
Besides using sums, like Hash Embeddings, the authors also propose element-wise multiplication 1 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 adds some more ﬂexibility in the indexing, including the ability for chunks to “wrap around” in the embedding table.
In our experiments ROBE was near indistinguishable from QR-concat 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 parts from the sketching framework, the ﬁrst hashing step of DHE is just sketching with a + dense random matrix H ∈ [ − 1 , 1] d 1 × 1024 .
While this is less efﬁcient than a sparse matrix, it can still be applied efﬁciently to sparse inputs, e i , and stored in small amounts of memory.
Unfortunately, in our experiments, DHE did not perform as well as the other methods, unless the MLP was taken to have just one layer, in which case it is just a linear transformation, e i HM , and we end up with a more expensive version of Hash Embeddings.
See ﬁg. 8b for the details.
Tensor Train (Yin et al., 2021) doesn’t use hashing, but like QR 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 ﬁrst step in reducing the input size is some kind of sketching.
See Appendix C for comparisons to a few more related methods.
Our algorithm is the ﬁrst that deviates from random sketching as the main vehicle of compression.
However, as we replace it by learned sketching, we can still express the embedding process as e i HM for a sparse matrix H and a small dense matrix M .
The difference is that we learn H from the data instead of using a random or ﬁxed matrix.
Experimentally and theoretically this allows better learning at the memory usage.
In section 4 we give the practical description of the algorithm, and in section 3 we analyze it mathematically in a simpliﬁed setting.  A NALYSIS
We consider a simpliﬁed model of the setup, where X ∈ R n × d 1 and Y ∈ R n × d 2 , and we want to ﬁnd the T ∈ R d 1 × d 2 that minimizes (cid:107) XT − Y (cid:107) 2 F .
In our setup n > d 1 >> d 2 .
(Here (cid:107) · (cid:107) F is the Frobenius norm, deﬁned by (cid:107) X (cid:107) 2 F = (cid:80) i,j X 2 i,j .) Obviously we can easily retrieve T from ( X, Y ) solving the least squares problem. Let’s call this optimal solution T ∗ ∈ R d 1 × d 2 .
However, we want to save memory, so we instead pick a sparse random matrix H 0 ∈ R d 1 × 2 k and solve
[Equation] [Figure] quickly solution sampled X ∼ N (0 , n d 1 , T ∼ N
0 , 1) d 1 × d let Y = XT . would be similar if Y noise added top. [Figure] [Figure]
[Figure]
[Figure]
[Figure]
{+[Figure] Then T 
 = H 
 M 0 is an approximation+} to T ∗ , in so far as H 0 preserves the column space of T ∗ .
Woodruff (2014) showed that this is indeed likely, if k > d 2 even if H 0 has just one non-zero per row.
For technical reasons we will let H 0 be the horizontal concatenation of two Count Sketch matrices, so it has 2 k columns with 2 non-zeros per row.
The idea is to improve over Count Sketch by analyzing T 0 . The classical k -means algorithm can be seen to ﬁnd the {+optimal factorization T 
 ≈ H (cid:48) M (cid:48) , where H (cid:48) ∈ { 0 , 1 } d 1 × k has just one 1 per row. So we have saved a bit of space.+}
We could stop here and declare victory, but we want to improve further. To do this we take H 1 = [ H (cid:48) | C (cid:48) ] – the concatenation of H (cid:48) with a new Count Sketch matrix. We can now repeat the optimization M 1 = arg min M ∈ R 2 k × d 2 (cid:107) XH 1 M − Y (cid:107) 2 F and rinse and repeat:
{+Multi-step CQR: Let H 
+}
{+= 0 ∈ R d 1 × 2 k , M 0 = 
 ∈ R 2 k × d 2 ; then repeat for i ≥ 0 : [Equation] Note that we only ever use the “assignment matrix” from k -means, not the actual cluster centers found by Lloyd’s algorithm.+}
One could use a different approximation to T i ’s column space instead, trading of the sparsity for convergence in fewer steps.
Indeed we will now modify the algorithm to a dense version that is more amendable to theoretical analysis.
Instead of k -means and Count Sketch, deﬁne H
{+i +
 = [ T i | N (0 , 1) d 1 × k ] . That is, H i is+} the previous
T i concatenated with a random normal matrix.
In appendix B we show experimentally that this algorithm (which uses a dense H i rather than a sparse, and thus is slower than what we propose above) behaves nearly identically to the multi-step CQR. We show that for this version of the algorithm: Theorem 1. [Equation]
where ρ = (cid:107) X (cid:107) 2 − 2 / (cid:107) X (cid:107) 2 F ≈ 1 /d 1 is the smallest singular value of X squared divided by the sum of singular values squared.
This means that after i
= O ( d 1 k log( (cid:107) XT ∗ (cid:107) F /ε )) iterations we have an ε 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 using less memory.
I MPLEMENTATION OF C LUSTERED C OMPOSITIONAL E MBEDDINGS (CQR)
We now describe the CQR method as used in our experiments.
The description here differs from ﬁg. 1 in two main aspects: 1) We build 4 separate tables, and concatenate the output.
And 2) We only do one iteration of the “learning, clustering, learning” loop. For illustration purposes, we use the setup of one of our experiments on the Kaggle dataset, where the embedding dimension is 16 and the number of parameters is 16 · 1000 .
Step 1: We start as in QR hybrid (see ﬁg. 4b).
with 2 times 4 tables T 1 , . . .
, T 4 and T (cid:48) 1 , . .
, T (cid:48) 4 , each of size 500 × 4 .
We also pick random hash functions h 1 , . .
, h 4 , h (cid:48) 1 , . .
, h (cid:48) 4 :
[ n ] → [500] .
Then the embedding of ID i ∈ [ n ] is deﬁned to be the concatenation  
T 1 [ h 1 ( i )] T 2 [ h 2 ( i )] T 3 [ h 3 ( i )]
{+T 4 [ h 4 ( i )] + + + + T (cid:48) 
 [ h (cid:48) 1 ( i )]+}
{+T (cid:48) 
 [ h (cid:48) 2 ( i )] T (cid:48) 
 [ h (cid:48) 3 ( i )]+}
T {+(cid:48) 
+} [ h (cid:48) 4 ( i )]   . We then train the embedding tables for either half an epoch or a full epoch.
Step 2:
Following ﬁg. 1 we now cluster the large implicit table for each column j = 1 . . .
We pick a sub-sample S ⊆
[ n ] of | S | = 256 · k IDs, and compute the vectors for each of them to obtain a table T cj of 1000 centroids.
We deﬁne h cj : [ n ] → [1000] to send i ∈ [ n ] to the nearest centroid of T j
[ h j ( i )] + T (cid:48)
j [ h (cid:48)
j ( i )] .
[-However, we want to save memory, so we instead pick a random matrix H 0 ∈ R d 1 × k and find [Equation] Then T 1 = H 
-]Step 3: We retrain the model with the new learned hash function.
The embedding vector of each i ∈ [ n ] is now assigned to
[Equation]
Note this is similar to the QR concat method (ﬁg. 4a) but with a learned hash function.
One may wonder about how best to store the learned hash functions h cj .
We give more details on this in appendix E.
E XPERIMENTS
The main experimental ﬁnding is shown in ﬁg. 6: The CQR method allows training a model with the baseline Binary Cross Entropy (BCE) using a third of the parameters needed with the hashing trick, and 1 . 5 times less than with the QR concat method.
If given the optimal amount of parameters and trained till convergence, it can give a measurably lower BCE.
In this section we explain the experimental setup we used, and show the results of various other experiments highlighting the strengths and weaknesses of the CQR method.
5. E XPERIMENTAL S ETUP
We follow the setup of the open-sourced DLRM code of Naumov et al.
Dataset.
We use Criteo’s Kaggle dataset and Criteo’s Terabyte dataset of click logs.
Both are public datasets and were used for benchmarking in DLRM (Naumov et al.
our backbone recommendation model.
Each dataset has 13 dense features and 26 categorical features.
[Figure]
• The Kaggle dataset has about 45 million samples over 7 days, where each categorical feature has at most 10 million possible values.
• The Terabyte dataset has about 4 billion samples over 24 days.
In our experiments, we sub-sampled one-eighth of the data and pre-hashed them with a simple modulus such that each categorical feature has at most 10 million possible values, as recommended by the default benchmarking setting in the DLRM code.
We use data from the days prior to the last day as the training set.
The data of the last data is split into a validation set and a test set.
Evaluation Metric.
We use Binary Cross-Entropy (BCE) to evaluate the performance of the recommendation model.
This metric is also used in Shi et al.
(2020) for evaluating the performance of the QR methods.
A small BCE means the prediction probability is close to the actual result.
Epochs.
Recommendation systems are typically just trained for one epoch since they are prone to over-ﬁtting.
It has however been observed (e.g. by Desai et al.
that compressed embedding tables beneﬁt from multi-epoch training.
We thus give results both for 1 epoch training, 2 epoch training and “best of 10 epochs” (in ﬁg. 6), making sure we get the actual best possible model performance.
Backbone Recommendation Models.
We adopt the deep learning recommendation model (DLRM) from Naumov et al.
(2019) as the backbone recommendation model; see ﬁg. 2.
The DLRM takes dense and sparse features as inputs, where a full embedding table (i.e., one with | V | many rows) is used for each categorical feature.
Since the model has an open-sourced implementation, we only have to modify the embedding table part of the code to implement our methods.
Number of runs of each model.
On the Kaggle dataset, PQ inference is run once, while all other models are run 3 times.
On the Terabyte dataset, each model is run once.
5. B ASELINE : F ULL EMBEDDING TABLE
Full embedding table means that for each categorical feature, we train an embedding table of the size of the number of categorical values.
Then a categorical value is mapped uniquely to a row in the table.
For example, if a feature has 10 million possible values, then its embedding table would have 10 million row of embedding vectors.
The embedding table is then trained for one epoch.
However, using a full embedding table is usually not practical. [Figure]
(a) Kaggle dataset, 1 epoch : Since CQR, as deﬁned, clusters the table after the ﬁrst epoch, we here suggest “CQR half” which clusters after only half an epoch. While this outperformed the previous methods (Hashing and QR concat) it was slightly worse than just keeping our own sketching method (QR hybrid) for the entire epoch. [Figure] is enough
[Figure]
to as methods previous papers, though not by as much as in ﬁg. 6 where we trained till convergence.
As described in the introduction, the CQR method mimics Product Quantization (PQ) on the full embedding table.
For that reason, PQ is a strong baseline which we compare with in ﬁg. 7a and ﬁg. 8a.
As it turns out we are mostly able to outperform this baseline, which while surprising may be explained by our PQ not being “ﬁnetuned” after clustering.
We tried training a full table, PQ clustering it and training this clustering for an epoch, but this approach immediately overﬁt and gave terrible results.
Interestingly this suggests CQR works as a regularization method for PQ. [Figure]
(a) Terabyte dataset, 1 epoch :
Like on the Kaggle dataset, running for just 1 epoch was not enough to show an improvement our sketching method, QR hybrid, but enough to beat QR concat and the Hashing method.
[Figure]
(b) DHE, Kaggle dataset, 1 epoch : Using
- gested setup from Kang et we were not able to get competitive performance when ﬁxing the number of layers to match the desired number of parameters. This is surprising, so we have opted to not include DHE on the other plots.  C ONCLUSION
We have shown the feasibility of compressing embedding tables at training time using clustering.
While there is still work to be done in expanding our theoretical understanding and testing the method in more situations, we believe this shows an exciting new paradigm for dynamic sparsity in neural networks and recommender systems in particular.
Some previous papers have made wildly different claims in terms of the possibility of compressing embedding tables.
We believe ﬁg. 7a, ﬁg. 7b and ﬁg. 6 help explain some of that confusion: With standard learning rates, it is not possible to improve over the DLRM baseline much, if at all, at 1 epoch of training.
This matches the results of Naumov et al.
However, if you train till convergence, you can match the baseline with 1000x times fewer parameters than the default.
Even if you just use the simple hashing trick.
This matches the results of Desai et al.
In all experiments, except ﬁg. 8a on the Terabyte dataset, our method, CQR, outperforms all other methods we tried.
We believe it likely would be better on the Terabyte dataset as well, if given time to train until convergence.
Training to convergence is however not very common in practical recommendation systems.
Our experiment in ﬁg. 6 may offer an explanation why: Most methods overﬁt if given too large embedding tables(!)
This is quite shocking, since the common knowledge is that bigger is always better and concepts should ideally have their own private embedding vectors.
We believe these experiments call for more investigation into overﬁtting in DLRM style models.
R EPRODUCIBILITY S TATEMENT
The backbone recommendation model, DLRM by Naumov et al.
(2019), has an open-sourced PyTorch implementation of the model available on github, with instructions on how to use them in the benchmarking folder.
It also has instructions on how to download the public datasets.
We follow very closely to their instructions, so reproducing the baseline result should be straightforward.
For the QR methods, we only need to change two functions in the code: create emb and apply emb .
We suggest using a class for each QR method; see ﬁg. 4.
For the random hash function, one could use a universal hash function or numpy.random.randint .
For the CQR method, apart from using the classes for the QR concat method and the QR hybrid method, we also need to add a section in the code for k -means clustering after some training.
We recommend using the open-sourced k -means clustering algorithm by Johnson et al.
(2019), since it is much faster than Scikit-learn’s k -means clustering algorithm (Pedregosa et al., 2011). 1 Introduction 2 Background and Related Work 2.
Embedding Tables as Linear Maps . .
3 Analysis 4 Implementation of Clustered Compositional Embeddings (CQR) 5 Experiments 5.
Experimental Setup . . Inference . . 6 Conclusion A
What
didn’t work
B Proof of the main theorem
B.1 Technical lemmas .
C Other Related Work
D
Hashing
E How to store the hash functions
F Graphs using AUC as the metric
G Table Collapse
G.1 What entropies are expected? . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
A W HAT 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 QR 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 G.
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 beneﬁts 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 CQR method combines Product Quantization (PQ) with the QR concat method.
We tried combining Residual vector quantization (RVQ) with the Hash
Embeddings method from Tito Svenstrup et al.
This method does not perform signiﬁcantly better than the Hash Embeddings method.
Seeding with PQ We ﬁrst 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 QR 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 CQR method:
Earlier clustering We currently have two versions of the CQR method: CQR half, where clustering happens at the middle of the ﬁrst epoch, and CQR, where clustering happens at the end of the ﬁrst epoch.
We observe that when we cluster earlier, the result is slightly worse.
Though in our case the CQR half method still outperforms the QR concat method.
More parameters before clustering The CQR method allows using two number of parameters, one in Step 1 where we follow the QR hybrid method to get a sketch, and one in Step 3 where we follow the QR 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 signiﬁcantly better.
B P ROOF OF THE MAIN THEOREM
Let’s remind ourselves of the “Dense CQR 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 ﬁnd a matrix T ∗ of size d 1 × dsuch that (cid:107) XT ∗ − Y (cid:107) F is minimized – the classical Least Squares problem.
However, we want to use memory less than the typical nd 21 .
We thus use this algorithm:
Dense CQR Algorithm: Let T {+0 = 
+}
∈ R d 1 × d 2 .
For i = 1 to m :
[Equation]
We will now argue that T m is a good approximation to T ∗ in the sense that (cid:107) XT m − Y (cid:107) 2 F is not much bigger than (cid:107) XT ∗ − Y (cid:107) 2 F .
Let’s consider a non-optimal choice of M i ﬁrst.
Suppose we set M i
= (cid:104) I
d 2 M (cid:48) i (cid:105)
where M (cid:48) i is chosen such that (cid:107) H
i M i − T ∗ (cid:107) F is minimized.
By direct multiplication, we have H i M i
=
T i − {+
+} + G i M (cid:48)
i .
Hence in this case minimizing (cid:107) H i M i − T ∗ (cid:107) F is equivalent to ﬁnding M (cid:48)
i at each step such that (cid:107) G i M (cid:48)
i − ( T ∗ − T i − 1 ) (cid:107) F is minimized.
In other words, we are trying to estimate T ∗ with (cid:80) i G i M (cid:48)
i , where each G i is random and each M (cid:48)
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: [Figure]
[Figure]
Theorem 1.
Given X ∈ R n × d 1 and Y ∈ R n × d 2 . Let T ∗ =
arg min T ∈ R
[Figure] (cid:107) XT − Y (cid:107) 2 F be an optimal solution to the least squares problem.
Then
[Equation]
where ρ = (cid:107) X (cid:107) 2 − 2 / (cid:107) X (cid:107) 2 F .
Here we use the notation that (cid:107) X (cid:107) − 2 is the smallest singular value of X .
Corollary 1.
In the setting of the theorem, if all singular values of X are equal, then
[Equation]
Proof of Corollary 1.
{+Note that (cid:107) X (cid:107) 
+}
F
is the sum of the d 1 singular values squared: (cid:107) X (cid:107) 2 F = (cid:80) i σ 2 i .
Since all singular values are equal, say to σ ∈ R , then (cid:107) X (cid:107) 2 F = d 1 σ 2 .
Similarly in this setting, (cid:107) X (cid:107) 2 − 2 = σ 2 so ρ = 1 /d 1 .
Using the inequality 1 − 1 /d 1 ≤ e − 1 /d 1 gives the corollary.
Proof of Theorem 1.
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 (cid:107) XT − Y (cid:107) F is the solution to the least squares problem.
By Pythagoras theorem we then have
[Equation]
so it sufﬁces to show
[Equation]
We will prove the theorem by induction over i .
{+In the case i = 
+}
we have T i = 0 , so E[ (cid:107) X
( T 0 −
T ∗ (cid:107) 2 F ] = E[ (cid:107) XT ∗ (cid:107) 2 F ] trivially.
For i ≥ 1 we insert T
i
=
H i M i and optimize over M (cid:48)
i :
E[
[Equation]
X
[Equation]
T
[Equation]
∗
[Equation]
where the last step followed by induction.
The critical step here was bounding
[Equation]
for a ﬁxed T .
We will do this in a series of lemmas below.
We show the lemma ﬁrst 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 1.
Let X ∈ R n × d be a matrix with singular values σ 1 ≥ · · · ≥ σ d ≥ 0 .
Deﬁne ρ
=
σ 2 d / (cid:80) i σ 2 i , then for any t ∈ R d ,
[Equation]
N
[Equation]
d
[Equation]
Proof.
{+Setting m = (cid:104) Xt, Xg (cid:105) / (cid:107) Xg (cid:107) 
+} we get [Equation]
[Equation]
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
[Equation]
where eq. (4) follows from U T U = I in the SVD.
We expand the upper sum to get [Equation] {+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 (cid:80) i σ 
+}
i t 2 i = 1 and deﬁne p i =
σ 2 i t 2 i .
Then the sum is just a convex combination:
[Equation]
Since σ i ≥ σ d and g i ’s are IID, by direct comparison we have
[Equation]
Hence
[Equation]
It remains to bound
[Equation] but this follows from a cute, but rather technical lemma, which we will postpone to the end of this section. (lemma 3.) 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 (cid:48) , where g (cid:48) ∼ N (0 , 1) d 1 × ( k − d 2 ) .
In that case we get
E (cid:18) (cid:104) Xt, Xg (cid:105)
[Equation]
So this way we recreate the ideal bound from corollary 1.
Note that (cid:107) X (cid:107) 2 − 2 (cid:107) X (cid:107) 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 ﬁg. 10 for the effect of this change experimentally. It’s an interesting problem how it might inspire a better CQR 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 2. Let X ∈ R n × d 1 be a matrix with singular values σ 1 ≥ · · · ≥ σ d 1 ≥ 0 . Deﬁne ρ =
σ 2 d 1 / (cid:80) i σ 2 i , then for any T ∈ R n × d 2 ,
[Equation]
Proof.
The case k = 1 , d 2 = 1 is proven above in lemma 1.
Case k
= 1 : We ﬁrst 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 matrix is simply the sum of the squared column l2 norms, we have
[Equation]
where in (10) we applied the single vector case.
[Figure] converges faster at least We also investigate how much we lose in the theorem by only considering M on the form [ I | M (cid:48) ] , rather than a general M that could take advantage of last rounds T i . The plots labeled “half noise” and “half smart noise” are limited in this way, while the two others are 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 to our analysis ����  �� 0. 1.0 p 1. 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 Tk 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
GM =
[Equation]
1 g i m Ti :
[Equation]
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. T ECHNICAL LEMMAS
It remains to show an interesting lemma used for proving the vector case in lemma 1.
Lemma 3.
Let a 1 . . . , a {+n ≥ 0 be IID random variables and assume some values p i ≥ 
+}
st. (cid:80) i p i
=and p n ≤ 1 /n .
Then
[Equation]
This completes the original proof with p i =
[Equation]
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 } , [Equation]
where eq. (13) uses Jensen’s inequality on the convex function 1 /x
[Equation]
Now deﬁne a = (cid:80) ni =1 a i . By permuting a n with the other variables, we get: [Equation]
where φ [Equation] is convex whenever 1 − p n n − 1 (cid:14) ( 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:
[Equation] which is what we wanted to show.
Recent parallel work by Ghaemmaghami et al. (2022) presents a different way to learn a clustering without having to learn the exact embedding tables ﬁrst. Instead of using hashing as a “pseudo table” for clustering to work on, they train a model with a lower dimensional table.
They run a special clustering algorithm designed for low dimensional data on this table, and then retrain the model with the clustering known. It is interesting future work to better compare the strengths and weaknesses of our two methods against each other. Other people have looked at “Learning CountSketch” such as Liu et al. (2020) who gave a formulation in the learning-based sketching paradigm proposed by Indyk et al.
However, these algorithms require repeatedly retraining the model, which for big recommendation systems would be much too slow. D H ASHING 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 (cid:107) x (cid:107)of any vector in R n is approximately preserved, (cid:107) Hx (cid:107) 
+}
≈ (cid:107) x (cid:107) 2 . This gives a simple theoretical way to think about the algorithms above: The learned matrix T (cid:48) = 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 (cid:48) .
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 & Wegman (1977), which can be as simple and fast as the “multiply shift” hash function by Dietzfelbinger et al.
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. 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 (cid:15) when looking for a 1 + (cid:15) 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.
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 ﬁrst 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 deﬁnitely 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 ﬁrst 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 G RAPHS USING AUC AS THE METRIC We also evaluate the models using AUC, another often employed metric for gauging the effectiveness of a recommendation model.
For example, it was used in (Kang et al., 2021). It provides the probability of getting a correct prediction when evaluating a test sample from a balanced dataset. Thus, a better model is implied by a larger AUC. In this section, we plot the graphs again using AUC.
[Figure] [Figure] [Figure]
G T ABLE C OLLAPSE
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 cj .
These functions can be thought as a table, where the ( i, j ) -entry is given by h cj ( i ) .
There are multiple failure modes we have to be aware of.
The ﬁrst one is column-wise [Figure]
[Figure] In this table the ﬁrst 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 deﬁne as follows:
For each column j , its column entropy is deﬁned to be the entropy of the probability distribution p
j : h cj ([ n ]) →
[0 , 1] deﬁned by
[Equation]
n .
Then we deﬁne H 1 to be the minimum entropy of the (here 3 ) column-entropies.
The second failure mode is pairwise collapse:
In this case the second column is just a permutation of the ﬁrst 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 deﬁned by the column entropy of h cj 1 ( · ) + max( h cj 1 ) h cj 2 ( · ) .
Then we deﬁne H 2 to be the minimum of such pair-entropies for all (cid:0) 32 (cid:1) 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.
G. W HAT 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 (cid:0) k 2 (cid:1)
≈ 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 QR 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:
Too high entropy: We are just doing QR more slowly.
Too low entropy: We have a table collapse.
Golden midpoint: Whatever entropy normal PQ gets.