A Challenge in Reweighting Data with Bilevel Optimization

In many practical learning scenarios, there is a discrepancy between the training and testing distribution. For instance, when training large language models, we may have access to a training set that contains many low-quality data points from different sources and want to train a model on this dataset to perform well on a testing set that contains a few high-quality points [7, 3, 18]. An appealing way to solve this problem is data reweighting [21, 23, 26], where one attributes one weight to each data point in the training set. The weight of a training sample should reflect how much this sample resembles the testing set and helps the model perform well on it. Figure 1 illustrates the general principle. Learning the optimal weights can be cast as a bilevel optimization problem [9], where the optimal weights are such that training the model with these weights leads to the smallest test loss possible. The weights are usually constrained to sum to one, leading to an optimization problem on the simplex, which is usually solved with mirror descent [19].