Fine-Grained Embedding Dimension Optimization During Training for Recommender Systems

Huge embedding tables in modern Deep Learning Recommender Models (DLRM) require prohibitively large memory during training and inference. Aiming to reduce the memory footprint of training, this paper proposes FIne-grained In-Training Embedding Dimension optimization (FIITED). Given the observation that embedding vectors are not equally important, FIITED adjusts the dimension of each individual embedding vector continuously during training, assigning longer dimensions to more important embeddings while adapting to dynamic changes in data. A novel embedding storage system based on virtually-hashed physically-indexed hash tables is designed to efficiently implement the embedding dimension adjustment and effectively enable memory saving. Experiments on two industry models show that FIITED is able to reduce the size of embeddings by more than 65% while maintaining the trained model's quality, saving significantly more memory than a state-of-the-art in-training embedding pruning method. On public click-through rate prediction datasets, FIITED is able to prune up to 93.75%-99.75% embeddings without significant accuracy loss. Huge embedding tables in modern Deep Learning Recommendation Models (DLRM) reach terabytes in size (Lian et al., 2022). Training DLRMs usually requires model parallelism (Ivchenko et al., 2022; Sethi et al., 2023), but even with embedding tables distributed over multiple compute nodes, memory still proves a scarce resource (Lian et al., 2022). Reducing the memory cost of embedding tables is crucial to enable efficient model training and deployment of DLRM and allow for sustainable model development. The size of an embedding table is determined by the number of rows (i.e., hash size), the number of columns (i.e., embedding dimension), and the size of each value in the embedding.