PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Yaoyiran Li, Haowen Ning, Mohamed Hammad

arXiv:2610.05559v1cs.LGcs.AIcs.ARcs.IR

TL;DR

Full-vocabulary BCE for multi-label recommendation can require materializing a large logits tensor, creating substantial HBM demands. CutBCE is an exact JAX/Pallas BCE operator with a fused loss formulation and custom backward kernel. On Yambda multi-label SASRec, it reduces HBM use and increases training speed with comparable accuracy.

Abstract

from arXiv · show

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.

I. INTRODUCTION AND MOTIVATION

Sequential recommenders increasingly predict multiple next items to capture multi-intent sessions, using BCE across the full item vocabulary. At industrial vocabulary sizes, materializing the resulting logits tensor consumes substantial HBM and can cause OOM failures.

I. INTRODUCTION AND MOTIVATION

  • ID-based sequential models remain the industry mainstream, while modeling user interaction trajectories over catalogs containing hundreds of thousands to tens of millions of item IDs.The passage cites submillisecond serving latency via ANN search and lightweight model architectures as reasons for their continued use.Source: §I · p. 1
  • Multi-label prediction of the next L items provides richer supervision and captures multi-intent sessions.The paper gives L = 8 as an example; BCE is evaluated over the entire item vocabulary.Source: §I · p. 1

A. The Logits Memory Wall in Standard BCE

  • At 876,939 items, the dense logits tensor creates a severe memory wall for full-vocabulary BCE.A single local tensor at B = 128, N = 256, and V = 876,939 requires about 13.38 GiB in 32-bit floats, leaving insufficient HBM for other training state.Source: A. The Logits Memory Wall in Standard… · p. 1

B. Research Gap and Key Contributions

CutBCE addresses the gap in large-vocabulary BCE optimization with an exact JAX/Pallas operator and hardware-aware kernels. Its contributions include a fused loss formulation, a custom backward path, distributed optimizations, and training metrics that avoid reconstructing full logits.

B. Research Gap and Key Contributions

  • Prior chunked-loss efforts target softmax cross-entropy, while large-vocabulary BCE optimization remains unexplored across deep-learning ecosystems.The paper introduces CutBCE as an exact, hardware-accelerated BCE loss operator implemented in JAX and Pallas.Source: B. Research Gap and Key Contributions · p. 2
  • CutBCE decomposes exact BCE into dense background loss and sparse target correction.This is the fused direct BCE formulation.Source: B. Research Gap and Key Contributions · p. 2
  • A custom VJP pairs a JAX/XLA forward scan with a Pallas TPU backward kernel, avoiding full-logit materialization in HBM.The backward kernel acts as an explicit rematerialization mechanism.Source: B. Research Gap and Key Contributions · p. 2
  • Dynamic VMEM budgeting, hierarchical sequence chunking, and sharding-aware collective hoisting support distributed execution across device meshes.Source: B. Research Gap and Key Contributions · p. 2
  • Cut Metrics compute confusion-matrix counts during the forward pass with zero auxiliary HBM overhead.These counts support binary classification metrics without reconstructing full logits.Source: B. Research Gap and Key Contributions · p. 2
  • On Yambda’s 876k-item catalog, CutBCE reduces HBM by 65.7% and increases training speed by 225.9% with comparable accuracy.The reported multi-label SASRec results are 65.7% HBM reduction and 225.9% speedup.Source: B. Research Gap and Key Contributions · p. 2

II. RELATED WORK

Prior large-vocabulary loss optimizations focus on softmax cross-entropy, including chunked LLM losses and GPU- or PyTorch-centric kernels. The paper positions exact BCE loss and gradient optimization for large vocabularies in JAX/TPU as an unaddressed area.

Loss Optimization for Large Vocabularies

  • Existing chunked-loss frameworks target softmax cross-entropy, whereas recommendation BCE evaluates independent Bernoulli losses over the full vocabulary with sparse positive targets.The passage names CCE, Liger Kernel, and Tokamax as prior approaches and describes CutBCE as the first to optimize exact large-vocabulary BCE in JAX/TPU and PyTorch/GPU.Source: §II · p. 2

JAX, Pallas, and Accelerator Kernels

  • Prior kernels are often GPU- or PyTorch-centric, while JAX and Pallas offer TPU-focused tools for accelerator memory and kernel control.The paper notes that Pallas kernels for large-vocabulary BCE had been absent.Source: §II · p. 2

III. METHODOLOGY

CutBCE reformulates BCE as dense background loss plus sparse positive-target correction, then computes vocabulary blocks without retaining the full logits tensor. This reduces peak intermediate activation memory while sequence tiling respects accelerator VMEM capacity.

III. METHODOLOGY

  • CutBCE expresses each class loss as a dense background term minus a sparse correction active only for positive targets.The formulation uses a numerically stable log-sum-exp identity to avoid overflow for large negative logits.Source: A. Mathematical Formulation of Fused… · p. 2
  • Batch and sequence dimensions are flattened and hierarchically tiled, with inner tiles of size BN = 128 to respect accelerator VMEM capacity.Source: A. Mathematical Formulation of Fused… · p. 2
  • Vocabulary blocks are projected and accumulated one at a time, reducing peak intermediate activation memory from O(BNV) to O(BNBV).Each block's logits are discarded immediately after its loss is accumulated.Source: A. Mathematical Formulation of Fused… · pp. 2–3

