Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

DimGrow: Memory-Efficient Field-level Embedding Dimension Search

Yihong Huang, Chen Chu

TL;DR

DimGrow presents a memory-efficient, SuperNet-free approach to field-level embedding dimension search by starting from a minimal $d_i=1$ per field and using a Shuffle Gate to generate importance signals for progressive expansion or reduction. The method extends Shuffle Gate from feature selection to dimension-level search, introducing dynamic embeddings, an adaptation layer, and a gating mechanism to drive dimension allocation during training. Empirical results on Aliccp, Avazu, and Criteo show that DimGrow achieves competitive AUC with significantly fewer embedding parameters and lower memory usage, particularly as embedding tables grow larger. This approach reduces search-time memory overhead and offers practical benefits for deploying efficient recommender systems without exhaustive SuperNet-based exploration.

Abstract

Key feature fields need bigger embedding dimensionality, others need smaller. This demands automated dimension allocation. Existing approaches, such as pruning or Neural Architecture Search (NAS), require training a memory-intensive SuperNet that enumerates all possible dimension combinations, which is infeasible for large feature spaces. We propose DimGrow, a lightweight approach that eliminates the SuperNet requirement. Starting training model from one dimension per feature field, DimGrow can progressively expand/shrink dimensions via importance scoring. Dimensions grow only when their importance consistently exceed a threshold, ensuring memory efficiency. Experiments on three recommendation datasets verify the effectiveness of DimGrow while it reduces training memory compared to SuperNet-based methods.

DimGrow: Memory-Efficient Field-level Embedding Dimension Search

TL;DR

DimGrow presents a memory-efficient, SuperNet-free approach to field-level embedding dimension search by starting from a minimal per field and using a Shuffle Gate to generate importance signals for progressive expansion or reduction. The method extends Shuffle Gate from feature selection to dimension-level search, introducing dynamic embeddings, an adaptation layer, and a gating mechanism to drive dimension allocation during training. Empirical results on Aliccp, Avazu, and Criteo show that DimGrow achieves competitive AUC with significantly fewer embedding parameters and lower memory usage, particularly as embedding tables grow larger. This approach reduces search-time memory overhead and offers practical benefits for deploying efficient recommender systems without exhaustive SuperNet-based exploration.

Abstract

Key feature fields need bigger embedding dimensionality, others need smaller. This demands automated dimension allocation. Existing approaches, such as pruning or Neural Architecture Search (NAS), require training a memory-intensive SuperNet that enumerates all possible dimension combinations, which is infeasible for large feature spaces. We propose DimGrow, a lightweight approach that eliminates the SuperNet requirement. Starting training model from one dimension per feature field, DimGrow can progressively expand/shrink dimensions via importance scoring. Dimensions grow only when their importance consistently exceed a threshold, ensuring memory efficiency. Experiments on three recommendation datasets verify the effectiveness of DimGrow while it reduces training memory compared to SuperNet-based methods.
Paper Structure (23 sections, 2 theorems, 30 equations, 4 figures, 1 table, 1 algorithm)

This paper contains 23 sections, 2 theorems, 30 equations, 4 figures, 1 table, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Embedding Compression
  4. Field-level EDS & Feature Selection
  5. PRELIMINARIES
  6. METHODOLOGY
  7. Dimension Importance Scoring via Shuffle Gate
  8. Generalized Shuffle Operation
  9. Application to Dimension-level Search
  10. Why Shuffle Gate?
  11. DimGrow: Progressive Dimension Expansion
  12. Dynamic Model Architecture
  13. Dimension Expansion and Reduction Strategy
  14. Experiments
  15. Setup
  16. ...and 8 more sections

Key Result

Theorem A.2

For an $\epsilon$-non-predictive dimension $k$ in feature field $i$, if $\alpha > \epsilon \cdot \mathbb{E}_{\mathbf{x} \sim \mathcal{X}}[\|\mathbf{e}_{i,k} - \tilde{\mathbf{e}}_{i,k}\|]$, gradient descent optimization will drive $g_{i,k}$ to 0.

Figures (4)

  • Figure 1: Shuffle Gate produces polarized importance scores with clear separation between important and unimportant components (right).
  • Figure 2: Framework of DimGrow
  • Figure 3: AUC vs. model size comparison of different embedding dimension search methods on Aliccp, Avazu, and Criteo datasets. The model size is measured in millions (M) of parameters.
  • Figure 4: Computational efficiency comparison: (a) search time (in seconds) and (b) GPU memory consumption (in MB) for different embedding dimension search methods.

Theorems & Definitions (5)

  • Definition A.1: $\epsilon$-non-predictive dimension
  • Theorem A.2
  • proof
  • Theorem A.3
  • proof