Title: Unified Embedding: Battle-Tested Feature Representations for Web-Scale ML Systems
Venue: NeurIPS 2023
Accepted: yes
Keywords: embedding learning; recommendation systems; representation learning
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

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, introducing 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 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 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.

1 Introduction

There have been many new and exciting breakthroughs across academia and industry towards largescale 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 et al., 2023; Wei et al., 2022a). Interestingly, large-scale model architectures are headed towards two extreme scenarios. In natural language understanding, speech, and computer vision, the transformer (Vaswani et al., 2017) dominates the Pareto frontier. Most transformer parameters are

Matthew Fahrbach

located in the hidden layers, with very small embedding tables (e.g., ~0.63B embedding parameters in the 175B GPT-3 (Brown et 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 in 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. 1 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 it is critical to learn high-quality representations.

Feature embedding learning has been a 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 a 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 × d matrix (i.e., an 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 a 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 a 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 a 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 a 'true' representation in a high-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-truth representation into a 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 is to reduce the dimension-reduction error. Early methods were based on the Johnson-Lindenstrauss lemma and sparse linear projections (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 are now state-of-the-art, but the focus on 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.

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.

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., multiplexing + 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.

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.

Battle-tested approach and industrial experiments at scale : Multi-size Unified Embedding has been launched in over a dozen web-scale SAR systems, significantly improving user engagement metrics and key business metrics. In 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 on 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 or more feature values v drawn from the vocabulary of the feature V . For a concrete example, consider the classic click-through rate (CTR) prediction task for ads systems, where we want to predict the probability of a user clicking on an advertisement. The ' ad_id ' feature may take on hundreds of billions of possible values-one for each unique ad served by the platform-while ' site_id ' describes the website where that ad was shown (a.k.a., the publisher).

Weembed the values of each feature into a space that is more suitable for modeling, e.g., R d or the unit ( d -1) -sphere. Given a vocabulary of N tokens V = { v 1 , v 2 , . . . , v N } , we learn a transformation that maps v ∈ V to an embedding vector e ∈ R d . 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 1 , y 1 ) , ( x 2 , y 2 ) , . . . , ( x | D | , y | D | ) } of examples with labels. The examples consist of values from T different categorical features with vocabularies { V 1 , V 2 , . . . , V T } , where feature t has vocabulary V t . Unless otherwise stated, we do not consider multivalent features and instead assume that each example x has one value for each categorical feature, i.e., x = [ v 1 , v 2 , . . . , v T ] where v i ∈ V i . However, this is for notational convenience only-the techniques in this paper extend to missing and multivalent feature values well.

Formally, the embedding table is a matrix E ∈ R M × d and the embedding function g ( x ; E ) transforms the categorical feature values into embeddings. We let h ( v ) : V → [ M ] be a 2-universal hash function (Vadhan et al., 2012) that assigns a feature value to a row index of the embedding table. We also consider a model function f ( e ; θ ) that transforms the concatenated embeddings into a model prediction, which is penalized by a loss glyph[lscript] . Using this decomposition, we define the joint feature learning problem as follows:

We use independent hash functions h t ( v ) for each feature t ∈ [ 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 e m to denote the m -th row of E , which means e h ( u ) is the embedding of u . Lastly, we let ✶ u,v be the indicator variable that denotes a hash collision between u and v , i.e., h ( u ) = h ( v ) .

Feature

1

Feature

1

1

Feature

V1

Feature

1

V2

V2

V2

V2

V3

V3

V3

2

V1

2

V1

2

V1

Feature 2

V1

Feature

V2

Feature

V2

Feature

V2

V2

V3

V3

V3

V3

V4

V4

V4

V4

V5

V5

V5

V5

Collisionless

Hash Table

Unified Table

Multisize Unified

3 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 in 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 on a multi-size (multi-probe) version of the multiplexed feature hashing trick (Figure 1), 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 to multiples of the table dimension, this is not a limiting constraint (and is needed by other methods such as ROBE-Z and product quantization).

Shared vocabulary It 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 semantically-similar 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 a brief analysis in the dimension-reduction 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., intra-feature ) and those occurring across features (i.e., inter-feature ).

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 a simplified but representative modelbinary logistic regression with learnable feature embeddings 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 proofs and detailed derivations in this section to Appendix A.

4.1 Parameter Efficiency in The Dimension-Reduction Framework

In the feature hashing problem, we are given two vectors x , y ∈ R N that we project using a linear map φ : R N → R M given by the sparse {± 1 } matrix in Weinberger et al. (2009). The goal is for φ to minimally distort inner product, i.e., 〈 φ ( x ) , φ ( y ) 〉 ≈ 〈 x , y 〉 . 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 at once in this framework, we consider the concatenations x = [ x 1 , x 2 ] and y = [ y 1 , y 2 ] . The components x 1 ∈ { 0 , 1 } N 1 and x 2 ∈ { 0 , 1 } N 2 are the one-hot encodings (or bag-of-words) for the two vocabularies V 1 and V 2 . The standard hashing trick approximates 〈 x , y 〉 by independently projecting x 1 to dimension M 1 and x 2 to dimension M 2 . The multiplexed version uses a single projection, but into dimension M 1 + M 2 . 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 is the signed sum of value counts within each hash bucket.

Definition 4.1. Let h : V →{ 1 , 2 , . . . , M } be a 2-universal hash function and ξ : V →{-1 , +1 } be a sign hash function. The function φ h,ξ : 2 V → R M is defined as φ h,ξ ( W ) = ∑ w ∈ W ξ ( w ) u h ( w ) , where u i is the i -th unit basis vector in R M .

The following result compares multiplexed and standard embeddings in the dimension reduction framework. Note that the randomness is over the choice of hash functions h ( v ) and ξ ( v ) . This result also directly extends to concatenating T feature embeddings.

Proposition 4.2. For any x 1 , y 1 ∈ { 0 , 1 } N 1 and x 2 , y 2 ∈ { 0 , 1 } N 2 , let x = [ x 1 , x 2 ] , y = [ y 1 , y 2 ] denote their concatenations. Let µ U , µ H , σ 2 U , and σ 2 H be the mean and variance of 〈 φ h,ξ ( x ) , φ h,ξ ( y ) 〉 for multiplexed and hash encodings, respectively. Then, µ U = µ H = 〈 x , y 〉 and

Observations Proposition 4.2 shows that multiplexed embeddings can do a good job balancing hash collisions across the parameter space. Consider the problem of estimating the inner product 〈 x , y 〉 . Suppose that ‖ x 1 ‖ 2 2 = ‖ y 1 ‖ 2 2 = k 1 and ‖ x 2 ‖ 2 2 = ‖ y 2 ‖ 2 2 = k 2 , i.e., features 1 and 2 are multivalent with k 1 and k 2 values, and assume x and y are orthogonal, i.e., 〈 x , y 〉 = 0 . Proposition 4.2 tells us that σ 2 U = ( k 1 + k 2 ) 2 / ( M 1 + M 2 ) and σ 2 H = k 2 1 /M 1 + k 2 2 /M 2 . The difference σ 2 H -σ 2 U factors as ( √ M 1 /M 2 k 1 -√ M 2 /M 1 k 2 ) 2 / ( M 1 + M 2 ) , which is non-negative, implying σ 2 U ≤ σ 2 H .

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 ∈ V 1 collides with y ∈ V 2 but no other values collide across vocabularies, then the embeddings from other tokens in V 1 and V 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 f ( z ; θ ) is the sigmoid function σ θ ( z ) = 1 / (1 + exp( -〈 z , θ 〉 )) , glyph[lscript] is the binary cross-entropy loss, and y ∈ { 0 , 1 } (e.g., click or non-click). We concatenate the embeddings of each feature, which means the input to f ( z ; θ ) is z = g ( x ; E ) = [ e h 1 ( x 1 ) , e h 2 ( x 2 ) , . . . , e h T ( x T ) ] , where e h t ( x t ) is

Criteo-EmbeddingNorms

Criteo-WeightOrthogonality

Feature

1

1.0

(六)

90

Hash Embedding

n

Eh(u)

01

Embedding Norm

0.8

Unified Embedding

Hash Embedding

Angle (Degrees)

08

Unified Embedding

0.6

60

b↑

2

Feature

Eh(v)

02

0.4

140

30

0.2

20

101

102

Table Size (M)

103

104

103

10

Total Number of Embeddings(M)

10?

104

103

10

the embedding for the t -th feature value in example x . We write the logistic regression weights as θ = [ θ 1 , θ 2 , . . . , θ T ] so that embedding e h t ( x t ) for feature t is projected via θ t ∈ R M . We illustrate this simple but representative model architecture in Figure 2.

The following analysis holds in general for T features, but we consider the problem for two categorical features with vocabularies V 1 and V 2 to simplify the presentation. We write the logistic regression objective as follows in Eq. (2), using C u,v, 1 and C u,v, 0 to denote the number of examples for which x = [ u, v ] and y = 1 and y = 0 , respectively. We present a full derivation in Appendix A.2.

To understand the effect of hash collisions on training dynamics, we take the gradient with respect to e h 1 ( u ) , i.e., the embedding parameters that represent u ∈ V 1 . By studying the gradients for collisionless, hash, and multiplexed embeddings, we quantify the effects of intra-feature and interfeature collisions. Interestingly, all three gradients can be written using the following three terms:

glyph[negationslash]

For collisionless embeddings, component (3) is the full gradient. For hash embeddings, the gradient is the sum of components (3) and (4). The multiplexed embedding gradient is the sum of all three terms.

Insights Our core observation is that embedding gradients can be decomposed into a collisionless (true) component (3), an intra-feature component (4), and an inter-feature component (5). The components from hash collisions act as a gradient bias. The intra-feature component is in the same direction as the true component, implying that intra-feature collisions are not resolvable by the model. This agrees with intuition since intra-feature hash collisions effectively act as value merging.

Inter-feature collisions, however, push the gradient in the direction of θ 2 , which creates an opportunity for model training to mitigate their effect. To see this, consider a situation where θ 1 and θ 2 are initialized to be orthogonal and do not change direction during training. 2 During SGD, the embedding e h 1 ( u ) is a linear combination of gradients over the training steps, which means e h 1 ( u ) can be decomposed into a true and intra-feature component, in the θ 1 direction, and an inter-feature component in the θ 2 direction. Since 〈 θ 1 , θ 2 〉 = 0 , the projection of θ glyph[latticetop] 1 e h 1 ( u ) effectively eliminates the inter-feature component .

Empirical study These observations allow us to form two testable hypotheses. First, we expect the projection vectors [ θ 1 , θ 2 , . . . , θ T ] to orthogonalize since this minimizes the effect of inter-feature collisions (and improves the loss). The orthogonalization effect should be stronger for embedding tables with fewer hash buckets since that is when inter-feature collisions are most prevalent. Second, intra-feature and inter-feature gradient contributions add approximately O ( N/M ) for each of their dimensions to the collisionless gradient since E [ ✶ u,v ] = 1 /M . Therefore, we expect the squared embedding norms to scale roughly as O ( N/M ) .

To investigate these hypotheses, we train a single-layer neural network on the categorical features of the Criteo click-through prediction task. We initialize θ = [ θ 1 , θ 2 , . . . , θ T ] with each θ t component in the same direction (i.e., the worst-case) to see if multiplexing encourages the weights to orthogonalize. We plot the results in Figure 2, which support our hypotheses.

4.3 Why Can Features Share a Single Table?

Returning to our parameter efficiency argument of Section 4.1, we can finally explain the full benefits of feature multiplexing. The dimension reduction argument shows that, given a parameter budget, the multiplexed embedding has the same overall number of hash collisions as a well-tuned hash embedding. Our gradient analysis, however, shows that not all collisions are equally problematic. Inter-feature collisions can be mitigated by a single-layer neural network because different features are processed by different model parameters. While our theoretical results are limited to single-layer neural networks, we expect deeper and more complicated network architectures to exhibit analogs of weight orthogonalization due to their relative overparameterization. However, we believe this is still compelling evidence that shared embedding tables work by load balancing unrecoverable collisions across a larger parameter space and effectively use one massive table for each feature.

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-valued continuous variables. We only embed the categorical features. Avazu is a 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 a traditional user-item collaborative filtering problem, which we convert to a binary prediction task by assigning '1's to all examples with rating ≥ 3 , and '0's to the rest. For more details, see the benchmarks in Zhu et al. (2022) for Criteo and Avazu and the review by Harper and Konstan (2015) for Movielens. For Criteo and Movielens, we apply the same continuous feature transformations, train-test split, and other preprocessing steps in Wang et al. (2021). For Avazu, we use the train-test split and preprocessing steps in Song et al. (2019).

Experiment design To find the Pareto frontier of the memory-AUC curve (Figure 3), we train a neural network with different embedding representation methods at 16 logarithmically-spaced memory budgets. Models can overfit on the second epoch, but compressed embeddings sometimes require multiple epochs. To showcase the best performance, we follow the guidelines of Tsang and Ahle (2022) and report the test performance of the best model found over three epochs. As baselines, we implement collisionless embeddings, the feature hashing trick, hash embeddings, HashedNets, ROBE-Z, compositional embeddings with element-wise product (QR), and compositional embeddings with concatenation (i.e., product quantization (PQ)).

For the multiplexed version of an existing embedding scheme, we use a single instance of the scheme to index all of the features. For example, we look up all feature values in a single multihash table, rather than using a separately-tuned multihash table for each feature. This is equivalent to salting the vocabulary of each feature with the feature ID and merging the salted vocabularies into a single massive vocabulary (which is then embedded as usual). We apply this procedure to each baseline to obtain a new multiplexed algorithm, reported below the horizontal bar in Table 1.

Implementation and hyperparameters All methods are implemented in Tensorflow 2.0 using the TensorFlow Recommenders (TFRS) framework. 3 For a fair comparison, all models are identical except for the embedding component. Following other studies (Naumov et al., 2019; Wang et al., 2021), we use the same embedding dimension for all features ( d ∈ { 39 , 32 , 30 } for Criteo, Avazu, and Movielens, respectively). This is followed by a stack of 1-2 DCN layers and 1-2 DNN feed-forward layers-we provide a full description of the architecture in Appendix B. We run a grid search to tune all embedding algorithm hyperparameters and conduct five runs per combination. When allocating embedding parameters to features, we size tables based on the cardinality of their vocabulary (e.g., a feature with 10% of the total vocabulary receives 10% of the memory budget).

Results Table 1 shows the results for the academic datasets. We give three points on the parameteraccuracy 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 parameteraccuracy tradeoff for all multiplexed and non-multiplexed methods. We observe that feature multiplexing is Pareto-optimal. Interestingly, the multiplexed hashing trick outperforms several embedding techniques (e.g., ROBE-Z) that were previously SOTA.

Collisionless embeddings are included in our benchmark evaluation only to provide an optimistic headroom reference point. Collisionless tables are often impossible to deploy in industrial recommendation systems since feature vocabularies can be massive and change dynamically due to churn. Multiplexing provides a significant improvement to methods that are fast, implementable, and orders of magnitude cheaper to train and serve than collisionless embeddings.

For significance tests comparing multiplexed methods to their corresponding non-multiplexed baseline (rather than the strongest baseline), see Appendix B. Appendix B also contains additional details about Criteo and Avazu, including an explanation for the collisionless-hash embedding performance gap on Criteo.

5.2 Industrial Deployments of Unified Embedding in Web-Scale Systems

Among the proposed approaches, Unified Embedding (multi-probe hashing trick + multiplexing) is uniquely positioned to address the infrastructure and latency constraints associated with training and serving web-scale systems. We adopt Unified Embedding in practice, using variable row lookups for each feature (typically 1~6) with concatenation to achieve the desired embedding width for our features. It should be noted that this process yields a similar algorithm to Multiplex PQ-one of the top performers from Table 1. We discuss how to adopt Unified Embedding in production SAR systems, where it has already served billions of users through launches in over a dozen models. We also present an experimental study of Unified Embedding in offline and online A/B experiments.

Criteo

Avazu

Movielens

0.774

0.89

0.806

0.773

0.88

0.804

0.772

0.87

%

0.771

0.86

C

0.802

AU

0.770

0.85

0.769

0.800

0.84

0.768

0.798

0.767

0.83

10-1

100

101

102

10-1

100

101

10-2

10-1

100

101

Memory (MB)

Memory (MB)

Memory (MB)

★

Collisionless

Standard

Multiplexed

Experiment and production setup To test the robustness and generalization capability of Unified Embedding, we conduct offline and online experiments with five key models from different SAR domains, including commerce, apps, and short-form videos. These systems represent a diverse set of production model architectures powered by state-of-the-art techniques, including two-tower retrieval models with LogQ correction (Yi et al., 2019), DCN-V2 (Wang et al., 2021), and MMOE (Zhao et al., 2019). These models handle a wide range of prediction tasks, including candidate retrieval, click-through rate (pCTR), conversion rate (pCVR) prediction (Chapelle, 2014), and multi-task user engagement modeling.

All of these production models are continuously trained and updated in an online streaming fashion, with the embedding tables being the dominant component (often >99%) in terms of model parameters (Anil et al., 2022; Fahrbach et al., 2023). Many of the baseline embedding methods are simple but quite strong (e.g., feature hashing and multiple hashing), and have been tuned with AutoML and human heuristics (Anil et al., 2022). In fact, these baseline embedding learning methods have roughly stayed the same for the past five years, with various efforts failing to improve them through novel techniques.

Practical benefits Unified Embedding tables offer solutions to several challenging and practical issues in large-scale embedding learning systems.

Simple configuration : With standard embedding methods, we need to independently tune the table size and dimension for each feature. Hyperparameter tuning is an arduous process due to the large number of categorical features used in production (hundreds if not thousands). Unified Embedding tables are far simpler to configure: the total table size is calculated based on the available memory budget, and we only need to tune the dimension on a per-feature basis-easily accomplished via AutoML (Bender et al., 2020) or attention-based search methods (Yasuda et al., 2023; Axiotis and Yasuda, 2023). This yields roughly a 50% reduction in the number of hyperparameters.

Adapt to dynamic feature distributions : In practice, the vocabulary size of each feature changes over time (new IDs enter the system on a daily basis, while stale items gradually vanish). For example, in a short-form video platform, we expect a highly dynamic video ID vocabulary with severe churn. In e-commerce, a significant number of new products are added during the holiday season compared to off-season. Thus, a fixed number of buckets for a given table can often lead to sub-optimal performance. Unified Embedding tables offer a large parameter space that is shared across all features, which can better accommodate fluctuations in feature distributions.

Hardware compatibility : Machine learning accelerators and frameworks have co-evolved with algorithms to support common use-cases, and the 'row embedding lookup' approach has been the standard technique for the past decade. While Unified Embedding tables are well-supported by the latest TPUs (Jouppi et al., 2023) and GPUs (Wei et al., 2022b; Mudigere et al., 2022), many of the recent and novel embedding methods require memory access patterns that are not as compatible with ML accelerators.

Offline and online results Table 2 displays the results of our offline and online experiments for Unified Embedding. It was applied to production models with vocabulary sizes spanning from 10M to 10B. We observe substantial improvements across all models. Note that the offline metrics and the key business metrics for different models can be different, and the gains are not directly comparable. For offline metrics, the models are evaluated using AUC (higher is better), Recall@1 (higher is better), and RMSE (lower is better). From these experiments, we found that feature multiplexing provides greater benefits when (i) the vocabulary sizes for features are larger, and (ii) the vocabularies are more dynamic (e.g., suffers from a high churn rate). For instance, an overwhelming number of new short-form videos are produced daily, creating a large and dynamic vocabulary.

6 Conclusion

In this work, we propose a highly generalizable Feature Multiplexing framework that allows many different features to share the same embedding space. One of our crucial observations is that models can learn to mitigate the effects of inter-feature collisions, enabling a single embedding vector to represent feature values from different vocabularies. Multiplexed versions of SOTA embedding methods provide a Pareto-optimal tradeoff, and Unified Embedding proves to be a particularly simple yet effective approach. We demonstrate that a Unified Embedding table performs significantly better in both offline and online experiments compared to highly-competitive baselines across multiple web-scale ML systems. We believe the interplay between feature multiplexing, dynamic vocabularies, power-law feature distributions, and DNNs provides interesting and important opportunities for future work.

Limitations While we expect deeper and more complicated network architectures to exhibit similar behavior, our theoretical analysis is limited to single-layer neural networks. The Pareto frontier from Section 5.1 is based on models that lag behind the current SOTA (we use 1-2 DCN layers + 1-2 DNN layers), though we provide evidence that feature multiplexing is also effective for SOTA models in Table 2. Unified embeddings impose limitations on the overall model architecture, and all embedding dimensions must share a common multiple. Because of the latency involved with looking up more than 6 components, we are often limited to 5-6 discrete choices of embedding width. Finally, unified embeddings sacrifice the ability to aggressively tune hyperparameters on a per-feature basis in exchange for flexibility, ease-of-use, and simplified feature engineering.

References

Achlioptas. Database-friendly random projections: Johnson-Lindenstrauss with binary coins. J. Comput. Syst. Sci. , 66(4):671-687, 2003.

Anil, S. Gadanho, D. Huang, N. Jacob, Z. Li, D. Lin, T. 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. In Proceedings of the 5th Workshop on Online Recommender Systems and User Modeling co-located with the 16th ACM Conference on Recommender Systems, ORSUM@RecSys , volume 3303 of CEUR Workshop Proceedings , 2022.

Axiotis and T. Yasuda. Performance of glyph[lscript] 1 regularization for sparse convex optimization. arXiv preprint arXiv:2307.07405 , 2023.

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 IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 14311-14320. Computer Vision Foundation / IEEE, 2020.

Bengio, R. Ducharme, and P. Vincent. A neural probabilistic language model. In Advances in Neural Information Processing Systems , pages 932-938, 2000.

B. Brown, B. Mann, N. Ryder, M. Subbiah, J. Kaplan, P. 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, S. Gray, B. Chess, J. Clark, C. Berner, S. McCandlish, A. Radford, I. Sutskever, and D. Amodei. Language models are few-shot learners. In Advances in Neural Information Processing Systems , volume 33, 2020.

Bubeck, V. Chandrasekaran, R. Eldan, J. Gehrke, E. Horvitz, E. Kamar, P. Lee, Y. T. Lee, Y. Li, S. Lundberg, et al. Sparks of artificial general intelligence: Early experiments with GPT-4. arXiv preprint arXiv:2303.12712 , 2023.

Chapelle. Modeling delayed feedback in display advertising. In The 20th ACM SIGKDD International Conference on Knowledge Discovery and Data Mining , pages 1097-1105. ACM, 2014.

Chen, J. Wilson, S. Tyree, K. Weinberger, and Y. Chen. Compressing neural networks with the hashing trick. In International Conference on Machine Learning , pages 2285-2294. 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.

Dahlgaard, M. Knudsen, and M. Thorup. Practical hash functions for similarity estimation and dimensionality reduction. Advances in Neural Information Processing Systems , 30, 2017.

Dasgupta and A. Gupta. An elementary proof of a theorem of Johnson and Lindenstrauss. Random Structures & Algorithms , 22(1):60-65, 2003.

Desai and A. Shrivastava. The trade-offs of model size in large recommendation models : 100GB to 10MB Criteo-tb DLRM model. In Advances in Neural Information Processing Systems , 2022.

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.

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 on Machine Learning , pages 9523-9546. PMLR, 2023.

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.

Hogan and H. W. Wind. Five Lessons: The Modern Fundamentals of Golf . 1985.

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 to embed categorical features without embedding tables for recommendation. In Proceedings of the 27th ACM SIGKDD Conference on Knowledge Discovery & Data Mining , pages 840-850, 2021.

Li, T. J. Hastie, and K. W. Church. Very sparse random projections. In Proceedings of the 12th ACM SIGKDD International Conference on Knowledge Discovery and Data Mining , pages 287-296, 2006.

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.

Mitzenmacher. The power of two choices in randomized load balancing. IEEE Transactions on Parallel and Distributed Systems , 12(10):1094-1104, 2001.

Moody. Fast learning in multi-resolution hierarchies. Advances in Neural Information Processing Systems , 1988.

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, P. Bhattacharya, P. Lapukhov, M. Naumov, A. Mathews, L. Qiao, M. Smelyanskiy, B. Jia, and V. Rao. Software-hardware 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.

Naumov, D. Mudigere, H.-J. M. Shi, J. Huang, N. Sundaraman, J. Park, X. Wang, U. Gupta, C.-J. Wu, A. G. Azzolini, et al. Deep learning recommendation model for personalization and recommendation systems. arXiv preprint arXiv:1906.00091 , 2019.

Papyan, X. Han, and D. L. Donoho. Prevalence of neural collapse during the terminal phase of deep learning training. Proceedings of the National Academy of Sciences , 117(40):24652-24663, 2020.

E. Rumelhart, G. E. Hinton, and R. J. Williams. Learning representations by back-propagating errors. Nature , 323(6088):533, 1986.

H.-J. M. Shi, D. Mudigere, M. Naumov, and J. Yang. Compositional embeddings using complementary partitions for memory-efficient recommendation systems. In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining , pages 165-175, 2020.

Song, C. Shi, Z. Xiao, Z. Duan, Y. Xu, M. Zhang, and J. Tang. Autoint: Automatic feature interaction learning via self-attentive neural networks. In Proceedings of the 28th ACM International Conference on Information and Knowledge Management , pages 1161-1170, 2019.

L.-H. Tsang and T. D. Ahle. Clustering embedding tables, without first learning them. arXiv preprint arXiv:2210.05974 , 2022.

P. Vadhan et al. Pseudorandomness. Foundations and Trends in Theoretical Computer Science , 7 (1-3):1-336, 2012.

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.

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.

Wei, X. Wang, D. Schuurmans, M. Bosma, E. Chi, Q. Le, and D. Zhou. Chain of thought prompting elicits reasoning in large language models. arXiv preprint arXiv:2201.11903 , 2022a.

Wei, M. Langer, F. Yu, M. Lee, J. Liu, J. Shi, and Z. Wang. A GPU-specialized inference parameter server for large-scale deep recommendation models. In Sixteenth ACM Conference on Recommender Systems , pages 408-419. ACM, 2022b.

Weinberger, A. Dasgupta, J. Langford, A. Smola, and J. Attenberg. Feature hashing for large scale multitask learning. In Proceedings of the 26th Annual International Conference on Machine Learning , pages 1113-1120, 2009.

Yasuda, M. Bateni, L. Chen, M. Fahrbach, G. Fu, and V. Mirrokni. Sequential attention for feature selection. In Proceedings of the 11th International Conference on Learning Representations (ICLR) , 2023.

Yi, J. Yang, L. Hong, D. Z. Cheng, L. Heldt, A. Kumthekar, Z. Zhao, L. Wei, and E. H. Chi. Sampling-bias-corrected neural modeling for large corpus item recommendations. In Proceedings of the 13th ACM Conference on Recommender Systems , pages 269-277. ACM, 2019.

Zhao, L. Hong, L. Wei, J. Chen, A. Nath, S. Andrews, A. Kumthekar, M. Sathiamoorthy, X. Yi, and E. H. Chi. Recommending what video to watch next: A multitask ranking system. In Proceedings of the 13th ACM Conference on Recommender Systems , pages 43-51. ACM, 2019.

Zhu, Q. Dai, L. Su, R. Ma, J. Liu, G. Cai, X. Xiao, and R. Zhang. Bars: Towards open benchmarking for recommender systems. In Proceedings of the 45th International ACM SIGIR Conference on Research and Development in Information Retrieval , pages 2912-2923, 2022.

A Theory

A.1 Missing Analysis from Section 4.1

Lemma A.1 (Weinberger et al. (2009, Lemma 2)) . For any x , y ∈ R n , E h,ξ [ 〈 φ ( x ) , φ ( y ) 〉 ] = 〈 x , y 〉 . Furthermore, we have glyph[negationslash]

where x ◦ y ∈ R n denotes the Hadamard product of x and y .

Proposition 4.2. For any x 1 , y 1 ∈ { 0 , 1 } N 1 and x 2 , y 2 ∈ { 0 , 1 } N 2 , let x = [ x 1 , x 2 ] , y = [ y 1 , y 2 ] denote their concatenations. Let µ U , µ H , σ 2 U , and σ 2 H be the mean and variance of 〈 φ h,ξ ( x ) , φ h,ξ ( y ) 〉 for multiplexed and hash encodings, respectively. Then, µ U = µ H = 〈 x , y 〉 and

Proof. The results for µ U , µ H and σ 2 U directly follow from Lemma A.1, and observing that

since x , y ∈ { 0 , 1 } N 1 + N 2 .

To compute σ 2 H , first decompose the inner product

Since the hash functions h 1 , ξ 1 , h 2 , ξ 2 are independent, it follows that σ 2 H is the sum of the variances of 〈 φ h 1 ,ξ 1 ( x 1 ) , φ h 1 ,ξ 1 ( y 1 ) 〉 and 〈 φ h 2 ,ξ 2 ( x 2 ) , φ h 2 ,ξ 2 ( y 2 ) 〉 , so we can again invoke Lemma A.1.

A.2 Derivation of Gradients

For easy reference, we re-introduce our logistic regression setting here. In Figure 4, we reproduce part of Figure 2 from the main text to further help clarify our notation.

Recall from the main text that we are interested in a logistic regression model with trainable embeddings for binary classification (i.e., a single-layer neural network with hashed one-hot encodings as input). This corresponds to our supervised learning framework (Eq. (1)), where f ( z ; θ ) is the sigmoid function σ θ ( z ) = 1 / (1 + exp( -〈 z , θ 〉 )) and glyph[lscript] is the binary cross-entropy loss.

We concatenate the embeddings of each feature, which means that the input to f ( z ; θ ) is z = g ( x ; E ) = [ e h 1 ( x 1 ) , e h 2 ( x 2 ) , . . . , e h T ( x T ) ] , where e h t ( x t ) is the embedding for the t -th feature value in example x .

We also partition the logistic regression weights θ = [ θ 1 , θ 2 , . . . , θ T ] , so that embedding e h t ( x t ) for feature t is projected via θ t ∈ R M .

For notational convenience we also partition the dataset by label. We write D 0 = { ( x i , y i ) ∈ D : y i = 0 } and D 1 = { ( x i , y i ) ∈ D : y i = 1 } . Using this notation, we can write the logistic regression objective as follows:

We proceed by writing the loss as a sum over vocabulary pairs rather than examples. We consider the problem for two categorical features with vocabularies V 1 and V 2 to simplify the presentation, and we use the notation C u,v, 0 to denote the number of co-occurences of ( u, v ) ∈ V 1 × V 2 in dataset D 0 .

Feature

1

Eh(u)

01

n

→0(eu,v· 0)

Feature 2

Eh(v)

02

θ

Likewise, C u,v, 1 denotes the corresponding quantity for D 1 . We also let e u,v = [ e h 1 ( u ) , e h 2 ( v ) ] so that the inner product 〈 [ θ 1 , θ 2 ] , [ e h 1 ( u ) , e h 2 ( v ) ] 〉 can be written concisely as θ glyph[latticetop] e u,v . It follows that

After combining the sigmoid functions, we get the following expression:

To get the gradients with respect to E h ( u ) (the embedding parameters used to represent value u ), we first take the gradient with respect to E m . These parameters accumulate gradient contributions from any value assigned to bucket m (i.e. for which ✶ u,m is one). It is then straightforward to express ∇ E h ( u ) L D ( E , θ ) in terms of the other values v that collide with u (i.e. for which u,v is one).

✶ Collisionless Embeddings: Because each row E m corresponds to a single value, we can take the gradient with respect to E h ( u ) directly.

Hash Embeddings: Here, row m corresponds to any value that is assigned to bucket m (i.e. for which ✶ u,m is one). We find the gradient for the embedding table corresponding to the first feature V 1 .

The interpretation of this equation is that with hash collisions, the gradient now contains contributions from any value assigned to the bucket. If we consider the bucket selected by a specfic value u ∈ V 1 , we can write this gradient as follows.

glyph[negationslash]

Unified Embeddings: Here, row m might correspond to any value - even ones from other features. The gradient with respect to E m now contains additional components.

V1

1

V1

Feature

V1

Feature

V21

Feature

V2

V2

4

V3-

V3

V4

V4

V5

V5

V5

V1、

V1

Feature 2

2

Feature 2

V2、

2

V2

Feature 2

V3

V4

V5

V67

V6

V7

V7

V7

Unified Table

Collisionless

Hashing Trick

Sharedmemory

Sharedmemory

Feature

V1

以以

Feature

1

V1

Feature

1

V1

V2

Note:notall

V2

Note:notall

V3

areshown

connections

areshown

connections

V5

V1

Feature 2

V1

Feature

2

V1

Weights

V2

V2

V2

34

V3

V3

e(v1)=

V4

V4

V5

ROBE / HashedNet

ROBE-Z

MultiHash

Feature

V1

以以

Feature

V1

以

V5

V5

e(v1)=

e(v1)=

Compositional (QR)

Compositional (PQ)

Rewriting the gradient to be in terms of E h ( u ) (for u ∈ V 1 ) reveals that in the unified embedding case, the gradient contains the collisionless component, an intra-feature component, and an inter-feature component:

glyph[negationslash]

B Experiments

Figure 5 illustrates our baselines. We describe the methods and their hyperparameters below.

Collisionless tables assign each value of each feature a unique d -dimensional representation. The only hyperparameter is d , the embedding dimension.

Hashing trick tables begin with a fixed number of representations and allow values to collide within features. Here, each table has two hyperparameters: d (the dimension) and M (the table size). The inspiration for this method is most commonly attributed to Weinberger et al. (2009), and it is widely used across industry.

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.

Compositional QR tables use multiple small tables for each feature, then perform elementwise multiplication of the resulting embeddings. The hyperparameters are k (the number of tables), M (the table size), and d (the embedding dimension). This method corresponds to the compositional embedding of Shi et al. (2020), using the 'concatenation' aggregation strategy.

Compositional PQ tables use multiple small tables for each feature, then perform concatenation. We call this 'product quantization' because it is similar to the vector compression technique. The hyperparameters are the same as for QR tables. This method corresponds to the compositional embedding of Shi et al. (2020), using the 'element-wise product' aggregation strategy.

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.

ROBE-Z tables look up contiguous blocks of Z dimensions in shared memory (we use Z = 2 in Figure 5). ROBE-Z has two extra hyperparameters: the size of the shared parameter array, and the number of blocks. Note that one can use a different Z for each block. For simplicity, we do not do this. This is the method used by Desai et al. (2022) and by Desai and Shrivastava (2022).

Hash embedding tables (or multihash , in Figure 5) perform multiple lookups into the same table, then combine the embeddings using a learned weighted combination. Multihash tables have several hyperparameters and are therefore trickier to tune well. The parameters are M 1 , the size of the embedding table, M 2 , the size of the importance weight table, k , the number of lookups and the dimension of the importance weight table, and d , the embedding dimension. This method was proposed by Tsang and Ahle (2022).

Hyperparameter tuning: For all methods, we consider sixteen logarithmically-spaced memory budgets: [0.001, 0.002, 0.005, 0.007, 0.01, 0.02, 0.05, 0.07, 0.1, 0.2, 0.5, 0.7, 1.0, 2.0, 5.0, 10.0] times the memory required for the collisionless table. We conduct 5 independent training runs for each hyperparameter configuration. For non-multiplexed representations, we allocate parameters to features based on vocabulary size. For example, a feature with 10% of the total vocabulary will be allocated 10% of the total parameter budget. We also tuned the network architecture and optimizer parameters. To avoid biasing our results in favor of any one method, we only used collisionless tables to tune these parameters. The embedding hyperparameters were tuned using grid search over the following choices.

Hash embeddings : For hash embeddings (and multiplexed hash embeddings), we chose the number of lookups k from [2, 3]. We also introduce an allocation parameter p that describes what fraction of the memory budget is used for the embedding table (e.g., M 1 ) versus the importance table (e.g., M 2 ). For example, p = 0 . 2 means that the embedding table are allocated 80% of the budget and the importance weight tables receives the remaining 20%. We chose p from [0.05, 0.1, 0.2].

ROBE-Z : We chose the number of lookups from [2, 4, 8, 16]. This corresponds to different values of Z depending on the embedding dimension of the dataset.

PQembeddings : For PQ embeddings (and multiplexed PQ), we selected the number of lookups k from [2, 3, 4, 8, 16]. We note that this parameter affects performance in different ways depending on the parameter budget - smaller budgets require a greater number of lookups for performance.

QR embeddings : For QR embeddings (and multiplexed QR), we selected the number of components k from [2, 3, 4].

Other methods : Our other methods did not require search over extra hyperparameters. For the special case of collisionless embeddings, we simply ran 5 replicates at a single parameter budget (1.0).

Taken together, we conducted 3205 training runs for each dataset to produce Figure 3. Note that in some cases (e.g., Avazu with a budget of 5 , 10 ), some runs with expensive configurations did not finish-we report only those results that completed training within 2 days. Finally, the total number of training runs (across all datasets) is somewhat higher, to account for our preliminary tuning over the architecture and optimizer parameters. We estimate that roughly 10,000 training runs were performed in total.

Criteo: We use the 7-day version of Criteo from the Kaggle competition. It is widely-known that additional filtering / merging of the vocabulary is necessary to achieve SOTA results on Criteo. We merge the least-frequent values, trimming the vocabulary down to the size listed in Table 4. For the network, we used a stack of two DCN cross layers followed by a 2-layer neural network with 748 nodes in each layer. This model gives a realistic picture of the tradeoff because near-SOTA performance on Criteo can be achieved with a well-tuned 3-layer feedforward network with 748 nodes and collisionless embeddings (Wang et al., 2021). We trained with a batch size of 512 for 300K steps using Adam with a learning rate of 0.0002. The full parameter-accuracy tradeoffs are presented in Figure 7.

Avazu: We use the version of Avazu from the Kaggle competition, but we use a different train-test split than the original. Following the practice of Song et al. (2019), we combine the train and test sets and do a 90-10 split after shuffling. We also apply vocabulary pruning to this dataset, with the results in Table 5. Note that we mod the 'hour' feature by 24 and do not use the 'id' feature. For the neural network, we use a stack of one DCN cross layer followed by two fully connected feedforward layers of 512 nodes each. We trained with a batch size of 512 for 300K steps using Adam with a learning rate of 0.0002. The full parameter-accuracy tradeoffs are presented in Figure 8.

Movielens: We use the Movielens-1M dataset with the same preprocessing used by Wang et al. (2021). We only use the 'userId,' 'movieId,' 'zipcode,' 'age,' 'occupation' and 'gender' features. We always use a collisionless table to embed the 'gender' feature due to small vocabulary size. We used a network composed of a single DCN cross layer followed by a single fully-connected layer with 192 nodes. We trained with a batch size of 128 for 300K steps using Adam with a learning rate of 0.0002. The full parameter-accuracy tradeoffs are presented in Figure 9.

Vocabulary distributions: Figure 6 presents the vocabulary distribution for each of our datasets. Criteo differs from Avazu and Movielens in terms of vocabulary distribution - it has a heavier-tailed vocabulary than Avazu or Movielens. This worsens the errors introduced by hash collisions because it is more likely for a colliding token to overwrite the shared embedding (all collisions are between heavy hitters). This is likely the cause of the performance gap between collisionless embeddings (where hash collisions do not occur) and the other methods on Criteo (Figure 3).

Statistical significance tests: We conducted additional statistical significance testing on Table 1. When we compare the multiplexed methods against the corresponding non-multiplexed baseline (rather than the strongest baseline), the following differences are significant with p < 0 . 05 (Welch's t -test).

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.

Avazu: All results in the 324 kB and 3.24 MB columns. Multiplex QR in the 32.4 MB column.

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.

Vocabulary Histograms

Criteo

Avazu

107

Movielens

log(Frequency)

105

103

101

100

101

102

103

104

105

Vocabulary (Log ID)

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 in job load / demand, these numbers have high variance ( σ > 5 ). However, our results still support the conclusion that multiplexing has a minimal effect on training time. We observed similar behavior in large-scale industrial experiments-multiplexing did not significantly increase the training or inference time.

B.1 Industry experiments

Hardware and training time: In several of our Unified Embedding deployments (Table 2), we use TPUv4 for training and/or inference. For models of practical interest, embedding table size rarely affects the model training time in models of practical interest. As long as there is enough memory to support the embedding tables (i.e., enough CPU RAM or total accelerator memory), the training time is mostly governed by the forwards and backwards passes on the rest of the (upstream) network. We do not observe a substantial increase in training / inference time on either CPU or TPUv4 (Jouppi et al., 2023).

Modeling tradeoffs: Unified embedding gives up some flexibility in setting the per-feature embedding dimension, and other (more complicated) multiplexed methods are better on the Pareto-frontier for offline experiments. In exchange, we get more per-feature embedding capacity, a much simpler

Criteo

Criteo

Criteo

0.805

0.805

0.806

0.800

AUC

560

Test AUC

0.800

Test AUC

0.804

0.790

56L0

0.802

0.785

0.790

0.780 -

Collisionless

Collisionless

0.800

Collisionless

0.775

0.770

Hashing Trick

0.785

HashEmbedding

MultiplexHashEmb

860

MultiplexHashedNet

HashedNet

105

Memory (MB)

106

10?

10#

0.780

105

Memory (MB)

106

10?

10a

105

Memory (MB)

106

10?

108

Criteo

Criteo

Criteo

0.8075

0.805

0.806

0.8050

0.800

Test AUC

0.804

AUC

0.8025

AUC

S60

0.8000

0.790

2080

Test

0.7975

S8/0

0.7950

0080

Collisionless

Collisionless

0840

Collisionless

ROBE-Z

0.7925

PQEmbedding

0.775

QREmbedding

0.798

MultiplexROBE-Z

0.7900

十

MultiplexQR

0.770

Memory (MB)

106

107

10#

105

Memory (MB)

106

107

10#

105

Memory (MB)

107

108

Criteo

Avazu

Avazu

0.775

508'0

0.800

0.772

0.765

AUC

56'0

Test AUC

Test AUC

0.771

0.760

0.790

TestA

02L0

0.785

0.755

0.769

0.780

0.775

Collisionless

Hashing Trick

0.745

0.750

Hash Embedding

Collisionless

0.768

Collisionless

HashedNet

0.770

0.740

十

MultiplexHashEmb

0.767

MultiplexHashedNet

105

Memory (MB)

106

10?

10#

105

Memory (MB)

106

107

10#

105

Memory (MB)

106

107

10a

Avazu

Avazu

Avazu

0.774

0.775

0.772

0.750

0.772

AUC

0.771

Test AUC

0.770

AUC

0.768

0.700

Test

0.770

0.766

Test/

0.675

0.769

0.764

0.768

....

Collisionless

ROBE-Z

0.762

Collisionless

PQ Embedding

0.625

059°0

Collisionless

QREmbedding

0.767

十

MultiplexROBE-Z

0.760

十

十

MultiplexQR

0.600

105

106

107

105

106

107

105

106

10?

Memory (MB)

Memory (MB)

Memory (MB)

hyperparameter configuration, hardware-friendly algorithm, and easy feature engineering. For example, to add new features to a model without multiplexing, we must specify and tune the table size and dimension for a new set of embedding tables. With multiplexing, we can simply add a new feature to an existing table at the cost of a few additional lookups. Another option is to use the conventional scheme for new features, then regularly consider a model update that migrates several new features into the unified embedding (typically resulting in performance gains).

Engineering tradeoffs: During our experience deploying Unified Embedding in a dozen applications, we have not had problems with observability, monitoring, or model health / stability. Instead, we find that Unified Embedding results in better load balancing and alleviate TPU hot-spot issues (as our

Movielens

Movielens

Movielens

0.90

0.89

0.85

0.875

0.88

AUC

0.80

AUC

0.850

Test AUC

0.87

0.75

0.825

0.86

Test

020

Test

0.800

0.85

0.65

0.775

0.60

Collisionless

Hashing Trick

0.750

Collisionless

Hash Embedding

0.84

Collisionless

HashedNet

0.55

MultiplexHashEmb

80

104

Memory (MB)

105

106

10

104

Memory (MB)

105

106

10

104

Memory (MB)

205

106

10?

Movielens

Movielens

Movielens

0.89

0.89 -

0.90

0.88

0.88

0.85

AUC

.80

AUC

0.87

AUC

0.75

0.80

0.86

0.86

Test

Test

0.70-

0.85

0.85

0.65

0.84

Collisionless

0.84

Collisionless

0.60

Collisionless

80

十

MultiplexROBE-Z

ROBE-Z

0.83

十

PQEmbedding

0.55

+

十

QREmbedding

MultiplexQR

0.50

104

Memory (MB)

105

106

107

104

Memory (MB)

105

106

107

104

Memory (MB)

105

106

107

embedding rows are distributed over cores). The high-level takeaway is that it is easier to balance one large table across accelerators than many small tables of different shapes.

Model health: All of our industry online experiments assume healthy models as a prerequisite. These models use online sequential training and adapt to constant variations in data distributions. The consistent improvements from unified embedding in these systems is strong evidence that our methods do not interfere with model health or add maintenance costs.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1) 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.\n\nWEAKNESSES:\n1) 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.\n\nQUESTIONS:\n1) 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?\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-17"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- Very important problem being addressed.\n- Good theoretical discussion of the method.\n- Good results shown in the production-level setup.\n\nWEAKNESSES:\n- In some places the explanations can be improved quite a lot.\n- The results are very mixed in some cases.\n\nQUESTIONS:\nPlease 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.\n\nLIMITATIONS:\nThe authors did not discuss the limitations, and it would be good to add a short paragraph.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-09"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\nPlease share your views on the aforementioned weaknesses.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-27"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n**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.\n\nWEAKNESSES:\n**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\nQUESTIONS:\nPlease refer to weakness.\n\nLIMITATIONS:\nThe authors did not discuss the limitations at all.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-08-02"}