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, we found that with the size of the pre-training dataset fixed, the best downstream performance comes with a balance on the intra-/inter-class diversity. To understand the underlying mechanism, we show theoretically that the downstream performance depends monotonically on both types of diversity. Notably, our theory reveals that the optimal class-to-sample ratio (#classes / #samples per class) is invariant to the size of the pre-training dataset, which motivates an application of predicting the optimal number of pre-training classes. We demonstrate the effectiveness of this application by an improvement of around 2 points on the downstream tasks when using ImageNet as the pre-training dataset.

翻译：预训练数据集对于构建最先进的机器学习模型至关重要，这促使人们对其在下游任务中的影响展开严谨研究。本文系统研究了监督预训练数据集中类内多样性（每类样本数）与类间多样性（类别数）之间的权衡关系。实验发现，在预训练数据集规模固定的情况下，平衡类内/类间多样性可获得最佳下游性能。为揭示其内在机制，我们从理论上证明下游性能对两类多样性均呈现单调依赖关系。特别值得注意的是，理论分析表明最优类样本比（类别数/每类样本数）与预训练数据集规模无关，这为预测最优预训练类别数提供了理论依据。以ImageNet作为预训练数据集的实验表明，该方法可使下游任务性能提升约2个百分点，有效验证了该应用的实际效果。