通过动态裁剪嵌入维度，大幅降低模型内存占用而不损性能。
Dimension Mask Layer: Optimizing Embedding Efficiency for Scalable ID-based Models
- 引入维度掩码层，自动裁剪嵌入向量前半部分维度。
- 在真实数据集上实现40%-50%的嵌入维度缩减，内存效率显著提升。
- 适合大规模推荐系统和社交平台优化部署资源的场景。
在Meta、TikTok、Instagram等现代推荐系统与社交媒体平台中，大规模基于ID的特征常需消耗大量内存的嵌入表。管理这些嵌入大小极具挑战性，导致模型臃肿，难以部署与维护。本文提出一种自动确定ID特征最优嵌入尺寸的方法，显著减少模型体积且保持性能。方法核心为自定义Keras层——维度掩码层，置于嵌入查找之后，仅允许前N维通过。该操作可将输入特征维度减半以上，同时模型性能几乎无损。实验表明，该方法在公开数据集上的离线测试及真实生产数据的在线A/B测试中，使有效嵌入维度降低40%-50%，极大改善内存效率。本方案为高密度ID特征平台提供了可扩展的资源与性能优化路径。
原文摘要 · Abstract (English)
In modern recommendation systems and social media platforms like Meta, TikTok, and Instagram, large-scale ID-based features often require embedding tables that consume significant memory. Managing these embedding sizes can be challenging, leading to bulky models that are harder to deploy and maintain. In this paper, we introduce a method to automatically determine the optimal embedding size for ID features, significantly reducing the model size while maintaining performance. Our approach involves defining a custom Keras layer called the dimension mask layer, which sits directly after the embedding lookup. This layer trims the embedding vector by allowing only the first N dimensions to pass through. By doing this, we can reduce the input feature dimension by more than half with minimal or no loss in model performance metrics. This reduction helps cut down the memory footprint of the model and lowers the risk of overfitting due to multicollinearity. Through offline experiments on public datasets and an online A/B test on a real production dataset, we demonstrate that using a dimension mask layer can shrink the effective embedding dimension by 40-50\%, leading to substantial improvements in memory efficiency. This method provides a scalable solution for platforms dealing with a high volume of ID features, optimizing both resource usage and model performance.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。