Reducing Per-Sample Interference in Stochastic Optimization
Apostolos Avranas
cs.LG, cs.AI, stat.ML
Submitted: 2026-06-28
Journal ref: Proceedings of the 43rd International Conference on Machine Learning (ICML 2026)
Code: https://github.com/avranasa/sample-harm-reduction
License: http://creativecommons.org/licenses/by-nc-nd/4.0/
The gist: Modern optimizers combine gradients from the current mini-batch with historical optimization state, such as momentum or adaptive moments.
Terminology
Abstract
Modern optimizers combine gradients from the current mini-batch with historical optimization state, such as momentum or adaptive moments. While effective, this standard practice can produce parameter updates that actively increase the loss of individual samples. We term this phenomenon per-sample interference and propose redefining the parameter update as an optimization problem that explicitly minimizes it. Because the exact formulation of the problem is computationally prohibitive, we introduce a highly efficient surrogate. By reducing the problem's dimensionality to the batch size and restricting the optimization to the last linear layer, we overcome memory and speed bottlenecks. This strategy hinges on our unexpected finding that this layer alone can reliably capture core second-order statistics of the full network. 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.
Sources
- Gradient Episodic Memory with a Soft Constraint for Continual Learning
- Adam: A Method for Stochastic Optimization
- Efficient Per-Example Gradient Computations in Convolutional Neural Networks
- Per-example gradients: a new frontier for understanding and improving optimizers
- Opacus: User-Friendly Differential Privacy Library in PyTorch
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: The Predictive Cost of ELBO Binding
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks