Unified Embedding: Battle-Tested Feature Representations for Web-Scale ML Systems
Authors: Benjamin Coleman
Wang-Cheng Kang
Matthew Fahrbach
Ruoxi Wang
Lichan Hong
Ed H. Chi
Derek Zhiyuan Cheng
Venue: NeurIPS 2023 spotlight (2023)
Keywords: embedding learning; recommendation systems; representation learning
TL;DR: We show that a single embedding space can support multiple vocabularies (multiplexing), driving significant practical improvements to recommendation systems.

Abstract: Learning high-quality feature embeddings efficiently and effectively is critical for the performance of web-scale machine learning systems. A typical model ingests hundreds of features with vocabularies on the order of millions to billions of tokens. The standard approach is to represent each feature value as a $d$-dimensional embedding, which introduces hundreds of billions of parameters for extremely high-cardinality features. This bottleneck has led to substantial progress in alternative embedding algorithms. Many of these methods, however, make the assumption that each feature uses an independent embedding table. This work introduces a simple yet highly effective framework, Feature Multiplexing, where one single representation space is used for many different categorical features. Our theoretical and empirical analysis reveals that multiplexed embeddings can be decomposed into components from each constituent feature, allowing models to distinguish between features. We show that multiplexed representations give Pareto-optimal space-accuracy tradeoffs for three public benchmark datasets. Further, we propose a highly practical approach called Unified Embedding with three major benefits: simplified feature configuration, strong adaptation to dynamic data distributions, and compatibility with modern hardware. Unified embedding gives significant improvements in offline and online metrics compared to highly competitive baselines across five web-scale search, ads, and recommender systems, where it serves billions of users across the world in industry-leading products.

Unified Embedding: Battle-Tested Feature Representations for Web-Scale ML Systems
Benjamin Coleman" Google DeepMind colemann eemmmen@@ooggl..commm com
Wang-Cheng Kang" Google DeepMind wckangdgoogle.com com
Matthew Fahrbach Google Research fahrbachOgoogle com
Ruoxi Wang Google DeepMind ruoxi@google. com
Lichan Hong Google DeepMind 1i Lichan@google. com
Ed H. Chi Google DeepMind
edchi@google, com
Derek Zhiyuan Cheng Google DeepMind zcheng@google. com
Abstract
Learning high-quality feature embeddings efficiently and effectively is critical for the performance of web-scale machine learning systems. A typical model ingests hundreds of features with vocabularies on the order of millions to bil- lions of tokens. The standard approach is to represent each feature value as a d-dimensional embedding, introducing hundreds of billions parameters for ex- tremely high-cardinality features. This bottleneck has led to substantial progress in alternative embedding algorithms. Many of these methods, however, make the assumption that each feature uses independent embedding table. This work introduces a simple yet highly effective framework, Feature Mulriplexing, where single representation space is used across many different categorical features. Our theoretical and empirical analysis reveals that multiplexed embeddings can be decomposed into components from each constituent feature, allowing models to distinguish between features. We show that multiplexed representations lead to Pareto-optimal parameter-accuracy tradeoff's for three public benchmark datasets. Further, we propose a highly practical approach called Unified Embedding with three major benefits: simplified feature configuration, strong adaptation to dynamic data distributions, and compatibility with modern hardware. Unified embedding gives significant improvements in offline and online metrics compared highly competitive baselines across five web-scale search, ads, and recommender systems, where serves billions of users across the world industry-leading products.
Introduction
There have been many and exciting breakthroughs across academia and industry towards large- scale models recently. The advent of larger datasets coupled with increasingly powerful accelerators has enabled machine learning (ML) researchers and practitioners to significantly scale up models. This scale has yielded surprising emergent capabilities, such as human-like conversation and reasoning (Bubeck al., 2023; Wei et al., 2022a). Interestingly, large-scale scale model architectures are headed towards two extreme scenarios. In natural language understanding, speech, and computer vision, the transformer (Vaswani al., 2017) dominates the Pareto frontier. Most transformer parameters are
Eqqual contribution.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

