[Paper Review] Unified Embedding: Battle-Tested Feature Representations for Web-Scale ML Systems

Benjamin Coleman, Wang-Cheng Kang+5 more|arXiv (Cornell University)|May 20, 2023

Machine Learning and Data Classification4 citations

TL;DR

This paper introduces Unified Embedding, a feature multiplexing framework that shares a single embedding space across multiple high-cardinality categorical features, enabling Pareto-optimal parameter-accuracy tradeoffs. It achieves significant improvements in offline and online metrics across five web-scale search, ads, and recommendation systems, with gains up to +2.2% AUC and +0.62% online click-through rate, while simplifying configuration and enhancing hardware compatibility.

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.

Motivation & Objective

• To address the parameter explosion and inefficiency in web-scale machine learning systems where each categorical feature uses a separate embedding table.
• To enable a unified representation space across multiple features without sacrificing model performance or interpretability.
• To improve parameter-accuracy tradeoffs in models with large, dynamic vocabularies common in search, ads, and recommendation systems.
• To design a system that is compatible with modern hardware accelerators and simplifies feature engineering in production environments.

Proposed method

• Proposes Feature Multiplexing, a framework where multiple categorical features share a single embedding space, reducing parameter count while preserving feature-specific representations.
• Uses a shared embedding table where each feature value is mapped to a unique vector in a common d-dimensional space, with model capacity to disambiguate features via learned interactions.
• Employs a unified lookup mechanism compatible with modern TPUs and GPUs, avoiding complex memory access patterns of alternative methods.
• Introduces Unified Embedding as a practical instantiation of feature multiplexing, enabling simplified configuration and dynamic adaptation to changing data distributions.
• Employs a shared embedding space where model training learns to distinguish between features despite shared vectors, leveraging deep neural network capacity to mitigate inter-feature collisions.
• Applies the method across diverse architectures (DCN-V2, Two-tower, MMOE) and domains (products, short-form videos, apps), demonstrating broad applicability.

Figure 1 : Embedding methods for two categorical features. We highlight the lookup process for the first value $v_{1}$ of each feature. Hash tables randomly share representations within each feature, while Unified Embedding shares representations across features. To implement Unified Embedding with

Experimental results

Research questions

• RQ1 Can a single shared embedding space effectively represent multiple high-cardinality categorical features without performance degradation?
• RQ2 Does feature multiplexing lead to Pareto-optimal tradeoffs between model parameters and accuracy across diverse web-scale benchmarks?
• RQ3 How does Unified Embedding perform in real-world production systems with dynamic, power-law distributed vocabularies?
• RQ4 To what extent can unified embeddings maintain or improve online business metrics (e.g., CTR, recall) compared to per-feature embedding baselines?
• RQ5 Can the unified approach maintain hardware efficiency and compatibility with modern accelerators like TPUs and GPUs?

Key findings

• Unified Embedding achieved a +2.2% improvement in AUC on a product recommendation task with a ~10B vocabulary, demonstrating strong offline performance gains.
• In a short-form video retrieval system, the method improved Recall@1 by +7.3%, indicating better ranking of relevant content.
• Online A/B tests showed a +0.62% increase in click-through rate (CTR) on a short-form video recommendation system, with +0.44% and +0.11% gains on other ad and app prediction tasks.
• The method maintained or reduced model size and training/serving costs across all production models, despite significant accuracy improvements.
• Feature multiplexing provided greater benefits in high-vocabulary and high-churn environments, such as dynamic video content systems.
• Unified Embedding demonstrated strong compatibility with modern hardware and simplified feature configuration, enabling broader deployment in production systems.

Figure 2 : Single-layer neural embedding model with per-feature weights $\boldsymbol{\theta}_{t}$ (left). Mean embedding $\ell^{2}$ -norm (middle) and mean angle between all pairs of weight vectors $\boldsymbol{\theta}_{t_{1}},\boldsymbol{\theta}_{t_{2}}$ (right) as a function of table size for Crit