TinyKG：面向知识图神经推荐系统的内存高效训练框架
信息检索
2022-12-12 v1
摘要
设计各类知识图神经网络（KGNNs）的兴趣激增，其在推荐中达到了最先进性能并提供了极佳的可解释性。其优异性能主要源于在知识图上捕获高阶邻近信息的能力。然而，大规模训练KGNNs因高内存使用而具挑战性。在前向传播中，自动微分引擎（如TensorFlow/PyTorch）通常需要缓存所有中间激活图以在反向传播中计算梯度，这导致巨大的GPU内存占用。现有工作利用多GPU分布式框架解决该问题。但这在将KGNNs部署于内存受限环境（尤其是工业级图）时带来实际挑战。本文提出TinyKG，一个基于GPU、面向推荐任务KGNNs的内存高效训练框架。具体而言，TinyKG在前向传播中使用精确激活值，同时在GPU缓冲中存储激活值的量化版本。在反向传播中，这些低精度激活值被反量化为全精度张量以计算梯度。为减小量化误差，TinyKG采用简单而有效的量化算法压缩激活值，确保无偏性与低方差。由此，KGNNs的训练内存占用大幅降低且精度损失可忽略。为评估TinyKG性能，我们在真实数据集上进行了全面实验。我们发现，采用INT2量化的TinyKG将激活图内存占用激进地降低，仅损失精度，使KGNNs可部署于内存受限设备。
引用
@article{arxiv.2212.04540,
title = {TinyKG: Memory-Efficient Training Framework for Knowledge Graph Neural Recommender Systems},
author = {Huiyuan Chen and Xiaoting Li and Kaixiong Zhou and Xia Hu and Chin-Chia Michael Yeh and Yan Zheng and Hao Yang},
journal= {arXiv preprint arXiv:2212.04540},
year = {2022}
}