SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Recall from the main text that we are interested in a logistic regression model with trainable embeddings for binary classification (i.e., a single-layer neural network with hashed one-hot encodings as input). This corresponds to our supervised learning framework (Eq. (1)), where  $f(\mathbf{z};\pmb {\theta})$  is the sigmoid function  $\sigma_{\pmb{\theta}}(\mathbf{z}) = 1 / (1 + \exp (-\langle \mathbf{z},\pmb {\theta}\rangle))$  and  $\ell$  is the binary cross-entropy loss.
We concatenate the embeddings of each feature, which means that the input to  $f(\mathbf{z};\pmb {\theta})$  is  $\mathbf{z} = g(\mathbf{x};\mathbf{E}) = [\mathbf{e}_{h_1(\mathbf{x}_1)},\mathbf{e}_{h_2(\mathbf{x}_2)},\dots ,\mathbf{e}_{h_T(\mathbf{x}_T)}]$ , where  $\mathbf{e}_{h_t(\mathbf{x}_t)}$  is the embedding for the  $t$ -th feature value in example  $\mathbf{x}$ .
We also partition the logistic regression weights  $\pmb{\theta} = [\pmb{\theta}_1, \pmb{\theta}_2, \dots, \pmb{\theta}_T]$ , so that embedding  $\mathbf{e}_{h_t(\mathbf{x}_t)}$  for feature  $t$  is projected via  $\pmb{\theta}_t \in \mathbb{R}^M$ .
For notational convenience we also partition the dataset by label. We write  $D_0 = \{(\mathbf{x}_i, y_i) \in D : y_i = 0\}$  and  $D_1 = \{(\mathbf{x}_i, y_i) \in D : y_i = 1\}$ . Using this notation, we can write the logistic regression objective as follows:
$$
\mathcal {L} _ {D} (\mathbf {E}, \pmb {\theta}) = - \sum_ {(\mathbf {x}, y) \in D _ {0}} \log \left(\frac {1}{1 + \exp (- \langle \pmb {\theta} , g (\mathbf {x} ; \mathbf {E}) \rangle)}\right) - \sum_ {(\mathbf {x}, y) \in D _ {1}} \log \left(\frac {1}{1 + \exp (\langle \pmb {\theta} , g (\mathbf {x} ; \mathbf {E}) \rangle)}\right).
$$
We proceed by writing the loss as a sum over vocabulary pairs rather than examples. We consider the problem for two categorical features with vocabularies  $V_{1}$  and  $V_{2}$  to simplify the presentation, and we use the notation  $C_{u,v,0}$  to denote the number of co-occurrences of  $(u,v)\in V_1\times V_2$  in dataset  $D_{0}$ .
Figure 4: Illustration of the logistic regression model, annotated with notation that corresponds to our theoretical analysis.
Likewise,  $C_{u,v,1}$  denotes the corresponding quantity for  $D_{1}$ . We also let  $\mathbf{e}_{u,v} = [\mathbf{e}_{h_1(u)},\mathbf{e}_{h_2(v)}]$  so that the inner product  $\langle [\pmb {\theta}_1,\pmb {\theta}_2],[\mathbf{e}_{h_1(u)},\mathbf{e}_{h_2(v)}]\rangle$  can be written concisely as  $\pmb{\theta}^{\top}\mathbf{e}_{u,v}$ . It follows that
$$
\mathcal {L} _ {D} (\mathbf {E}, \boldsymbol {\theta}) = - \sum_ {u \in V _ {1}} \sum_ {v \in V _ {2}} C _ {u, v, 0} \log \left(\frac {1}{1 + \exp (- \boldsymbol {\theta} ^ {\top} \mathbf {e} _ {u , v})}\right) + C _ {u, v, 1} \log \left(\frac {1}{1 + \exp (\boldsymbol {\theta} ^ {\top} \mathbf {e} _ {u , v})}\right).
$$
After combining the sigmoid functions, we get the following expression:
$$
\mathcal {L} _ {D} (\mathbf {E}, \boldsymbol {\theta}) = - \sum_ {u \in V _ {1}} \sum_ {v \in V _ {2}} C _ {u, v, 0} \log \exp (\boldsymbol {\theta} ^ {\top} \mathbf {e} _ {u, v}) - (C _ {u, v, 0} + C _ {u, v, 1}) \log (1 + \exp (\boldsymbol {\theta} ^ {\top} \mathbf {e} _ {u, v})).
$$
To get the gradients with respect to  $\mathbf{E}_{h(u)}$  (the embedding parameters used to represent value  $u$ ), we first take the gradient with respect to  $\mathbf{E}_m$ . These parameters accumulate gradient contributions from any value assigned to bucket  $m$  (i.e. for which  $\mathbb{1}_{u,m}$  is one). It is then straightforward to express  $\nabla_{\mathbf{E}_{h(u)}}\mathcal{L}_D(\mathbf{E},\theta)$  in terms of the other values  $v$  that collide with  $u$  (i.e. for which  $\mathbb{1}_{u,v}$  is one).
Collisionless Embeddings: Because each row  $\mathbf{E}_m$  corresponds to a single value, we can take the gradient with respect to  $\mathbf{E}_{h(u)}$  directly.
$$
\nabla_ {\mathbf {E} _ {h (u)}} \mathcal {L} _ {D} (\mathbf {E}, \theta) = \theta_ {1} \sum_ {v \in V _ {2}} C _ {u, v, 0} - (C _ {u, v, 0} + C _ {u, v, 1}) \sigma_ {\theta} (\mathbf {e} _ {u, v})
$$
Hash Embeddings: Here, row  $m$  corresponds to any value that is assigned to bucket  $m$  (i.e. for which  $\mathbb{1}_{u,m}$  is one). We find the gradient for the embedding table corresponding to the first feature  $V_{1}$ .
$$
\nabla_ {\mathbf {E} _ {m}} \mathcal {L} _ {D} (\mathbf {E}, \theta) = \sum_ {u \in V _ {1}} \mathbb {1} _ {u, m} \theta_ {1} \sum_ {v \in V _ {2}} C _ {u, v, 0} - (C _ {u, v, 0} + C _ {u, v, 1}) \sigma_ {\theta} (\mathbf {e} _ {u, v})
$$
The interpretation of this equation is that with hash collisions, the gradient now contains contributions from any value assigned to the bucket. If we consider the bucket selected by a specific value  $u \in V_1$ , we can write this gradient as follows.
$$
\begin{array}{l} \nabla_ {\mathbf {E} _ {h (u)}} \mathcal {L} _ {D} (\mathbf {E}, \theta) = \theta_ {1} \sum_ {v \in V _ {2}} C _ {u, v, 0} - (C _ {u, v, 0} + C _ {u, v, 1}) \sigma_ {\theta} (\mathbf {e} _ {u, v}) \\ +\theta_{1}\sum_{\substack{w\in V_{1}\\ w\neq u}}\mathbb{1}_{u,w}\sum_{v\in V_{2}}C_{u,w,0} - (C_{u,w,0} + C_{u,w,1})\sigma_{\theta}(\mathbf{e}_{u,v}) \\ \end{array}
$$
Unified Embeddings: Here, row  $m$  might correspond to any value - even ones from other features. The gradient with respect to  $\mathbf{E}_m$  now contains additional components.
$$
\begin{array}{l} \nabla_ {\mathbf {E} _ {m}} \mathcal {L} _ {D} (\mathbf {E}, \theta) = \sum_ {u \in V _ {1}} \mathbb {1} _ {u, m} \theta_ {1} \sum_ {v \in V _ {2}} C _ {u, v, 0} - (C _ {u, v, 0} + C _ {u, v, 1}) \sigma_ {\theta} (\mathbf {e} _ {u, v}) \\ + \sum_ {v \in V _ {2}} \mathbb {1} _ {v, m} \theta_ {2} \sum_ {u \in V _ {1}} C _ {u, v, 0} - \left(C _ {u, v, 0} + C _ {u, v, 1}\right) \sigma_ {\theta} (\mathbf {e} _ {u, v}) \\ \end{array}
$$
Unified Table
Collisionless
Hashing Trick
ROBE / HashedNet
ROBE-Z
MultiHash
Compositional (QR)
Compositional (PQ)
Figure 5: Illustration of baseline embedding methods, all which are compatible with feature multiplexing.
Rewriting the gradient to be in terms of  $\mathbf{E}_{h(u)}$  (for  $u\in V_1$ ) reveals that in the unified embedding case, the gradient contains the collisionless component, an intra-feature component, and an inter-feature component:
$$
\begin{array}{l} \nabla_ {\mathbf {E} _ {h (u)}} \mathcal {L} _ {D} (\mathbf {E}, \theta) = \theta_ {1} \sum_ {v \in V _ {2}} C _ {u, v, 0} - (C _ {u, v, 0} + C _ {u, v, 1}) \sigma_ {\theta} (\mathbf {e} _ {u, v}) \\ +\theta_{1}\sum_{\substack{w\in V_{1}\\ w\neq u}}\mathbb{1}_{u,w}\sum_{v\in V_{2}}C_{w,v,0} - (C_{w,v,0} + C_{w,v,1})\sigma_{\theta}(\mathbf{e}_{u,v}) \\ + \theta_ {2} \sum_ {v \in V _ {2}} \mathbb {1} _ {u, v} \sum_ {w \in V _ {1}} C _ {w, v, 0} - (C _ {w, v, 0} + C _ {w, v, 1}) \sigma_ {\theta} (\mathbf {e} _ {w, u}). \\ \end{array}
$$
Figure 5 illustrates our baselines. We describe the methods and their hyperparameters below.
1. Collisionless tables assign each value of each feature a unique  $d$ -dimensional representation. The only hyperparameter is  $d$ , the embedding dimension.
2. Hashing trick tables begin with a fixed number of representations and allow values to collide within features. Here, each table has two hyperparameters:  $d$  (the dimension) and  $M$  (the table size). The inspiration for this method is most commonly attributed to Weinberger et al. (2009), and it is widely used across industry.
3. Unified embeddings use a similar process but allow values to collide within and across features. Unified embeddings have the same hyperparameters as hash tables, but there is now only one table size to tune.
4. Compositional QR tables use multiple small tables for each feature, then perform elementwise multiplication of the resulting embeddings. The hyperparameters are  $k$  (the number of tables),  $M$  (the table size), and  $d$  (the embedding dimension). This method corresponds to the compositional embedding of Shi et al. (2020), using the "concatenation" aggregation strategy.
5. Compositional PQ tables use multiple small tables for each feature, then perform concatenation. We call this "product quantization" because it is similar to the vector compression technique. The hyperparameters are the same as for QR tables. This method corresponds to the compositional embedding of Shi et al. (2020), using the "element-wise product" aggregation strategy.
6. HashedNet tables look up each dimension of each embedding separately in a flat, shared memory. The only extra hyperparameter is the size of the shared parameter array. This is the method of Chen et al. (2015), but with the HashedNet algorithm applied only to the embedding tables.
7. ROBE-Z tables look up contiguous blocks of  $Z$  dimensions in shared memory (we use  $Z = 2$  in Figure 5). ROBE-Z has two extra hyperparameters: the size of the shared parameter array, and the number of blocks. Note that one can use a different  $Z$  for each block. For simplicity, we do not do this. This is the method used by Desai et al. (2022) and by Desai and Shrivastava (2022).
8. Hash embedding tables (or multihash, in Figure 5) perform multiple lookups into the same table, then combine the embeddings using a learned weighted combination. Multihash tables have several hyperparameters and are therefore trickier to tune well. The parameters are  $M_{1}$ , the size of the embedding table,  $M_{2}$ , the size of the importance weight table,  $k$ , the number of lookups and the dimension of the importance weight table, and  $d$ , the embedding dimension. This method was proposed by Tsang and Ahle (2022).
Hyperparameter tuning: For all methods, we consider sixteen logarithmically-spaced memory budgets: [0.001, 0.002, 0.005, 0.007, 0.01, 0.02, 0.05, 0.07, 0.1, 0.2, 0.5, 0.7, 1.0, 2.0, 5.0, 10.0] times the memory required for the collisionless table. We conduct 5 independent training runs for each hyperparameter configuration. For non-multiplexed representations, we allocate parameters to features based on vocabulary size. For example, a feature with  $10\%$  of the total vocabulary will be allocated  $10\%$  of the total parameter budget. We also tuned the network architecture and optimizer parameters. To avoid biasing our results in favor of any one method, we only used collisionless tables to tune these parameters. The embedding hyperparameters were tuned using grid search over the following choices.
- Hash embeddings: For hash embeddings (and multiplexed hash embeddings), we chose the number of lookups  $k$  from [2, 3]. We also introduce an allocation parameter  $p$  that describes what fraction of the memory budget is used for the embedding table (e.g.,  $M_1$ ) versus the importance table (e.g.,  $M_2$ ). For example,  $p = 0.2$  means that the embedding table are allocated  $80\%$  of the budget and the importance weight tables receive the remaining  $20\%$ . We chose  $p$  from [0.05, 0.1, 0.2].
- ROBE-Z: We chose the number of lookups from [2, 4, 8, 16]. This corresponds to different values of  $Z$  depending on the embedding dimension of the dataset.
- PQ embeddings: For PQ embeddings (and multiplexed PQ), we selected the number of lookups  $k$  from [2, 3, 4, 8, 16]. We note that this parameter affects performance in different ways depending on the parameter budget – smaller budgets require a greater number of lookups for performance.
- QR embeddings: For QR embeddings (and multiplexed QR), we selected the number of components  $k$  from [2, 3, 4].
- Other methods: Our other methods did not require search over extra hyperparameters. For the special case of collisionless embeddings, we simply ran 5 replicates at a single parameter budget (1.0).
Taken together, we conducted 3205 training runs for each dataset to produce Figure 3. Note that in some cases (e.g., Avazu with a budget of 5, 10), some runs with expensive configurations did not finish—we report only those results that completed training within 2 days. Finally, the total number of training runs (across all datasets) is somewhat higher, to account for our preliminary tuning over the architecture and optimizer parameters. We estimate that roughly 10,000 training runs were performed in total.
Criteo: We use the 7-day version of Criteo from the Kaggle competition. It is widely-known that additional filtering / merging of the vocabulary is necessary to achieve SOTA results on Criteo. We merge the least-frequent values, trimming the vocabulary down to the size listed in Table 4. For the network, we used a stack of two DCN cross layers followed by a 2-layer neural network with 748 nodes in each layer. This model gives a realistic picture of the tradeoff because near-SOTA performance on Criteo can be achieved with a well-tuned 3-layer feedforward network with 748 nodes and collisionless embeddings (Wang et al., 2021). We trained with a batch size of 512 for 300K steps using Adam with a learning rate of 0.0002. The full parameter-accuracy tradeoffs are presented in Figure 7.
Avazu: We use the version of Avazu from the Kaggle competition, but we use a different train-test split than the original. Following the practice of Song et al. (2019), we combine the train and test sets and do a 90-10 split after shuffling. We also apply vocabulary pruning to this dataset, with the results in Table 5. Note that we mod the "hour" feature by 24 and do not use the "id" feature. For the neural network, we use a stack of one DCN cross layer followed by two fully connected feedforward layers of 512 nodes each. We trained with a batch size of 512 for 300K steps using Adam with a learning rate of 0.0002. The full parameter-accuracy tradeoffs are presented in Figure 8.
Movielens: We use the Movielens-1M dataset with the same preprocessing used by Wang et al. (2021). We only use the "UserID," "movieId," "zipcode," "age," "occupation" and "gender" features. We always use a collisionless table to embed the "gender" feature due to small vocabulary size. We used a network composed of a single DCN cross layer followed by a single fully-connected layer with 192 nodes. We trained with a batch size of 128 for 300K steps using Adam with a learning rate of 0.0002. The full parameter-accuracy tradeoffs are presented in Figure 9.
Vocabulary distributions: Figure 6 presents the vocabulary distribution for each of our datasets. Criteo differs from Avazu and Movielens in terms of vocabulary distribution - it has a heavier-tailed vocabulary than Avazu or Movielens. This worsens the errors introduced by hash collisions because it is more likely for a colliding token to overwrite the shared embedding (all collisions are between heavy hitters). This is likely the cause of the performance gap between collisionless embeddings (where hash collisions do not occur) and the other methods on Criteo (Figure 3).
Statistical significance tests: We conducted additional statistical significance testing on Table 1. When we compare the multiplexed methods against the corresponding non-multiplexed baseline (rather than the strongest baseline), the following differences are significant with  $p < 0.05$  (Welch's  $t$ -test).
1. Criteo: All results in the 2.5 MB column. Multiplex Hashing Trick and Multiplex PQ in the 12.5 MB column. Multiplex Hashing Trick, Multiplex HashEmb, Multiplex PQ, and Multiplex QR in the 25MB column.
2. Avazu: All results in the  $324\mathrm{kB}$  and 3.24 MB columns. Multiplex QR in the 32.4 MB column.
3. Movielens: All results in the  $158\mathrm{~kB}$  column. All results in the  $791\mathrm{~kB}$  column, except for Multiplex ROBE-Z. Multiplex HashEmb, Multiplex PQ, and Multiplex QR in the 1.6 MB column.
Hardware and training time: We used CPU training for our academic experiments. End-to-end CPU training on a shared cluster took approximately 3-4 hours for Criteo, 4-5 hours for Avazu, and 30 minutes for Movielens. The choice of embedding method did not significantly affect the training time.
Table 3: Criteo training time (steps / second)
Method	Standard	Multiplexed
Collisionless	32.9	NA
Hashing Trick	29.2	29.4
(Multiple) Hash Embedding	33.5	31.3
HashedNet	44.1	39.8
ROBE-Z Embedding	30.0	34.4
PQ Embedding	37.6	40.3
QR Embedding	37.4	31.3
Vocabulary Histograms
Figure 6: Vocabulary distribution for all categorical feature values (merged) in Criteo, Avazu, and Movielens. All of these datasets display power law behavior, but Criteo has a much heavier tail.
In the Table 3, we report the average number of steps / second for several methods on Criteo for the  $25\mathrm{MB}$  table size with CPU training (higher is better). This is representative of the total training time because the number of steps to convergence is fairly stable across methods (250k batches of size 512 for Criteo, 200K of size 512 for Avazu, and 50K of size 128 for Movielens). Note that because our shared cluster has large variations in job load / demand, these numbers have high variance  $(\sigma > 5)$ . However, our results still support the conclusion that multiplexing has a minimal effect on training time. We observed similar behavior in large-scale industrial experiments—multiplexing did not significantly increase the training or inference time.
B.1 Industry experiments
Hardware and training time: In several of our Unified Embedding deployments (Table 2), we use TPUv4 for training and/or inference. For models of practical interest, embedding table size rarely affects the model training time in models of practical interest. As long as there is enough memory to support the embedding tables (i.e., enough CPU RAM or total accelerator memory), the training time is mostly governed by the forwards and backwards passes on the rest of the (upstream) network. We do not observe a substantial increase in training / inference time on either CPU or TPUv4 (Jouppi et al., 2023).
Modeling tradeoffs: Unified embedding gives up some flexibility in setting the per-feature embedding dimension, and other (more complicated) multiplexed methods are better on the Pareto-frontier for offline experiments. In exchange, we get more per-feature embedding capacity, a much simpler
Figure 7: Full parameter-accuracy tradeoffs for all methods on the Criteo dataset. Naming convention is the same as in Table 1. Note the log scale and that plots do not share the y-axis scale. The error bar width denotes the standard deviation over 5 model training runs.
Figure 8: Full parameter-accuracy tradeoffs for all methods on the Avazu dataset. Naming convention is the same as in Table 1. Note the log scale and that plots do not share the y-axis scale. The error bar width denotes the standard deviation over 5 model training runs.
hyperparameter configuration, hardware-friendly algorithm, and easy feature engineering. For example, to add new features to a model without multiplexing, we must specify and tune the table size and dimension for a new set of embedding tables. With multiplexing, we can simply add a new feature to an existing table at the cost of a few additional lookups. Another option is to use the conventional scheme for new features, then regularly consider a model update that migrates several new features into the unified embedding (typically resulting in performance gains).
Engineering tradeoffs: During our experience deploying Unified Embedding in a dozen applications, we have not had problems with observability, monitoring, or model health / stability. Instead, we find that Unified Embedding results in better load balancing and alleviate TPU hot-spot issues (as our
Figure 9: Full parameter-accuracy tradeoffs for all methods on the Movielens-1M dataset. Naming convention is the same as in Table 1. Note the log scale and that plots do not share the y-axis scale. The error bar width denotes the standard deviation over 5 model training runs.
embedding rows are distributed over cores). The high-level takeaway is that it is easier to balance one large table across accelerators than many small tables of different shapes.
Model health: All of our industry online experiments assume healthy models as a prerequisite. These models use online sequential training and adapt to constant variations in data distributions. The consistent improvements from unified embedding in these systems is strong evidence that our methods do not interfere with model health or add maintenance costs.
Table 4: Criteo vocabulary pruning.
Feature	Vocabulary Size
14	676
15	533
16	17447
17	19995
18	180
19	13
20	9693
21	337
22	3
23	14637
24	4378
25	17795
26	3067
27	26
28	6504
29	18679
30	10
31	3102
32	1557
33	3
34	18230
35	10
36	14
37	13079
38	56
39	10581
Table 5: Avazu vocabulary pruning.
Feature	Vocabulary Size
C1	8
C14	2309
C15	9
C16	10
C17	405
C18	5
C19	66
C20	167
C21	56
app_category	29
app_domain	277
app_id	4438
banner_pos	8
device_connect_type	5
device_id	67767
device_ip	163804
device_model	6217
device_type	6
site_category	24
site_domain	3887
site_id	3317
hour	24