Adaptive Training Distributions with Scalable Online Bilevel Optimization
This work addresses distribution mismatch in pretraining for practitioners using large neural networks, though it is incremental with mixed empirical results.
The paper tackles the problem of mismatched distributions between large-scale pretraining data and target application domains by proposing an algorithm that adapts pretraining distributions using a small sample of target data. The approach shows benefits over existing domain adaptation strategies in some cases but not others, with the authors providing a test to predict when it will work well.
Large neural networks pretrained on web-scale corpora are central to modern machine learning. In this paradigm, the distribution of the large, heterogeneous pretraining data rarely matches that of the application domain. This work considers modifying the pretraining distribution in the case where one has a small sample of data reflecting the targeted test conditions. We propose an algorithm motivated by a recent formulation of this setting as an online, bilevel optimization problem. With scalability in mind, our algorithm prioritizes computing gradients at training points which are likely to most improve the loss on the targeted distribution. Empirically, we show that in some cases this approach is beneficial over existing strategies from the domain adaptation literature but may not succeed in other cases. We propose a simple test to evaluate when our approach can be expected to work well and point towards further research to address current limitations.