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

arXiv: https://arxiv.org/abs/2305.12102 · HTML: https://mlsanity.com/paper/2305.12102/

## Key facts

- Cited by 14 papers in the map (data through 14 Oct 2026).
- Builds on 23 of its 48 references that are in the map.
- First posted on arXiv 20 May 2023.

## Builds on (most cited first, 23 of 23)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Chain-of-Thought Prompting Elicits Reasoning in Large Language Models](https://mlsanity.com/paper/2201.11903/) (2022)
- [Distributed Representations of Words and Phrases and their Compositionality](https://mlsanity.com/paper/1310.4546/) (2013)
- [Sparks of Artificial General Intelligence: Early experiments with GPT-4](https://mlsanity.com/paper/2303.12712/) (2023)
- [Prevalence of Neural Collapse during the terminal phase of deep learning training](https://mlsanity.com/paper/2008.08186/) (2020)
- [Deep Learning Recommendation Model for Personalization and Recommendation Systems](https://mlsanity.com/paper/1906.00091/) (2019)
- [DCN V2: Improved Deep & Cross Network and Practical Lessons for Web-scale Learning to Rank Systems](https://mlsanity.com/paper/2008.13535/) (2020)
- [AutoInt: Automatic Feature Interaction Learning via Self-Attentive Neural Networks](https://mlsanity.com/paper/1810.11921/) (2018)
- [TPU v4: An Optically Reconfigurable Supercomputer for Machine Learning with Hardware Support for Embeddings](https://mlsanity.com/paper/2304.01433/) (2023)
- [Compressing Neural Networks with the Hashing Trick](https://mlsanity.com/paper/1504.04788/) (2015)
- [Software-Hardware Co-design for Fast and Scalable Training of Deep Learning Recommendation Models](https://mlsanity.com/paper/2104.05158/) (2021)
- [Compositional Embeddings Using Complementary Partitions for Memory-Efficient Recommendation Systems](https://mlsanity.com/paper/1909.02107/) (2019)
- [Can weight sharing outperform random architecture search? An investigation with TuNAS](https://mlsanity.com/paper/2008.06120/) (2020)
- [Hash Embeddings for Efficient Word Representations](https://mlsanity.com/paper/1709.03933/) (2017)
- [BARS: Towards Open Benchmarking for Recommender Systems](https://mlsanity.com/paper/2205.09626/) (2022)
- [Learning to Embed Categorical Features without Embedding Tables for Recommendation](https://mlsanity.com/paper/2010.10784/) (2020)
- [On the Factory Floor: ML Engineering for Industrial-Scale Ads Recommendation Models](https://mlsanity.com/paper/2209.05310/) (2022)
- [Fully Understanding the Hashing Trick](https://mlsanity.com/paper/1805.08539/) (2018)
- [Sequential Attention for Feature Selection](https://mlsanity.com/paper/2209.14881/) (2022)
- [A GPU-specialized Inference Parameter Server for Large-Scale Deep Recommendation Models](https://mlsanity.com/paper/2210.08804/) (2022)
- [Practical Hash Functions for Similarity Estimation and Dimensionality Reduction](https://mlsanity.com/paper/1711.08797/) (2017)
- [Learning Rate Schedules in the Presence of Distribution Shift](https://mlsanity.com/paper/2303.15634/) (2023)

## Built on by (most cited first, 14 of 14)

- [Conditional Memory via Scalable Lookup: A New Axis of Sparsity for Large Language Models](https://mlsanity.com/paper/2601.07372/) (2026)
- [PLUM: Adapting Pre-trained Language Models for Industrial-scale Generative Recommendations](https://mlsanity.com/paper/2510.07784/) (2025)
- [A Thorough Performance Benchmarking on Lightweight Embedding-based Recommender Systems](https://mlsanity.com/paper/2406.17335/) (2024)
- [PriorBoost: An Adaptive Algorithm for Learning from Aggregate Responses](https://mlsanity.com/paper/2402.04987/) (2024)
- [BayesCNS: A Unified Bayesian Approach to Address Cold Start and Non-Stationarity in Search Systems at Scale](https://mlsanity.com/paper/2410.02126/) (2024)
- [On-device Content-based Recommendation with Single-shot Embedding Pruning: A Cooperative Game Perspective](https://mlsanity.com/paper/2411.13052/) (2024)
- [Two-dimensional Sparse Parallelism for Large Scale Deep Learning Recommendation Model Training](https://mlsanity.com/paper/2508.03854/) (2025)
- [Probabilistic Hash Embeddings for Online Learning of Categorical Features](https://mlsanity.com/paper/2511.20893/) (2025)
- [TokenMinds: Pretrained User Tokens and Embeddings for User Understanding in Large Recommender Systems](https://mlsanity.com/paper/2606.25147/) (2026)
- [MMbeddings: Parameter-Efficient, Low-Overfitting Probabilistic Embeddings Inspired by Nonlinear Mixed Models](https://mlsanity.com/paper/2510.22198/) (2025)
- [Efficient Hyperparameter Search for Non-Stationary Model Training](https://mlsanity.com/paper/2512.01258/) (2025)
- [tensorFM: Low-Rank Approximations of Cross-Order Feature Interactions](https://mlsanity.com/paper/2602.15229/) (2026)
- [Balanced Co-Clustering of Users and Items for Embedding Table Compression in Recommender Systems](https://mlsanity.com/paper/2604.18351/) (2026)
- [JourneyFormer: Encoding Airbnb Guest Journey with Sequence Modeling](https://mlsanity.com/paper/2606.19108/) (2026)

## Abstract (arXiv)

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.
