中文

TT-Rec:面向深度学习推荐模型的张量列压缩

机器学习 2021-01-29 v1

摘要

深度学习推荐模型(DLRMs)中嵌入表的内存容量在整个行业正从数十 GB 急剧增长到 TB 级。鉴于 DLRMs 的快速增长,迫切需要有新颖的解决方案,以实现快速且高效的 DLRM 创新。同时,必须在不使基础设施容量需求呈指数级增长的情况下完成这一目标。在本文中,我们展示了张量列分解用于 DLRMs(TT-Rec)的 promising 潜力,这是一个重要但未被充分研究的场景。我们设计并实现了优化内核(TT-EmbeddingBag)来评估所提出的 TT-Rec 设计。TT-EmbeddingBag 比当前最优的 TT 实现快 3 倍。TT-Rec 的性能通过批处理矩阵乘法和用于嵌入向量查找操作的缓存策略得到进一步优化。此外,我们从数学和实证上展示了权重初始化分布对 DLRM 精度的影响,并提议遵循采样高斯分布初始化 TT-Rec 的张量核。我们使用 Criteo 的 Kaggle 和 Terabyte 数据集训练 MLPerf-DLRM,从内存容量、精度和时序性能三个重要设计空间维度评估 TT-Rec。TT-Rec 对 Kaggle 和 Terabyte 分别实现了 117 倍和 112 倍的模型尺寸压缩。与未压缩基线相比,这种显著的模型尺寸缩减可以带来零精度损失且零训练时间开销。

关键词

引用

@article{arxiv.2101.11714,
  title  = {TT-Rec: Tensor Train Compression for Deep Learning Recommendation Models},
  author = {Chunxing Yin and Bilge Acun and Xing Liu and Carole-Jean Wu},
  journal= {arXiv preprint arXiv:2101.11714},
  year   = {2021}
}

备注

To appear in Conference on Machine Learning and Systems (MlSys 2021, https://mlsys.org/)