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

• URL: http://arxiv.org/abs/2510.15308v1
• Date: Fri, 17 Oct 2025 04:46:10 GMT
• Title: Dimension Mask Layer: Optimizing Embedding Efficiency for Scalable ID-based Models
• Authors: Srijan Saket, Ikuhiro Ihara, Vaibhav Sharma, Danish Kalim,
• Abstract summary: In recommendation systems and social media platforms like Meta, TikTok, and Instagram, large-scale ID-based features often require embedding tables that consume significant memory.<n>This paper introduces a method to automatically determine the optimal embedding size for ID features.<n>We show that using a dimension mask layer can shrink the effective embedding dimension by 40-50%, leading to substantial improvements in memory efficiency.
• Score: 1.2599533416395765
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• LESA: Learnable LLM Layer Scaling-Up [57.0510934286449]
  Training Large Language Models (LLMs) from scratch requires immense computational resources, making it prohibitively expensive.<n>Model scaling-up offers a promising solution by leveraging the parameters of smaller models to create larger ones.<n>We propose textbfLESA, a novel learnable method for depth scaling-up.
  arXiv  Detail & Related papers  (2025-02-19T14:58:48Z)
• Sparse Gradient Compression for Fine-Tuning Large Language Models [58.44973963468691]
  Fine-tuning large language models (LLMs) for downstream tasks has become increasingly crucial due to their widespread use and the growing availability of open-source models.<n>High memory costs associated with fine-tuning remain a significant challenge, especially as models increase in size.<n>We propose sparse compression gradient (SGC) to address these limitations.
  arXiv  Detail & Related papers  (2025-02-01T04:18:28Z)
• Instruction-Following Pruning for Large Language Models [58.329978053711024]
  We move beyond the traditional static pruning approach of determining a fixed pruning mask for a model.<n>In our method, the pruning mask is input-dependent and adapts dynamically based on the information described in a user instruction.<n>Our approach, termed "instruction-following pruning", introduces a sparse mask predictor that takes the user instruction as input and dynamically selects the most relevant model parameters for the given task.
  arXiv  Detail & Related papers  (2025-01-03T20:19:14Z)
• Bypass Back-propagation: Optimization-based Structural Pruning for Large Language Models via Policy Gradient [57.9629676017527]
  We propose an optimization-based structural pruning that learns the pruning masks in a probabilistic space directly by optimizing the loss of the pruned model.<n>We achieve this by learning an underlying Bernoulli distribution to sample binary pruning masks.<n>Experiments conducted on LLaMA, LLaMA-2, LLaMA-3, Vicuna, and Mistral models demonstrate the promising performance of our method in efficiency and effectiveness.
  arXiv  Detail & Related papers  (2024-06-15T09:31:03Z)
• Fine-Grained Embedding Dimension Optimization During Training for Recommender Systems [17.602059421895856]
  FIITED is a system to automatically reduce the memory footprint via FIne-grained In-Training Embedding Dimension pruning. We show that FIITED can reduce DLRM embedding size by more than 65% while preserving model quality. On public datasets, FIITED can reduce the size of embedding tables by 2.1x to 800x with negligible accuracy drop.
  arXiv  Detail & Related papers  (2024-01-09T08:04:11Z)
• Vocabulary-level Memory Efficiency for Language Model Fine-tuning [36.1039389951318]
  We show that a significant proportion of the vocabulary remains unused during fine-tuning.<n>We propose a simple yet effective approach that leverages this finding to minimize memory usage.<n>Our approach does not impact downstream task performance, while allowing more efficient use of computational resources.
  arXiv  Detail & Related papers  (2023-09-15T19:00:00Z)
• Dynamic Embedding Size Search with Minimum Regret for Streaming Recommender System [39.78277554870799]
  We show that setting an identical and static embedding size is sub-optimal in terms of recommendation performance and memory cost. We propose a method to minimize the embedding size selection regret on both user and item sides in a non-stationary manner.
  arXiv  Detail & Related papers  (2023-08-15T13:27:18Z)
• Continuous Input Embedding Size Search For Recommender Systems [60.89189829112067]
  Continuous input embedding size search (CIESS) is a novel RL-based method that operates on a continuous search space with arbitrary embedding sizes to choose from.<n> CIESS is also model-agnostic and hence generalizable to a variety of latent factor RSs.<n> experiments on two real-world datasets have shown state-of-the-art performance of CIESS under different memory budgets.
  arXiv  Detail & Related papers  (2023-04-07T06:46:37Z)
• Masked Autoencoding for Scalable and Generalizable Decision Making [93.84855114717062]
  MaskDP is a simple and scalable self-supervised pretraining method for reinforcement learning and behavioral cloning. We find that a MaskDP model gains the capability of zero-shot transfer to new BC tasks, such as single and multiple goal reaching.
  arXiv  Detail & Related papers  (2022-11-23T07:04:41Z)
• A Model or 603 Exemplars: Towards Memory-Efficient Class-Incremental Learning [56.450090618578]
  Class-Incremental Learning (CIL) aims to train a model with limited memory size to meet this requirement. We show that when counting the model size into the total budget and comparing methods with aligned memory size, saving models do not consistently work. We propose a simple yet effective baseline, denoted as MEMO for Memory-efficient Expandable MOdel.
  arXiv  Detail & Related papers  (2022-05-26T08:24:01Z)
• Binary Code based Hash Embedding for Web-scale Applications [12.851057275052506]
  Deep learning models are widely adopted in web-scale applications such as recommender systems, and online advertising. In these applications, embedding learning of categorical features is crucial to the success of deep learning models. We propose a binary code based hash embedding method which allows the size of the embedding table to be reduced in arbitrary scale without compromising too much performance.
  arXiv  Detail & Related papers  (2021-08-24T11:51:15Z)
• Learning Effective and Efficient Embedding via an Adaptively-Masked Twins-based Layer [15.403616481651383]
  We propose an Adaptively-Masked Twins-based Layer (AMTL) behind the standard embedding layer. AMTL generates a mask vector to mask the undesired dimensions for each embedding vector. The mask vector brings flexibility in selecting the dimensions and the proposed layer can be easily added to either untrained or trained DLRMs.
  arXiv  Detail & Related papers  (2021-08-24T11:50:49Z)
• Training with Multi-Layer Embeddings for Model Reduction [0.9046327456472286]
  We introduce a multi-layer embedding training architecture that trains embeddings via a sequence of linear layers. We show that it allows reducing d by 4-8X, with a corresponding improvement in memory footprint, at given model accuracy.
  arXiv  Detail & Related papers  (2020-06-10T02:47:40Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.