TT-Rec: Tensor Train Compression for Deep Learning Recommendation Models

By Chunxing Yin, Bilge Acun, Xing Liu, Carole-Jean Wu. 2021

122 citations recorded by the source corpus.

The memory capacity of embedding tables in deep learning recommendation models (DLRMs) is increasing dramatically from tens of GBs to TBs across the industry. Given the fast growth in DLRMs, novel solutions are urgently needed, in order to enable fast and efficient DLRM innovations. At the same time, this must be done without having to exponentially increase infrastructure capacity demands. In this paper, we demonstrate the promising potential of Tensor Train decomposition for DLRMs (TT-Rec), an important yet under-investigated context. We design and implement optimized kernels (TT-EmbeddingBag) to evaluate the proposed TT-Rec design. TT-EmbeddingBag is 3 times faster than the SOTA TT implementation. The performance of TT-Rec is further optimized with the batched matrix multiplication and caching strategies for embedding vector lookup operations. In addition, we present mathematically and empirically the effect of weight initialization distribution on DLRM accuracy and propose to initialize the tensor cores of TT-Rec following the sampled Gaussian distribution. We evaluate TT-Rec across three important design space dimensions -- memory capacity, accuracy, and timing performance -- by training MLPerf-DLRM with Criteo's Kaggle and Terabyte data sets. TT-Rec achieves 117 times and 112 times model size compression, for Kaggle and Terabyte, respectively. This impressive model size reduction can come with no accuracy nor training time overhead as compared to the uncompressed baseline.

Subfields: Artificial Intelligence, Computational Mathematics, Information Systems, Efficiency & Compression

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computational Mathematics

In Information Systems

In Efficiency & Compression

Citations in the corpus

Sustainable AI: Environmental Implications, Challenges and Opportunities (2021, methodology): “For example, the Tensor-Train compression technique (TT-Rec) achieves more than 100× memory capacity reduction with negligible training time and accuracy trade-off (Yin et al., 2021).”
RecShard: statistical feature-based memory optimization for industry-scale neural recommendation (2022, background): “MLP layers (both top and bottom) are replicated across all trainers (GPUs in figure) in a data-parallel manner, while EMBs are sharded across trainers to exploit model-parallelism [17, 18, 31, 44].”
High-performance, Distributed Training of Large-scale Deep Learning Recommendation Models (2021, methodology): “…large embedding table size, we also applied different embedding compression techniques like row-wise sparse AdaGrad; low/mixed-precision training us-ing a high-precision cache backed by low precision embedding tables [57]; and advanced factorization techniques like Tensor Train compression [59].”