SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Feiyang Kang*
Bradley Department of ECE
fyk@vt.edu
Hoang Anh Just*
just@vt.edu
Anit Kumar Sahu
Alexa AI
Seattle, WA 98121
anit.sahu@gmail.com
Traditionally, data selection has been studied in settings where all samples from prospective sources are fully revealed to a machine learning developer. However, in practical data exchange scenarios, data providers often reveal only a limited subset of samples before an acquisition decision is made. Recently, there have been efforts to fit scaling functions that predict model performance at any size and data source composition using the limited available samples. However, these scaling functions are usually black-box, computationally expensive to fit, highly susceptible to overfitting, or/and difficult to optimize for data selection. This paper proposes a framework called projector, which predicts model performance and supports data selection decisions based on partial samples of prospective data sources. Our approach distinguishes itself from existing work by introducing a novel two-stage performance inference process. In the first stage, we leverage the Optimal Transport distance to predict the model's performance for any data mixture ratio within the range of disclosed data sizes. In the second stage, we extrapolate the performance to larger undisclosed data sizes based on a novel parameter-free mapping technique inspired by neural scaling laws. We further derive an efficient gradient-based method to select data sources based on the projected model performance. Evaluation over a diverse range of applications (e.g., vision, text, fine-tuning, noisy data sources, etc.) demonstrates that projector significantly improves existing performance scaling approaches in terms of both the accuracy of performance inference and the computation costs associated with constructing the performance predictor. Also, projector outperforms by a wide margin in data selection effectiveness compared to a range of other off-the-shelf solutions. We provide projector as an open-source toolkit.
The choice of training data is one of the most crucial components when it comes to extracting the best performance out of a model. Since data is typically acquired from various sources, such as different organizations or vendors, machine learning practitioners often encounter a central question: how to select and combine samples from these data sources?
*Equal contribution. This work also appears at Data-centric Machine Learning Research (DMLR) Workshop@ICML 2023. Code Repository: https://github.com/ruoxi-jia-group/projektor.
Figure 1: (a) Overview of projector, which take as inputs the public pilot data from each source, a selection budget, a target model, a validation set representing the test distribution, and return the optimal combination of data sources as well as the prediction of the resulting model performance. (b) Optimal data source composition and performance projection in a practical autonomous driving data acquisition scenario [15]. projector achieves accurate model performance projection from 1K pilot samples and effective selection. Please see Evaluation Metrics (Section 5) for details on mAP.
Although data selection has been extensively studied in the literature related to active learning [1], coreset selection [2], and data valuation [3-8], most techniques are designed for a fully-observable setting where all data sources are fully revealed to the model developer. The core ideas behind these techniques are to compare the relative importance of different data points or enumerate possible combinations of data points, all of which require complete knowledge of the entire collection of data points. While these methods have shown promising results, their practical applications in real-world scenarios are limited due to a significant gap: the acquisition decision-making processes require knowledge of the entire data sets, while data owners may only reveal limited samples before an acquisition decision is made (e.g., [9, 10] provide the examples in real-world data markets).
To bridge the gap, this paper explores strategic data selection in partially observable settings, where only limited samples of data sources (referred to as pilot datasets) are accessible. The goal is to determine an optimal allocation of the selection budget to each source, only based on the pilot datasets, such that the model trained on the mixture of collected data achieves the best performance at some given objectives.
Technical challenges. In the fully-observable setting, the evaluation and eventually ranking the candidate data selection decisions, including the number of samples to be selected and the ratio of samples from each source ("mixing ratios"), can be determined directly on the available datasets [11, 12]. However, the partially observable setting presents considerable challenges for evaluating a selection decision as one can no longer directly evaluate model performance on the entire data. With limited samples from each data source, the best possible evaluation is the resulting model performance for any combination of the pilot datasets. Then, to make an informed selection decision, it is necessary to understand the model's performance when trained on potentially larger datasets (target scales) at various mixing ratios. In other words, there is a need for prediction and projection of model performance onto larger data scales at different mixing ratios.
A recent study [13] proposes a performance scaling law that takes into account the data size and mixing ratio to predict model performance. Providing a preliminary exploration of this problem, though, this approach faces two major limitations: (1) The numerical instability of its high-order form for scaling functions causes significant difficulty in fitting its parameters, rendering the fitted function susceptible to overfitting and often fails to extrapolate model performance on unseen data mixtures. (2) It hypothesizes on the separability of model performance scaling with data composition and with data size, which is generally untrue as evidenced by latest research [14] and leads to unsatisfactory performance prediction results in empirical observations. Besides, this method requires parameters that grow quadratically with the number of data sources, demanding many (mixing ratio, resulting performance) pairs to fit the function, resulting in substantial computational overhead. Thus, there remains a considerable lack of effective and practical approaches to this problem.
Contributions. The paper investigates two fundamental building blocks for strategic data selection in the partially observable setting:
(Q1) How to provide an accurate projection of model performance trained on any combination of data sources based on limited samples? and (Q2) How to determine the optimal selection strategy?
Towards that end, the paper makes the following contributions.
- Parameter-Efficient Performance Prediction based on Optimal Transport (Addressing Q1). In contrast to existing model performance scaling methods that feature a one-shot fitting of a noninformative parametric model ("surrogate"), our approach is a novel two-stage performance inference process, where the first stage addresses the dependence of the model performance on the mixing ratio by fitting a parameter-efficient model between model performance and Optimal Transport [16] distance between the mixtures of training data and the validation data (Section 4.2). Then, for stage two, we propose a parameter-free mapping that directly projects model performance onto larger data scales, achieving remarkable accuracy as it fully preserves this the dependency of model performance scaling with data sizes and data distributions (Section 4.3).
- Determining optimal data selection strategies (Addressing Q2). We consider the typical data selection goal: maximizing the resulting model performance with fixed data acquisition budgets (data quantity). With model performance predicted by the proposed tools, these problems translate into convex losses that are optimized effectively via gradient-based methods (Section 4.4). We also provide in Appendix.B how it similarly applies to alternative objectives such as minimizing data acquisition costs for the resulting model performance to reach a given level.
- Experiments. We experiment on a variety of applications (vision, natural language processing (NLP) etc., with simple to complex models) with a rich diversity of tasks and real-world scenarios (e.g., autonomous driving, fine-tuning, noisy data sources, unlabeled data). The proposed approach is highly effective in performance prediction, demonstrating superior prediction precision to many baselines while being much more efficient to be constructed (Section 5). We test the performance of data selection by optimizing the performance predictor and show that it improves over existing methods by  $3\%$  on ImageNet-100.
The recent line of research on Data valuation aims to assess the relative importance of each data source ("value") to machine learning applications [3-8]. While originally designed for data pricing, these values are frequently used to inform data selection [4, 17]: in more detail, one can rank the data sources based on their values and select the data points with the highest values. While value-based selection shows some promising results, data values are not directly related to model performance and hence cannot inform the prediction of model performance resulting from the selected data. Besides, values for different data sources typically cannot be combined to measure the value of their compositions [18, 8]. Notably, distributional distances including Optimal Transport have seen a major presence in data valuation as an implicit proxy for model performance [8, 19], but no connection has been made to directly relate data distance to model performance. Our work bridges this gap and directly addresses this long-standing problem. On another line, Coreset selection attempts to find a representative subset of data points to speed up learning [2]. Coreset selection methods have been studied for different learning algorithms [20, 21]. For example, a popular coreset selection approach for neural networks is to cast data selection as a bilevel optimization problem that selects the optimal subset to maximize the performance evaluated on a validation dataset [22]. However, coreset selection techniques rely on access to all the data samples to be chosen from, which limits their use in the partially observable setting. Besides, Predicting the resulting model performance associated with a dataset without performing actual training on it has attracted a lot of attention in different use cases, such as interpreting the model learning process [23, 24]–which leverages surrogate functions to model the black-box relationships between model performance and training data, or predicting performance under the distributional shift [25]. Our work resembles the idea of predicting model performance from data but differs in the technique of leveraging the data distance in the performance predictor. Scaling laws, predicting how the model performance changes with the scale of training data, model parameters, and computation budget [26], have seen increasingly successful in a variety of tasks pertaining to vision and text processing [27]. The performance of machine-learning models generally adheres to a power law relationship with the scale of these variables, which allows for predicting the model performance on larger scales with high precision [26] and provides a viable approach to predicting the potential usefulness of target data from only a small proportion of the set. [14] shows that data from different distributions generally scale at different rates. Our work provides a novel approach that materializes this dependency of scaling relationships with data distributions and achieves remarkable empirical results.
Data provider. Suppose that there are  $m$  prospective data providers. Datasets (data sources) held by these providers are denoted by  $D_{1}^{\mathrm{all}}$ , ...,  $D_{m}^{\mathrm{all}}$ , respectively. We focus on the case that only partial data (samples) from these sources are made available to the public, replicating practical data exchange scenarios [9]. We refer to the public subset of each data source as a pilot dataset and denote it by  $D_{i}^{\mathrm{pi}}$ , where  $D_{i}^{\mathrm{pi}} \subseteq D_{i}^{\mathrm{all}}$  and  $|D_{i}^{\mathrm{pi}}| = n_{i} \ll \bar{N}_{i} = |D_{i}^{\mathrm{all}}|$  for all  $i$ . Each provider  $i$ , upon accepting the purchasing order for acquiring  $n_i$  samples ( $n_i \leq \bar{N}_i$ ), will randomly sample a subset  $S_i$  of size  $n_i$  from  $D_{i}^{\mathrm{all}}$  and return the subset to the requester.[2]
Data collector (or requester, machine learning practitioner). Now, consider a data collector who would like to acquire samples from the providers to train a model. Notably, the collector's acquisition decisions must be made based only on the pilot datasets. We assume the collector has a validation set  $D^{\mathrm{val}}$ , representing the desired target data distribution. For ease of exposition, we assume the collector has a target learning algorithm  $\mathcal{A}^3$  that is going to be applied to the collected data as well as a target performance metric  $\mathcal{L}$  which takes the input of a trained model and a validation set and returns a performance score. The model performance resulting from training on any dataset  $S$  can be thus expressed as  $\mathcal{L}(\mathcal{A}(S), D^{\mathrm{val}})$ .
Given a selection budget of  $N$  samples, a mixing ratio of data sources  $\mathbf{p} = \{p_1,\dots ,p_m\}$  such that  $\forall i,0\leq p_i\leq 1$  and  $\sum_{i = 1}^{m}p_{i} = 1$ , and  $m$  datasets  $D_{1},\ldots ,D_{m}$  to be mixed, we denote the selected dataset by  $\mathcal{D}(N,\mathbf{p}) = S_1\cup \dots \cup S_m$ , where each  $S_{i}$  is a random subset of  $D_{i}^{\mathrm{all}}$  and  $|S_{i}| = p_{i}N$ . Using these notations, we now describe the typical acquisition goals that can be accommodated by our approach:
- (Primary) Fixed-budget selection for maximal performance: The collector seeks to maximize the resulting model performance by strategically choosing the mixing ratio  $\mathbf{p}$  of  $m$  data sources at a pre-specified selection budget  $N_{s} \leq \sum_{i=1}^{m} \bar{N}_{i}$ . The objective can be formalized as  $\max_{\mathbf{p}} \mathcal{L}(\mathcal{A}(\mathcal{D}(N_{s}, \mathbf{p})), D^{\mathrm{val}})$ .
- (Alternative) Flexible-budget selection for reaching performance threshold with minimal costs: The collector seeks to attain a target model performance  $u^{\mathrm{tar}}$  by choosing both the mixing ratio  $\mathbf{p}$  as well as the selection budget  $N$ . More formally, the objective can be expressed as  $\min_{N,\mathbf{p}} N$  s.t.  $\max_{\mathbf{p}} \mathcal{L}(\mathcal{A}(\mathcal{D}(N,\mathbf{p}), D^{\mathrm{val}}) \geq u^{\mathrm{tar}})$ .
The alternative objective can be treated as a direct extension of the primary, where one solves the "performance maximization" problem for different data quantities  $N$  and performs a line search for minimal data quantity  $N$  that meets the performance requirement. We defer to Appendix.B for its complete solution procedure due to the similarity.
Design challenge and key idea. The primary challenge is that the collector cannot access  $D_1^{\mathrm{all}}, \ldots, D_m^{\mathrm{all}}$  for decision making and hence cannot directly evaluate the two optimization objectives for every  $N$ . Yet, as the pilot datasets are public, the collector can evaluate and observe the model performance associated with various mixtures of the pilot datasets for  $N \in \{N : Np_i \leq |D_i^{\mathrm{pi}}|, i = 1, \ldots, m\}$  project the evaluations onto larger data scales. Our high-level idea to tackle the challenge is to first predict the model performance associated with any mixture of prospective unrevealed data sources based on observations on pilot datasets and project the predictions onto different data scales using scaling laws, then determine the data selection strategy by optimizing the predicted performance at the target scales.
4 Methodology of projector: prediction, projection, and selection
4.1 Preliminaries on Optimal Transport
Optimal Transport (OT) is a metric for measuring the discrepancy between probability distributions [16]. Compared to other measures such as the Kullback-Leibler Divergence [29] or Maxi-
This paper will assume each provider honestly provides requested samples and leave an in-depth study of potential security risks, such as malicious data manipulation [28] to future work.
3The proposed data selection approach can support multiple learning algorithms by simply applying it to different choices of algorithms/metrics and picking the best one.
mum Mean Discrepancies [30], OT enjoys advantageous analytical properties (is a valid metric; compatible with sparse-support distributions; stable with respect to deformations of the distributions' supports [31, 32]). Given probability measures  $\mu_t, \mu_v$  over the space  $\mathcal{Z}$ , the OT distance is defined as [33]  $\mathrm{OT}(\mu_t, \mu_v) := \min_{\pi \in \Pi(\mu_t, \mu_v)} \int_{\mathcal{Z}^2} \mathcal{C}(z, z') d\pi(z, z')$ , where  $\Pi(\mu_t, \mu_v) := \{\pi \in \mathcal{P}(\mathcal{Z} \times \mathcal{Z}) \mid \int_{\mathcal{Z}} \pi(z, z') dz = \mu_t, \int_{\mathcal{Z}} \pi(z, z') dz' = \mu_v\}$  denotes a collection of couplings between two distributions  $\mu_t$  and  $\mu_v$ ,  $\mathcal{C}: \mathcal{Z} \times \mathcal{Z} \to \mathbb{R}^+$  is a symmetric positive-definite cost function (with  $\mathcal{C}(z, z) = 0$ ), respectively. A popular choice of  $\mathcal{C}$  is given in [34] by considering  $z$  as the feature-label pair  $(x, y)$ . The computation of the OT distance usually relies on the Sinkhorn algorithm [35], which attains almost linear time complexity and memory overhead with the state-of-the-art implementations and applies to large scales with parallel computing [31, 32]. Given  $D_t = \{(x_i, y_i)\}_{i=1}^N$  of size  $N$ , and  $D_v = \{(x_i', y_i')\}_{i=1}^T$  of size  $T$ , one can construct discrete probability measures  $\mu_t(x, y) := \frac{1}{N} \sum_{i=1}^{N} \delta_{(x_i, y_i)}$  and  $\mu_v(x, y) := \frac{1}{T} \sum_{i=1}^{T} \delta_{(x_i', y_i')}$ , where  $\delta$  is the Dirac delta function. With slight abuse of notation, we use  $\mathrm{OT}(D_t, D_v)$  to denote the OT distance between their corresponding discrete measures  $\mathrm{OT}(\mu_t(x, y), \mu_v(x, y))$ .
Extensive theoretical studies show that the OT distance between two distributions provides an upper bound on the difference of a model's performance when evaluated on the two distributions [36, 37, 8]. Largely built upon Kantorovich-Rubinstein Duality [38], existing theoretical results require assumptions on the Lipschitz constant of the model with respect to the input space. However, the constant is rarely known in practice, nor can it be bounded tightly for complex models such as deep neural networks. As a result, despite its widespread popularity as a performance proxy [8, 36, 39], one cannot directly apply the existing theoretical results to directly estimate the model performance based on the OT distance, posing an important gap.
4.2 Aligning data distance with performance predictions
Inspired by the theoretical results that the upper bound on the difference between training loss and validation loss can be tightly bounded by an affine transformation of the OT distance [33, 8], our first proposed approach is to directly estimate this transformation by empirically fitting data distances to model performance and then the estimated transformation can be used for predicting the model performance for different data mixtures. Formally, we consider the following performance estimator:
$$
\hat {\mathcal {L}} \left(\mathcal {A} (\mathcal {D} (N, \mathbf {p})), D ^ {\text {v a l}}\right) = a _ {1} \cdot \operatorname {O T} \left(\mathcal {D} (N, \mathbf {p}), D ^ {\text {v a l}}\right) + a _ {0}, \tag {1}
$$
where scaling parameter  $a_1$  and centering parameter  $a_0$  define the affine transformation. These two parameters can be estimated through least-square fitting. In particular, consider collecting the "training data" by forming the set of tuples  $\{(N_j,\mathbf{p}_j,\mathcal{L}(\mathcal{A}(\mathcal{D}(N_j,\mathbf{p}_j)),D^{\mathrm{val}}))\}_{j = 1}^{\ell}$ , where  $N_{j}$  is randomly sampled from  $\{1,\ldots ,\sum_{i = 1}^{m}|D_i^{\mathrm{pi}}|\}$  and  $\mathbf{p}_j$  is sampled from a probability simplex. Then, these parameters can be estimated as
$$
\left(\hat {a _ {1}}, \hat {a _ {0}}\right) = \arg \min  _ {a _ {1}, a _ {0}} \sum_ {j = 1} ^ {\ell} \left(\hat {\mathcal {L}} \left(\mathcal {A} (\mathcal {D} (N _ {j}, \mathbf {p} _ {j})), D ^ {\mathrm {v a l}}\right) - \mathcal {L} \left(\mathcal {A} (\mathcal {D} (N _ {j}, \mathbf {p} _ {j})), D ^ {\mathrm {v a l}}\right)\right) ^ {2}.
$$
We refer to this method as center-scaling (projektor/CS).  $a_1$  can be considered an empirical estimate of the Lipschitz constant, and this treatment has been informally adopted in various works under different names [40-42]. projector/CS has only two parameters that need to be estimated, which brings an important benefit of efficiency: we only need to have a few training iterations to get the "training data" for the above least-square fitting.
While proposed projector/CS is sufficient to provide reliable performance predictions in most circumstances, we found it possible to further improve the prediction accuracy by making the scaling parameter and centering parameter a function of the mixing ratio. The intuition is that for samples from different data sources (i.e., data lying in different manifolds of the input space), the Lipschitz constant of the model along the combined manifold may vary with the mixing ratio. Hence, we supplement projector/CS with simple nonlinear terms to characterize the dependence on each data source, leading to the pseudo-quadratic (projector/PQ) method, which is given as
$$
\hat {\mathcal {L}} (\mathcal {A} (\mathcal {D} (N, \mathbf {p})), D ^ {\text {v a l}}) = \sum_ {i = 1} ^ {m} \left(b _ {2} ^ {i} \cdot p _ {i} ^ {2} + b _ {1} ^ {i} \cdot p _ {i} + b _ {0}\right) \cdot \mathrm {O T} (\mathcal {D} (N, \mathbf {p}), D ^ {\text {v a l}}) + \sum_ {i = 1} ^ {m} \left(c _ {2} ^ {i} \cdot p _ {i} ^ {2} + c _ {1} ^ {i} \cdot p _ {i} + c _ {0}\right), \tag {2}
$$
where  $(\mathbf{b}_2, \mathbf{b}_1, \mathbf{b}_0, \mathbf{c}_2, \mathbf{c}_1, \mathbf{c}_0)$  are pseudo-quadratic parameters where the fitting process is similar to Eq. (1). Projektor/PQ has  $\mathcal{O}(m)$  parameters ( $m$  is the number of data sources). So its fitting
process will require more re-training than projector/CS. However, as we will show in Section 5, it still significantly improves the efficiency over the existing baselines [13]. This pseudo-quadratic form is chosen for it contains the simplest nonlinear terms and we want to preserve the convexity for numerical stability. We trimmed off the cross terms in the quadratic function as they often do not contribute much, resulting in the number of parameters growing linearly rather than quadratically with the number of data sources as in [13], greatly easing the computation burden in parameter fitting.
4.3 Parameter-free performance projection onto larger data scales
Once the parameters are learned, the performance predictors (1) and (2) can be used to predict the validation performance associated with a training set by calculating the OT distance between the training and the validation set and plugging it as input to the predictors. Then, we need to project these predictions onto the target data scales. Neural scaling laws showcase the predictability of empirical performance with respect to the size of the training dataset, where it typically follows a log-linear scaling relationship as
$$
\mathbb {E} _ {V} \left[ \mathcal {L} \left(\mathcal {A} \left(\mathcal {D} (N, \mathbf {p})\right); D ^ {\mathrm {v a l}}\right) \right] \approx - \alpha \log (N) + C
$$