DimGrow: Memory-Efficient Field-level Embedding Dimension Search

TL;DR

DimGrow is a memory-efficient method for automatically adjusting feature embedding dimensions in recommendation models, avoiding large SuperNets and enabling scalable, importance-driven dimension growth or shrinkage.

Contribution

It introduces a lightweight, importance-based approach for dynamic embedding dimension search that eliminates the need for memory-intensive SuperNets.

Findings

Reduces training memory compared to SuperNet-based methods.

Effectively adjusts embedding dimensions based on feature importance.

Validated on three recommendation datasets.

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.