Dimension Mask Layer: Optimizing Embedding Efficiency for Scalable ID-based Models

2025-10-17

AI Summary

To address the high memory footprint and deployment challenges of recommendation models caused by large-scale ID feature embeddings, this paper proposes a learnable Dimensional Masking Layer that dynamically prunes individual dimensions of embedding vectors, enabling fine-grained, adaptive compression of embedding dimensionality. Integrated into the Keras embedding lookup backend, the layer is trained end-to-end to jointly optimize masking parameters and model weights, thereby mitigating multicollinearity and suppressing overfitting. Extensive evaluations on multiple public benchmarks and online A/B tests demonstrate an average 40–50% reduction in embedding dimensionality, yielding significant decreases in model memory consumption and inference latency, while preserving core metrics—including CTR and AUC—at stable levels. To our knowledge, this is the first work to introduce a differentiable, dimension-level selection mechanism for industrial-scale ID embedding compression, achieving a favorable trade-off between model lightweighting and performance robustness.

Technology Category

• Machine Learning: Dimensionality Reduction/Feature Selection
• Data Mining & Knowledge Management: Data Compression
• Search and Optimization: Learning to Search

Application Category

• User Modeling, Personalization and Recommendation: Attacks and countermeasures in recommendation systems
• Graph Algorithms and Modeling for the Web: Graph embeddings and representation learning for Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

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.

Problem

• Optimizing embedding size for ID features in recommendation systems
• Reducing memory consumption of large-scale embedding tables
• Maintaining model performance while shrinking embedding dimensions

Innovation

• Dimension mask layer optimizes embedding vector size
• Automatically determines optimal embedding dimensions for IDs
• Reduces embedding dimensions by half with minimal loss