Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Fine-Grained Embedding Dimension Optimization During Training for Recommender Systems

Qinyi Luo, Penghan Wang, Wei Zhang, Fan Lai, Jiachen Mao, Xiaohan Wei, Jun Song, Wei-Yu Tsai, Shuai Yang, Yuxi Hu, Xuehai Qian

TL;DR

Experiments show that FIITED can reduce DLRM embedding size by more than 65% while preserving model quality, outperforming state-of-the-art in-training embedding pruning methods and increasing the reduction ratio by 1.67, improvement in reduction ratio compared to baselines while improving model throughput.

Abstract

Huge embedding tables in modern deep learning recommender models (DLRM) require prohibitively large memory during training and inference. This paper proposes FIITED, a system to automatically reduce the memory footprint via FIne-grained In-Training Embedding Dimension pruning. By leveraging the key insight that embedding vectors are not equally important, FIITED adaptively adjusts the dimension of each individual embedding vector during model training, assigning larger dimensions to more important embeddings while adapting to dynamic changes in data. We prioritize embedding dimensions with higher frequencies and gradients as more important. To enable efficient pruning of embeddings and their dimensions during model training, we propose an embedding storage system based on virtually-hashed physically-indexed hash tables. Experiments on two industry models and months of realistic datasets show that FIITED can reduce DLRM embedding size by more than 65% while preserving model quality, outperforming state-of-the-art in-training embedding pruning methods. On public datasets, FIITED can reduce the size of embedding tables by 2.1x to 800x with negligible accuracy drop, while improving model throughput.

Fine-Grained Embedding Dimension Optimization During Training for Recommender Systems

TL;DR

Experiments show that FIITED can reduce DLRM embedding size by more than 65% while preserving model quality, outperforming state-of-the-art in-training embedding pruning methods and increasing the reduction ratio by 1.67, improvement in reduction ratio compared to baselines while improving model throughput.

Abstract

Huge embedding tables in modern deep learning recommender models (DLRM) require prohibitively large memory during training and inference. This paper proposes FIITED, a system to automatically reduce the memory footprint via FIne-grained In-Training Embedding Dimension pruning. By leveraging the key insight that embedding vectors are not equally important, FIITED adaptively adjusts the dimension of each individual embedding vector during model training, assigning larger dimensions to more important embeddings while adapting to dynamic changes in data. We prioritize embedding dimensions with higher frequencies and gradients as more important. To enable efficient pruning of embeddings and their dimensions during model training, we propose an embedding storage system based on virtually-hashed physically-indexed hash tables. Experiments on two industry models and months of realistic datasets show that FIITED can reduce DLRM embedding size by more than 65% while preserving model quality, outperforming state-of-the-art in-training embedding pruning methods. On public datasets, FIITED can reduce the size of embedding tables by 2.1x to 800x with negligible accuracy drop, while improving model throughput.
Paper Structure (18 sections, 1 equation, 15 figures, 4 tables, 1 algorithm)

This paper contains 18 sections, 1 equation, 15 figures, 4 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Background and Motivations
  3. Deep Learning Recommendation Models
  4. Embedding Table Characteristics
  5. In-Training Embedding Dimension Optimization
  6. FIITED Overview
  7. Dimension Optimization Algorithm
  8. System design
  9. Overhead optimization
  10. Experimental Results
  11. Methodology
  12. Memory Saving Results
  13. Overhead analysis
  14. Ablation Study
  15. Related Work
  16. ...and 3 more sections

Figures (15)

  • Figure 1: Existing Embedding Dimension Search (EDS) methods vs. FIITED.
  • Figure 2: Embedding characteristics. (a) Histogram of the number of ranks needed to preserve a certain percentage of squared SVD values for embedding tables. (b) Normalized row access and gradient squared sum of randomly sampled rows from 10 biggest tables. Both axes are log-scale. (c) Importance rankings of the top 10 important features change over time. (d) SVD characteristics change over time. The average absolute change is 4.4.
  • Figure 3: Change of dimension importance on two DLRM systems with different dimension sizes. The deeper color indicates a larger importance value.
  • Figure 4: In-training embedding dimension pruning.
  • Figure 5: Zero padding vs. linear projection given the same per-feature embedding dimensions.
  • ...and 10 more figures