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.

翻译：在现代推荐系统及Meta、TikTok、Instagram等社交媒体平台中，基于ID的大规模特征通常需要消耗大量内存的嵌入表。管理这些嵌入维度具有挑战性，往往导致模型体积庞大，难以部署和维护。本文提出一种自动确定ID特征最优嵌入维度的方法，能在保持模型性能的同时显著减小模型规模。该方法通过定义一个名为维度掩码层的自定义Keras层实现，该层直接置于嵌入查找层之后，通过仅允许前N个维度通过的方式对嵌入向量进行裁剪。采用此方法，我们可将输入特征维度减少一半以上，且模型性能指标仅出现极小损失甚至无损失。这种维度缩减有助于降低模型的内存占用，并减少因多重共线性导致的过拟合风险。通过在公开数据集上的离线实验及真实生产数据集上的在线A/B测试，我们证明使用维度掩码层可将有效嵌入维度压缩40-50%，从而显著提升内存效率。该方法为处理海量ID特征的平台提供了可扩展的解决方案，同时优化了资源利用与模型性能。