[论文解读] 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 Classification被引用 4

一句话总结

本文提出统一嵌入（Unified Embedding），一种特征多路复用框架，通过在多个高基数类别特征之间共享单一嵌入空间，实现参数-精度权衡的帕累托最优。该方法在五个大规模网络搜索、广告和推荐系统中显著提升了离线与在线指标，AUC最高提升+2.2%，在线点击率最高提升+0.62%，同时简化了配置并增强了硬件兼容性。

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.

研究动机与目标

• 解决大规模网络机器学习系统中每个类别特征使用独立嵌入表所导致的参数爆炸与效率低下问题。
• 在不牺牲模型性能或可解释性的前提下，实现多个特征间的统一表征空间。
• 改善搜索、广告和推荐系统中常见且大规模、动态变化的词汇表模型的参数-精度权衡。
• 设计一种与现代硬件加速器兼容、并简化生产环境中特征工程的系统。

提出的方法

• 提出特征多路复用（Feature Multiplexing），一种多个类别特征共享单一嵌入空间的框架，降低参数量的同时保留特征特异性表征。
• 使用共享嵌入表，将每个特征值映射到一个统一的d维空间中的唯一向量，通过模型容量学习特征间的交互作用以实现区分。
• 采用与现代TPU和GPU兼容的统一查找机制，避免了其他方法中复杂的内存访问模式。
• 提出统一嵌入作为特征多路复用的实际实现，支持简化配置与对动态数据分布的动态适应。
• 采用共享嵌入空间，通过模型训练学习在共享向量下区分不同特征，利用深度神经网络的容量缓解特征间冲突。
• 在多种架构（DCN-V2、双塔、MMOE）和领域（商品、短视频、应用）中应用该方法，证明其广泛适用性。

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

实验结果

研究问题

• RQ1 单一共享嵌入空间是否能在不降低性能的前提下有效表征多个高基数类别特征？
• RQ2 特征多路复用是否能在多种大规模网络基准上实现模型参数与精度之间的帕累托最优权衡？
• RQ3 统一嵌入在具有动态、幂律分布词汇表的真实生产系统中表现如何？
• RQ4 与每个特征独立嵌入的基线相比，统一嵌入在保持或提升在线业务指标（如CTR、召回率）方面能达到何种程度？
• RQ5 统一方法是否能保持硬件效率并兼容现代加速器（如TPU和GPU）？

主要发现

• 在词汇表规模约100亿的物品推荐任务中，统一嵌入实现了AUC +2.2%的提升，证明了其强大的离线性能增益。
• 在短视频检索系统中，该方法将Recall@1提升了+7.3%，表明对相关结果的排序能力显著增强。
• 在线A/B测试显示，在短视频推荐系统中点击率（CTR）提升了+0.62%；在其他广告和应用预测任务中，分别实现了+0.44%和+0.11%的提升。
• 尽管精度显著提升，该方法在所有生产模型中均保持或降低了模型大小以及训练/推理成本。
• 特征多路复用在高词汇量和高变化率的环境中（如动态视频内容系统）带来了更大收益。
• 统一嵌入在现代硬件上表现出强兼容性，并简化了特征配置，支持在生产系统中更广泛部署。

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