Reducing Per-Sample Harm in Stochastic Optimization

Modern optimizers combine gradients from the current mini-batch with historical optimization state, such as momentum or adaptive moments. While highly effective, aggregating across the batch and incorporating this history can produce parameter updates that increase the loss of individual samples. We term this effect harm and formalize the parameter update as an optimization problem that explicitly minimizes the conflicting impact of both batch averaging and past optimization state on current data. Because the exact formulation is intractable, we introduce a highly efficient proxy. We first reduce the problem's dimensionality to the batch size, and then drastically cut memory and speed bottlenecks by successfully restricting the optimization to the last linear layer. This hinges on the unexpected finding that this layer alone reliably captures the second-order statistics of the per-sample gradients. The resulting surrogate problem integrates readily into standard optimizers like SGD and AdamW, and can be solved using a small number of GPU-friendly iterations. Crucially, the method exhibits favorable scaling properties, as the relative computational overhead shrinks as the model size or input grows. Experiments on image classification benchmarks confirm reduced per-sample interference and improved generalization.

Modern AI models learn through iterative updates. At each iteration, a group of examples is used together to update the model. While effective, this approach can hide a problem: an update that looks good on average may still harm some individual examples in the group. We introduce a method that adjusts training updates to reduce this hidden harm. Designing such updates is challenging because training must remain efficient. To address this, we introduce efficient approximations that make the method practical for modern AI models. Across several image classification experiments, our approach improves accuracy with little extra training cost. This work shows that AI systems can learn more reliably by paying attention not only to the average behavior of a group of examples, but also to how each update affects individual examples.