Faster and Memory-Efficient Training of Sequential Recommendation Models for Large Catalogs

TL;DR

<3-5 sentence high-level summary> Sequential recommender models over large catalogs face prohibitive memory costs when using cross-entropy loss. The authors introduce Cut-Cross-Entropy (CCE) and Cut-Cross-Entropy with Negative Sampling (CCE-)—GPU-optimized losses implemented via Triton kernels that fuse logit/softmax and exploit gradient sparsity to dramatically reduce memory and accelerate training. They show memory reductions of orders of magnitude and speedups across diverse datasets, and demonstrate that reallocating saved memory to larger batches, longer sequences, or more negatives improves accuracy on large catalogs. The results extend to Bert4Rec, illustrating broad applicability, and they provide practical guidance on memory-accuracy trade-offs and sampling strategies for industrial SR training.

Abstract

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.

Figures (9)

• Figure 1: Here, we compare several methods CE, $\mathrm{CE}^{-}$, CCE, $\mathrm{CCE}^{-}$, and SCE in terms of (1) NDCG@10, (2) Training time per epoch in seconds, and (3) Memory consumption in Gb for the SASRec model across six datasets. We demonstrate the best performance achieved through optimized hyperparameters. We can see that CE, $\mathrm{CE}^{-}$ require much more time and memory usage.
• Figure 2: Illustration of CE, CCE and $\mathrm{CCE}^{-}$ implementation. CE: Materializes full logits in GPU memory, matmul and softmax are computed separately $\rightarrow$ high GPU memory usage, full $C \times E$ matmul, extra data transfers between HBM and SRAM lead to higher time delay. CCE : No logits materialization, fused matmul + softmax $\rightarrow$ store only positive logits and LSE vector, reduced HBM $\leftarrow \rightarrow$ SRAM transfers, but still full CxE matmul. $\mathrm{CCE}^{-}$: No logits materialization, fused matmul + softmax on sampled items $\rightarrow$ store only positive logits and LSE vector, reduced HBM $\leftarrow \rightarrow$ SRAM transfers, matmul only for sampled columns of C and corresponding E.
• Figure 3: An illustration of computing gradients for the weights of the final layer.
• Figure 4: Histograms of the gradient distributions at the final layer at the beginning (left) and end (right) of the training, the Megamarket dataset.
• Figure 5: Memory scaling across three dimensions: (1) marker size ($sl$: sequence length), (2) color ($bs$: batch size), and (3) marker shape ($ns$: number of negative samples). Asterisks denote points using the CE loss. For example, brown markers (larger batch sizes) generally yield better average performance, though the largest batch size does not consistently dominate. Higher NDCG occurs with 128–512 negative samples (squares-hexagons); only the Zvuk dataset shows improved performance with asterisks (full $\mathrm{CE}^{-}$ loss).