Dimension Mask Layer: Optimizing Embedding Efficiency for Scalable ID-based Models
Srijan Saket, Ikuhiro Ihara, Vaibhav Sharma, Danish Kalim
TL;DR
The paper tackles the memory burden of large embedding tables in ID-based systems by introducing Dimension Mask Layer (DML), a train-time mechanism that learns an effective embedding dimension after the initial embedding lookup. DML uses a continuous control $effective\_dim$ and a regularization term, selecting the final dimension as $\lceil effective\_dim \rceil$ and gating dimensions with a pseudo-random multiplier $z$ defined by $z = 1 / (1 + e^{-alpha (2x - y - 0.5)})$ where $y ~ U(0,1)$. Offline and online experiments on ShareChat production data show embedding-dimension reductions of over 50% with substantial memory and compute savings, while maintaining or improving key metrics; Avazu experiments corroborate the method’s generality and reveal parameter sensitivities between L1/L2 regularization and initial embedding sizes. Overall, DML provides a scalable, real-time approach to automatically optimize embedding sizes for high-cardinality ID features, enabling memory-efficient deployments without compromising performance.
Abstract
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.