Changing the Training Data Distribution to Reduce Simplicity Bias Improves In-distribution Generalization
This addresses the issue of poor generalization due to simplicity bias for practitioners training neural networks, offering an incremental improvement by combining with existing methods like SAM and data augmentation.
The paper tackles the problem of simplicity bias in neural networks by proposing a method to modify the training data distribution, which improves in-distribution generalization; it achieves state-of-the-art performance on multiple datasets like CIFAR10, STL10, CINIC10, Tiny-ImageNet, CIFAR100, and with models such as ResNet18, ResNet34, VGG19, and DenseNet121.
Can we modify the training data distribution to encourage the underlying optimization method toward finding solutions with superior generalization performance on in-distribution data? In this work, we approach this question for the first time by comparing the inductive bias of gradient descent (GD) with that of sharpness-aware minimization (SAM). By studying a two-layer CNN, we rigorously prove that SAM learns different features more uniformly, particularly in early epochs. That is, SAM is less susceptible to simplicity bias compared to GD. We also show that examples containing features that are learned early are separable from the rest based on the model's output. Based on this observation, we propose a method that (i) clusters examples based on the network output early in training, (ii) identifies a cluster of examples with similar network output, and (iii) upsamples the rest of examples only once to alleviate the simplicity bias. We show empirically that USEFUL effectively improves the generalization performance on the original data distribution when training with various gradient methods, including (S)GD and SAM. Notably, we demonstrate that our method can be combined with SAM variants and existing data augmentation strategies to achieve, to the best of our knowledge, state-of-the-art performance for training ResNet18 on CIFAR10, STL10, CINIC10, Tiny-ImageNet; ResNet34 on CIFAR100; and VGG19 and DenseNet121 on CIFAR10.