Beyond Loss Gaps: Stability and Geometry in Dynamic Sample Reweighting

Abstract

Dynamic sample reweighting can emphasize informative examples, but it also changes update geometry. Under interpolation, the usual loss-gap view is sufficient; outside it, minibatch-dependent weights can amplify the mean-squared deviation of the weighted gradient from the full gradient by a factor quadratic in the minibatch size relative to uniform averaging. We construct a well-conditioned deterministic full-batch problem where the best local linear complexity of entropy-smoothed loss reweighting grows with component-gradient variation, while uniform gradient descent is unaffected by that variation. Motivated by this failure, we propose Geometry-Aware Reweighting (GAR), which retains the quadratic update norm in the same descent calculation. GAR preserves the usual minibatch-variance scale of uniform SGD, improves its convergence upper bound by an explicit nonnegative margin, and locally damps the directions in which loss-only weighting creates additional curvature. Experiments on language models from 120M to 300M parameters show improved average log-perplexity over uniform AdamW and loss-only reweighting.

Chulhee Yun

Associate Professor

I am an Associate Professor at KAIST AI. I am interested in optimization and machine learning theory.