Batch Size or Negatives? A Selection Rule for Memory-Constrained Recommender Training

Large-scale neural recommender systems are typically trained with a softmax cross-entropy objective over the full item vocabulary. For a typical large number of possible items K, the final classification layer dominates memory, requiring O(nK) logits and gradients to materialize for a batch of n examples. Sampled softmax reduces this cost by restricting the objective to only k ≪ K candidate negative items, resulting in an O(nk) memory. However, for a fixed budget B = nk, it remains unclear whether one should prioritize larger batches or the inclusion of more negative items.

We address this question by analyzing sampled-softmax training under a fixed memory constraint. Under standard smoothness and variance assumptions, our theoretical evidence suggests that the fastest convergence arises from an n ∼ B, k ∼ 1 allocation. So, an actionable rule is to include as many objects as possible given computational constraints. Our theory is supported by controlled synthetic and four real sequential recommendation benchmarks, including MovieLens-20M. The suggested configuration achieves faster convergence and better final recommendation quality than imbalanced alternatives within the same memory constraint. These findings provide a theoretical and empirical foundation for configuring memory during the training of recommender systems.

The paper formalizes this trade-off as a constrained stochastic optimization problem with B = n · k, where n controls mini-batch noise and k controls class-sampling noise. The key contribution is a theoretical analysis showing that these two noise sources play different roles. Under standard smoothness and variance assumptions, we demonstrate that we should maximize the number of objects in batch with sampling small number of classes. This allocation yields near-optimal convergence behavior, providing an actionable guideline.

The theoretical analysis focuses on the final softmax layer, which can be viewed as a multi-class logistic regression model on top of learned sequence representations. The analysis shows that the gradient variance decomposes into a mini-batch component controlled by the number of objects n and a class-sampling component controlled by the number of sampled negative classes k. Balancing these two components leads to maximization of the number of considered objects: n ≍ B, k ≍ 1. However, in practice computations become efficient for k ≥ n and we use the rule n = k = √B.

The contributions are: (1) a framework for convergence analysis for stochastic first order methods, formulating sampled-softmax recommender training as stochastic optimization under a fixed memory constraint nk ≤ B, leading to a theoretical framework for analyzing stochastic gradients with two sources of randomness: mini-batch sampling and class sampling; (2) a theoretically optimal allocation rule, deriving a practical allocation suggesting that we should prioritize sampling examples to sampling classes in typical applied settings, giving a guideline for configuring sampled-softmax training without an expensive grid search; (3) empirical validation of the rule on sequential recommendation benchmarks, including MovieLens-1M and Gowalla, where proposed configurations converge faster and often achieve better final recommendation quality than strongly imbalanced alternatives under the same memory budget.

Extensive experiments on MovieLens-1M and Gowalla with various memory budgets and optimizers confirm the theoretical findings. There, configurations with minimum possible values of k consistently outperform imbalanced ones in both convergence speed and final recommendation quality across the optimizers considered, including pure SGD and Adam. The empirical results show that in the vast majority of scenarios, especially for a better performing Adam optimizer, a configuration with less k achieves the minimum AUL (Area Under the Validation Loss Curve). The empirical evidence shows that the AUL is a strictly decreasing function of the n/k ratio, confirming that given a rigid hardware memory budget B, it is consistently more advantageous to maximize the batch size n while scaling down the number of negative classes k.

A negative result obtained during evaluation concerns the choice of the objective function. When comparing the vanilla cross-entropy loss against the unbiased cross-entropy loss with correction terms, no statistically significant difference was observed in terms of optimization speed or final recommendation metrics. The corrected objective neither accelerates nor slows the training process, and its loss curves completely overlap with the standard baseline. This suggests that while class-imbalance correction is theoretically sound, its practical impact on optimization dynamics within sequential recommendation architectures (like SASRec) remains negligible. The accelerated convergence is purely driven by the structural optimization benefits of the (n, k) trade-off rather than modifications to the loss landscape.

Improvements to AI Systems:

1. Memory-Aware Training Scheduler for Recommenders: Implement an adaptive configuration rule that, given a fixed memory budget B, automatically sets batch size n about B and negative samples k about 1 (or n = k = sqrt B for practical efficiency). This eliminates manual grid search and ensures near-optimal convergence under hardware constraints.

2. Noise-Source Decomposition for Gradient Variance Reduction: Modify the optimizer to explicitly separate and balance two stochastic noise sources—mini-batch sampling noise (controlled by n) and class-sampling noise (controlled by k). The system can dynamically adjust n and k during training to minimize total gradient variance, leading to faster convergence and better final accuracy.

3. Budget-Optimal Negative Sampling Strategy: Replace fixed negative sampling ratios with a rule that prioritizes maximizing the number of training examples per step over the number of negatives. This yields a more stable and efficient training process, especially for large-scale item catalogs, without requiring loss-function corrections.

4. Loss-Function Simplification: Drop unbiased cross-entropy correction terms (e.g., importance weighting) from the objective, as they show no measurable benefit in optimization speed or final quality for sequential recommenders. This reduces computational overhead and model complexity without sacrificing performance.

5. Convergence-Predictive Validation Metric: Use the Area Under the Validation Loss Curve (AUL) as a primary early-stopping and hyperparameter selection criterion, since it is a strictly decreasing function of the n/k ratio. This allows automated tuning of memory allocation to achieve faster convergence and better final recommendation quality.

---

What the Improved AI System Can Do:

• Train large-scale recommender systems (e.g., SASRec, BERT4Rec) with up to 10–100× less memory for the softmax layer, enabling deployment on smaller GPUs or larger batch sizes on existing hardware.

• Achieve faster convergence (fewer steps to target accuracy) and higher final recommendation quality (e.g., Recall@10, NDCG@10) under the same memory budget, by simply reallocating memory from negatives to batch examples.

• Automatically configure training without manual hyperparameter tuning for n and k, saving days of experimentation.

• Maintain robust performance across different optimizers (SGD, Adam) and memory budgets, making it suitable for both research and production environments.

• Simplify the loss function, reducing code complexity and computational cost, while retaining identical optimization dynamics.