Batch Size or Negatives? A Selection Rule for Memory-Constrained Recommender Training
Artyom Sabitov, Daniil Volkov, Alexey Zaytsev
cs.LG
2026-08-11
Theory and four real benchmarks show sampled-softmax training converges fastest by maximizing batch size under a fixed memory budget, not by piling on negatives.
The last layer of a large recommender is a massive multiclass problem: the item vocabulary runs to tens of millions, and every example must compute logits and gradients over all of them, costing O(nK) memory. The standard fix is sampled softmax, drawing only k negative items per batch (k much smaller than K) to bring memory down to O(nk).
But once the memory budget is fixed at B = n x k, how do you split it? Spend it on a bigger batch (large n) or more negatives (large k)? It is a concrete engineering question with no clear answer until now.
The authors decompose the sampled-softmax gradient variance and, under standard assumptions (logits approximately normal, variance bounded, positive-class gradient much smaller than negative), derive an upper bound:
variance <= Cy/(Kn) + Cphat/(Bk) + const
The first term falls with batch size n, the second with negative count k. The conclusion: to minimize variance (and thus converge fastest), push n as large as possible. The theoretical optimum is n near B, k near 1; given that distinct positive classes within a batch require k >= sqrt(B), the practical optimum is n roughly equal to k roughly equal to sqrt(B).
One-line rule: given the memory, pack in as many examples as you can; negatives just need to be enough.
On synthetic data and four real sequential-recommendation datasets (MovieLens-1M, MovieLens-20M, Gowalla, Netflix), three same-budget configs were compared: (32,512), (64,256), (128,128):
| Dataset | AUL (32,512) -> (128,128) |
| MovieLens-20M | 233.1 -> 61.5 |
| Gowalla | 158.0 -> 43.7 |
| Netflix | best at (128,128) |
AUL (area under the loss curve; smaller means faster convergence) dropped as the n/k ratio rose on every dataset, with the large-batch config converging clearly faster and final quality (NDCG@10) matching or beating the rest. The authors also tried an unbiased gradient correction derived from their theorem; it showed no statistically significant difference from vanilla cross-entropy. The correction did nothing.
This is a counterintuitive but immediately usable result. Many engineers assume more negatives is better and pump up k. Under memory pressure the opposite holds: give the budget to batch size and you converge faster with no quality loss. The lesson is not limited to recommendation; any large-classification head trained with sampled softmax (language models, retrieval) benefits.
The authors admit the scale is modest, using a single strong architecture and common datasets, leaving larger-scale validation for later. The theory covers only the last layer (multinomial logistic regression), not the full deep network. The theoretically "more correct" unbiased gradient correction was a wash in practice, a negative result the authors record honestly but do not explain.