🤖 AI Summary
This work addresses the single-sample interference problem in stochastic optimization—where parameter updates inadvertently increase the loss of individual samples due to momentum or other historical state dependencies—by reformulating the update process as a constrained optimization problem that explicitly minimizes such interference. The key insight is that the final layer alone suffices to accurately capture the network-wide second-order statistical properties essential for this mitigation. Leveraging this observation, the authors construct a computationally efficient proxy problem whose dimensionality scales with batch size rather than model size, thereby reducing relative computational overhead as models grow larger. The proposed method is compatible with mainstream optimizers such as SGD and AdamW and employs a GPU-friendly iterative solver. Experiments demonstrate that it effectively suppresses single-sample interference and consistently improves generalization performance on standard image classification benchmarks.
📝 Abstract
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.