located in the hidden layers, with very small embedding tables (e.g., -0.63B embedding parameters in the 175B GPT-3 (Brown al., 2020)). On the other hand, search, ads, and recommendation (SAR) systems require staggering scale for state-of-the-art performance (e.g., 12T parameters in Mudigere et al. (2022)), but here, most of the parameters are the embedding tables. Typical downstream hidden layers are 3 to 4 orders of magnitude smaller than the embedding tables. Recent work blurs the boundary between these extremes- large transformers started adopting bigger embedding tables for new types of tokens, and new SAR models leverage deeper and ever more-complicated networks. Therefore, embedding learning is a core technique that we only expect to become more critical and relevant to large-scale models in the future.
Motivated by web-scale ML for SAR systems, we study the problem of learning embeddings for categorical (sparse) features. The weekend golf warriors in the ML community may appreciate our analogy between embeddings and the golf club grip (Hogan and Wind, 1985). In golf, the quality of your grip (feature embeddings) dictates the quality of your swing (model performance). Feature embeddings are the only point of contact between the data and the model, so is critical to learn high-quality representations.
Feature embedding learning has been highly active research area for nearly three decades (Rumelhart et al., 1986; Bengio et al., 2000; Mikolov et al., 2013; Papyan et al., 2020). In the linear model era, sparse features were represented as one-hot vectors. Linear sketches for dimension reduction became popular way to avoid learning coefficients for each coordinate ( (Weinberger et al., 2009). In deep learning models, sparse features are not sketched but rather represented directly as embeddings. Learned feature embeddings carry a significant amount of semantic information, providing avenues to meaningful model interactions in the high-dimensional latent space.
In practice The simplest transformation is to represent each feature value by a row in an N x d matrix (i.e., embedding table). This approach is sufficient for many language tasks, but SAR features often have a massive vocabulary (e.g., tens of billions of product IDs in commerce) and highly dynamic power-law distributions. It is impractical to store and serve the full embedding table.
The "hashing trick" (Moody, 1988; Weinberger et al., 2009) provides a workaround by randomly assigning feature values to one of M rows using hash function. Hash embeddings (?) are a variation on the hashing trick where each feature value is assigned to multiple rows in the table. The final embedding is weighted sum of rows, similar to a learned form of two-choice hashing (Mitzenmacher, 2001). Compositional embeddings (Shi et al., 2020) look up the feature in several independent tables and construct the final embedding from the components by concatenation, addition, or element-wise product. HashedNet embeddings (Chen et al., 2015) independently look up each dimension of the embedding in a flattened parameter space. ROBE embeddings (Desai et al., 2022) look up chunks, instead of single dimensions, to improve cache efficiency. Tsang and Ahle (2022) learn sparse matrix to describe the row/dimension lookups, and deep hash embeddings (Kang et al., 2021) use a neural network to directly output the embedding.
In theory The analysis of embedding algorithms is heavily influenced by linear models and the work of Weinberger et al. (2009) on the hashing trick. This line of reasoning starts with the assumption that all inputs have "true" representation in a nigh-dimensional space that preserves all of the properties needed for learning (e.g., the one-hot encoding). The next step is to perform linear dimension reduction to compress the ground-trute truth representation into smaller embedding, while approximately preserving inner products and norms. The final step is to extend the inner product argument to models that only interact with their inputs through inner products and norms (e.g., kernel methods).
In this framework, a key opportunity for improving performance error. Early methods were based on the JohnsonnLindennstrauss lemma and sparse linear projec- tions (Achlioptas, 2003; Li et al., 2006; Weinberger et al., 2009). Since then, a substantial effort has gone into improving the hash functions (Dahlgaard et al., 2017) and tightening the theoretical analysis of feature hashing (Dasgupta and Gupta, 2003; Freksen et al., 2018). Learned embeddings now ofate-of-the-ars but the focus dimension reduction remains: most recent algorithms are motivated by dimension reduction arguments. Previous attempts to understand hash embeddings do not study the full give-and-take relationship between feature encoding and the downstream learning problem. In this work, we carry the analysis forward to better understand learned embeddings.
reduce the inneension-reductinn
For simplicity, we call this problem "feature embedding learning." but we are primarily concerned with representing sparse features such as ZIP codes, word tokens, and product IDs.

---

Contributions We present an embedding learning framework called Feature Multiplexing, where multiple features share one representation space. Our theoretical analysis and experimental studies on public benchmark datasets show that Feature Multiplexing is theoretically sound and generalizable to many popular embedding methods. We also propose a highly-efficient multiplexed approach called Unified Embedding, which is battle-tested and serves billions of users in industry-leading products via various web-scale search, ads, and recommendation models.
1. Feature multiplexing: We propose the Feature Multiplexing framework, where a single representation space (embedding table) is used to support different sparse features. All of the SOTA benchmark embedding algorithms in our experiments show better results when adopting Feature Multiplexing. Furthermore, the Unified Embedding approach (i.e., multi- plexing + hashing trick) outperforms many SOTA embedding methods that are noticeably more complicated to implement and impractical for low-latency serving systems on modern day hardware.
2. Analysis beyond collision counts:
We analyze hash embeddings in the standard framework as well as in a supervised learning setting to capture interactions between feature embedding and downstream learning. Our analysis gives insights and testable predictions about unified embeddings that cannot be obtained by dimension reduction arguments alone.
3. Battle-tested approach and industrial experiments at scale: Multi-size Unified Embed- ding has been launched in over dozen web-scale SAR systems, significantly improving user engagement metrics and key business metrics. addition, Unified Embedding presents major practical benefits: simplified feature configuration, strong adaptation to dynamic data distributions, and compatibility with ML accelerators. In Section 5.2, we share our insights, challenges, and practical lessons adopting Unified Embedding.
2 Preliminaries
SAR applications often present supervised learning problems over categorical (sparse) features that represent properties of users, queries, and items. Each categorical feature consists of one more feature values v drawn from the vocabulary of the feature V. For concrete example, consider the classic click-through rate (CTR) prediction task for ads systems, where we want predict the probability of a user clicking on an advertisement. The "ad_ _ad_id" feature may take on hundreds of billions of possible values-ont for each unique ad served by the platform- while "site_ id" describes the website where that ad was shown (a.k.a,, the publisher).
We embed the values of each feature into space that is more suitable for modeling, c.g., Ed or the unit (d- -sshere Given a vocabulary of N tokens V = {01, 00, U2,.... UN }. we learn a transformation that maps e E V to an embedding vector e E Idd After repeating this process for each categorical feature, we get a set of embeddings that are then concatenated and fed as input to a neural network (Covington et al., 2016; Anil et al., 2022).
Problem statement We are given a dataset D = {(x)) y1), (æ2, v2),..., (ar D), UID }} of exam- ples with labels. The examples consist of values from T different categorical features with vocabular- ies (Vi, V2.... }. where feature has vocabulary V,, Unless otherwise stated, we do not consider multivalent features and instead assume that each example æ has one value for each categorical fea- ture, i.e., [u), u2..... 2T] where E V.. However, this is for notational convenience only techniques in this paper extend to missing and multivalent feature values well.
Formally, the embedding table is a matrix E E MM xd and the embedding function gxx; E) transforms the categorical feature values into embeddings. We let h(v)) : V - [M] be a 2-universal hash function (Vadhan et al., 2012 2) that assigns feature value to row index of the embedding table. We also consider a model function f(e 0) that transforms the concatenated embeddings into a model prediction, which is penalized by loss f. Using this decomposition, we define the joint feature learning problem as follows:
arg min (E, where (D( = Σ E); E,0
(1)
(D D
We use independent hash functions (v()) for each feature t E [T] so that tokens appearing in multiple vocabularies have the chance to get mapped to different locations, depending on the feature. We use the notation em to denote the m-th row of E, which means ea(x) is the embedding of u. Lastly, let 1n be the indicator variable that denotes a hash collision between and v, i.c. h(u) h(v).
3

---

Feature 1
Feature Feature
Feature Feature
Feature 2 Feature
V11
%g
Kg
Feature 2
V/2
V2.
V2
V9
V5
V4
V4
Collisionless
Hash Table
Unified Table
Multisize Unified
Figure I: Embedding methods for two categorical features. We highlight the lookup process for the first value en of each feature. Hash tables randomly share representations within each feature, while Unified Embedding shares representations across features. To implement Unified Embedding with different dimensions (multi-size or variable-length), we perform multiple lookups and concatenate the results.
Feature Multiplexing: The Inter-Feature Hashing Trick
It is standard practice to use a different hash table for each categorical feature. There are a few approaches that deviate from this common practice-for example, HashedNet and ROBE-Z models use parameter sharing across all features, depending on the implementation. However, conventional wisdom suggests that each categorical vocabulary benefits from having an independent representation, so shared representations are not well-studied.
Parameter reuse We present a Feature Multiplexing framework that only uses one embedding table to represent all of the feature vocabularies. In contrast, typical models in SAR systems have hundreds of embedding tables, each representing one input feature. Multiplexed embeddings allow models to reuse parameters and lookup operations, improving efficiency terms of both space and time. This is critical for models that must fit into limited GPU/TPU memory or satisfy latency and serving cost constraints. In principle, any feature embedding scheme (e.g., hashing trick, multihash, compositional, or ROBE-Z) can be used as the shared representation for feature multiplexing. In Section 5, we experiment with multiplexed versions of six different embedding algorithms. Here, we focus multi-size (multi-probe) version of the multiplexed feature hashing trick (Figure I), which we call Unified Embedding. Unified Embedding presents unique practical advantages that we describe in Section 5.2.
Tunable embedding width By using a unified table, we make the architectural assumption that all features share the same embedding dimension. However, this is often sub-optimal in practice since independent tuning of individual embedding dimensions can significantly improve model performance. For use-cases where different features require different embedding dimensions, we concatenate the results of a variable number of table lookups to obtain the final embedding. While this limits the output dimension multiples of the table dimension, this a limiting constraint (and is needed by other methods such as ROBE-Z and product quantization).
Shared vocabulary is common for categorical feature vocabularies to share tokens. For example, two features that represent character bigrams of different text fields (e.g., query and document) might have substantial overlap. Our default approach is to use a different hash function for each feature by using a different hash seed for each feature. However, in cases where many cemantically-simile features are present, the performance may slightly improve by using the same hash function.
4 Analysis of Feature Multiplexing
In this section, we present theoretical arguments about feature multiplexing applied to the "multiplexed feature hashing trick" setting. We start with brief analysis in the dimension- reduucioon framework, suggesting that multiplexing can optimally load balance feature vocabularies across buckets in the table. This classic framework, however, cannot be used to distinguish between collisions that happen within the vocabulary of a feature (i.e., fetra-feature) and those occurring across features (i.e., feterrfeature).

---

To understand these interactions, we also need to consider the learning task and model training. Thus, we analyze the gradient updates applied to each bucket under simplified but representative model- binary logistic regression with learnable feature emheddings for CTR prediction. Our analysis demonstrates a novel relationship between the embedding parameters and model parameters during gradient descent. Specifically, we find that the effect of inter-feature collisions can be mitigated if the model projects different features using orthogonal weight vectors. In Section 4.2, we empirically verify this by training models on public benchmark on click-through prediction (Criteo) and finding that the weight vectors do, in fact, orthogonalize. We defer the proof' and detailed derivations in this section to Appendix A.
4.1 Parameter Efficiency in The Dimension-Reduciio Framework
In the feature hashing problem, we are given two vectors x, y E EN that we project using linear map o: R. -> RM given by the sparse {+1} matrix in Weinberger al. (2009). The goal is for to minimally distort inner product, i.e., (o(x), o(y)) = (x. Since the estimation error directly depends on the variance, it is sufficient to compare hashing schemes based on the moments of their inner product estimators, We examine the variance of the hashing trick and multiplexing trick.
To analyze multiple features once in this framework, we consider the concatenations x =[x1, X2] and y = [21. The components X1 E {00 13.2i and x² E {0, 1).2 are the one-hot encodings (or bag-of-words) for the two vocabularies V1 and V2. The standard hashing trick approximates (x. by independently projecting
to dimension M1 and X2 to dimension My. The multiplexed version uses a single projection, but into dimension M1 + M2. Now, we formalize the hashing trick of Weinberger et al. (2009), but with slightly revised notation. Given a set of tokens from the vocabulary. the embedding the signed sum of value counts within each hash bucket.
Definition 4.1. Let h: V + (1, .M be 2-universal hash function and : V - +1, be a sign hash function. The function dn,. 2V + AM is defined as on.e (W)== -wew ((w) where u, is the i-th unit basis vector in AR..
The following result compares multiplexed and standard embeddings in the dimension reduction framework. Note that the randomness is over the choice of hash functions li(v) and E(v). This result also directly extends to concatenating T feature embeddings.
Proposition 4.2. For any X1. E 10,1 andx2, E [0,1]N², lerx= = x1, X2), y [2] denote their concatenations. Let μu, PH. of and JH be the mean and variance of (ón.e (x), on.g (y)) for multiplexed and hash encodings, respectively. Then, Pe = μh = (x,y) and
x
yy + (x, - y(x,
o
M1
M2
x
y1
2(x1, + (x1. -
X2 1 + (x), 2(x2, y2)
0
Mi
M2
Observations Proposition 4.2 shows that multiplexed embeddings can do good job balancing hash collisions across the parameter space. Consider the problem of estimating the inner product (x, y)- Suppose that ||x, [/ = ||y: 1 kn and ||x2|l ||y> = l2 i.e., features 1 and 2 are multivalent with k and ky values, and assume x and y are orthogonal, i.e., (x,y) = 0. Proposition 4.2 tells us that o (k1 k2)² (M, + M2) and oH kī/ /M1 + k3 M2/ The difference o oi factors as (v Mi / Maki / M2/ MMaMMKa)2/(MM + M2). which is neganvga,vve implying or < oî
2
4.2 Trajectories of Embeddings During SGD for Single-Layer Neural Networks
Intuition suggests that inter-feature collisions are less problematic than intra-feature collisions. For example, if x E Vi collides with y E V/2 but no other values collide across vocabularies, then the embeddings from other tokens in 112 and 1/2 are free to rotate around the shared embedding for x, y to preserve their optimal similarity relationships.
We analyze 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 Eq. (1) where (2; 0) is the sigmoid function o (z) 1/(1 + exp(-(z, (z, 0))). é is the binary cross-entropy loss, and y E 1) (e.g., click or non-click). We concatenate the embeddings of each feature,
which means the input to fzz 0) is g( E) (en, (ha(w2): ...., (hr(wr) where (a.) is
5

---

Critea Embeddling Norms
Criteo Weight Orthogonality
Heas Emoedting Unitted Embedding
Feature
En(u)
1
Entoodding Norm
fidd Erocoding
+
Feature 2
E(())
02
Cuv
Table Size
Total Number Embeddings (M)
Figure 2: Single-layer neural embedding model with per-feature weights Or (left). Mean embedding fo-norm (middle) and mean angle between all pairs of weight vectors O O1 (right) as a function of table size for Criteo across all 26 categorical features. Note that the horizontal axes are in log scale.
the embedding for the t-th feature value in example æ. We write the logistic regression weights as [01. ....... , 0y] so that embedding (m,) for feature t is projected via 0, E RMM We illustrate this simple but representative model architecture in Figure 2.
The following analysis holds in general for features, but we consider the problem for two categorical features with vocabularies V1 and V2 to simplify the presentation. We write the logistic regression objective as follows in Eq. (2), using Cn...11 and Ca to denote the number of examples for which [u, uj and y and y O, respectively. We present full derivation in Appendix A.2.
(D - Σ
Σ
Cu.o. log exp(O' e...)) - (Cu,v,0 + Cu.v.) log(1 + exp((0 eu,.)) (2)
WEV, VEV2
To understand the effect of hash collisions on training dynamics, we take the gradient with respect to (u), i.e., the embedding parameters that represent u E V1. By studying the gradients for collisionless, hash, and multiplexed embeddings, we quantify the effects of intra-feature and inter- feature collisions. Interestingly, all three gradients can be written using the following three terms:
V (mn Lp (E 01 Σ Cu,v,D (Cu,e,0 + Cu,v,1)oo(Cu.)))
VEVs
u.w Cu,o,o (Cu, + C (o,,,))o((e..))
weV ui
DEV2
+ O2
Cm, - (Cm,v.0 + Cm.,, 100
(5)
vEV2
WEV,
For collisionless embeddings, component (3) is the full gradient. For hash embeddings, the gradient the sum of components (3) and (4). The multiplexed embedding gradient is the sum of all three terms.
Insights Our core observation 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 the true component, implying that intra-feature collisions This agrees with intuition since intra-feature hash collisions effectively act as value merging.
resolvable by the model.
Inter-feature collisions, however, push the gradient in the direction of O2, which creates opportunity for model training to mitigate their effect. To see this, consider situation where 0, and Oz are initialized to be orthogonal and do not change direction during training. During SGD, the embedding eb is linear combination of gradients over the training steps, which means e (w) can be decomposed into true and mntra-fcature component, in the O2 direction, and inter-feat inter-feature component the 0₂ direction. Since (01, 02) 0, the projection of 0 Ch,( effectively eliminates the inter-feature component.
"This restriction does not affect the representation capacity because learned embeddings can rotate around O, i.c., for any 0, and nonzero constant we can learn some e' that satisfies (e', u) (e. 0).

---

Empirical study These observations allow us to form two testable hypotheses. First, we expect the projection vectors [0,, 02...................... ................................ 0r] to orthogonalize since this minimizes the effect of inter-feature collisions (and improves the loss). The orthogonalizationn effect should be stronger for embedding tables with fewer hash buckets since that 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 [[[,,]] = 1/M. Therefore, we expect the squared embedding norms to scale roughly as O(N/M).
To investigate these hypotheses, we train single-layer neural network on the categorical features of the Criteo click-through prediction task. We initialize 0 = [01. 02.... 0yl with each 0, com- ponent 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 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 single-layer neural network because different features are processed by different model parameters. While our theoretical results are limited 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.
5 Experiments
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 -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- real-valued continuous variables. We only embed the categorical features. Avazu is similar advertisement click dataset but with -36 million examples (11 days of data). Each example here contains 23 categorical features and no continuous features. Movielens is traditional user-item collaborative filtering problem, which we convert to a binary prediction task by assigning "1"s to all examples with rating 2 3, and "0"s to the rest. For more details, see the benchmarks 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 neural network with different embedding representation methods at 16 logarithnically-saced 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 basclines, 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 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 I.
7

---

Implementation and hyperparameters All methods are implemented in Tensorflow 2.0 using the TensorFlow Recommenders (TFRS) framework." For 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 E {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., feature with 10% of the total vocabulary receives 10% of the memory budget).
Table 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 con-multiplexed methods. We observe that feature multi- plexing is Pareto-optimal. Interestingly, the multiplexed hashing trick outperforms several embedding techniques (e.g.. ROBE-Z) that were previously SOTA.
Results
Collisionless embeddings are included in our benchmark evaluation only to provide an optimistic headroom reference point. Collisionless tables often impossible deploy in industrial recom- mendation systems since feature vocabularies can be massive and change dynamically due to churn. Multiplexing provides significant improvement 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 В also contains additional details about Criteo and Avazu, including explanation for the collisionless-hat embedding performance gap on Criteo.
Table I: AUC of embedding methods on click-through datasets. We give results for several points on the parameter-accurage 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 T* and 0.05* significance levels of Welch's 4-test.
Criteo
Avazu
Movielens-
Method
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
ROBB-Z 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 Multiplex PQ-one 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.
Phttps://gitaan htw.://aw..mm//eeoooeo////////ee

---

Criteo
Avazu
Movielens
6.744
0.32
0.773
0.03
0.722
0.07
AUC (%)
0.711
0.03
0.70
0.302
1154
0.79
0.747
32"
1P
14'1
Memory (MB)
Memory (MB)
Memory (MB)
Collisionless
Standard
Multiplexed
Figure 3: Pareto frontier of multiplexed and original (non-multiplexcd) 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-ar 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 parame- ters (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 AutoMI. 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.
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 basis- accomplished via AutoML (Bender et al., 2020) attention- based search methods (Yasuda et al., 2023; Axiotis and Yasuda, 2023). This yields roughly a 50% reduction in the number of hyperparameters.
Simple configuration:
Adapt to dynamic feature distributions: In practice, the vocabulary size of each feature changes over time (new IDs enter the system on 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, significant number of new products are added during the holiday season compared to off-season. Thus, a fixed number of buckets for 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. was applied 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 metries 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 ge.g,, suffers from high churn rate). For instance, an overwhelming number of new short-form videos are produced daily, creating 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 DCN-V2
Products
Short-form videos MMOE -1B AUC +0.39% (task 1) RMSE -0.53% (task 2) +0.62%
Apps DCN-V2 -10M
Apps DCN-V2
Architecture
Two-tower -10B
Vocabulary Size
10BB
-10M
Offine Metric
AUC +2.2%
Recall@1 1 +7.3%
AUC +0.25%
AUC +0.17%
Online Metric
Positive
Positive
+0.44%
+0.11%
Conclusion
In this work, we propose highly generalizable Feature Mulriplexing framework that allows many different features to share the same embedding space. One of our crucial observations is that models can learn 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 Pareto-optimal tradeoff, and Unified Embedding proves 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 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
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 often limited to 5-6 discrete choices of embedding width. Finally, unified embeddings sacrifice the ability to aggressively tune hyperparameters on per-feature basis in exchange for flexibility, ease-of-use, and simplified feature engineering.
Acknowledgements
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 leong, 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.
10

---

References
D. Achlioptas. Database-friendly random projections: Johnson-LLindeenttaan with binary coins. J. Compud. Syst. Sci., 66(4):671-687, 2003.
R. Anil, S. Gadanho, D. Huang, N. Jacob, Z. Li, D. Lin, Phillips, C. Pop. K. Regan, G. I. Shamir, R. Shivanna, and Q. Yan. On the factory floor: ML engineering for industrial-scale ads recommendation models. Systems and User Modeling co-located with the 16th ACM Conference Recommender Systems, ORSUM@ RecSys, volume 3303 of CEUR Workshop Proceedings, 2022.
Proceedings of the Sth Workshop on Online Recommender
K. Axiotis and Yasuda. Performance of regularization for sparse convex optimization. arXiv preprint S.tv:2337.00400000 2023.
G. Bender, H. Liu, B. Chen, G. Chu, S. Cheng, P. Kindermans, and Q. V. Le. Can weight sharing outperform random architecture search? An investigation with TuNAS. In 2020 IEEF/CVF Conference on Computer Vision and Pattern Recognition, pages 14311-14320. Computer Vision Foundation IEEE, 2020.
Y. Bengio, R. Ducharme, and P. Vincent. A neural probabilistic language model. In Advances in Neural Information Processing Systems, pages 932- 998, 2000.
T. B. Brown, Mann, N. Ryder, M. Subbiah, Kaplan, Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, S. Agarwal, A. Herbert- Voss, G. Krueger, T. Henighan, R. Child, A. Ramesh, D. M. Ziegler, J. Wu, C. Winter, C. Hesse, M. Chen, E. Sigler, M. Litwin, Gray, Chess, J. Clark, C. Berner, S. McCandlish, A. Radford, I. Sutskever, and D. Amodei. Language models are few-shot learners. In Advances Neural Information Processing Systems, volume 33, 2020.
S. Bubeck, V. Chandrasekaran, R. Eldan, J. Gehrke, Horvitz, E. Kamar, Lee, Y. T. Lee, Y. Li, S. Lundberg, et al. Sparks of artificial general intelligence: Early experiments with GPT-4. arXiv preprint arXiv:2333.12712 2023.
O. Chapelle. Modeling delayed feedback in display advertising. In The 20th ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, pages 1097-1 105. ACM, 2014.
W. Chen, J. Wilson, S. Tyree, K. Weinberger, and Y. Chen. Compressing neural networks with the hashing trick. International Conference on Machine Learning, pages
PMLR, 2015.
Covington, J. Adams, and E. Sargin. Deep neural networks for YouTube recommendations. In Proceedings of the 10th ACM Conference on Recommender Systems, pages 191-198 2016.
S. Dahlgaard, M. Knudsen, and M. Thorup. Practical hash functions for similarity estimation and dimensionality reduction. Advances in Neural Information Processing Systems, 30, 2017.
S. Dasgupta and A. Gupta. An elementary proof of a theorem of Johnson and Lindenstrauss. Random Structures & Algorithmss, 22(1):60-65 2003.
A. Desai and A. Shrivastava. The trade-offs of model size in large recommendation models: 100GB to 1OMB Criteo-tb DLRM model. In Advances in Neural Information Processing Systems, 2022.
A. Desai, L. Chou, and A. Shrivastava. Random offset block embedding (ROBE) for compressed embedding tables in deep learning recommendation systems. Proceedings of Machine Learning and Systems, 4:762-778, 2022.
M. Fahrbach, A. Javanmard, V. Mirrokni, and P. Worah. Learning rate schedules in the presence of distribution shift. In Proceedings of the 40th Annual International Conference Machine Learning, pages 9523-9546. PMLR, 2023.
C. B. Freksen, L. Kamma, and K. Green Larsen. Fully understanding the hashing trick. Advances in Neural Information Processing Systems, 31, 2018.
M. Harper and J. A. Konstan. The movielens datasets: History and context. ACM Transactions on Interactive Intelligent Systems, 5(4):1-19, 2015.
B. Hogan and H. W. Wind. Five Lessons: The Modern Fundamentals of Golf. 1985.
11

---

N. P. Jouppi, G. Kurian, S. Li, P. Ma, R. Nagarajan, L. Nai. N. Patil, S. Subramanian, A. Swing, B. Towles, et al. Tpu v4: An optically reconfigurable supercomputer for machine learning with hardware support for embeddings. Proceedings of the 50th International Symposium of Computer Architecture, 2023.
W.-C. Kang, D. Z. Cheng, T. Yao, X. Yi, T. Chen, L. Hong, and E. H. Chi. Learning embed categorical features without embedding tables for recommendation. In Proceedings of the 27th ACM SIGKDD Conference on Knowledge Discovery Dala Mining, pages 840-850, 2021.
P. Li, J. Hastie, and K. W. Church. Very sparse random projections. In Proceedings of the ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, pages 287-296, 2006.
T. Mikolov, I. Sutskever, K. Chen, G. S. Corrado, and J. Dean. Distributed representations of words and phrases and their compositionality. In Advances in Neural Information Processing Systems, pages 3111-3119, 2013.
M. Mitzenmacher. The power of two choices in randomized load balancing. IEEE Transactions on Parallel and Distributed Systems, 12(10):1094-11) 1104, 2001.
J. Moody. Fast learning in multi-resolution hierarchies. Advances in Neural Information Processing Systems, 1988.
D. Mudigere, Y. Hao, J. Huang, Z. Jia, A. Tulloch, S. Sridharan, X. Liu, M. Ozdal, J. Nie. J. Park, L. Luo, J. A. Yang, L. Gao, D. Ivchenko, A. Basant, Y. Hu, J. Yang. E. K. Ardestani, X. Wang, R. Komuravelli, C. Chu, S. Yilmaz, H. Li, J. Qian, Z. Feng, Y. Ma, J. Yang. E. Wen, H. Li, L. Yang, C. Sun, W. Zhao, D. Melts, K. Dhulipala, K. R. Kishore, T. Graf, A. Eisenman, K. K. Matam, A. Gangidi, G. J. Chen, M. Krishnan, A. Nayak, K. Nair, B. Muthiah, M. khorashadi, tacharya, P. Lapukhov, M. Naumov, A. Mathews, L. Qiao, M. Smelyanskiy, B. Jia, and V. Rao. horrware co-design for fast and scalable training of deep learning recommendation models. In The 49th Annual International Symposium on Computer Architecture, pages 993-1011. ACM, 2022.
Bhat-
M. Naumov, D. Mudigere, H.-J. M. Shi, J. Huang, N. Sundaraman, J. Park, X. Wang, U. Gupta, C.-. Wu, A. G. Azzolini, et al. Deep learning recommendation model for personalization and recommendation systems. arXiv preprint arXiv: /906. 1906.0001, 2019.
V. Papyan, X. Han, and D. L. Donoho. Prevalence of neural collapse during the terminal phase of deep learning training. Proceedings of the National Acudemy of Sciences, 117(40))4652-24463,, 2020.
D. E. Rumelhart, G. E. Hinton, and R. J. Williams. Learning representations by back- proopppppggg errors. Nature, 323(6088):533, 1986.
H.J. M. Shi, D. Mudigere, M. Naumov, and J. Yang. Compositional embeddings using comple- mentary partitions for memory-efficient recommendation systems. Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, pages 165-175, 2020.
W. Song, Shi, Z. Xiao, Z. Duan, Xu, M. Zhang, and J. Tang. Autoint: Automatic feature inter- action learning via self-attentive neural networks. In Proceedings of the 28th ACM International Conference on Information and Knowledge Management, pages 1161-1170, 2019.
H. L.-H. Tsang and T. D. Ahle. Clustering embedding tables, without first learingg them. arXiv preprint arXiv:2210.05974 2022.
S. P. Vadhan al. Pseudorandomness. Foundations and Trends in Theoretical Computer Science, (1-3):1-336, 2012.
A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, ł.. Kaiser, and I. Polosukhin. Attention is all you need. Advances in Neural Information Processing Systems, 2017.
R. Wang, R. Shivanna, D. Z. Cheng, S. Jain, D. Lin, L. Hong, and E. H. Chi. DCN V2: Improved deep & cross network and practical lessons for web-scale learning to rank systems. In The Web Conference, pages 1785-1797. ACM/ IW3C2, 2021.
12

---

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 arXi::220.111003, 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 019.419.oo ACM, 2022b.
K. Weinberger, A. Dasgupta, Langford, A. Smola, and 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. Proceedings of the 11th International Conference on Learning Representations (ICLR), 2023.
X. Yi, Yang. L. Hong. D. Z. Cheng, L. Heldt, A. Kumthekar, 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. We Chen, A. Nath, S. Andrews, A. Kumthekar, M. Sathiamoorthy, X. Yi, and E. H. Chi. Recommending what video watch next: multitask ranking system. Proceedings of the 13th ACM Conference on Recommender Systems, pages 43-51. ACM, 2019.
J. Zhu, Q. Dai, L. Su, R. Ma, 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.
13

---

so that embedding (x,(x.) for A Theory
A.1 Missing Analysis from Section 4.1
Lemma A.1 (Weinberger et al. (2009, Lemma 2)). For any x, y IR., (n,e((((x), (x,y). Furthermore, we have
Var ((o,,(((()))
1 2t + ,,
Σ
m
1 (xx, (y.y) + (x. y) y, y,
m
where x y E R denotes the Hadamard product of x and y-
Proposition 4.2. For any X1, S1 E 10,1 and X2, Y2 E (0.1)% 13 letx (x1,X), X2), [y1. 82) denote their concatenations. Let yuu M.H of and on be the mean and variance of ( (x)x 2h. (y)) for multiplexed and hash encodings, respectively. Then, Pu = Pn = (x, and
ot = x 1 + (x, y)2 y(x,
Mi + M2
| 1 (x). v1x
|xa|B1lye212 + (x2, y2) 22)
o =
M1
M2
The results for pu. pa and of directly follow from Lemma A. and observing that Proof.
(xo = yx),, = (x.
since x, E (0.1)N
To compute o first decompose the inner product
(0(x), (by)) (on &i (x1), &E (vr)) (x2), On (y2)).
Since the hash functions f1j E1, f22 &2 are independent, follows that of is the sum of the variances of (0h, (x1), & (yr)) and (oha (x2), &h (y2)), so we can again invoke Lemma A.I. O
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 further help clarify notation.
Recall from the main text that we are interested a logistic regression model with trainable embed- dings for binary classification (i.e., single-layer neural network with hashed one-hot encodings as input). This corresponds to our supervised learning framework (Eq. (1)), where f(z; 0) is the sigmoid function (z) = 1/(1 + exp( - (z, 0))) and is the binary cross-entropy loss.
We concatenate the embeddings of each feature, which means that the input to f(z: 0) is g(x; E) (x1), (kz())))____.._..))))))))>)))))) (xr)) where (x)())) is the embedding for the (-th feature value example x.
We also partition the logistic regression weights 0 = =0,, feature t is projected via 0, E R.M.
For notational convenience we also partition the dataset by label. We write Do {(x,, vi) E D 0} and Dj = ((x,, v:) E D y 1}. Using this notation, we can write the logistic regression
objective follows:
(D.
Σ Ex
log
1
log
I exp( g(:: B))) (
[+exp(((00
( g(,
(x.))E
We proceed by writing the loss sum over vocabulary pairs rather than examples. We consider the problem for two categorical features with vocabularies Vi and V/2 simplify the presentation, and we use the notation Cu,v. to denote the number of co-occurences of (u, v) Vi x V2 in dataset Do.
14

---

Feature Feature 1
En(u))
01
dleuv' 0)
Ey(v)
0z
u 0
Figure 4: Illustration of the logistic regression model, annotated with notation that corresponds to our theoretical analysis.
Likewise, Ca,e,1 denotes the corresponding quantity for D₁. We also let Cu,e (w) Cha(u)) so that the inner product ([01, 02), (u), (hg(w))) can be written concisely as 0 eu,v. It follows that
() - - Σ o log
1 exp( 1opp--0aa "Te.,.)) + Ca,v,1 log
+ exp(O em..))
WEVi VEVS
After combining the sigmoid functions, we get the following expression:
(D( =
- Σ Σ
Cu,u, log exp(0" eu,v) (Cu,v,0 + log, + exp(0 ew,))
VE VEV,
To get the gradients with respect to Ex(w) (the embedding parameters used to represent value u), we first take the gradient with respect to Em These parameters accumulate gradient contributions from any value assigned to bucket m
(i.e. for which 1um is one). then straightforward to express VB,co E, 8) in terms of the other values that collide with u (i.e. for which La, is one).
Collisionless Embeddings: Because each row Em corresponds to single value, we can take the gradient with respect En(w) directly.
V E (D(E, 0) 01
Σ Cu,,o (Cu,v,o + Cu.,))ao(eu.o) veV/
Hash Embeddings: Here, row m corresponds to any value that is assigned to bucket m (i.e. for which Lam is one). We find the gradient for the embedding table corresponding to the first feature Vi.
V (E(
Σ 1
In.m O2
Cu,v,0 - (Cu,v.o + Cu,,, (oo1)))oaee)) (eu,))
WEV,
JEVS
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 speefic value Vi, we can write this gradient as follows.
V V (E, 01
Σ
Cu, - (Cu,v,0 + Cu.,.))oo(eu..)
vEVs
0,
Σ WEV,
I
u,w
Cunro (Cu,w,, + Cu,,,1)o0(een))
DEV.
wyFu
Unified Embeddings: Here, row might correspond to any value even ones from other features. The gradient with respect to Em now contains additional components.
E., CD (E,
Σ
Or
Σ
Ca,v, o (Cu,e,o + (o,,
NEV
"eVa
Σ
L LumO0
Σ
Cu,,oo - (Cu,v.0 + Ca.e.))o((en.))
V
WEV,
15

---

Feature
Feature
Feature
Feature 2
Feature 2
Feature
Collisionless
Hashing Trick
Unified Table
memory
Shared memory
V1
Feature
Feature
V1
Feature 2 Feature
Vi
Noter not ad contnccioons oe chown
Nori: comnectine no show
V1
Weights
Feature
V2.
etv,)
ROBE/ / NashedNet
ROBE-Z
MultiHash
Feature
V1
Feature
V6
e(vy) -
e(v;)
Compositional (QR)
Compositional (PQ)
Figure 5: Illustration of baseline embedding methods, all which are compatible with feature multi- plexing.
Rewriting the gradient be in terms of En(u) (for Vi) reveals that in the unified embedding case, the gradient contains the collisionless component, an intra-feature component, and an inter-feature component:
VEN CD E 0) = 01 Σ
Cu.v,0 - (Cu,v,0 + (o,,,)))p(eu,)))
VVV2
+ 01
Σ
Σ
C (Cur,v,0 + Doo (en.v)
u,w
WEV,
vEV2
+ 02
1
Cm,o.0 (Cm,v,0 + Ca, (00)))
.EV.
WEV,
B Experiments
Figure illustrates our baselines. We describe the methods and their hyperparameters below.
Collisionless tables assign each value of each feature unique d-dimensiinal representation. The only hyperparameter is d, the embedding dimension.
16

---

2. Hashing trick tables begin with fixed number of representations and allow values to collide within features. Here, each table has two hyperparameters: (the dimension) and M (the table size). The inspiration for this method is most commonly attributed to Weinberger et al. (2009), and 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 element- wise multiplication of the resulting embeddings. The hyperparameters are & (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 concate- nation. 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 L dimensions in shared memory (we use = 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, 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 learned weighted combination. Multihash tables have several hyperparameters and are therefore trickier to tune well. The parameters are M1, the size of the embedding table, M2, the size of the importance weight table, ks, 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 logaritmmically-spaeee 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 independent training runs for each hyperparameter configuration. For non-multiplexed representations, we allocate parameters to features based on vocabulary size. For example, 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., M1) versus the importance table (e.g., Mz). For example, p = 0.2 means that the embedding table are allocated 80% of the budget and the importance weight tables receives the remaining 20%. We chose 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 compo- nents & from [2, 3, 4].
17

---

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 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. 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 Criteo 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 presented in Figure 7.
be achieved with a well-tuned 3-layer feedforward network with 748
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 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 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 learning rate of 0.0002. The full parameter-accurage tradeoffs are presented in Figure
Movielens: We use the Movielens- 1M dataset with the same preprocessing used by Wang al. (2021). We only use the "userld," "movield," " "Zipcode,"" " "age," "occupation" and "gender" features. We always use a collisionless table embed the "gender" feature due 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 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. Critco differs from Avazu and Movielens in terms of vocabulary distribution it has heavier-taile vocabulary than Avazu or Moviclens. This worsens the errors introduced by hash collisions because 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 mon-multiplexed baseline (rather than the strongest baseline), the following differences are significant with 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 kB and 3.24 MB columns. Multiplex QR in the 32.4 MB column.
3. Movielens: All results in the 158 kB column. All results in the 791 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.
18

---

Table 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
Criteo
Avazu
16*
Movielens
log(Frequency)
100
Heavier Tail
Vocabulary (Log ID)
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 much heavier tail.
In the Table 3, we report the average number of steps / second for several methods on Criteo for the 25 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 job load demand, these numbers have high variance (a >> 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 experimntss-muliilexiiin did not significantly increase the training inference time.
B.1 Industry experiments
Hardware and training time: In several of our Unified Embedding deployments (Table 2), 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 there enough memory to support the embedding tables (i.c., enough CPU RAM or total accelcrator memory), the training time mostly governed by the forwards and backwards passes the rest of the (upstream) network. We do not observe a substantial increase training / inference time on either CPU or TPUv4 (Jouppi al., 2023).
Modeling tradeoffs: Unified embedding gives up some flexibility in setting the per-feature embed- ding dimension, and other (more complicated) multiplexed methods are better on the Pareto-frontier for offline experiments. In exchange, we get more per- per-feature embedding capacity, much simpler
19

---

Criteo
Criteo
Criteo
3.0
Test AUC
Test AUC
2.00
Collistanless Hashing
Collitiniass
Collisioness
1.75
Emoedding
Muttplex Hashing trik
Muttppix HashEmb
Mutiplex HashedNet
Memory (MB)
Memory (MB)
Memory (MB)
Criteo
Criteo
Criteo
Test AUC
Test AUC
Test
Collisionles
Collision ess Embedding
Collsioniess OR Embedding
ROBE
Muttilex ROBE-Z
Multipiex PQ
Multipex QR
Memory (MB)
Memory (MB)
Memory (MB)
Figure 7: Full parameter-accuracc tradeoffs for all methods on the Criteo dataset. Naming convention the same as in Table I, 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.
Criteo
Avazu
Avazu
4.74
Test AUC
Test AUC
Collisionless Hashing trick
Collisinnlass
Collisionless
Hash Embedkiing
HashedNet
6.1
Mutiplex Hashing Trick
378
Muttipixx HachEmb
k20
Muttplex HashodNet
Memory (MB)
Memory (MB)
Memory (MB)
Avazu
Avazu
Avazu
4.72
Test AUC
Test AUC
Test AUC
Collisionless ROBE
Collsionless PQ Embedding
Collssioness
QR Embedding
Muttiplox ROBE-2
Multiplex
Multipexx OR
Memory (MB)
Memory (MB)
Memory (MB)
Figure 8: Full parameter-accuracy tradeoffs for all methods on the Avazu dataset. Naming convention is the same 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 model training runs.
hyperparameter configuration, hardware-friendly algorithm. and easy feature engineering. For exam- plc, to add new features to a model without multiplexing, we must specify and tune the table size and dimension for new set of embedding tables. With multiplexing, 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 model update that migrates several new features into the unified embedding (typically resulting in performance gains).
can simply add new feature to
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
Engineering tradeoffs:
20

---

Movielens
Movielens
Movielens
Test Test AUC
120
Collisioniess
Callisoriess
Collisionless
Heshing Trick
178
Hosh Embeddinng
MashedlNet
850
Multiplex Hashing Trick
J.0
Multiplex HashEmb
Multiplex HashedNett
Memory (MB)
Memory (MB)
Memory (MB)
Movielens
Movielens
Movielens
0.53
Test AUC
Test
Collisionless
Collisionless Embedcing
Collisionless
QR Embedding
Multipiex ROBE Z
Multiplex PQ
Multiplex QR
6.4
Memory (MB)
Memory (MB)
Memory (MB)
Figure 9: Full parameter-accuracy tradeoffs for all methods on the Movielens- M 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 eгrоr bar width denotes the standard deviation over model training runs.
embedding rows are distributed over cores). The high-level takeaway is that 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 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.
21

---

Table 4: Criteo vocabulary pruning.
Feature
Vocahulary Size
Table Avazu vocabulary pruning.
14
676
Feature
Vocabulary Size
533
16
17447
C1
8
17
19995
C14
2309
18
180
C15
9
C16
9693
C17
405
21
337
C18
5
22
C19
66
23
14637
C20
167
C21
25
17795
app_category
29
26
3067
app_domain
277
27
26
app_id
4438
6504
banner_po pos
8
18679
device_conn_type
5
30
10
device_id
67767
31
3102
devicc_ip
163804
32
1557
device model
62.7
device type
34
18230
site category
24
35
10
stt__domain
3887
36
14
site_id
3317
13079
hour
56
39
10581
22


Reviews:
[{"review_id": "8XVe7LEYbN", "reviewer": "Reviewer_fvbA", "summary": "The paper introduces a novel \"Feature Multiplexing\" framework which uses a shared representation space (embedding table) for multiple sparse features. This approach aims to find a balance between model size and accuracy for industrial level recommender system. Besides, the authors provide a theoretical analysis, highlighting that inter-feature collisions can be alleviated if features are projected using orthogonal weight vectors. Further gradient analysis reveals that these collisions are not uniformly detrimental; the adverse effects can be mitigated when features are processed by distinct parameters in a single-layer neural network.", "questions": "Please refer to weakness. ", "limitations": "The authors did not discuss the limitations at all.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "**Pros**:\n\n- The paper presents the a novel \"Feature Multiplexing\" framework, offering a straightforward and effective method to optimize the trade-off between model size and accuracy.\n\n- This framework promises considerable practical advantages, especially in the context of large-scale recommendation systems.\n\n- A  theoretical analysis is provided, addressing the advantage of the Feature Multiplexing system from rigorous theoretical analysis . It explains how inter-feature collisions can be reduced when the model uses orthogonal weight vectors to project distinct features. This analysis provide good theoretical ground for a fairly practical design. ", "weaknesses": "**Cons**:\n\n- Certain claims in the paper, such as \"0.02% increase in test AUC is considered significant in Avazu and Criteo\" and \"+0.1% is considered significant in online systems,\" raise eyebrows. These claims appear to be based on subjective opinions rather than objective facts.\n\n- The authors have not provided source code for their work. This factor makes it difficult to reproduce and verify the claims made in the paper, which is a foundational principle of the NeurIPS community. If this work is industry-driven and the code cannot be released easily, perhaps the authors should consider venues more suited to applied data mining.\n\n- The content and focus of the paper lean heavily towards data mining and address real-world, industrial-scale problems. As such, it might be better suited for venues like KDD, WWW, or SIGIR.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review. We address the weaknesses below.\n\n**Regarding significance of results:** It is completely fair to be skeptical about the significance of a +0.1% improvement. However, seemingly small improvements can translate to huge numbers in large online systems. For example, Anil et al. state that “accuracy improvements above 0.1% are considered significant” in their [paper](https://arxiv.org/abs/2209.05310) “On the Factory Floor: ML Engineering for Industrial-Scale Ads Recommendation Models'' when discussing Google Ads models. This is an industry with $200B+ annual revenue. In “BarsCTR: Open Benchmarking for Click-Through Rate Prediction” ([CIKM 2021](https://arxiv.org/abs/2009.05794)), Zhu et al. summarize the literature: “existing studies from Google [8, 47] and Microsoft [28] reveal that an absolute improvement of 0.1% in logloss (or AUC) is considered as practically significant in real CTR prediction problems.” All of the results in Table 2 are statistically significant with $p < 0.05$, and 0.1% is very practically significant in industrial systems.\n\nThis is also true in academic evaluations: the differences between logistic regression and SOTA on Avazu and Criteo are +1.7% and +2.1% AUC, respectively (see BarsCTR). The [DCNv2 paper](https://arxiv.org/abs/2008.13535) states, “For Criteo, a 0.001-level improvement is considered significant (see [13, 46, 50]).” The [AutoInt paper](https://arxiv.org/abs/1810.11921) claims that “a slightly higher AUC or lower LogLoss at 0.001-level is regarded significant for CTR prediction task.” Our +0.02% AUC number from Table 1 is based on the 95% confidence interval surrounding the mean AUCs, with $\\sigma = 0.00022$ (approximate population standard deviation, estimated via ~20 runs). However, we really should have made this much clearer, so thank you for making this point in your review. In the revision, we have performed an unpaired t-test between the values in Table 1, finding most differences to be statistically significant. You can find the revised table in the rebuttal PDF.\n\n**Regarding code:** We are in the process of releasing an open source implementation for the academic experiments to reproduce Table 1. Due to our internal publication review process, it wasn’t possible to release this by the paper submission deadline, but it will be present in the revision.\n\n**Regarding venue:** Embedding learning (or representation learning) is the heart of deep learning, and is a problem that occurs in NLP, vision, and bioinformatics as well as in data mining. Like many recent papers (e.g. ROBE-Z at NeurIPS 2022), we focus on SAR (search, ads, and recommendation) systems because they are the most compelling application of categorical representation learning. However, feature multiplexing is a general technique with many potential applications. Applications of embedding learning include scaling up the vocabulary size in transformers and improving the compute-memorization tradeoff of embedding retrieval-augmented transformers and LLMs. Finally, we note that the top baseline methods for this paper recently appeared at NeurIPS (Multihash Hash Embeddings in 2017, ROBE-Z in 2022) and similar venues (HashedNets and Hashing Trick at ICML).\n\n**Regarding limitations:** Due to space constraints, we discussed limitations throughout the text. Limitations of the theory are addressed in Section 4.3, modeling limitations (due to embedding width constraints, etc) in Section 3 (around line 150), and the constraints of our experimental setup in Section 5. \n\nHowever, after reading the reviews, it seems that this was easy to miss. We agree that it would be clearer to consolidate this in a separate section. We’re happy to add the following paragraph to the revision:\n\n*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 do 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 effectively 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.\n\nOnce again, thank you for taking the time to review our paper. If we have addressed the questions and weaknesses from your review, we hope that you will consider raising the score.\n"}, {"review_id": "BjoHZFGM5f", "reviewer": "Reviewer_Azh8", "summary": "This paper proposes that in web-scale machine learning systems, features from different fileds can share the embedding matrix without significantly affecting the model's performance. The insight lies in that different feature fields are processed by different model parameters. Therefore, compared to inter-feature collisions, in the case of intra-feature collisions, the embeddings of different features can tend to be orthogonal, which is beneficial to the final model performance. This point is confirmed by the empirical study based on logistic regression. The experimental results show that the multiplexed embedding scheme is more effective than the existing schemes that only consider inter-feature embedding sharing.\n\n------\nAFTER REBUTTAL:\n\nI've read and appreciated the author’s rebuttal. I understand that collisionless embeddings are considered the upper bound. However, it would be helpful if the advantages of feature multiplexing could be mentioned and demonstrated in the discussion and experiments. Thus I would like to keep the score.", "questions": "Please share your views on the aforementioned weaknesses.", "limitations": "Yes.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The overall organization and writing of the paper are excellent.\n- The proposed multiplexed feature embedding scheme is novel and its feasibility is verified both theoretically and experimentally.\n- The paper conducts experiments on multiple public datasets, and the overall results are promising.", "weaknesses": "- In the experiments, although the performance is better compared to the inter-feature embedding sharing scheme, it is worse than the Collisionless scheme on Criteo, very close on Avazu, but no relevant discussion is provided.\n- Compared to the Collisionless scheme, the advantages of using feature multiplexing do not seem to be fully discussed, nor are they reflected in the experimental results.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review. We address the questions raised in review below.\n\n**Regarding experiment results:** There seems to be a misunderstanding about collisionless embeddings. This method is not usually a feasible baseline, and it is included in our benchmark evaluation only as a headroom reference point. In industrial recommendation systems, where feature vocabularies can be massive and change dynamically due to churn, the hashing trick is the standard approach (e.g. Google [1], Facebook [2], Twitter [3]). Because of memory limits and dynamic vocabularies that are not fully known up-front, collisionless tables are often impossible to deploy. \n\nFor example, our pCTR application in Table 2 would require > 1 TB of memory for a 32-dimensional collisionless table. Facebook’s DLRM model [4] requires a 40 TB table even after hashing. While Facebook does not disclose their total vocabulary size, a collisionless table would likely require space that is 1-2 orders of magnitude larger (> 100 TB in their case). We conducted our academic experiments on the 7-day Criteo dataset, where the collisionless table only requires 25 MB. When more realistic time periods are considered, the size of the vocabulary grows by several orders of magnitude (e.g., the collisionless table for the 30-day Criteo dataset is larger than 100 GB [5]).\n\nHence, the hashing-based methods are the only realistic baselines in our evaluation. Collisionless embedding performance effectively represents an upper bound on what is practically attainable. We include them in our (small-scale) experiments as an optimistic reference point, to investigate the Pareto frontier and estimate the ideal model capacity when we are not limited by hash collisions. The experiments show that multiplexing achieves close to full capacity on Movielens / Avazu and provides the best performance on Criteo versus all baselines. \n\n**Criteo vs. Avazu / Movielens:** The performance gap on Criteo is likely a result of the vocabulary distribution. Criteo 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). We have included a plot of the vocabulary distribution in the rebuttal PDF, which we hope will shed some light on this point.\n\n**Regarding our advantages:** We cover the algorithm-level advantages of feature multiplexing in Section 3 and the practical advantages in Section 5.2. Compared to collisionless embedding, these approaches are fast, implementable, and orders of magnitude cheaper to train and serve.\n\nWe will add this discussion to the revision. Once again, thanks for reviewing our work. If we have addressed your questions and the weaknesses identified in the review, we hope that you will consider raising the score.\n\n**References**\n- [1] Learning to Embed Categorical Features without Embedding Tables for Recommendation (https://arxiv.org/pdf/2010.10784.pdf)\n- [2] Compositional Embeddings Using Complementary Partitions for Memory-Efficient Recommendation Systems (https://arxiv.org/pdf/1909.02107.pdf)\n- [3] \"Model Size Reduction Using Frequency Based Double Hashing for Recommender Systems\" [RecSys 2020](https://arxiv.org/pdf/2007.14523.pdf)\n- [4] \"Software-Hardware Co-design for Fast and Scalable Training of Deep Learning Recommendation Models\" [ISCA'22](https://arxiv.org/abs/2104.05158)\n- [5] \"The trade-offs of model size in large recommendation models: 100GB to 10MB Criteo-tb DLRM model\" [NeurIPS 22](https://arxiv.org/abs/2207.10731)\n\n"}, {"review_id": "StB5w2H1i7", "reviewer": "Reviewer_fp86", "summary": "The authors present a method for multiplexing embeddings of various features in the recommender and similar applications, i.e., sharing the feature embeddings in order to save space and improve performance, especially at lower memory budgets. They provide a detailed overview of the relevant prior work, and give strong both theoretical and empirical analysis of the proposed method. They show the benefits of the method on three public data sets, and also show how the method helped in large-scale production setting.", "questions": "Please find the detailed comments below:\n- The data set details should be explained better. E.g., in Fig 2 the authors mention 26 Criteo features, and it is unclear which ones they are referring to. This is adding to some confusion.\n- This holds also for other data sets. that are considered. Adding more details there would make the experimental section much more readable.\n- The results in Table 1 show that the method is only outperforming other baselines on low-budget Criteo, while everywhere else it is worse. This is quite a weak result, and it is unclear why would someone use their method with such mixed performance.\n- Multiplexed versions of the other baselines should be better explained, at least briefly. The methods are just given, without much discussion.\n\n*************** UPDATE AFTER REBUTTAL **************\nI would like to thank the authors for their responses. They do address my concerns, and I am happy to increase the score.", "limitations": "The authors did not discuss the limitations, and it would be good to add a short paragraph.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Very important problem being addressed.\n- Good theoretical discussion of the method.\n- Good results shown in the production-level setup.", "weaknesses": "- In some places the explanations can be improved quite a lot.\n- The results are very mixed in some cases.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review. We address the questions and weaknesses below. In particular, we think there is a misunderstanding about our experiment results (which we would like to clarify).\n\n**Regarding datasets:** This is a good point - we agree that it would be valuable to have more detailed explanations here. We will elaborate in the revision. To clarify the confusion about Figure 2, Criteo is an online advertisement dataset of about 45 million examples (7 days of data), where each example corresponds to a user interaction (click/no-click decision) and contains an additional 39 features. 26 of these features are categorical and 13 of them are real-valued continuous variables. Figure 2 refers to the 26 categorical features in Criteo (the continuous variables are typically fed alongside the embedded categorical features, but do not use an embedding table - see the [DLRM paper](https://arxiv.org/abs/1906.00091) and the related literature for details).\n\nFor the other datasets: Avazu is a collection of 11 days of advertisement click data (approximately 36 million examples), where each example contains 23 categorical features and no continuous features. Movielens is a traditional user-item recommendation problem, very similar to the well-known Netflix Prize task. These are all highly popular datasets that are widely used to evaluate recommendation algorithms. We plan to refer the reader to \"BarsCTR: Open Benchmarking for Click-Through Rate Prediction,\" ([CIKM 2021](https://arxiv.org/abs/2009.05794)) for a very detailed description of Avazu and Criteo and to \"The Movielens Datasets: History and Context\" (ACM Transactions on Interactive Intelligent Systems, 2015) for a thorough description of Movielens.\n\n**Regarding experiments:** We think there is a misunderstanding here. Table 1 shows that multiplexed embeddings (our proposal) outperform all of the baselines at all memory budgets on all datasets (except for collisionless embeddings, which are infeasible in practice and included only as a headroom reference point - see the response to Reviewer Azh8 for a detailed discussion of this matter). Note that all of the methods below the horizontal bar are newly proposed, while existing baselines are listed above the bar. Figure 3 shows the top baseline (blue) against the top multiplexed method (red) and is another view of the same data behind Table 1.\n\nIn practice, we use Multisize Unified Embeddings (see line 306), which are very similar to Multiplex PQ - one of the top performers from Table 1. We would also like to highlight that in large-scale industrial applications, the extreme low-budget case is standard practice because vocabulary sizes regularly exceed 10 billion, making collisionless schemes infeasible. Our performance improvement in this memory-constrained setting is therefore a major advantage. To summarize, all of the methods below the bar in Table 1 are ours (and are new), and we show very strong results in real-world settings where the vocabulary is on the order of millions to billions (Table 2).\n\nTo remove this confusion, we plan to revise the names to be clearer and more descriptive:\n- What was previously called “Multisize Unified Embeddings” (in Section 5.2) will be referred to simply as “Unified Embedding.”\n- What was previously called “Unified Embedding” (in Table 1) will be referred to as “Multiplexed Hashing Trick.”\n\nSee the rebuttal PDF for a draft revision of Table 1.\n\n**Regarding other multiplexed methods:** We describe the last 20 years of SOTA embedding methods in the appendix, but we agree that it would be valuable to have a detailed description of how to construct a multiplexed method given an existing embedding strategy. We will add the following description:\n\n“To construct a multiplexed version of an existing embedding scheme, we use a single instance of the scheme to index all of the features. For example, we can look up all feature values in a single multihash table, rather than using a separately-tuned multihash table for each feature. Practically, this is equivalent to salting the vocabulary of each feature with the feature ID and merging the salted vocabularies into a single, massive categorical vocabulary (which is then embedded as usual).”\n\n**Regarding limitations:** Due to space constraints, we discussed limitations throughout the text. Limitations of the theory are addressed in Section 4.3, modeling limitations (due to embedding width constraints, etc) in Section 3 (around line 150), and the constraints of our experimental setup in Section 5.\n\nHowever, after reading the reviews, it seems that this was very easy to miss. We plan to add the following paragraph to the revision:\n\n*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 do 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 effectively 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.\n\nFinally, thank you for taking the time to review our paper. If you feel that we have addressed your concerns, we hope that you will consider raising the score. "}, {"review_id": "6Yy4MXiTxC", "reviewer": "Reviewer_5UVX", "summary": "This paper introduces a novel approach known as Feature Multiplexing, which allows multiple features to share a single representation space in machine learning systems. This is significant for web-scale systems which handle hundreds of features with vocabularies of up to billions of tokens, where the standard embedding approaches introduce a vast number of parameters. The authors propose a new solution called Unified Embedding, which simplifies feature configuration, adapts to dynamic data distributions, and is compatible with modern hardware. The empirical results from multiple web-scale search, ads, and recommender systems show superior performance compared to highly competitive baselines.", "questions": "1) Could the authors provide more information on the impact of Feature Multiplexing and Unified Embedding on the ML infrastructure, particularly in terms of computational costs, training time, and hardware usage?\n\n2) Can the authors elaborate on the specific behaviors of the Criteo dataset compared to Avazu and Movielens datasets?\n\n3) Could the authors discuss the limitations and potential trade-offs of the proposed method? How might it affect the ease of extending the model and conducting future R&D?", "limitations": "The authors have not adequately addressed the limitations and potential trade-offs of their proposed technique. Future work may be constrained or impacted by these unaddressed issues. For instance, the authors have not discussed the ease (or lack thereof) of extending the model with new features or conducting new R&D with the proposed method. They also have not explored the potential maintenance costs and impacts on model health and observability, which can be crucial for deploying such systems in real-world applications. Further information on these aspects could greatly benefit the audience's understanding of the practical viability and potential challenges of implementing the proposed technique.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1) The paper addresses a crucial problem in large-scale machine learning systems related to efficient and effective learning of feature embeddings. The proposed framework, Feature Multiplexing, is innovative, allowing multiple features to share the same representation space.\n\n2) The authors provide a well-written and clear explanation of the concepts and the proposed solution. The paper is well-structured, with a good balance of theory, experimentation, and discussion.\n\n3) The problem this paper addresses is of substantial significance, considering the scale at which modern machine learning systems operate. The introduction of Unified Embedding could lead to substantial advancements in web-scale ML systems, serving billions of users globally.", "weaknesses": "1) The paper lacks details regarding the computational benefits of the proposed technique, specifically in terms of infrastructure gains, parameter size, hardware usage, and training time. Providing such details would make the comparison to the baseline more comprehensive and persuasive. (particularly the large scale experiments explained at the end)\n\n2) Some specific analysis and explanations are missing. For example, why the Criteo dataset behaves differently from Avazu and Movielens is not explained. A more in-depth exploration would strengthen the understanding of the behavior of the proposed technique across datasets.\n\n3) The authors could have provided more insights into why online deployment results are providing gains. A detailed explanation could better support the claim of real-world applicability and insights into future users. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review. We address the weaknesses and questions below.\n\n**W1 and Q1:**\n- **Regarding training time:** 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.\n\n- In the following table, we report the average number of steps / second for several methods on Criteo for the 25 MB table size with CPU training (higher is better). This is representative of total training because the number of steps to convergence is fairly stable across methods (250k batches for Criteo, 200K for Avazu, and 50K for Movielens).\n|Method|Standard|Multiplexed|\n|--|--|--|\n|Collisionless|32.9|NA|\n|Hashing Trick|29.2|29.4|\n|(Multiple) Hash Embedding|33.5|31.3|\n|HashedNet|44.1|39.8|\n|ROBE-Z Embedding|30.0|34.4|\n|PQ Embedding|37.6|40.3|\n|QR Embedding|37.4|31.3|\n\n- Note that because our shared cluster has large variations in job load / demand, these numbers have high variance ($\\sigma > 5$). However, these results still support the conclusion that multiplexing has a minimal effect on training time. We observed similar behavior in large scale experiments (see caption of Table 2).\n\n- **Hardware:** We use proprietary ML accelerators with dedicated embedding support - see [TPUv4](https://arxiv.org/pdf/2304.01433.pdf) and [MITA v1](https://ai.meta.com/blog/meta-training-inference-accelerator-AI-MTIA/) for examples of this kind of system (to preserve anonymity during review, we will disclose the full details after the review process). For our academic evaluations, end-to-end CPU training took approximately 3-4 hours for Criteo, 4-5 hours for Avazu, and 30 minutes for Movielens.\n\n**W2 and Q2:** This is a good question. 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. See the rebuttal PDF for a plot of the vocabulary distributions.\n\n**W3:** In short, online systems in industry are regularly memory-bound. We can almost always improve model quality with larger embedding tables, but this requires more CPU/GPU/accelerator memory and is not worth the resource cost/revenue trade-off. Improving model quality for the same memory budget is an immediate win.\nOur offline experiments (from Table 1 / Figure 3) show that multiplexed embeddings lift the Pareto frontier, allowing us to either have better performance at a fixed size or equal performance with a smaller size. The following evidence suggests that this tradeoff drives our online performance gains:\n- Increases in model capacity often lead to better online performance (see [“Scaling Law for Recommendation Models”](https://arxiv.org/pdf/2111.11294.pdf)). Our theory and offline experiments show that multiplexing increases the effective capacity of a model, allowing a smaller model to behave like a larger one.\n- We observe (line 353- 355) that multiplexing provides the greatest benefits for problems with large and/or dynamic vocabulary. These are exactly the applications known to require higher model capacity (see [“Learning to Embed Categorical Features without Embedding Tables for Recommendation”](https://arxiv.org/abs/2010.10784)).\n\n**Q3:** While we discuss some limitations in the text (inline, due to the page limit), we agree that a description of our engineering and modeling tradeoffs would be beneficial to include. In the revision, we plan to summarize these points in a clearly-marked limitations section and provide a full discussion in the appendix. Thank you for raising these valuable questions.\n\n**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 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).\n\n**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 embeddings have better load balancing and alleviate hot-spot issues (our 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.\n\n**Model health:** All of the industry online experiments in Table 2 assume healthy models as a prerequisite. These models train online 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. \n\nFinally, thank you for taking the time to do a detailed review of our paper. If you feel that we have addressed your concerns, we hope that you will consider raising the score."}]

Author rebuttal:
We'd like to thank all of the reviewers for their efforts. We have included a rebuttal PDF that contains:
- Vocabulary distributions for Critoe, Avazu, and Movielens. This plot helps to explain the differences in behavior for embedding algorithms on Criteo vs Avazu and Movielens.
- A revised version of Table 1 that we hope will clear up some common confusions about our results. We would like to highlight that multiplexed embedding algorithms (our proposal) show improvements across the full memory-accuracy tradeoff for each of our three datasets. The revision of Table 1 also includes results for t-tests at the $p < 0.01$ and $p < 0.05$ significance levels, which demonstrate statistically significant differences between multiplexed vs. non-multiplexed embeddings.

Please see the individual rebuttals for responses to specific reviewer questions. 

Meta-review: The paper provides an approach to the problem of feature embedding, which is a crucial part in several ML designs, and in particular in search ads and recommendation systems (called SAR in the paper). The core idea, called feature multiplexing, is sensible and effective. The reviews agree that the application is novel, and appreciated the theoretical analysis provided. Furthermore, the experiments are overall agreed to show the effectiveness of the proposed method and its advantage over existing work. The reviews pointed out a need for a better discussion or explanation of the results, but these seem to be focused on particular points and the authors’ response mitigated these concerns. I think that after editing based on the discussions here, the paper can be made ready to be published in NeurIPS. The novelty of the work and wide impact of the technique would make it a welcome addition to the venue.
Decision: Accept (spotlight)