B. Gradient Formulation and Custom VJP Rematerialization

The BCE logit gradient is proportional to the difference between predicted probability and binary target. A custom VJP retains only activations, embeddings, and sparse target IDs, recomputing logits by chunk during backward.

B. Gradient Formulation and Custom VJP Rematerialization

C. Hybrid Architecture and Pallas TPU Backward Kernel

CutBCE pairs a JAX/XLA forward pass with a dedicated Pallas TPU backward kernel. The backward implementation manages VMEM use and sharded execution across distributed meshes.

C. Hybrid Architecture and Pallas TPU Backward Kernel

D. Zero-Overhead Cut Metrics for Training

Cut Metrics compute classification counts during the forward pass without reconstructing full logits. They support common classification metrics during training, while ranking evaluation remains separate.

D. Zero-Overhead Cut Metrics for Training

E. Multi-Label SASRec-BCE Application Pipeline

The pipeline uses a causal, non-autoregressive SASRec model with a GPT-2 backbone to predict eight next items at once. Experiments use TPU v5e and v6e hardware with JAX and Keras 3.

E. Multi-Label SASRec-BCE Application Pipeline

  • The SASRec-based pipeline uses a GPT-2 backbone to predict the next 8 items simultaneously.Every non-padding training position predicts its next 8 items using CutBCE.Source: E. Multi-Label SASRec-BCE Application… · p. 3
  • Evaluation projects the hidden state at the last non-padding position to full-vocabulary logits and reports multi-label Hit@8.A user is counted as a hit when at least one of the eight held-out items appears among the top eight ranked items.Source: E. Multi-Label SASRec-BCE Application… · p. 3

IV. EXPERIMENTAL EVALUATION

  • Experiments run on Google Cloud TPU v5e and v6e using JAX and Keras 3.The v5e has 16 GiB HBM and the v6e has 32 GiB HBM.Source: §IV · p. 3

A. Mini-Benchmark Experiments

Single-chip TPU mini-benchmarks compare CutBCE with standard Keras BCE, including numerical equivalence and performance at large workloads. The tests measure isolated loss and gradient computation after compilation.

A. Mini-Benchmark Experiments

  • The benchmark measures isolated loss and backward-gradient computation on one TPU v6e/v5e chip across 500 steady-state iterations after JIT compilation, excluding transformer-backbone and optimizer overhead.Source: A. Mini-Benchmark Experiments · p. 3
  • At B = N = 32 and V = 100k, CutBCE matches Keras BCE within 1.1×10−7 relative loss error and 7.1×10−5 relative ℓ2 error in ∇AL and ∇EL.Unit tests also check unaligned N and V, padded and duplicate target IDs, and Pallas and pure-JAX backends.Source: A. Mini-Benchmark Experiments · p. 4
  • At B = 128, N = 128, V = 100k, standard Keras BCE crashes with OOM, while CutBCE runs in 6.93 ms on TPU v6e and 12.79 ms on TPU v5e.At B = 32, N = 32, speedups are 53.8%–91.9% on v6e and 69.4%–83.3% on v5e; speedup improves as vocabulary grows.Source: A. Mini-Benchmark Experiments · p. 4, Table II · p. 4

B. Multi-Label SASRec Experiments on Yambda Dataset

On Yambda-50M, CutBCE reduces HBM use and increases training speed for multi-label SASRec while achieving similar Hit@8 to standard BCE. The conclusion presents it as an exact, memory-efficient operator and reports that it supports large-batch, long-sequence training on a large item catalog.

B. Multi-Label SASRec Experiments on Yambda Dataset

  • The evaluation filters Yambda-50M to 9,074 users, 46.46M interactions, and 876,938 unique items, and trains for 300,000 steps on an 8-chip TPU v6e slice.The model uses a 3-layer GPT-2 backbone and AdamW; validation and test run on independent TPU slices.Source: B. Multi-Label SASRec Experiments on… · p. 4
  • At global B = 64 and N = 256, CutBCE reduces per-chip HBM from 22.37 GiB to 7.67 GiB (65.7%) and raises training speed from 6.14 to 20.01 steps/second (225.9%).CutBCE achieves 0.183 Hit@8 versus 0.186 for standard BCE; standard BCE OOMs at global B = 128 or N = 512, while CutBCE runs at 7.89–8.94 GiB HBM.Source: B. Multi-Label SASRec Experiments on… · p. 4, §V · p. 4

V. CONCLUSION AND OPEN SOURCE

  • CutBCE is presented as an exact, memory-efficient BCE operator with fused loss, custom VJP rematerialization, Pallas TPU backward kernels, sharding-aware collective hoisting, and count-based metrics.The conclusion says it is open-sourced as a modular JAX operator and handles tile sizes, sharding, and Pallas kernel dispatch automatically.Source: §V · p. 4
Loading 2610.05559v1…