Making Robust Generalizers Less Rigid with Loss Concentration
This work tackles the problem of achieving robust generalization on rare or difficult data points for machine learning practitioners, offering a more flexible alternative to existing methods, though it appears incremental as it builds on prior sharpness-aware and loss transformation techniques.
The paper addresses the breakdown of sharpness-aware minimization for robust generalization in simpler models with extreme difficulty gaps, proposing a training criterion that penalizes poor loss concentration and can be combined with loss transformations like exponential tilting, CVaR, or DRO to control tail emphasis.
While the traditional formulation of machine learning tasks is in terms of performance on average, in practice we are often interested in how well a trained model performs on rare or difficult data points at test time. To achieve more robust and balanced generalization, methods applying sharpness-aware minimization to a subset of worst-case examples have proven successful for image classification tasks, but only using overparameterized neural networks under which the relative difference between "easy" and "hard" data points becomes negligible. In this work, we show how such a strategy can dramatically break down under simpler models where the difficulty gap becomes more extreme. As a more flexible alternative, instead of typical sharpness, we propose and evaluate a training criterion which penalizes poor loss concentration, which can be easily combined with loss transformations such exponential tilting, conditional value-at-risk (CVaR), or distributionally robust optimization (DRO) that control tail emphasis.