Batch Size or Negatives? A Selection Rule for Memory-Constrained Recommender Training
Artyom Sabitov, Daniil Volkov, Alexey Zaytsev
cs.LG
2026-08-11
理论加四个真实数据集证明,sampled softmax 在固定内存下应尽量做大 batch,负样本够用就行,收敛更快。
大型推荐系统的最后一层是个超大的多分类:词表里有上千万个商品,每个样本都要对全部商品算一遍 logit 和梯度,显存开销是 O(nK)。常见解法是 sampled softmax,每批只采 k 个负样本(k 远小于 K),显存降到 O(nk)。
但固定显存预算 B = n×k 之后,该怎么分配?是把预算花在更大的 batch(n 大),还是更多的负样本(k 大)?这是个很实际的工程问题,之前没有清楚的答案。
作者把 sampled softmax 的梯度方差拆开分析,在「logit 近似正态」「方差有界」「正类梯度远小于负类」几个标准假设下,得到一个上界:
方差 ≤ Cy/(Kn) + Cp̂/(Bk) + 常数
第一项随 batch size n 下降,第二项随负样本数 k 下降。结论是:要最小化方差(也就是最快收敛),应该尽量做大 n。理论上最优是 n 接近 B、k 接近 1;考虑到每批样本的正类不能重复采(要求 k ≥ √B),实际可行的最优是 n ≈ k ≈ √B。
一句话规则:内存够的情况下,尽量多塞样本,负样本够用就行。
在合成数据和四个真实序列推荐数据集(MovieLens-1M、MovieLens-20M、Gowalla、Netflix)上,对比三组同预算配置 (32,512)、(64,256)、(128,128):
| 数据集 | AUL (32,512) → (128,128) |
| MovieLens-20M | 233.1 → 61.5 |
| Gowalla | 158.0 → 43.7 |
| Netflix | 最优在 (128,128) |
AUL(损失曲线下面积,越小收敛越快)在所有数据集上都随 n/k 比例上升而下降,大 batch 组收敛明显更快,最终推荐质量(NDCG@10)也更好或持平。作者还试了用定理推出的无偏梯度修正项,结果跟普通交叉熵没有统计显著差异,修正项没帮上忙。
这是个反直觉但能直接用的结论。很多工程师默认「负样本越多越好」,会去把 k 撑大。这篇说明在显存吃紧时反过来才对:把预算让给 batch size,收敛更快、质量不差。结论不限于推荐,任何用 sampled softmax 训大分类头的场景(语言模型、检索)都适用。
作者承认实验规模偏小,只用了一个强架构和几个常见数据集,更大规模的验证留待后续。理论分析只针对最后一层(多类逻辑回归),没有覆盖完整深度网络。无偏梯度修正项这个理论上「更对」的做法实测没用,作者诚实记录了这个负面结果,但没给出解释。