Making Robust Generalizers Less Rigid with Loss Concentration

TL;DR

This paper introduces a new training criterion that penalizes poor loss concentration to improve robustness in models, especially where the difficulty gap between data points is significant, surpassing traditional sharpness-aware methods.

Contribution

It proposes a flexible loss concentration penalty that can be combined with various loss transformations to enhance robustness beyond overparameterized neural networks.

Findings

Loss concentration penalty improves robustness in simpler models.

Combining the criterion with loss transformations like CVaR enhances tail performance.

Traditional sharpness-aware methods fail under models with large difficulty gaps.

Abstract

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.