Keywords: generalized geodesic, transfer learning
Abstract: Data for pretraining machine learning models often consists of collections of heterogeneous datasets. Although training on their union is reasonable in agnostic settings, it might be suboptimal when the target domain ---where the model will ultimately be used--- is known in advance. In that case, one would ideally pretrain only the dataset(s) most similar to the target one. Instead of limiting this choice to those datasets already present in the pretraining collection, here we explore extending this search to all datasets that can be synthesized as `combinations' of them. We define such combinations as multi-dataset interpolations, formalized through the notion of generalized geodesics from optimal transport (OT) theory. We compute these geodesics using a recent notion of distance between labeled datasets, and derive alternative interpolation schemes based on it: using either barycentric projections or optimal transport maps, the latter computed using recent neural OT methods. These methods are scalable, efficient, and ---notably--- can be used to interpolate even between datasets with distinct and unrelated label sets. Through various experiments in transfer learning, we demonstrate this promising new approach to targeted on-demand dataset synthesis.