Efficient loss computation speeds up large-scale sequential recommendations
Cut Binary Cross Entropy: Efficient Large-Vocabulary Loss and Gradient Kernels for Sequential Recommendation
Machine LearningArtificial IntelligenceHardware ArchitectureInformation Retrieval
Summary
Big recommendation systems suggest items to users from huge catalogs with millions of products. Training these systems takes a lot of memory and sometimes crashes because of how the math is done. The authors created CutBCE, a faster and more memory-efficient way to calculate the loss used during training, specifically for multi-label recommendation tasks with very large item lists. Their approach reduces memory use and speeds up training on specialized hardware while keeping recommendation accuracy high.
What this means in practice
- •For recommender system engineers: Train multi-label sequential recommendation models on massive catalogs faster and with less memory using hardware-accelerated loss operators.
- •For machine learning platform teams: Integrate efficient BCE loss kernels to improve training speed and reduce memory footprint on TPU-based distributed training workloads.
Authors
Yaoyiran Li, Haowen Ning, Mohamed Hammad
Abstract
Industrial sequential recommender systems operate over massive item catalogs (e.g., 10^5--10^7 items). Multi-label recommendation models are trained with Binary Cross-Entropy (BCE) loss over the full vocabulary, but standard BCE materializes a dense [B, N, V] logits tensor in High Bandwidth Memory (HBM), incurring prohibitive $O(BNV)$ memory and fatal Out-Of-Memory (OOM) errors. While chunked loss optimizations exist for Softmax Cross-Entropy in LLMs, large-scale multi-label BCE optimization remains unexplored across deep learning ecosystems. We propose CutBCE, an exact, hardware-accelerated BCE loss and gradient operator implemented in JAX and Pallas for large-vocabulary workloads. CutBCE introduces (1) an exact fused reformulation evaluating dense background loss and sparse target corrections; (2) a custom Vector-Jacobian Product (VJP) with a dedicated Pallas TPU backward kernel computing logit tiles on-chip in both passes so logits and their gradients never reside in HBM; (3) dynamic VMEM budgeting and sharding-aware collective hoisting for distributed meshes; and (4) count-based zero-overhead training metrics. On single-chip TPU v5e/v6e mini-benchmarks, CutBCE eliminates OOM errors with up to 91.9% speedup. On 8-chip TPU slice training for multi-label SASRec with 876k items (Yambda-50M), CutBCE reduces peak HBM by 65.7% (>14 GiB saved per chip) and increases training speed by 225.9% with comparable accuracy. CutBCE is open-sourced at https://github.com/AI-Hypercomputer/RecML/blob/main/recml/core/ops/binary_cross_entropy_ops.py.