Learning Effective and Efficient Embedding via an Adaptively-Masked Twins-based Layer

TL;DR

This paper introduces an adaptively-masked twins-based layer (AMTL) for embedding size selection in deep learning recommendation models, improving efficiency and flexibility while reducing memory usage by 60%.

Contribution

It proposes a novel AMTL layer that adaptively masks embedding dimensions, enabling flexible and memory-efficient embedding size selection in DLRMs.

Findings

Outperforms baseline methods on benchmark tasks.

Reduces memory usage by 60%.

Maintains performance metrics with improved efficiency.

Abstract

Embedding learning for categorical features is crucial for the deep learning-based recommendation models (DLRMs). Each feature value is mapped to an embedding vector via an embedding learning process. Conventional methods configure a fixed and uniform embedding size to all feature values from the same feature field. However, such a configuration is not only sub-optimal for embedding learning but also memory costly. Existing methods that attempt to resolve these problems, either rule-based or neural architecture search (NAS)-based, need extensive efforts on the human design or network training. They are also not flexible in embedding size selection or in warm-start-based applications. In this paper, we propose a novel and effective embedding size selection scheme. Specifically, we design 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. Extensive experimental evaluations show that the proposed scheme outperforms competitive baselines on all the benchmark tasks, and is also memory-efficient, saving 60\% memory usage without compromising any performance metrics.