Changing the Training Data Distribution to Reduce Simplicity Bias Improves In-distribution Generalization

TL;DR

This paper explores how modifying training data distribution can reduce simplicity bias, leading to better in-distribution generalization, by comparing gradient descent and sharpness-aware minimization, and proposing a clustering-based upsampling method.

Contribution

It introduces a novel data reweighting technique that alleviates simplicity bias and improves generalization across various models and datasets.

Findings

SAM learns features more uniformly than GD in early training.

The proposed method improves generalization performance across multiple datasets.

Combining the method with existing strategies achieves state-of-the-art results.

Abstract

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. Our code is available at https://github.com/BigML-CS-UCLA/TADA.