[论文解读] TT-Rec: Tensor Train Compression for Deep Learning Recommendation Models

Chunxing Yin, Bilge Acun+2 more|arXiv (Cornell University)|Jan 25, 2021

Tensor decomposition and applications参考文献 32被引用 13

一句话总结

TT-Rec 提出了一种基于张量列车分解的方法，用于压缩深度学习推荐模型（DLRMs）中的嵌入表，实现高达117倍的模型尺寸压缩，且无精度损失，仅增加13.9%的训练时间开销。该方法采用优化的TT-EmbeddingBag内核、自适应缓存机制，并提出一种基于高斯分布的新初始化策略，以在大规模场景下保持性能与精度。

ABSTRACT

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.

研究动机与目标

• 解决深度学习推荐模型（DLRMs）日益增长的内存需求，这些模型如今需要TB量级的嵌入表。
• 在不牺牲精度或产生不可接受的训练时间开销的前提下，减少模型内存容量。
• 探索张量化作为一种系统化、参数化的嵌入压缩方法，以在保持模型质量的同时压缩 DLRM 嵌入。
• 设计一种计算高效且硬件感知的压缩框架，适用于工业规模的 DLRMs。
• 通过优化内核设计与缓存策略，实现大规模推荐模型的实用化、低开销训练。

提出的方法

• 使用张量列车（TT）分解将 DLRMs 中的大规模嵌入表替换为一系列低秩矩阵乘积。
• 实现一种自定义的、优化的 TT-EmbeddingBag 内核，加速基于 TT 的嵌入查找，相比最先进 TT 实现提升3倍速度。
• 引入一种缓存机制，将频繁访问的嵌入向量以未压缩形式存储，以减少重复计算并恢复精度。
• 提出一种新颖的权重初始化策略，从高斯分布中采样，以稳定训练并提升模型精度。
• 利用批量矩阵乘法进一步优化 GPU 加速系统上的推理与训练性能。
• 在多个设计维度（内存、精度、速度）上调整 TT-秩与压缩参数，以识别帕累托最优配置。

实验结果

研究问题

• RQ1 张量列车分解是否能在不降低精度的前提下，实现 DLRMs 中显著的模型压缩？
• RQ2 权重初始化的选择如何影响 TT-Rec 在推荐模型中的性能与收敛性？
• RQ3 缓存能在多大程度上缓解训练期间解压 TT 结构嵌入所带来的性能开销？
• RQ4 在不同的 TT-秩与嵌入维度设置下，内存压缩、训练时间与模型精度之间的权衡关系如何？
• RQ5 TT-Rec 是否能实现数量级的压缩，同时保持与工业规模训练流水线及硬件加速器的兼容性？

主要发现

• TT-Rec 在 Criteo Kaggle 和 Terabyte 数据集上分别实现了117倍和112倍的模型尺寸压缩，与未压缩基线相比无精度下降。
• 所提出的 TT-EmbeddingBag 内核相比最先进 TT 实现快3倍，显著提升了计算效率。
• 在90%缓存命中率下，TT-Rec 的训练时间性能优于标准 EmbeddingBag，验证了缓存策略的有效性。
• 即使在极端压缩条件下，模型仍保持与基线相同的精度，验证了所提出初始化与压缩方案的鲁棒性。
• 与基线相比，训练时间开销仅增加13.9%，使 TT-Rec 在大规模在线训练中具备实际可行性。
• 该方法实现了内存容量、精度与训练速度之间的帕累托最优权衡，且可通过可调参数适配不同硬件与部署约束。