Cut Binary Cross Entropy: Efficient Large-Vocabulary Loss and Gradient Kernels for Sequential Recommendation
CutBCE：大词表序列推荐的高效 BCE 损失与梯度内核
🧭 大词表 BCE 训练优化 · TPU Kernel ★ 深度解读
Google Cloud
CutBCE 通过分块重算与稀疏修正避免 [B,N,V] logits 驻留 HBM，在 876k 词表上将训练显存降 65.7%、速度提 225.9%
📌 工作参考价值（4 点）
- 训练多标签召回/序列模型时，用分块 BCE：把 BCE 写成 max(x,0)+log(1+e^{-|x|}) 减去 y*x，按 vocab chunk 计算并立即释放，避免 [B,N,V] logits；PyTorch 里可实现为 chunked loss 或自定义 autograd.Function。
- 反向不要存完整 logits：自定义 VJP/rematerialization 只保留 activations、embedding 和稀疏 target IDs，逐块重算 logits 和梯度；对电商百万级 catalog，这是时间换空间，在 TPU 上因减少 HBM 往返甚至更快。
- 如果部署在 TPU/JAX，可以用 Pallas 把 embedding 梯度累积在 VMEM，降低 HBM 读写；动态查询 VMEM 预算并做序列 chunking，分布式训练用 shard_map 显式控制 psum 和 vocab_offset。
- 训练监控用 count-based 指标：前向直接统计 TP/FP/FN/TN 获得 accuracy/precision/recall/Fβ，不要为监控重建全量 logits，保持低显存。