SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Henry Ling-Hei Tsang*
henrylhtsang@meta.com
Thomas Dybdahl Ahle*
Meta†
Normal Computing
thomas@ahle.dk
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.
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
$\dagger$  Work done mostly at Probability at Meta.
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<sup>3</sup>
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
$^3$ Reducing the precision to 16-bit floats is often feasible during training, but this work aims for memory reductions much larger than that.
focused only on updating code words already assigned to the concept. Our goal is more ambitious: We want 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.
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 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 \in \mathbb{R}^{d_1 \times d_2}$ , where each ID  $i \in [d_1]$  is mapped to the  $i$ -th row,  $T[i]$ . Alternatively,  $i$  can be expressed as a one-hot row-vector  $e_i \in \{0, 1\}^{d_1}$  in which case  $T[i] = e_i T \in \mathbb{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 \in \{0,1\}^{d_1 \times k}$  and a dense matrix  $M \in \mathbb{R}^{k \times d_2}$ , where  $k \ll d_1$ , and take  $T = HM$ . In other words, to compute  $T[i]$
Figure 2: Typical Recommendation System Architecture: The DLRM model Naumov et al. [2019] embeds each categorical feature separately and combines the resulting vectors with pair-wise dot products. Other architectures use different interaction layers or a single embedding table for all categorical features, but the central role of the embedding table is universal. (Picture credit: Nvidia).
we compute  $(e_iH)M$ . Since  $H$  and  $e_i$  are both sparse, this requires very little memory and takes only constant time. The vector  $e_iH \in \mathbb{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] \to [k]$ , such that  $i$  is given the vector  $M[h(i)]$ , where  $M$  is a table with just  $k \ll d_1$  rows. Alternatively, we can think of this trick as multiplying  $e_i$  with a random matrix  $H \in \{0,1\}^{d_1 \times k}$  which has exactly one 1 in each row. Then the embedding of  $i$  is  $M[h(i)] = e_iHM$ , where  $HM \in \mathbb{R}^{d_1 \times d_2}$ .
Hash Embeddings [Tito Svenstrup et al., 2017] map each ID  $i \in 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 \in \{0,1\}^{m \times 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) = \lfloor i / p \rfloor$  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<sup>4</sup> and concatenation. Concatenation  $[T[h_1(i)], T[h_2(i)]]$  can again be described with a matrix  $H \in \{0, 1\}^{d_1 \times 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, \ldots, h_{1024} : [d_1] \to [-1,1]$  and feed the vector  $(h_1(i), \ldots, 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 \in [-1,1]^{d_1 \times 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_iHM$ , 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.
4While combining vectors with element-wise multiplication is not a linear operation, from personal communication with the authors, it is unfortunately hard to train such embeddings in practice. Hence we focus on the two linear variants.
(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 embedding vector stored at the location. Many IDs are likely to share the same vector.
(b) Hash Embeddings: Each ID is hashed to two rows, one per table, and its embedding vector is assigned to be the sum of those two vectors. Here, we use two separate tables unlike in Tito Svenstrup et al. [2017].
(c) ROBE: Similar to CE with concatenation, but instead of using separate columns, it uses a continuous array from which the (in this case  $\dim = 4$ ) pieces are picked. The pieces may even overlap.
(d) Deep Hash Embeddings:Computes a weighted sum of all the table entries, where the weights come from hashing the input 1000 times to  $[-1,1]$  .The produced (hidden) vector is then sent through an MLP for further refinement.
(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.
Figure 3: The evolution of hashing-based methods for embedding tables. The Hashing Trick and Hash Embeddings shown at the top, side by side with an equal amount of parameters. Next we introduce the idea of splitting the space into multiple concatenated subspaces. This is a classic idea from product quantization and reminiscent of multi-head attention in transformers. Finally in CCE we combine both methods in a way that allows iterative improvement using clustering.
3 Sparse and Dense CCE
The goal 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 \in \mathbb{R}^{n \times d_1}$  and we want to find a matrix  $T$  that minimizes the sum of squares,  $\| XT - Y\|_F^2 = \sum_{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.
Algorithm 1 Dense CCE for Least Squares
Require: X ∈ Rn×d1, Y ∈ Rn×d2, k ∈ N
1: H0 = 0 ∈ Rd1×2k
2: M0 = 0 ∈ R2k×d2
3: for i = 0,1,... do
4: Ti = HiMi
5: N ~ N(0,1)d1×k
6: Hi+1 = [Ti | N]
7: Mi+1 = arg minM ||XHi+1M - Y||F
8: end for
Algorithm 2 Sparse CCE for Least Squares
Require: X ∈ Rn×d1, Y ∈ Rn×d2, k ∈ N
1: H0 = 0 ∈ Rd1×2k	▷ Initialize assignm
2: M0 = 0 ∈ R2k×d2	▷ Initialize code
3: for i = 0,1,... do
4: Ti = HiMi ∈ Rd1×d2
5: A = kmeans(Ti) ∈ {0,1}d1×k
6: C ~ countsketch() ∈ {-1,0,1}d1
7: Hi+1 = [A | C] ∈ Rd1×2k
8: Mi+1 = arg minM ||XHi+1M - Y
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^{*} \in \mathbb{R}^{d_{1} \times d_{2}}$  be the matrix that minimizes  $\| X T^{*} - Y \|_{F}^{2}$ , then  $T_{i} = H_{i}M_{i}$  from Algorithm 1 exponentially approaches the optimal loss in the sense
$$
E \left[ \| X T _ {i} - Y \| _ {F} ^ {2} \right] \leq (1 - \rho) ^ {i k} \| X T ^ {*} \| _ {F} ^ {2} + \| X T ^ {*} - Y \| _ {F} ^ {2},
$$
where  $\rho = \| X\|_{-2}^2 /\| X\| _F^2\approx 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\left(\frac{d_1}{k}\log (1 / \varepsilon)\right)$  iterations we have a  $1 + \varepsilon$  approximation to the optimal solution. Note that the standard least squares problem can be solved in  $O(nd_1d_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 \in \mathbb{R}^{d_1 \times 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\times k}$  matrices with one non-zero per row. A matrix  $C$  is sampled based on hash functions  $h_i:[d_1]\to [k]$  and  $s_i:[d_1]\to \{-1,1\}$  such that  $C_{j,\ell} = s_i(j)$  if  $h_i(j) = \ell$  and 0 otherwise. We see that with  $A$  and  $C$  sparse,  $XHM$  can be computed efficiently.
<sup>5</sup>We can ignore the sign,  $\pm 1$  in most models if  $M$  has mean 0. See Charikar et al. [2002].
Figure 4: CCE also improves models even when trained for just one epoch. This shows useful information is available from clustering even when only parts of the data have been seen.
Table 1: Memory Reduction Rates Across all Datasets. For each algorithm, dataset and epoch limit we measured the necessary number of parameters to reach baseline BCE. The compression ratios with ranges are estimated using degree 1 and 2 polynomial extrapolation. We can compare these results with reported compression rates from Desai et al. [2022] (ROBE) which gets  $1000 \times$  with multi epoch training; [Yin et al., 2021] (Tensor Train) which reports a  $112 \times$  reduction at 1 epoch on Kaggle; and Yin et al. [2021] which reports a  $16 \times$  reduction at 1 epoch with "Mixed Dimension methods" on Kaggle.
Method	Dataset	Epochs	Embedding Compression
CCE (This Paper)	Criteo Kaggle	≤ 10	8,500×
CE with Concatenation	Criteo Kaggle	≤ 10	3,800×
The Hashing Trick	Criteo Kaggle	≤ 10	4,600×
Deep Hash Embeddings	Criteo Kaggle	≤ 10	1,300×
CCE (This Paper)	Criteo Kaggle	1	212×
CE with Concatenation	Criteo Kaggle	1	127 – 155×
The Hashing Trick	Criteo Kaggle	1	78 – 122×
Deep Hash Embeddings	Criteo Kaggle	1	7 – 25×
CCE (This Paper)	Criteo TB	1	101×
CE with Concatenation	Criteo TB	1	25 – 48×
The Hashing Trick	Criteo TB	1	23 – 32×
Deep Hash Embeddings	Criteo TB	1	2 – 6×
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.
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  $2kd_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.