Modern optimizers combine gradients from the current mini-batch with historical optimization states, such as momentum or adaptive moments. While effective, this aggregation can lead to increased loss for individual samples, a phenomenon we term 'harm'. We formalize the parameter update as an optimization problem aimed at explicitly minimizing the conflicting impacts of batch averaging and past optimization states on current data. Given the intractability of the exact formulation, we introduce a highly efficient proxy.
We first reduce the problem's dimensionality to the batch size and drastically cut memory and speed bottlenecks by restricting the optimization to the last linear layer. This unexpected finding reveals that this layer alone reliably captures the second-order statistics of per-sample gradients. The resulting surrogate problem integrates seamlessly into standard optimizers like SGD and AdamW and can be solved with a small number of GPU-friendly iterations. Importantly, the method exhibits favorable scaling properties, as the relative computational overhead shrinks with growing model size or input. Experiments on image classification benchmarks confirm the reduced per-sample interference and improved generalization.
Blogger's Review: This paper presents an innovative optimization method that highlights the importance of considering inter-sample interference in stochastic optimization. By effectively reducing dimensions and focusing on critical layers, it significantly enhances optimization, especially in large-scale models, demonstrating promising scalability and broad applicability. I look forward to seeing more in-depth research on optimization issues in the future.