SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Jieyu Zhang $^{1*}$ , Bohan Wang $^{2*}$ , Zhengyu Hu $^{3}$ , Pang Wei Koh $^{1}$ , Alexander Ratner $^{1,4}$
<sup>1</sup> University of Washington <sup>2</sup> USTC <sup>3</sup> HKUST(GZ) <sup>4</sup> Snorkel AI, Inc.
{jieyuz2,pangwei,ajratner}@cs.washington.edu
zhu021@connect.hkust-gz.edu.cn
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  $\left(\frac{\text{#classes}}{\text{#samples per class}}\right)$ , 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.
Many state-of-the-art deep neural network models are pre-trained on large datasets before being finetuned for downstream tasks [13, 17, 23, 1]. While the composition of their pre-training dataset has been shown to be a key factor in the performance of these models [7, 9, 14, 8, 12, 26], how best to design these pre-training datasets still remains underexplored. In this work, we focus on supervised pre-training, one of the most popular pre-training paradigms, and study two key quantities of a supervised pre-training dataset: intra-class diversity (the number of different samples within each pre-training class) and inter-class diversity (the number of different pre-training classes). Intuitively, both diversities are beneficial for supervised pre-training [13]. Yet when the size of the pre-training dataset is fixed, these diversities trade off, since increasing one will decrease the other. Our work studies the impact of this dataset diversity trade-off on downstream performance, as well as how to balance them to design a supervised pre-training dataset with the best downstream performance.
Empirically, with ImageNet [24] as the pre-training dataset and the pre-training dataset size fixed, we show that the optimal performance on the downstream tasks occurs when a balance on the intra-/inter-class diversity is achieved. We then offer a theoretical explanation for this effect by first modeling the dataset generation process through a two-step sampling framework, and then demonstrating that the test error of the downstream task displays a rational relationship with respect to the class-to-sample ratio, i.e., the ratio of the number of pre-training classes to the number of samples per class, or, in other words, the ratio between inter-/intra-class diversity. The established analytical relationship
between downstream performance and the class-to-sample ratio can serve as a guiding principle in designing a supervised pre-training dataset by estimating the optimal class-to-sample ratio rather than the grid search.
Notably, our theory shows that given a source of a pre-training dataset and a downstream task, the optimal class-to-sample ratio is invariant to the size of the pre-training dataset. Based on such an invariance, one could estimate the optimal class-to-sample ratio with small pre-training datasets and then leverage it to build a large-scale pre-training dataset. In particular, the optimal number of pretraining classes  $\bar{K}$  and the number of examples per class  $n$  are proportional to the square root of the size of the pre-training dataset  $N$ , i.e.,  $\bar{K} \propto \sqrt{N}$ , which leads to an invariant optimal class-to-sample ratio. We empirically verify our theoretical findings on ImageNet [24] and present the effectiveness of its application in predicting the optimal number of classes for pre-training datasets with different sizes. In addition, we conducted experiments with different pre-trained datasets, different model backbones, and downstream tasks of different domains to demonstrate that our findings are consistent across many scenarios.
Our major findings and contributions are as follows:
- In supervised pre-training, we observe that with a fixed pre-training dataset size, there exists a trade-off between intra-class and inter-class diversities. This balance between diversities plays a crucial role in shaping the downstream performance, underscoring the significance of considering both aspects when designing the pre-training dataset;
- We then theoretically explain this effect by first modeling the dataset generation process through a two-step sampling framework and then showing that the test error of the downstream task displays a convex relationship with respect to the class-to-sample ratio, serving as a guiding principle in designing a supervised pre-training dataset.;
- Our theory also uncovers the invariance of the optimal class-to-sample ratio with respect to the size of the pre-training dataset, allowing us to predict the optimal number of classes with a small number of pre-training data before building a larger pre-training dataset for a downstream task.
2 Empirical Observations
The goal of this work is to study the trade-off of intra-/inter-class diversity in a supervised pre-training dataset and its impact on the pre-trained model's performance on downstream tasks. Specifically, the inter-class diversity refers to the diversity of classes in pre-training dataset, i.e., how many different classes we have  $(K)$ ; while the intra-class diversity refers to the diversity of samples within each class, i.e., how many different samples in each class  $(n)$ . When the size of the pre-training dataset is fixed, increasing either type of diversity will by definition decrease the other, leading to a dataset diversity trade-off. To study the impact of such dataset diversity trade-off, we experiment with pre-training datasets with varying numbers of classes and number of samples per class. The experimental details can be found in Appendix A.2.
Evaluation protocol. Following common practice [13], we use the ImageNet [24] as the dataset for supervised pre-training. In this work, we mainly use ResNet-18 [10] as the backbone model. For evaluating the performance of the pre-trained model on downstream tasks, we perform linear probing (tuning the head but freezing the lower layers). We repeat each individual experiment five times and report the averaged top-1 accuracy.
Downstream tasks. We adopt the following six datasets as the downstream classification tasks: Stanford40 dataset [25] for action recognition, StanfordDogs [15] for fine-grained object recognition, MIT67 [22] for scene classification, CIFAR10 [16] for image recognition datasets, Flowers102 [20] for image classification dataset, FGVCAircraft [18] for aircraft classification dataset.
While having all the other configurations fixed, during pre-training, we vary the number of classes and the number of samples per class. Specifically, given  $N$  as the size of the pre-training dataset and  $K$  as the number of classes, we randomly sample  $K$  classes from ImageNet and then uniformly sample  $n = \frac{N}{K}$  samples from each class to compose the dataset. We experiment with the following  $N$  and  $K$  values: {1K, 2K, 5K, 10K, 20K, 50K, 100K} and {2, 5, 10, 20, 50, 100, 200, 500, 1000} respectively. Note that with larger  $N$  (e.g.,  $N = 10\mathrm{K}$ ), we cannot evaluate smaller values of  $K$  (e.g.,  $K = 2$ ), since in ImageNet each class has at most 1300 samples.
(a) CIFAR10
(d) MIT67
Figure 1: Test error rate (the darker, the better) as a function of intra-class diversity on the y-axis  $(\log n)$  and inter-class diversity on the x-axis  $(\log K)$ . Each plot represents a different dataset. The red dashed anti-diagonal lines indicate fixed pre-training dataset sizes  $(\log K + \log n = \log N$  is constant), from which we can see that obtaining the best downstream performance given a fixed pre-training dataset size requires balancing both diversities.
(b) FGVCAircraft
(e) Stanford40
(c) Flowers102
(f) StanfordDogs
Results and observations. We visualize the results in Figure 1. In the contour plot, the z-value is the error rate on the test set, thus lower is better. The x-axis and y-axis are inter-class diversity  $(\log K)$  and intra-class diversity  $(\log n)$  in the log space, respectively. The values on anti-diagonal lines  $(y = -x + c)$  share the same pre-training dataset size as  $\log N = \log K + \log n$ . From the results, we have two observations: 1) Both intra-/inter-class diversity are beneficial for downstream tasks: We can see that increasing either inter-class diversity  $(\log K)$  or intra-class diversity  $(\log n)$ , given the other is fixed, would lead to a better test error rate. This is intuitive and as expected, since the size of the pre-training dataset  $N$  would increase accordingly, which is known to be beneficial for downstream tasks [11, 7, 13]. 2) A trade-off of intra-/inter-class diversity on downstream task performance: More importantly, by looking at the anti-diagonal lines where  $\log N$  is fixed and equals  $\log K + \log n$ , we can see a trade-off between intra-/inter-class diversity on the test error rate of downstream tasks: either cases of 1) high inter-class diversity, low intra-class diversity and 2) low inter-class diversity and high intra-class diversity would not render the best performance. Instead, some point in the middle of the anti-diagonal line leads to the lowest test error rate.
3 Theoretical Understanding
In this section, we first present the theoretical setup and notations and then provide a theory on the impact of the pre-training dataset diversity on downstream performance. We also show that the optimal class-to-sample ratio  $\left(\frac{K}{n}\right)$  is invariant to the size of the pre-training dataset; such a property can be leveraged to predict the optimal number of classes for building a large pre-training dataset with a small number of data samples first.
3.1 Setup and notations
Dataset. To be consistent with our experimental setup, we consider the supervised pre-training task. Specifically, we can access two datasets, one for the pre-training task (denoted as  $S^p$ ) and another for the downstream task (denoted as  $S^d$ ). Each example in the pre-training dataset consists of input features  $x \in \mathcal{X} = \mathbb{R}^{d_1}$  (where  $d_1$  is the dimension of data) and a label  $y \in [K]$  (where  $K$  is the number of classes). Specifically, we denote  $S^p = \{(x_1, y_1), \dots, (x_N, y_N)\}$ , where  $N$  is the size of  $S^p$ , and assume that  $S^p$  is sampled according to some underlying distribution  $\mathcal{P}$  (we do
not specify  $\mathcal{P}$  here because we will analyze cases with different  $\mathcal{P}$  latter). Every example in the downstream dataset consists of input features  $\tilde{x} \in \mathcal{X}$  and a label  $\tilde{y} \in [\tilde{K}]$  (note that  $\tilde{K}$  does not necessarily equal-to  $K$ ), and is sampled i.i.d. according to an underlying distribution  $\hat{\mathcal{P}}$ . We denote  $S^d = \{(\tilde{x}_1, \tilde{y}_1), \dots, (\tilde{x}_N, \tilde{y}_{\tilde{N}})\}$  and thus  $S^d \sim \tilde{\mathcal{P}}^{\tilde{N}}$ .
Model. The models for both pre-training and downstream tasks consist of two components: the feature extractor and the classifier. Specifically, the model for the pre-training task is given as  $f_{S^p} \circ h_{S^p}$ , where  $f_{S^p}: \mathbb{R}^{d_2} \to \mathbb{R}^K$  is the pre-training classifier ( $d_2$  is the dimension of feature) and  $h_{S^p}: \mathbb{R}^{d_1} \to \mathbb{R}^{d_2}$  is the feature extractor. We denote the set of all possible  $f_{S^p}$  as  $\mathcal{F}$ , and the set of all possible  $h$  as  $\mathcal{H}$ . The model for the downstream task is given as  $f_{S^d} \circ h_{S^d}$ , where  $f_{S^d}: \mathbb{R}^{d_2} \to \mathbb{R}^{\tilde{K}}$  is the downstream classifier and  $h$  is the feature extractor shared with the pre-training task. We set all possible  $f_{S^p}$  as  $\tilde{\mathcal{F}}$ .
Loss. To measure the correctness of model's predictions, we use the cross-entropy loss. Specifically, given an example  $(x,y)$  and pre-training/downstream model  $f\circ h$ , the corresponding cross-entropy loss is defined as
$$
\ell (f \circ h (x), y) = - \log \frac {e ^ {f _ {y} \circ h (x)}}{\sum_ {i} e ^ {f _ {i} \circ h (x)}},
$$
where  $f_{i} \circ h(x)$  is the  $i$ -th coordinate of  $f \circ h(x)$ . We make the following assumption about the complexity of the model.
Assumption 1. There exist a positive constant  $M_{\ell}$ , such that  $\forall i, \forall f \in \mathcal{F} \cup \tilde{\mathcal{F}}, h \in \mathcal{H}, x \in \mathcal{X}$
$$
\ell (f \circ h (x), i) \leq M _ {\ell}.
$$
Furthermore, for any distribution  $Q$  over the feature space  $\mathbb{R}^{d_1}$ , any  $m \in \mathbb{N}$ , and any  $\mathcal{F}' \in \{\mathcal{F}, \tilde{\mathcal{F}}\}$ , define the Gaussian complexity of function class  $\mathcal{F}' \circ \mathcal{H}$  over the marginal distribution  $Q$  with sample number  $m$  as
$$
G _ {m} ^ {Q} (\mathcal {F} ^ {\prime} \circ \mathcal {H}) \triangleq \mathbb {E} _ {(x _ {i}) _ {i = 1} ^ {m} \sim Q ^ {m}} \mathbb {E} _ {(\sigma_ {i, j}) _ {i \in [ N ], j \in K ^ {\prime}} \sim \mathcal {N} (0, \mathbb {1} _ {n \times K ^ {\prime}})} \sup  _ {f \in \mathcal {F}, h \in \mathcal {H}} \sum_ {i = 1} ^ {m} \sum_ {j = 1} ^ {K ^ {\prime}} \sigma_ {i, j} f _ {j} \circ h (x _ {i}).
$$
Here  $K^{\prime} = \dim (\mathcal{F}^{\prime})$ . We assume that  $G_{m}^{Q}(\mathcal{F}\circ \mathcal{H})\leq G\sqrt{m}$ , where  $G$  is independent of  $Q$  and  $m$ .
Assumption 1 constrains the dependence of model complexity over the number of sample  $N$ , which is a standard assumption in generalization analysis [19].
Risks and corresponding minimizers. We first define the empirical risk  $\mathcal{R}_p(f\circ h,S^p)$  and population risk  $\mathcal{R}_p(f\circ h,\mathcal{P})$  over the pre-training task as
$$
\bar {\mathcal {R}} _ {p} (f \circ h, S ^ {p}) \triangleq \frac {1}{N} \sum_ {i = 1} ^ {N} [ \ell (f \circ h (x _ {i}), y _ {i}) ], \mathcal {R} _ {p} (f \circ h, \mathcal {P}) \triangleq \mathbb {E} _ {S ^ {p} \sim \mathcal {P}} \bar {\mathcal {R}} _ {p} (f \circ h, S ^ {p}).
$$
The corresponding feature extractor and classifier of the empirical risk minimizer over the pre-training task as
$$
h _ {S ^ {p}} \triangleq \arg \min  _ {h \in \mathcal {H}} \left(\min  _ {f \in \mathcal {F}} \bar {\mathcal {R}} _ {p} (f \circ h, S ^ {p})\right), f _ {S ^ {p}} \triangleq \arg \min  _ {f \in \mathcal {F}} \left(\min  _ {h \in \mathcal {H}} \bar {\mathcal {R}} _ {p} (f \circ h, S ^ {p})\right).
$$
Given a feature extractor  $h \in \mathcal{H}$ , we measure its performance over the pre-training task through a classifier agnostic approach by considering
$$
\mathcal {R} _ {p} (h, \mathcal {P}) \triangleq \min  _ {f \in \mathcal {F}} \mathcal {R} _ {p} (f \circ h, \mathcal {P}).
$$
Note here we slightly abuse the notation of  $\mathcal{R}_p$  without causing confusion as  $f\circ h$  and  $h$  have different image spaces. The representation error of  $h$  over  $\mathcal{P}$  is then defined as the gap between the risk of  $h$  and the smallest possible risk
$$
\mathcal {E} _ {p} (h, \mathcal {P}) \triangleq \mathcal {R} _ {p} (h, \mathcal {P}) - \min _ {\tilde {h} \in \mathcal {H}} \mathcal {R} _ {p} (\tilde {h}, \mathcal {P}).
$$
2This inequality holds for a wide range of models, including deep neural networks [2].
Similarly, the empirical risk  $\bar{\mathcal{R}}_d(f\circ h,S^d)$  and population risk  $\mathcal{R}_d(f\circ h,\tilde{\mathcal{P}})$  over the downstream task are defined as
$$
\bar {\mathcal {R}} _ {d} (f \circ h, S ^ {d}) \triangleq \sum_ {i = 1} ^ {\tilde {N}} \ell (\tilde {f} \circ h (\tilde {x} _ {i}), \tilde {y} _ {i}), \mathcal {R} _ {d} (f \circ h, \tilde {\mathcal {P}}) \triangleq \mathbb {E} _ {S ^ {d} \sim \tilde {\mathcal {P}} ^ {\tilde {N}}} \bar {\mathcal {R}} _ {d} (f \circ h, S ^ {d}).
$$
The learned classifier from the downstream task can then be defined as
$$
f _ {S ^ {d}} \triangleq \arg \min  _ {f \in \bar {\mathcal {F}}} \left(\min  _ {h \in \mathcal {H}} \bar {\mathcal {R}} _ {d} (f \circ h _ {S ^ {p}}, S ^ {d})\right).
$$
The corresponding performance and excess risk of  $h$  over the downstream task can then be defined as
$$
\mathcal {R} _ {d} (h, \tilde {\mathcal {P}}) \triangleq \min _ {\tilde {f} \in \tilde {\mathcal {F}}} \mathcal {R} _ {d} (\tilde {f} \circ h, \tilde {\mathcal {P}}), \mathcal {E} _ {d} (h, \tilde {\mathcal {P}}) \triangleq \mathcal {R} _ {d} (h, \tilde {\mathcal {P}}) - \min _ {\tilde {h} \in \mathcal {H}} \mathcal {R} _ {d} (\tilde {h}, \tilde {\mathcal {P}}).
$$
Finally, we are interested in the excess risk of the obtained model  $f_{S^d} \circ h_{Sp}$  over the downstream task, i.e.,
$$
\mathcal {E} _ {d} (f _ {S ^ {d}} \circ h _ {S _ {p}}, \tilde {\mathcal {P}}) \triangleq \mathcal {R} _ {d} (f _ {S ^ {d}} \circ h _ {S _ {p}}, \tilde {\mathcal {P}}) - \min _ {\tilde {h} \in \mathcal {H}} \mathcal {R} _ {d} (\tilde {h}, \tilde {\mathcal {P}}).
$$
3.2 A theory on the impact of intra-/inter-class diversity trade-off
In this work, we focus on a common practice of collecting the pre-training dataset: one first sample  $K$  classes and then collect  $n$  samples for each class. We start with a detailed characterization of such data generation process, followed by assumptions and the main result.
Data generation process 1. We assume the pre-training data is generated through a two-step sampling process. Specifically, suppose that there is a distribution  $\mathcal{D}$  over  $\Delta (\mathcal{X})$  (the set consisting of all distributions on  $\mathcal{X}$ ). Then,  $S^p$  is generated by the following procedure (and  $\mathcal{P}$  is naturally induced)
- Sample  $K$  classes by i.i.d. sampling  $K$  distributions  $\{\mathcal{P}_i\}_{i=1}^K$  according to  $\mathcal{D}$ . These are respectively the underlying distributions of  $K$  classes;
- For each  $i \in [K]$ , i.i.d. sample  $n$  data  $\{x_{i,1}, \dots, x_{i,n}\}$  according to  $\mathcal{P}_i$  and denote  $S_i = \{(x_{i,1}, i), \dots, (x_{i,n}, i)\}$ . Note here  $x_{i,j}$  does not contain the information of label, as its label information is already contained in  $i$ . The whole dataset is obtained by putting all  $S_i$  together, i.e.,  $S^p = \{S_1, \dots, S_K\}$ .
We make the following assumption on the correlation between the representation powers of the pre-training and the downstream task.
Assumption 2. Given  $\mathcal{P}$ , there exists non-negative coefficients  $\nu_0^{\vec{\mathcal{P}}}(\mathcal{P})$  and  $\nu_{1}^{\vec{\mathcal{P}}}(\mathcal{P})$ , such that  $\forall h\in \mathcal{H}$
$$
\mathcal {E} _ {d} (h, \tilde {\mathcal {P}}) \leq \nu_ {1} ^ {\tilde {\mathcal {P}}} (\mathcal {P}) \mathcal {E} _ {p} (h, \mathcal {P}) + \nu_ {0} ^ {\tilde {\mathcal {P}}} (\mathcal {P}).
$$
We further assume that  $\nu_{0}^{\tilde{P}}(\mathcal{P})$  and  $\nu_{1}^{\tilde{P}}(\mathcal{P})$  are stable, that is, there exist two  $\tilde{\mathcal{P}}$ -dependent positive constants  $M_0^{\tilde{P}}$  and  $M_1^{\tilde{P}}$ , such that for any  $\mathcal{P} = \Pi_{i=1}^{K}(\mathcal{P}_i, i)^n$  and  $\mathcal{P}' = \Pi_{i=1}^{K-1}(\mathcal{P}_i, i)^n \times (\mathcal{P}_K', K)^n$  which differ by only one component, we have that  $|\nu_{0}^{\tilde{P}}(\mathcal{P}) - \nu_{0}^{\tilde{P}}(\mathcal{P}')| \leq \frac{M_0^{\tilde{P}}}{K}$  and  $|\nu_{1}^{\tilde{P}}(\mathcal{P}) - \nu_{1}^{\tilde{P}}(\mathcal{P}')| \leq \frac{M_1^{\tilde{P}}}{K}$ . Moreover, we assume that  $\nu_{0}^{\tilde{P}}(\mathcal{P})$  and  $\nu_{1}^{\tilde{P}}(\mathcal{P})$  concentrate around their means, i.e., there exist  $\nu_{0}^{\tilde{P}}(\mathcal{D})$ ,  $\nu_{1}^{\tilde{P}}(\mathcal{D})$ ,  $C_0^{\tilde{P}}$  and  $C_1^{\tilde{P}}$ , such that  $|\mathbb{E}_{\{\mathcal{P}_i\}_{i=1}^{K} \sim \mathcal{D}^{K}} \nu_{0}^{\tilde{P}}(\mathcal{P}) - \nu_{0}^{\tilde{P}}(\mathcal{D})| \leq \frac{C_0^{\tilde{P}}}{\sqrt{K}}$  and  $|\mathbb{E}_{\{\mathcal{P}_i\}_{i=1}^{K} \sim \mathcal{D}^{K}} \nu_{1}^{\tilde{P}}(\mathcal{P}) - \nu_{1}^{\tilde{P}}(\mathcal{D})| \leq \frac{C_1^{\tilde{P}}}{\sqrt{K}}$ .
Assumption 2 assumes that the pre-training representation error can bound the downstream representation error, which is a common assumption in existing works [6, 26, 3]. Also, as  $\mathcal{P}$  is derived by sampling  $K$  distributions according to  $\mathcal{D}$ , we make mild assumptions that the coefficients  $\nu_0^{\tilde{\mathcal{P}}}(\mathcal{P})$  and  $\nu_{1}^{\tilde{\mathcal{P}}}(\mathcal{P})$  is robust when changing the underlying distribution of only one class, and when  $K$  grows, the expectation of  $\nu_0^{\tilde{\mathcal{P}}}(\mathcal{P})$  and  $\nu_{1}^{\tilde{\mathcal{P}}}(\mathcal{P})$  converge to some limits.
Theorem 3.1. Let Assumptions 1 and 2 hold. For data generation process 1, with probability over the sampling of the datasets at least  $1 - \delta$ , we have
$$
\begin{array}{l} \mathcal {E} _ {d} \left(f _ {S ^ {d}} \circ h _ {S _ {p}}, \tilde {\mathcal {P}}\right) \leq \left(\nu_ {1} ^ {\tilde {\mathcal {P}}} (\mathcal {D}) + M _ {1} \sqrt {\frac {\log \frac {4}{\delta}}{2 K}} + \frac {C _ {1}}{\sqrt {K}}\right) \left(5 M _ {\ell} \sqrt {\frac {\log \frac {6}{\delta}}{2 n}} + \frac {2 G \sqrt {2}}{\sqrt {n}}\right) + \nu_ {0} ^ {\tilde {\mathcal {P}}} (\mathcal {D}) \\ + M _ {0} \sqrt {\frac {\log \frac {6}{\delta}}{2 K}} + \frac {C _ {0}}{\sqrt {K}} + 5 M _ {\ell} \sqrt {\frac {\log \frac {6}{\delta}}{2 \tilde {N}}} + 2 \sqrt {2} G \frac {1}{\sqrt {\tilde {N}}}. \tag {1} \\ \end{array}
$$
The detailed proof is deferred to Appendix A.8.1. Below we simplify the right-hand-side of Equation 1 and show that the empirically observed downstream performance trade-off can be explained by such a result.
Simplifying the Theorem 3.1. Denote the RHS of Equation 1 as  $U$ , we have
$$
\begin{array}{l} U = \nu_ {1} ^ {\bar {\mathcal {P}}} (\mathcal {D}) \left(5 M _ {\ell} \sqrt {\frac {\log \frac {6}{\delta}}{2}} + 2 G \sqrt {2}\right) \frac {1}{\sqrt {n}} + \left(M _ {0} \sqrt {\frac {\log \frac {6}{\delta}}{2}} + C _ {0}\right) \frac {1}{\sqrt {K}} + \left(M _ {1} \sqrt {\frac {\log \frac {6}{\delta}}{2}} + C _ {1}\right) \\ \times \left(5 M _ {\ell} \sqrt {\frac {\log \frac {6}{\delta}}{2}} + 2 G\right) \frac {1}{\sqrt {N}} + \nu_ {0} ^ {\tilde {\mathcal {P}}} (\mathcal {D}) + 5 M _ {\ell} \sqrt {\frac {\log \frac {6}{\delta}}{2 \tilde {N}}} + 2 \sqrt {2} G \frac {1}{\sqrt {\tilde {N}}}. \\ \end{array}
$$
We can see that the above equation can be simplified as
$$
U = \frac {A}{\sqrt {n}} + \frac {B}{\sqrt {K}} + \frac {C}{\sqrt {N}} + D, \tag {2}
$$
where  $A, B, C, D$  do not depend on  $N, K, n$ , but instead only depend on the properties of the underlying pre-training and the downstream task data distribution.
Explaining downstream performance trade-off given a fixed  $N$ . From Equation 2 we can see that the performance on the target task would increase when we increase 1) intra-class diversity  $n$ , 2) inter-class diversity  $K$ , and 3) the size of pre-training dataset  $N$ . When  $N$  is fixed, however, increasing either intra-class diversity or inter-class diversity would decrease the other (since  $N = n \times K$ ) and therefore eventually lead to a performance drop. Another way to see this is to parameterize  $U$  as a function of  $K$  without  $n$ :
$$
U (K) = \frac {A \sqrt {K}}{\sqrt {N}} + \frac {B}{\sqrt {K}} + \frac {C}{\sqrt {N}} + D, \tag {3}
$$
From this we can clearly see that both extremes of  $K$  (too large or too small) would not lead to optimal performance. A similar conclusion can be drawn regarding  $n$  when parametrizing  $U$  as a function of  $n$  without  $K$ .