On the Trade-off of Intra-/Inter-class Diversity for Supervised Pre-training

Pre-training datasets are critical for building state-of-the-art machine learning models, motivating rigorous study on their impact on downstream tasks. In this work, we study the impact of the trade-off between the intra-class diversity (the number of samples per class) and the inter-class diversity (the number of classes) of a supervised pre-training dataset. Empirically, given a fixed pre-training dataset size, we find that the best downstream performance comes with a balance on the intra-/inter-class diversity. To understand the underlying mechanism, we show theoretically that downstream performance depends monotonically on both types of diversity. Notably, our theory reveals that the optimal class-to-sample ratio (#classes/#samples per class), i.e., the ratio of the number of pre-training classes to the number of samples per class, is invariant to the size of the pre-training dataset, enabling the prediction of the optimal number of pre-training classes. We demonstrate the effectiveness of this application by an improvement of approximately 2 points on average on downstream tasks when pre-training on ImageNet.