PEACH: Personalized Negative Caching and Sampling for Unbiased Optimization of Extreme Multiclass Classification

Extreme multiclass classification with the cross-entropy loss is computationally challenging because evaluating its normalization term requires summing over a huge number of classes. Approximating this term with a small subset of classes, as in sampled softmax, yields biased stochastic gradients that in general do not help stochastic optimization converge to a stationary point of the full objective. Hard-negative mining can reduce this approximation error, but does not remove the bias and incurs additional computational cost. Recent work removes the bias by reformulating cross-entropy minimization as an equivalent min-min problem, in which an auxiliary variable learns the normalization term and unbiased stochastic gradients can be computed efficiently. We find, however, that uniform sampling within this framework produces high-variance gradient estimates and slow convergence in practice. To address this limitation, we propose $\text{PEACH}$, an unbiased optimization method based on $\text{efficient, personalized negative caching and sampling}$. PEACH maintains a small cache of hard negative classes for each example, evaluates their contributions exactly, and uses simple random sampling to estimate the sum over the remaining classes. Our theoretical analysis shows that this strategy preserves unbiasedness while substantially reducing the variance of stochastic gradients for min-min optimization. Across datasets with hundreds of thousands of classes, PEACH substantially accelerates optimization in wall-clock time compared with existing methods. At matched training loss, PEACH achieves up to $2\text{textendash}10\times$ speedup compared with biased and unbiased optimization methods, with the advantage becoming pronounced at larger scales.