SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Nicholas Roberts
Jue Wang
Together AI
Together AI, University of Chicago
The quality of training data impacts the performance of pre-trained large language models (LMs). Given a fixed budget of tokens, we study how to best select data that leads to good downstream model performance across tasks. We develop a new framework based on a simple hypothesis: just as humans acquire interdependent skills in a deliberate order, language models also follow a natural order when learning a set of skills from their training data. If such an order exists, it can be utilized for improved understanding of LMs and for data-efficient training. Using this intuition, our framework formalizes the notion of a skill and of an ordered set of skills in terms of the associated data. First, using both synthetic and real data, we demonstrate that these ordered skill sets exist, and that their existence enables more advanced skills to be learned with less data when we train on their prerequisite skills. Second, using our proposed framework, we introduce an online data sampling algorithm, SKILL-IT, over mixtures of skills for both continual pre-training and fine-tuning regimes, where the objective is to efficiently learn multiple skills in the former and an individual skill in the latter. On the LEGO synthetic in the continual pre-training setting, SKILL-IT obtains 37.5 points higher accuracy than random sampling. On the Natural Instructions dataset in the fine-tuning setting, SKILL-IT reduces the validation loss on the target skill by  $13.6\%$  versus training on data associated with the target skill itself. We apply our skills framework on the RedPajama dataset to continually pre-train a 3B-parameter LM, achieving higher accuracy on the LM Evaluation Harness with 1B tokens than the baseline approach of sampling uniformly over data sources with 3B tokens.
Large language models (LMs) exhibit remarkable capabilities, including producing creative content [55], writing source code [9], and chatting with users [8]. A key ingredient in enabling models to perform such tasks is the data on which the models are trained [18, 20, 59]. A natural way to unlock particular capabilities is to improve this training data. However, it is unclear how to select data from a large corpus for these capabilities given a fixed budget of training tokens, as data selection methods for current state-of-the-art LMs mostly rely on heuristics for filtering and mixing together different datasets [33, 59]. We lack a formal framework for capturing how data influences the model's capabilities and how to utilize this data effectively for improving LM performance.
*Correspondence to mfchen@cs.stanford.edu.
Ordered skill set
Figure 1: Inspired by how humans acquire knowledge, we hypothesize that LMs best learn skills in a particular order and that this can help improve our understanding and training of LMs. We show that these ordered skill sets exist in real data, which enables skills to be learned with less data given that we train on their prerequisite skills. We then propose SKILL-IT, an online data selection algorithm that learns skills quickly by exploiting their ordering.
To develop such a framework, we take inspiration from how humans acquire knowledge. A classic idea in education literature is the concept of skills that form a learning hierarchy [65]. For example, one study found that students learned mathematical and scientific skills most quickly when these skills were presented in a particular order [12]. We seek to understand the extent that similar skill-based orderings characterize LM training. Such orderings, if they exist, may provide a better understanding of LMs as well as a mechanism for data-efficient training. For instance, to train an LM for Spanish question generation, we wish to know if training first on related but simpler tasks, such as Spanish grammar and English question generation, helps.
We study if the idea of skill orderings can help us build a framework that relates data to LM training and behavior. This requires addressing two challenges revolving around the connection between skills and data. First, in order to show that there exist sets of skills that the LM learns most efficiently in some particular order, an operational definition of LM skill and skill ordering must be developed and validated on data. In initial experiments, we investigated if semantic groupings of data, such as metadata attributes or embedding clusters, were sufficient to represent a skill and characterize how models learn. For instance, we partitioned the Alpaca dataset [56] by instruction type—a technique used to capture dataset diversity [62]—but we found that sampling based on instruction types and random sampling resulted in similar model performance, suggesting that not just any existing notion of data groups can characterize skills.
Second, these definitions of skills must be used to construct sampling distributions to actually improve model training. To develop criteria for a data selection algorithm that learns skills efficiently, we identify challenges that naive selection approaches face. The standard approach of random uniform sampling over data fails to learn skills optimally due to not accounting for skill imbalance and ordering. Skills can be distributed unevenly in the data, with more complex skills being rare—for instance, Spanish and question generation (QG) are  $5\%$  and  $4\%$  of the Natural Instructions dataset [63], respectively, but Spanish QG is only  $0.2\%$ . Random sampling also provides no mechanism for taking into account a particular training order and dependency structure on skills. More sophisticated techniques like curriculum learning account for sample-level ordering, but not skills or their dependencies. Our goal framework must account for these issues of imbalance and ordering.
Skill-based framework We define a skill as a unit of behavior that a model can learn using an associated slice of data (Definition 1). An ordered skill set is a collection of skills with a directed skills graph that is neither complete nor empty, where an edge from a prerequisite skill to a skill exists if the amount of training it takes to learn the skill can be reduced if the prerequisite skill is also learned (Definition 2, Figure 1 left, center). We show that ordered skill sets exist in synthetic and real datasets using this operational definition. Interestingly, the existence of these ordered skill sets unveils that one can learn a skill quickly not by training solely on that skill, but on a mixture of that skill and prerequisite skills. For instance, in Figure 3 we observe that Spanish QG can be learned more efficiently when the model also learns English QG and Spanish—we can achieve  $4\%$  lower validation loss than training on only Spanish QG over a fixed budget of overall training steps.
Next, given an ordered skill set to train on, we use our framework to propose methods for how to select data so that the LM learn skills faster: skill-stratified sampling and an online generalization, SKILL-IT. We address the issue of unevenly distributed skills in datasets by proposing skill-stratified sampling, a
simple approach that allows us to explicitly optimize for learning skills by uniformly sampling relevant skills (such as a target skill and its prerequisite skills in fine-tuning). Skill-stratified sampling uses the construction of the ordered skill set but is static, which does not incorporate the ordering as training proceeds and results in oversampling skills that may be already learned early on in training. We address this issue by proposing an online data selection algorithm, SKILL-IT, for selecting mixtures of training skills that allocates more weight towards learning skills that are not yet learned or towards prerequisite influential skills (Figure 1 right). SKILL-IT is derived from an online optimization problem over the training skills for minimizing loss on a set of evaluation skills given a fixed budget of data and the skills graph. SKILL-IT is inspired by online mirror descent and can be adapted for continual pre-training, fine-tuning, or out-of-domain evaluation depending on the relationship between the evaluation skill set and the training skill set.
We evaluate SKILL-IT on synthetic and real datasets at two model scales, 125M and 1.3B parameters. For the continual pre-training setting, we show on the LEGO synthetic [72] that we obtain a 37.2 point improvement in accuracy over randomly selecting training data and a 9.7 point improvement over curriculum learning [4]. For the fine-tuning setting, we show that on the widely-used Natural Instructions dataset [41, 64], our algorithm over a mixture of skills is able to achieve up to  $13.6\%$  lower loss on that skill than solely training on that skill, given the same overall training budget. For the out-of-domain setting when our training skills do not align perfectly with evaluation skills, our algorithm is able to achieve the lowest loss on 11 out of 12 evaluation skills corresponding to task categories in the Natural Instructions test tasks dataset over random and skill-stratified sampling on the training data. We finally apply our framework to a case study on the recent RedPajama 1.2 trillion token dataset [57]. We use the data mixture produced by SKILL-IT to continually pre-train a 3B parameter model. We find that SKILL-IT achieves higher accuracy with 1B tokens than uniform sampling over data sources with 3B tokens.
An extended related work can be found in Appendix B. Existing work on data selection for LMs has generally ranged from more computationally expensive methods for dataset condensation on smaller datasets [47, 58, 48] to broader dedduplication and filtering techniques for web-scale datasets [1, 33, 69]. Another way of improving model performance through choice of data is via curriculum learning [4], which also draws inspiration from how humans learn and arranges data in order from easiest to hardest over samples or groups [60]. In contrast to existing works in both curriculum learning and data selection, our work focuses on selecting data for learning an ordered set of skills more efficiently. How LMs learn is also a topic of growing interest; one framework posits that models learn over quanta, discrete units of computation [38], and another proposes that scaling laws for LMs can be understood in terms of learning combinations of skills [2]. Lastly, the notion of skill has been studied in education, ranging from classical research on learning hierarchies [66] to methods for decision-making over lesson sequences [49].
3 Skills framework
First, we propose definitions of skills and ordered skill sets in order to formalize our intuition around how models learn skills, and we demonstrate that not just any existing notion of data groups can characterize an ordered skill set in the dataset. Then, we demonstrate the existence of ordered skill sets on synthetic and real data, which show how viewing data through a skills-based framework can help with training and understanding model performance. Finally, we explore unsupervised skill recovery from data, finding that embedding-based approaches do not adequately recover synthetic skills.
We first present a definition of an individual skill. Let the input space of all possible text data be  $\mathcal{X}$ , where  $x\in \mathcal{X}$  is an individual text sample that a next-token-prediction LM  $f\in \mathcal{F}:\mathcal{X}\to \mathcal{X}$  is trained on. We quantify learning via a metric  $L:\mathcal{F}\times \mathcal{X}\rightarrow \mathbb{R}$ , which maps from a model and evaluation data to a scalar quantity. In our setup, we use the cross-entropy validation loss applied over next-token predictions as our metric  $L$ .
Definition 1 (Skill) A skill  $s$  is a unit of behavior with associated data  $\mathcal{X}_s \subseteq \mathcal{X}$  such that if  $f$  is trained on an dataset  $\mathcal{D}_s \subset \mathcal{X}_s$ ,  $f$  has improved metric  $L$  afterwards on samples belonging to  $\mathcal{X}_s \setminus \mathcal{D}_s$  on average.
This definition of a skill is flexible—it simply means that given a training dataset associated with the skill, a model  $f$  has an improved metric (e.g., decreasing validation loss) when evaluated on validation
Figure 2: Heatmaps of adjacency matrices we compute for skill graphs for Alpaca, Pile of Law, and Natural Instructions. Negative elements and diagonals are thresholded to 0 for clarity. See Appendix D.2 for descriptions of how they were constructed and larger versions.
data associated with this skill. Under this definition, a skill could be a granular task, such as Spanish question generation for a subset of Wikipedia articles, or can be defined over a data source, such as next-token prediction of legal data from tax court rulings. However, our next definition, the ordered skill set, has a more specific construction and provides a framework for how models learn across dependent skills.
Definition 2 (Ordered skill set, skills graph) An ordered skill set for  $f$  is a collection of skills  $S = \{s_1, \dots, s_k\}$  over which there is a directed skills graph  $G = (S, E)$  on the skill set that is neither complete or empty, where  $(s_i, s_j) \in E$  if the amount of data needed to learn  $s_j$  when uniformly sampling from  $\mathcal{D}_{s_i} \cup \mathcal{D}_{s_j}$  is no more than the amount of data needed when sampling only from  $\mathcal{D}_{s_j}$ . We equate learning a skill  $s_j$  to  $f$  attaining a certain value of  $L$  or lower on average over  $\chi_{s_j} \backslash \mathcal{D}_{s_j}$ .
This definition isolates complete and empty graphs as extrema that do not capture meaningful sets of skills. We discuss the three types of skill graphs—complete, empty, intermediate—and their implications for data selection. In particular, we discuss how several initial attempts of defining skills over datasets via semantic groupings resulted in the extrema cases (see Appendix D.2 for full results):
- The complete graph demonstrates that all skills influence each other. A random partition is an example of a skill set that yields a complete graph. This graph suggests that the best approach for learning any skill or set of skills is random sampling on the dataset. This is not a setting where we can gain much with skill-based sampling. For example, using instruction types as skills on the Alpaca dataset results in a nearly complete estimated skills graph (97.4% dense), and we find that stratified sampling on these skills only improves validation loss per skill by 0.007 points over random sampling on average (Figure 2 left), suggesting that utilizing skills does not improve model performance in this case.
- The empty graph demonstrates that each skill is independent. This can occur if skills are too granular; for instance, learning Spanish math problems is unlikely to help with English poem generation. This graph suggests that the best approach for learning an individual skill is to train on the skill itself. We see that empty graphs exist in real data; in Figure 2 (center), using data sources as skills on the Pile of Law [22] results in a nearly empty skills graph (3.9% dense).
- Graphs that are neither empty nor complete thus suggest a nontrivial order of how skill influence each other. This is the setting in which we expect that identifying skills and exploiting their ordering will help the most. In Figure 2 right, we use task categories, which capture broader reasoning patterns, as skills on Natural Instructions and find that the estimated graph has intermediate density (42.7% dense). We show concrete examples of how skills can be learned more efficiently on Natural Instructions in Section 3.2.
While these intuitive groupings result in ordered skill sets on some datasets (e.g., task categories on NI), this is not always the case (e.g., instruction types on Alpaca and sources on Pile of Law). Even though these groupings capture some notion of diversity in the dataset, our findings suggest that not just any semantic grouping induces an ordered skill set. We now empirically demonstrate that our definition of ordered skill sets aligns with how models learn and can be exploited for more data-efficient training.
3.2 Examples of skills and ordered skill sets
We provide examples of ordered skill sets on the LEGO synthetic dataset, an addition synthetic dataset, and subsets of the Natural Instructions dataset. On these datasets, we find that certain skills are better learned when trained along with their prerequisite skills rather than in isolation.
Figure 3: On the LEGO synthetic, 3-digit addition, and Natural Instructions, we identify examples of ordered skill sets in which training on a mixture of skills helps learn an individual skill faster than just training on that skill itself, given a fixed training budget.
LEGO skills The LEGO synthetic, first introduced in [72], can evaluate a model's ability to follow a chain of reasoning. In this synthetic, the letters of the alphabet,  $\mathcal{A}$ , are variables each with some binary label in  $\{0,1\}$ . An individual sample consists of  $k$  clauses for some fixed  $k$  across the dataset, each of the form  $a = g x$  where  $a, x \in \mathcal{A}$  and  $g$  is either a negation ("not") or assertion ("val"), e.g. we assign  $a$  to the value of  $x$ , or we assign  $a$  to the opposite label. At the end of the sentence, we prompt the model for what the value of one of these variables is. Two samples  $x \in \mathcal{X}$  are given below for  $k = 5$ :
Input:  $\mathrm{b} =$  not  $\mathrm{y},\mathrm{r} =$  val 1,  $\mathrm{m} =$  val  $\mathrm{b},\mathrm{q} =$  val  $\mathrm{m},\mathrm{y} =$  not  $\mathrm{r}$  . Output:  $\mathrm{b} = 1$  .
Input:  $c = \text{val} x, p = \text{val} f, x = \text{val} k, f = \text{not} c, k = \text{val} 0$ . Output:  $k = 0$ .
These samples each correspond to a chain of reasoning; for instance the first sample has the chain  $r, y, b, m, q$ , where knowing  $q$ 's label requires the most reasoning steps. We define the  $i$ th skill  $s_i$  as the model's ability to know the  $i$ th variable of the chain. From our example above, the first sample belongs to  $\mathcal{X}_{s_3}$  and the second sample belongs to  $\mathcal{X}_{s_1}$ . To demonstrate the existence of ordered skill sets, we continually pre-train the 125M parameter GPT-Neo model [14, 6] over various mixtures of LEGO skills with  $k = 5$ . In Figure 3 (left), we find that in  $35.9\%$  fewer training steps, training on a balanced mixture of  $\mathcal{X}_{s_1}, \mathcal{X}_{s_2}$ , and  $\mathcal{X}_{s_3}$  resulted in the same validation loss of 0.01 as training solely on  $\mathcal{X}_{s_3}$ . This suggests that  $s_1, s_2$  helped unlock performance on  $s_3$  and that there exist edges from  $s_1$  or  $s_2$  to  $s_3$  in the skill graph. Additional observations are available in Appendix E.1, where we examine other edges as well as more complex reasoning chains, and the full skills graph corresponding to the ordered skill set for LEGO with  $k = 5$  is in Figure 11.
Addition skills We consider a variant of a synthetic 5-digit addition dataset analyzed in [45]. We show the existence of ordered skill sets for a simplified 3-digit addition dataset where we treat each digit prediction as a skill—the outputs, in this case, are the integers  $\{0,1,\dots,9\}$ . Examples are of the following form:
Input:  $\mathrm{A} = 106 + 071$ ,  $\mathrm{A}0 = ?$  Output: 7 Input:  $\mathrm{A} = 606 + 879$ ,  $\mathrm{A}2 = ?$  Output: 4
where 'A 0' refers to the ones digit of the output  $(s_1)$  and 'A 2' refers to the hundreds digit  $(s_3)$ . In Figure 3 (center), we find that in  $32\%$  fewer training steps, training on a balanced mixture of  $\mathcal{X}_{s_1}$  and  $\mathcal{X}_{s_2}$  resulted in the same validation loss of 0.01 as training solely on  $\mathcal{X}_{s_1}$ . That is, the ones digit addition skill can be improved by simultaneously learning the tens digit addition skill, even though the former should not require information from the latter—this is in line with observations from prior work that models do not always learn the ones digit addition first [45]. The full skills graph corresponding to the ordered skill set over 3-digit addition is in Figure 12.
Natural Instructions (NI) skills We show that ordered skill sets exist in NI [63] when we treat task categories as skills.
- In Figure 3 (top right), we show that ordered skill sets exist over crosslingual task categories. Training on Spanish question generation (QG) along with equal parts of English QG, Spanish question answering (QA), and English QA results in  $4.1\%$  lower validation loss than training only on Spanish QG. Remarkably, the former only uses  $25\%$  of the latter's Spanish QG data. This suggests that there are edges from Spanish QA, English QA, and English QG to Spanish QG.
- In Figure 3 (bottom right), we see that training on the task category Text Matching along with Stance Detection helps decrease the loss on Stance Detection by  $11\%$ . This suggests that these categories, which both involve understanding the relationship between two input texts, share an edge.
Table 1: Summary of three settings—continual pre-training, fine-tuning, and out-of-domain. These settings are determined by how  $S_{\mathrm{eval}}$  is defined and result in different skills graphs used for our sampling methods.
Setting	Seval	Skills graph
Continual pre-training	Seval=Strain	A∈Rk×k, edges among all Strain
Fine-tuning	Seval⊂Strain	A∈Rk×m, edges from training skills to target subset
Out-of-domain	Seval∩Strain=∅	A∈Rk×m, edges from training skills to evaluation skill set
The full skills graphs corresponding to the ordered skill sets over these task categories are in Figure 14. While equating task categories to skills may be noisy, these examples suggest that there is signal within real data that suggests that ordered skill sets can improve data efficiency.
3.3 Skill recovery
A final component of characterizing skills is unsupervised recovery of ordered skill sets. We consider embedding-based clustering approaches and a loss-based clustering approach for recovering LEGO skills. When clustering data using various trained and pre-trained embeddings, we find that they were unable to achieve above  $39\%$  accuracy on LEGO. Instead, we find that taking 10 random training runs and clustering data by their loss per timestep per run recovers the skills with  $61\%$  accuracy (Table 3). The intuition behind this method is that the validation losses on points from the same skill have similar trajectories as models learn. We discuss this approach more in Appendix E.2.
4 Skills-based data selection
Now that we have established the existence of ordered skill sets, we discuss how to use them for data selection. We state the data selection problem for learning across skills in Section 4.1. We discuss how to learn the skills graph that will be exploited in our data selection methods in Section 4.2. We then introduce two sampling methods that utilize the graph, a simple skill-stratified sampling method and the online sampling method SKILL-IT, in Section 4.3.
4.1 Problem statement
We are given an ordered training skill set  $S_{\mathrm{train}} = \{s_{\mathrm{train},1},\dots,s_{\mathrm{train},k}\}$  on the training data, each with associated support set  $\mathcal{X}_{s_{\mathrm{train},1}},\ldots,\mathcal{X}_{s_{\mathrm{train},k}}$ , and an ordered evaluation skill set  $S_{\mathrm{eval}} = \{s_{\mathrm{eval},1},\dots,s_{\mathrm{eval},m}\}$  of  $m$  evaluation skills on a separate evaluation dataset. We aim to select  $n$  samples from  $S_{\mathrm{train}}$  via a mixture of training skills,  $p \in \Delta^{k-1}$ , to achieve three goals depending on how  $S_{\mathrm{eval}}$  is constructed: