# Cut Binary Cross Entropy: Efficient Large-Vocabulary Loss and Gradient Kernels for Sequential Recommendation

arXiv: https://arxiv.org/abs/2610.05559 · HTML: https://mlsanity.com/paper/2610.05559/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 8 of its 10 references that are in the map.
- First posted on arXiv 4 Oct 2026.

## Builds on (most cited first, 8 of 8)

- [Self-Attentive Sequential Recommendation](https://mlsanity.com/paper/1808.09781/) (2018)
- [BERT4Rec: Sequential Recommendation with Bidirectional Encoder Representations from Transformer](https://mlsanity.com/paper/1904.06690/) (2019)
- [Liger Kernel: Efficient Triton Kernels for LLM Training](https://mlsanity.com/paper/2410.10989/) (2024)
- [Cut Your Losses in Large-Vocabulary Language Models](https://mlsanity.com/paper/2411.09009/) (2024)
- [CALRec: Contrastive Alignment of Generative LLMs for Sequential Recommendation](https://mlsanity.com/paper/2405.02429/) (2024)
- [Yambda-5B -- A Large-Scale Multi-modal Dataset for Ranking And Retrieval](https://mlsanity.com/paper/2505.22238/) (2025)
- [SessionRec: Next Session Prediction Paradigm For Generative Sequential Recommendation](https://mlsanity.com/paper/2502.10157/) (2025)
- [FANS: Fast Non-Autoregressive Sequence Generation for Item List Continuation](https://mlsanity.com/paper/2304.00545/) (2023)

## Abstract (arXiv)

Industrial sequential recommender systems operate over massive item catalogs (e.g., 105--107 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/binarycrossentropyops.py.