Unified Embedding: Battle-Tested Feature Representations for Web-Scale ML Systems

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.

Paper

References (48)

Scroll for more · 36 remaining

Similar papers

Peer review

Reviewer 5UVX7/10 · confidence 4/52023-06-17

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.

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. 2) 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. 3) 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) 2) 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. 3) 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.

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? 2) Can the authors elaborate on the specific behaviors of the Criteo dataset compared to Avazu and Movielens datasets? 3) 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?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

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.

Reviewer fp867/10 · confidence 3/52023-07-09

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.

Strengths

- Very important problem being addressed. - Good theoretical discussion of the method. - Good results shown in the production-level setup.

Weaknesses

- In some places the explanations can be improved quite a lot. - The results are very mixed in some cases.

Questions

Please find the detailed comments below: - 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. - This holds also for other data sets. that are considered. Adding more details there would make the experimental section much more readable. - 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. - Multiplexed versions of the other baselines should be better explained, at least briefly. The methods are just given, without much discussion. *************** UPDATE AFTER REBUTTAL ************** I would like to thank the authors for their responses. They do address my concerns, and I am happy to increase the score.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The authors did not discuss the limitations, and it would be good to add a short paragraph.

Reviewer Azh86/10 · confidence 3/52023-07-27

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. ------ AFTER REBUTTAL: I'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.

Strengths

- The overall organization and writing of the paper are excellent. - The proposed multiplexed feature embedding scheme is novel and its feasibility is verified both theoretically and experimentally. - 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. - 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.

Questions

Please share your views on the aforementioned weaknesses.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

4 excellent

Contribution

3 good

Limitations

Yes.

Reviewer fvbA5/10 · confidence 4/52023-08-02

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.

Strengths

**Pros**: - 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. - This framework promises considerable practical advantages, especially in the context of large-scale recommendation systems. - 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**: - 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. - 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. - 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.

Questions

Please refer to weakness.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The authors did not discuss the limitations at all.

Reviewer 5UVX2023-08-18

Thanks to the authors for detailed response. I think crucial and adding a lot of practical value to the paper and impact. I strongly recommend these comments to be embeded into the final version. I am happy to raise my score as well. I recommend accept.

Reviewer fp862023-08-21

I would like to thank the authors for their responses. They do address my concerns, and I have increased my score.

Program Chairsdecision2023-09-21

Decision

Accept (spotlight)

© 2026 NYSGPT2525 LLC