提出CCE-方法,让推荐模型训练更快更省内存。
Faster and Memory-Efficient Training of Sequential Recommendation Models for Large Catalogs
- 用改进的交叉熵损失配合负采样,降低显存占用。
- 训练速度提升2倍,内存减少10倍以上。
- 适合大商品库场景,尤其适合工业级推荐系统。
基于Transformer的序列推荐模型在实际应用中广泛使用,需频繁重训练以适应用户偏好的动态变化。然而,训练这类模型常因评分大规模商品库(通常超过数千项)导致高计算成本,主要源于交叉熵损失使峰值显存随商品库大小、批次大小和序列长度成比例增长。当前业界通常通过交叉熵损失结合负采样来缓解显存压力,但少量负样本会损害模型性能。本文发现:增加负样本数和批次大小能进一步提升性能,但很快超出工业级GPU(约40GB)的显存容量。为此,我们提出CCE-方法,实现一种高效的交叉熵损失与负采样结合的GPU实现,训练速度最高提升2倍,显存消耗降低10倍以上。利用此节省的显存,可在大商品库数据集上显著提升模型精度。我们还分析了关键显存超参数,指出应平衡负样本数量与批次大小,而非单一最大化。为促进落地,我们开源了基于Triton的高效实现内核。
原文摘要 · Abstract (English)
Sequential recommendations (SR) with transformer-based architectures are widely adopted in real-world applications, where SR models require frequent retraining to adapt to ever-changing user preferences. However, training transformer-based SR models often encounters a high computational cost associated with scoring extensive item catalogs, often exceeding thousands of items. This occurs mainly due to the use of cross-entropy loss, where peak memory scales proportionally to catalog size, batch size, and sequence length. Recognizing this, practitioners in the field of recommendation systems typically address memory consumption by integrating the cross-entropy (CE) loss with negative sampling, thereby reducing the explicit memory demands of the final layer. However, a small number of negative samples would degrade model performance, and as we demonstrate in our work, increasing the number of negative samples and the batch size further improves the model's performance, but rapidly starts to exceed industrial GPUs' size (~40Gb). In this work, we introduce the CCE- method, which offers a GPU-efficient implementation of the CE loss with negative sampling. Our method accelerates training by up to two times while reducing memory consumption by more than 10 times. Leveraging the memory savings afforded by using CCE- for model training, it becomes feasible to enhance its accuracy on datasets with a large item catalog compared to those trained with original PyTorch-implemented loss functions. Finally, we perform an analysis of key memory-related hyperparameters and highlight the necessity of a delicate balance among these factors. We demonstrate that scaling both the number of negative samples and batch size leads to better results rather than maximizing only one of them. To facilitate further adoption of CCE-, we release a Triton kernel that efficiently implements the proposed method.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。