SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\mathcal {L} _ {D} (\mathbf {E}, \boldsymbol {\theta}) = - \sum_ {u \in V _ {1}} \sum_ {v \in V _ {2}} C _ {u, v, 0} \log \exp \left(\boldsymbol {\theta} ^ {\top} \mathbf {e} _ {u, v}\right) - \left(C _ {u, v, 0} + C _ {u, v, 1}\right) \log \left(1 + \exp \left(\boldsymbol {\theta} ^ {\top} \mathbf {e} _ {u, v}\right)\right) \tag {2}
$$
To understand the effect of hash collisions on training dynamics, we take the gradient with respect to  $\mathbf{e}_{h_1(u)}$ , i.e., the embedding parameters that represent  $u \in V_1$ . By studying the gradients for collisionless, hash, and multiplexed embeddings, we quantify the effects of intra-feature and interfeature collisions. Interestingly, all three gradients can be written using the following three terms:
$$
\begin{array}{l} \nabla_ {\mathbf {e} _ {h _ {1} (u)}} \mathcal {L} _ {D} (\mathbf {E}, \boldsymbol {\theta}) = \boldsymbol {\theta} _ {1} \sum_ {v \in V _ {2}} C _ {u, v, 0} - \left(C _ {u, v, 0} + C _ {u, v, 1}\right) \sigma_ {\boldsymbol {\theta}} \left(\mathbf {e} _ {u, v}\right) (3) \\ + \boldsymbol {\theta} _ {1} \sum_ { \begin{array}{c} w \in V _ {1} \\ w \neq u \end{array} } \mathbb {1} _ {u, w} \sum_ {v \in V _ {2}} C _ {w, v, 0} - \left(C _ {w, v, 0} + C _ {w, v, 1}\right) \sigma_ {\boldsymbol {\theta}} \left(\mathbf {e} _ {u, v}\right) (4) \\ + \boldsymbol {\theta} _ {2} \sum_ {v \in V _ {2}} \mathbb {1} _ {u, v} \sum_ {w \in V _ {1}} C _ {w, v, 0} - \left(C _ {w, v, 0} + C _ {w, v, 1}\right) \sigma_ {\boldsymbol {\theta}} \left(\mathbf {e} _ {w, v}\right). (5) \\ \end{array}
$$
For collisionless embeddings, component (3) is the full gradient. For hash embeddings, the gradient is the sum of components (3) and (4). The multiplexed embedding gradient is the sum of all three terms.
Insights Our core observation is that embedding gradients can be decomposed into a collisionless (true) component (3), an intra-feature component (4), and an inter-feature component (5). The components from hash collisions act as a gradient bias. The intra-feature component is in the same direction as the true component, implying that intra-feature collisions are not resolvable by the model. This agrees with intuition since intra-feature hash collisions effectively act as value merging.
Inter-feature collisions, however, push the gradient in the direction of  $\theta_{2}$ , which creates an opportunity for model training to mitigate their effect. To see this, consider a situation where  $\theta_{1}$  and  $\theta_{2}$  are initialized to be orthogonal and do not change direction during training. During SGD, the embedding  $\mathbf{e}_{h_1(u)}$  is a linear combination of gradients over the training steps, which means  $\mathbf{e}_{h_1(u)}$  can be decomposed into a true and intra-feature component, in the  $\theta_{1}$  direction, and an inter-feature component in the  $\theta_{2}$  direction. Since  $\langle \theta_{1},\theta_{2}\rangle = 0$ , the projection of  $\theta_{1}^{\top}\mathbf{e}_{h_{1}(u)}$  effectively eliminates the inter-feature component.
2This restriction does not affect the representation capacity because learned embeddings can rotate around  $\theta$ , i.e., for any  $\theta, \mathbf{e}$ , and nonzero constant  $\mathbf{u}$ , we can learn some  $\mathbf{e}'$  that satisfies  $\langle \mathbf{e}', \mathbf{u} \rangle = \langle \mathbf{e}, \theta \rangle$ .
Empirical study These observations allow us to form two testable hypotheses. First, we expect the projection vectors  $[\theta_1,\theta_2,\dots ,\theta_T]$  to orthogonalize since this minimizes the effect of inter-feature collisions (and improves the loss). The orthogonalization effect should be stronger for embedding tables with fewer hash buckets since that is when inter-feature collisions are most prevalent. Second, intra-feature and inter-feature gradient contributions add approximately  $O(N / M)$  for each of their dimensions to the collisionless gradient since  $\mathbb{E}[\mathbb{1}_{u,v}] = 1 / M$  . Therefore, we expect the squared embedding norms to scale roughly as  $O(N / M)$
To investigate these hypotheses, we train a single-layer neural network on the categorical features of the Criteo click-through prediction task. We initialize  $\pmb{\theta} = [\pmb{\theta}_1,\pmb{\theta}_2,\dots,\pmb{\theta}_T]$  with each  $\pmb{\theta}_t$  component in the same direction (i.e., the worst-case) to see if multiplexing encourages the weights to orthogonalize. We plot the results in Figure 2, which support our hypotheses.
4.3 Why Can Features Share a Single Table?
Returning to our parameter efficiency argument of Section 4.1, we can finally explain the full benefits of feature multiplexing. The dimension reduction argument shows that, given a parameter budget, the multiplexed embedding has the same overall number of hash collisions as a well-tuned hash embedding. Our gradient analysis, however, shows that not all collisions are equally problematic. Inter-feature collisions can be mitigated by a single-layer neural network because different features are processed by different model parameters. While our theoretical results are limited to single-layer neural networks, we expect deeper and more complicated network architectures to exhibit analogs of weight orthogonalization due to their relative overparameterization. However, we believe this is still compelling evidence that shared embedding tables work by load balancing unrecoverable collisions across a larger parameter space and effectively use one massive table for each feature.
Finally, we present experimental results on three public benchmark datasets, followed by anonymized, aggregate results from our industrial deployments across multiple web-scale systems. We give more details in Appendix B, including full parameter-accuracy tradeoffs, additional experiments, and the neural network architectures used.
5.1 Experiments with Public Benchmark Datasets
Datasets Criteo is an online advertisement dataset with  $\sim 45$  million examples (7 days of data). Each example corresponds to a user interaction (click/no-click decision) and has 39 corresponding features. 26 of these features are categorical, and 13 are real-valued continuous variables. We only embed the categorical features. Avazu is a similar advertisement click dataset but with  $\sim 36$  million examples (11 days of data). Each example here contains 23 categorical features and no continuous features. Movielens is a traditional user-item collaborative filtering problem, which we convert to a binary prediction task by assigning "1"s to all examples with rating  $\geq 3$ , and "0"s to the rest. For more details, see the benchmarks in Zhu et al. (2022) for Criteo and Avazu and the review by Harper and Konstan (2015) for Movielens. For Criteo and Movielens, we apply the same continuous feature transformations, train-test split, and other preprocessing steps in Wang et al. (2021). For Avazu, we use the train-test split and preprocessing steps in Song et al. (2019).
Experiment design To find the Pareto frontier of the memory-AUC curve (Figure 3), we train a neural network with different embedding representation methods at 16 logarithmically-spaced memory budgets. Models can overfit on the second epoch, but compressed embeddings sometimes require multiple epochs. To showcase the best performance, we follow the guidelines of Tsang and Ahle (2022) and report the test performance of the best model found over three epochs. As baselines, we implement collisionless embeddings, the feature hashing trick, hash embeddings, HashedNets, ROBE-Z, compositional embeddings with element-wise product (QR), and compositional embeddings with concatenation (i.e., product quantization (PQ)).
For the multiplexed version of an existing embedding scheme, we use a single instance of the scheme to index all of the features. For example, we look up all feature values in a single multihash table, rather than using a separately-tuned multihash table for each feature. This is equivalent to salting the vocabulary of each feature with the feature ID and merging the salted vocabularies into a single massive vocabulary (which is then embedded as usual). We apply this procedure to each baseline to obtain a new multiplexed algorithm, reported below the horizontal bar in Table 1.
Implementation and hyperparameters All methods are implemented in TensorFlow 2.0 using the TensorFlow Recommenders (TFRS) framework. For a fair comparison, all models are identical except for the embedding component. Following other studies (Naumov et al., 2019; Wang et al., 2021), we use the same embedding dimension for all features ( $d \in \{39, 32, 30\}$  for Criteo, Avazu, and Movielens, respectively). This is followed by a stack of 1-2 DCN layers and 1-2 DNN feed-forward layers—we provide a full description of the architecture in Appendix B. We run a grid search to tune all embedding algorithm hyperparameters and conduct five runs per combination. When allocating embedding parameters to features, we size tables based on the cardinality of their vocabulary (e.g., a feature with  $10\%$  of the total vocabulary receives  $10\%$  of the memory budget).
Results Table 1 shows the results for the academic datasets. We give three points on the parameter-accuracy tradeoff and highlight the Pareto-optimal choices, but present the full tradeoff (including standard deviations across runs) in Appendix B. Figure 3 shows the Pareto frontier of the parameter-accuracy tradeoff for all multiplexed and non-multiplexed methods. We observe that feature multiplexing is Pareto-optimal. Interestingly, the multiplexed hashing trick outperforms several embedding techniques (e.g., ROBE-Z) that were previously SOTA.
Collisionless embeddings are included in our benchmark evaluation only to provide an optimistic headroom reference point. Collisionless tables are often impossible to deploy in industrial recommendation systems since feature vocabularies can be massive and change dynamically due to churn. Multiplexing provides a significant improvement to methods that are fast, implementable, and orders of magnitude cheaper to train and serve than collisionless embeddings.
For significance tests comparing multiplexed methods to their corresponding non-multiplexed baseline (rather than the strongest baseline), see Appendix B. Appendix B also contains additional details about Criteo and Avazu, including an explanation for the collisionless-cache embedding performance gap on Criteo.
Table 1: AUC of embedding methods on click-through datasets. We give results for several points on the parameter-accuracy tradeoff, listed as sub-headings under the dataset that describe the total embedding table size. Our multiplexed methods (listed below the bar) outperform the strongest baseline at  $0.01^{**}$  and  $0.05^{*}$  significance levels of Welch's  $t$ -test.
Method	Criteo	Avazu	Movielens-1M
25MB	12.5MB	2.5MB	32.4MB	3.24MB	324kB	1.6MB	791kB	158kB
Collisionless	80.70	-	-	77.35	-	-	88.72	-	-
Hashing Trick	80.21	79.98	79.44	77.24	76.71	75.10	85.37	83.00	77.07
Hash Embedding	79.82	80.34	80.40	77.29	76.89	76.86	88.28	87.71	86.05
HashedNet	80.54	80.50	80.42	77.30	77.19	76.89	87.92	87.92	87.31
ROBE-Z Embedding	80.55	80.52	80.41	77.32	77.18	76.88	88.05	87.92	87.40
PQ Embedding	80.29	80.25	79.82	77.39	77.11	76.48	88.49	88.38	87.32
QR Embedding	80.29	80.21	79.74	76.98	76.61	75.86	88.49	88.63	88.22
Multiplex Hash Trick	80.43	80.47	80.49**	77.35	77.18	76.86	87.74	86.93	82.00
Multiplex HashEmb	80.52	80.42	80.31	77.28	77.16	76.82	88.79**	88.74*	87.66
Multiplex HashedNet	80.53	80.55	80.48**	77.30	77.23*	76.94**	87.95	87.93	87.57
Multiplex ROBE-Z	80.57	80.56	80.50*	77.32	77.24**	76.94**	88.00	88.03	87.62
Multiplex PQ	80.43	80.49	80.54**	77.39	77.29**	76.96**	88.81**	88.69*	88.09
Multiplex QR	80.22	80.28	80.48**	77.29	77.16	76.90	88.83**	88.53	88.27**
5.2 Industrial Deployments of Unified Embedding in Web-Scale Systems
Among the proposed approaches, Unified Embedding (multi-probe hashing trick + multiplexing) is uniquely positioned to address the infrastructure and latency constraints associated with training and serving web-scale systems. We adopt Unified Embedding in practice, using variable row lookups for each feature (typically 1~6) with concatenation to achieve the desired embedding width for our features. It should be noted that this process yields a similar algorithm to Multiplex PQ—one of the top performers from Table 1. We discuss how to adopt Unified Embedding in production SAR systems, where it has already served billions of users through launches in over a dozen models. We also present an experimental study of Unified Embedding in offline and online A/B experiments.
<sup>3</sup>https://github.com/tensorflow/recommenders
Figure 3: Pareto frontier of multiplexed and original (non-multiplexed) methods. Top-left is better.
Experiment and production setup To test the robustness and generalization capability of Unified Embedding, we conduct offline and online experiments with five key models from different SAR domains, including commerce, apps, and short-form videos. These systems represent a diverse set of production model architectures powered by state-of-the-art techniques, including two-tower retrieval models with LogQ correction (Yi et al., 2019), DCN-V2 (Wang et al., 2021), and MMOE (Zhao et al., 2019). These models handle a wide range of prediction tasks, including candidate retrieval, click-through rate (pCTR), conversion rate (pCVR) prediction (Chapelle, 2014), and multi-task user engagement modeling.
All of these production models are continuously trained and updated in an online streaming fashion, with the embedding tables being the dominant component (often  $>99\%$ ) in terms of model parameters (Anil et al., 2022; Fahrbach et al., 2023). Many of the baseline embedding methods are simple but quite strong (e.g., feature hashing and multiple hashing), and have been tuned with AutoML and human heuristics (Anil et al., 2022). In fact, these baseline embedding learning methods have roughly stayed the same for the past five years, with various efforts failing to improve them through novel techniques.
Practical benefits Unified Embedding tables offer solutions to several challenging and practical issues in large-scale embedding learning systems.
- Simple configuration: With standard embedding methods, we need to independently tune the table size and dimension for each feature. Hyperparameter tuning is an arduous process due to the large number of categorical features used in production (hundreds if not thousands). Unified Embedding tables are far simpler to configure: the total table size is calculated based on the available memory budget, and we only need to tune the dimension on a per-feature basis—easily accomplished via AutoML (Bender et al., 2020) or attention-based search methods (Yasuda et al., 2023; Axiotis and Yasuda, 2023). This yields roughly a  $50\%$  reduction in the number of hyperparameters.
- Adapt to dynamic feature distributions: In practice, the vocabulary size of each feature changes over time (new IDs enter the system on a daily basis, while stale items gradually vanish). For example, in a short-form video platform, we expect a highly dynamic video ID vocabulary with severe churn. In e-commerce, a significant number of new products are added during the holiday season compared to off-season. Thus, a fixed number of buckets for a given table can often lead to sub-optimal performance. Unified Embedding tables offer a large parameter space that is shared across all features, which can better accommodate fluctuations in feature distributions.
- Hardware compatibility: Machine learning accelerators and frameworks have co-evolved with algorithms to support common use-cases, and the "row embedding lookup" approach has been the standard technique for the past decade. While Unified Embedding tables are well-supported by the latest TPUs (Jouppi et al., 2023) and GPUs (Wei et al., 2022b; Mudigere et al., 2022), many of the recent and novel embedding methods require memory access patterns that are not as compatible with ML accelerators.
Offline and online results Table 2 displays the results of our offline and online experiments for Unified Embedding. It was applied to production models with vocabulary sizes spanning from 10M to 10B. We observe substantial improvements across all models. Note that the offline metrics and the key business metrics for different models can be different, and the gains are not directly comparable. For offline metrics, the models are evaluated using AUC (higher is better), Recall@1 (higher is better), and RMSE (lower is better). From these experiments, we found that feature multiplexing provides greater benefits when (i) the vocabulary sizes for features are larger, and (ii) the vocabularies are more dynamic (e.g., suffers from a high churn rate). For instance, an overwhelming number of new short-form videos are produced daily, creating a large and dynamic vocabulary.
Table 2: Unified Embedding applied to different production model architectures and prediction tasks across various domains. Note that  $+0.1\%$  is considered significant due to the large amount of traffic served in online A/B experiments. All multiplexed models are neutral or better in terms of model size and training/serving costs.
Task	pCTR	Retrieval	Multi-Task Ranking	pCVR	pCTR
Domain	Products	Products	Short-form videos	Apps	Apps
Architecture	DCN-V2	Two-tower	MMOE	DCN-V2	DCN-V2
Vocabulary Size	~10B	~10B	~1B	~10M	~10M
Offline Metric	AUC +2.2%	Recall@1 +7.3%	AUC +0.39% (task 1)	AUC +0.25%	AUC +0.17%
RMSE -0.53% (task 2)
Online Metric	Positive	Positive	+0.62%	+0.44%	+0.11%
In this work, we propose a highly generalizable Feature Multiplexing framework that allows many different features to share the same embedding space. One of our crucial observations is that models can learn to mitigate the effects of inter-feature collisions, enabling a single embedding vector to represent feature values from different vocabularies. Multiplexed versions of SOTA embedding methods provide a Pareto-optimal tradeoff, and Unified Embedding proves to be a particularly simple yet effective approach. We demonstrate that a Unified Embedding table performs significantly better in both offline and online experiments compared to highly-competitive baselines across multiple web-scale ML systems. We believe the interplay between feature multiplexing, dynamic vocabularies, power-law feature distributions, and DNNs provides interesting and important opportunities for future work.
Limitations While we expect deeper and more complicated network architectures to exhibit similar behavior, our theoretical analysis is limited to single-layer neural networks. The Pareto frontier from Section 5.1 is based on models that lag behind the current SOTA (we use 1-2 DCN layers + 1-2 DNN layers), though we provide evidence that feature multiplexing is also effective for SOTA models in Table 2. Unified embeddings impose limitations on the overall model architecture, and all embedding dimensions must share a common multiple. Because of the latency involved with looking up more than 6 components, we are often limited to 5-6 discrete choices of embedding width. Finally, unified embeddings sacrifice the ability to aggressively tune hyperparameters on a per-feature basis in exchange for flexibility, ease-of-use, and simplified feature engineering.
The authors would like to thank our amazing collaborators in helping develop the research, integrate our approaches into production, and reviewing our paper. Here is the list of contributors in alphabetical order by last name: Shuchao Bi, Bo Chang, Kevin Chang, Shuo Chen, Chen Cheng, Chris Collins, Evan Ettinger, Minjie Fan, Gang Fu, Xing Gao, Eu-jin Goh, Matthew Gonzalgo, Cristos Goodrow, Wenshuo Guo, Zhongze Hu, Gui Huan, Da Huang, Yanhe Huang, Samuel Ieong, Arpith Jacob, Maciej Kula, George Kurian, Duo Li, Yang Li, Yueling Li, Silva Lin, Liang Liu, Wenjing Ma, Lifeng Nai, Abdul Saleh, Mahesh Sathiamoorthy, Edward Shen, Rakesh Shivanna, Jaideep Singh, Ryan Smith, Myles Sussman, Jiaxi Tang, Alice Xu, Luning Yang, Tiansheng Yao, Wenxing Ye, Xinyang Yi, Angel Yu, David Zats, Fan Zhang, Jerry Zhang, Tianyin Zhou, Yun Zhou, Qinyun Zhu.
J. Wei, X. Wang, D. Schuurmans, M. Bosma, E. Chi, Q. Le, and D. Zhou. Chain of thought prompting elicits reasoning in large language models. arXiv preprint arXiv:2201.11903, 2022a.
Y. Wei, M. Langer, F. Yu, M. Lee, J. Liu, J. Shi, and Z. Wang. A GPU-specialized inference parameter server for large-scale deep recommendation models. In Sixteenth ACM Conference on Recommender Systems, pages 408-419. ACM, 2022b.
K. Weinberger, A. Dasgupta, J. Langford, A. Smola, and J. Attenberg. Feature hashing for large scale multitask learning. In Proceedings of the 26th Annual International Conference on Machine Learning, pages 1113-1120, 2009.
T. Yasuda, M. Bateni, L. Chen, M. Fahrbach, G. Fu, and V. Mirrokni. Sequential attention for feature selection. In Proceedings of the 11th International Conference on Learning Representations (ICLR), 2023.
X. Yi, J. Yang, L. Hong, D. Z. Cheng, L. Heldt, A. Kumthekar, Z. Zhao, L. Wei, and E. H. Chi. Sampling-bias-corrected neural modeling for large corpus item recommendations. In Proceedings of the 13th ACM Conference on Recommender Systems, pages 269-277. ACM, 2019.
Z. Zhao, L. Hong, L. Wei, J. Chen, A. Nath, S. Andrews, A. Kumthekar, M. Sathiamoorthy, X. Yi, and E. H. Chi. Recommending what video to watch next: A multitask ranking system. In Proceedings of the 13th ACM Conference on Recommender Systems, pages 43-51. ACM, 2019.
J. Zhu, Q. Dai, L. Su, R. Ma, J. Liu, G. Cai, X. Xiao, and R. Zhang. Bars: Towards open benchmarking for recommender systems. In Proceedings of the 45th International ACM SIGIR Conference on Research and Development in Information Retrieval, pages 2912-2923, 2022.
A.1 Missing Analysis from Section 4.1
Lemma A.1 (Weinberger et al. (2009, Lemma 2)). For any  $\mathbf{x},\mathbf{y}\in \mathbb{R}^n$ $\mathbb{E}_{h,\xi}[\langle \phi (\mathbf{x}),\phi (\mathbf{y})\rangle ] = \langle \mathbf{x},\mathbf{y}\rangle$  Furthermore, we have
$$
\begin{array}{l} \operatorname {V a r} _ {h, \xi} (\langle \phi (\mathbf {x}), \phi (\mathbf {y}) \rangle) = \frac {1}{m} \sum_ {i \neq j} x _ {i} ^ {2} y _ {j} ^ {2} + x _ {i} y _ {i} x _ {j} y _ {j} \\ = \frac {1}{m} (\langle \mathbf {x}, \mathbf {x} \rangle \langle \mathbf {y}, \mathbf {y} \rangle + \langle \mathbf {x}, \mathbf {y} \rangle \langle \mathbf {x}, \mathbf {y} \rangle - 2 \langle \mathbf {x} \circ \mathbf {y}, \mathbf {x} \circ \mathbf {y} \rangle), \\ \end{array}
$$
where  $\mathbf{x} \circ \mathbf{y} \in \mathbb{R}^n$  denotes the Hadamard product of  $\mathbf{x}$  and  $\mathbf{y}$ .
Proof. The results for  $\mu_U, \mu_H$  and  $\sigma_U^2$  directly follow from Lemma A.1, and observing that
$$
\langle \mathbf {x} \circ \mathbf {y}, \mathbf {x} \circ \mathbf {y} \rangle = \langle \mathbf {x} \circ \mathbf {x}, \mathbf {y} \circ \mathbf {y} \rangle = \langle \mathbf {x}, \mathbf {y} \rangle ,
$$
since  $\mathbf{x},\mathbf{y}\in \{0,1\}^{N_1 + N_2}$
To compute  $\sigma_H^2$ , first decompose the inner product
$$
\langle \phi (\mathbf {x}), \phi (\mathbf {y}) \rangle = \langle \phi_ {h _ {1}, \xi_ {1}} (\mathbf {x} _ {1}), \phi_ {h _ {1}, \xi_ {1}} (\mathbf {y} _ {1}) \rangle + \langle \phi_ {h _ {2}, \xi_ {2}} (\mathbf {x} _ {2}), \phi_ {h _ {2}, \xi_ {2}} (\mathbf {y} _ {2}) \rangle .
$$
Since the hash functions  $h_1, \xi_1, h_2, \xi_2$  are independent, it follows that  $\sigma_H^2$  is the sum of the variances of  $\langle \phi_{h_1, \xi_1}(\mathbf{x}_1), \phi_{h_1, \xi_1}(\mathbf{y}_1) \rangle$  and  $\langle \phi_{h_2, \xi_2}(\mathbf{x}_2), \phi_{h_2, \xi_2}(\mathbf{y}_2) \rangle$ , so we can again invoke Lemma A.1.
A.2 Derivation of Gradients
For easy reference, we re-introduce our logistic regression setting here. In Figure 4, we reproduce part of Figure 2 from the main text to further help clarify our notation.