SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Data selection for LMs There have been several studies of large-scale data selection for LMs. Data dedduplication [33, 1, 23], in which identical or nearly identical samples are removed, is a method that enables LMs to be trained on smaller, cleaned datasets and has been increasingly used as a pre-processing step for training data [59, 5, 71]. Other methods applied at scale involve ensuring high quality of data by explicitly filtering out samples or comparing the training dataset with a cleaned reference dataset [8, 59, 32]. Importance reweighting approaches have also been proposed for identifying training data from a large corpus that best approximates a smaller target distribution [69], and influence functions have been used to select a subset of training data to improve performance on downstream tasks [61]. These approaches can identify data pertaining to a particular target distribution or filter out low quality data according to some heuristic, while our work aims to understand how the choice of data is related to the numerous skills that LMs learn.
Recent development of LMs has shifted focus from emphasizing the scale of the model to prioritizing the training data utilized. For example, models like Alpaca [56], Vicuna [10], and Koala [16] are all based on the LLaMA model combined with instruction data generated by an existing LM. Palm 2's technical report states that the data mixture was a critical component of the final model [18], and Mosaic ML's recent MPT model was trained on a hand-engineered mixture of the RedPajama dataset [43]. However, these works lack rigorous explanation for why their training datasets were constructed in this way.
Finally, perhaps most related to our approach is the contemporary work DoReMi [68], which uses group distributionally robust optimization on a smaller LM to select data source mixtures for training a larger LM. Their approach focuses on selecting data at the data source level for optimizing worst-case performance across the training data sources, rather than at the more general skills level for a variety of target skill sets. Furthermore, we focus on understanding how skills are related to each other and induce some order in how LMs learn by explicitly modeling skill graph structure, which we find to be important for data-efficient LM training (see ablations in Appendix E.5).
Data selection methods Many data selection methods have been proposed for supervised, task-specific settings. In this setting, the most typical objective is dataset condensation, which aims to identify a small subset of data that captures the larger dataset's properties with respect to the model.
Some approaches include constructing coresets [31, 48], identifying samples that the model forgets during training [58]; identifying samples with the largest gradients [47] or gradients that approximate the overall gradient [40]; clustering in embedding space and selecting points farthest from cluster centers [53]; and selecting samples with the highest uncertainty or entropy [34]. These approaches have also been shown to transfer from smaller models to larger models [11]. Unlike these methods, we study how to select data for learning one or many skills at the mixture level for LMs instead of the instance level.
Another area of interest is data selection for domain adaptation and multitask learning. For domain adaptation, there are a wide range of methods that select data to best match the target distribution. For example, the Moore-Lewis method matches data based on the difference in cross-entropy using a model trained on the target versus a model trained on the source data [42]. Several other approaches suggest training a model to distinguish between source and target and selecting points with high uncertainty [50], or selecting points based on some divergence in an embedding space [51]. In comparison to these approaches, our work focuses on learning one or many skills and also finds that embedding-based heuristics do not fully identify skills.
Data attribution Another perspective on understanding training data is data attribution, which seeks to identify what data is responsible for particular model behaviors. Influence functions [29] and shapley values [17] are two ways to quantify the role of individual samples. Datamodels [24] fit a model to predict behavior given a subset of training data, providing a framework for understanding individual samples as well as dataset counterfactuals. Simfluence [21] fits a Markov process to a set of training trajectories for finer-grained understanding of how data impacts training. We focus on understanding how groups of data associated with skills elicit broader model capabilities, and utilize this understanding to select data for more efficient training.
Curriculum learning Curriculum learning [4] proposes to show the model data in order from easy samples to hard ones. Various criteria have been used to determine hardness, and anticurriculum as well as various pacing functions and mixing rates have been explored [54]. Curriculum learning can also be performed at the group level [60]. More sophisticated approaches include parametrizing each sample with a dynamic importance [52], and also accounting for irrelevant and noisy data [39]. Our approach similarly utilizes a curriculum, but it is defined over a skills graph and does not necessarily align with training on easiest to hardest skills.
How LMs learn Many different explanations for how LMs learn from data have been proposed. One hypothesis is that there exist discrete, universal building blocks of LM knowledge called quanta, and power law scaling emerges from a learning over a particular distribution of quanta in the right order [38]. Another work proposes that scaling laws for LMs can be understood in terms of learning combinations of skills [2]. Others have provided theoretical analysis of how transformers learn topics by studying co-occurrences of words in the training data [35]. Empirically, how models learn is still a mystery—for instance, models trained on code are found to perform fairly well at commensense reasoning [37]. Our work initiates a study on how LMs learn various skills and how to exploit this for better data selection.
Task selection In multitask auxiliary learning, the goal is to train a model to perform well on a target task(s) by selecting the most beneficial source tasks to train on. One can use feature similarity to select tasks [30], but we find in our synthetics that feature similarity does not always recover skills. In Taskonomy [70], a hypergraph over a set of tasks is learned and used to select tasks. The methods used to develop the taxonomy can be applied to further expand our graph learning (e.g., studying transitive and higher-order properties). However, their focus is on task selection in computer vision rather than data selection for LMs to learn skills. Lastly, the contemporary work of TaskWeb [25] builds a graph among 22 common NLP tasks in order to determine what the best source tasks are for a target task. Their definition of an edge in the task graph is less strict than ours (their comparison is on if training on additional data from  $s_i$  helps with  $s_j$ , while we fix the overall amount of training data over both  $s_i$  and  $s_j$ ). Overall, our approach is similar in use of the skills graph, but we incorporate it into a dynamic sampling algorithm. Furthermore, we look more broadly at skills, rather than tasks, and characterize when we expect using the skills graph to improve model performance.
Education The notion of skill has been studied in education. Classical research on learning hierarchies [66] identify sets of skills that make up subordinate capabilities for students. For instance, [13] identified that in order for students to solve linear equations, there were many prerequisites
skills, ranging from the simplest being symbol recognition to the most complex being the ability to add, subtract, multiple, and divide from both sides of the equation. More recently, decision-making over lesson sequences based on skills, e.g., what the student already knows versus what the lesson teaches, has become an area of interest in personalized learning [49].
C Additional Algorithmic Details
C.1 Derivation of SKILL-IT Update Rule
First, we provide the derivation of our update rule from online mirror descent using the proximal point view [19]. We restate our optimization problem from (3):
$$
\underset {p _ {1}, \dots , p _ {T} \in \Delta^ {k - 1}} {\text {m i n i m i z e}} \frac {1}{m} \sum_ {j = 1} ^ {m} L _ {\text {e v a l}, j} \left(f _ {T}\right) \tag {5}
$$
$$
s. t L _ {\text {e v a l}, j} \left(f _ {t}\right) = L _ {\text {e v a l}, j} \left(f _ {t - 1}\right) \left(1 - \alpha A _ {:, j} ^ {\top} p _ {t - 1}\right) \forall j \in [ m ], t = 1, \dots , T
$$
$$
f _ {t} = \Phi (f _ {t - 1}, p _ {t - 1}) \forall t = 1, \dots T
$$
Let  $\bar{L}_t(p) = \frac{1}{m}\sum_{i = j}^{m}L_{\mathrm{eval},j}(f_{t + 1}) = \frac{1}{m}\sum_{i = j}^{m}L_{\mathrm{eval},j}(\Phi (f_t,p))$ ; that is,  $p$  is the mixture we must choose at time  $t$  and  $\bar{L}_t$  is the average loss per skill of the model after it is trained on  $p$  at round  $t$ . A greedy approximation of (5) is minimize  $\bar{L}_t(p)$ , given the model and mixtures at previous rounds. A linear approximation of  $\bar{L}_t(p)$  is
$$
\bar {L} _ {t} (p) \approx \bar {L} _ {t} \left(p _ {t - 1}\right) + \left\langle \nabla \bar {L} _ {t - 1} \left(p _ {t - 1}\right), p - p _ {t - 1} \right\rangle \tag {6}
$$
Then, the problem of minimizing  $\bar{L}_t(p)$  becomes
$$
\operatorname {a r g m i n} _ {p \in \Delta^ {k - 1}} \left\langle \eta \nabla \bar {L} _ {t - 1} \left(p _ {t - 1}\right), p \right\rangle \tag {7}
$$
after we drop terms from (6) that do not depend on  $p$ . Note that the  $\eta$  is a constant and does not impact the solution. The optimal solution to this problem is selecting the  $p$  that has the most weight on the slice with the largest gradient. To improve stability and prevent overfitting, we introduce regularization via a Bregman divergence  $D_{h}(p||p_{t - 1}) = h(p) - h(p_{t - 1}) - \langle \nabla h(p_{t - 1}),p - p_{t - 1}\rangle$ . After dropping terms that do not contain  $p$ , our problem is now
$$
\operatorname {a r g m i n} _ {p \in \Delta^ {k - 1}} \left\langle \eta \nabla \bar {L} _ {t - 1} \left(p _ {t - 1}\right), p \right\rangle + h (p) - \left\langle \nabla h \left(p _ {t - 1}\right), p \right\rangle \tag {8}
$$
Taking the gradient and setting it equal to 0 gives us
$$
\eta \nabla \bar {L} _ {t - 1} \left(p _ {t - 1}\right) + \nabla h (p) - \nabla h \left(p _ {t - 1}\right) = 0 \tag {9}
$$
Similar to in standard multiplicative weights, we set  $h(p) = \sum_{i} p_{i} \ln p_{i}$  and  $\nabla h(p) = [\ln p_{i} + 1]_{i}$ . Then,
$$
\begin{array}{l} \ln p ^ {i} = \ln p _ {t - 1} ^ {i} - \eta \nabla_ {i} L _ {t - 1} \left(p _ {t - 1}\right) \\ \Rightarrow p _ {t + 1} ^ {i} = p _ {t} ^ {i} \exp (- \eta \nabla_ {i} \bar {L} _ {t} (p _ {t})) \tag {10} \\ \end{array}
$$
where  $\nabla_{i}$  is the  $i$ th element of the gradient. Now we wish to compute  $\nabla_{i}\bar{L}_{t}(p_{t}) = \frac{1}{m}\sum_{j=1}^{m}\nabla_{i}[L_{\mathrm{eval},j}(f_{t+1})] = \frac{1}{m}\sum_{j=1}^{m}\nabla_{i}[L_{\mathrm{eval},j}(\Phi(f_{t},p_{t}))]$ . Recall the dynamics model for  $L_{\mathrm{eval}}$ :
$$
L _ {\text {e v a l}, j} \left(f _ {t + 1}\right) = L _ {\text {e v a l}, j} \left(f _ {t}\right) \left(1 - A _ {:, j} ^ {\top} p _ {t}\right), \tag {11}
$$
The gradient of this model with respect to each training skill  $s_i$  is
$$
\begin{array}{l} \nabla_ {i} L _ {\text {e v a l}, j} \left(f _ {t + 1}\right) = - A _ {i j} L _ {\text {e v a l}, j} \left(f _ {t}\right) \tag {12} \\ \Rightarrow \nabla_ {i} \bar {L} _ {t} \left(p _ {t}\right) = \frac {1}{m} \sum_ {j = 1} ^ {m} - A _ {i j} L _ {\text {e v a l}, j} \left(f _ {t}\right) \\ \end{array}
$$
Plugging this back into (10),
$$
p _ {t + 1} ^ {i} = p _ {t} ^ {i} \exp \left(\eta \sum_ {j = 1} ^ {m} A _ {i j} L _ {\text {e v a l}, j} \left(f _ {t}\right)\right), \tag {13}
$$
where we can absorb the  $\frac{1}{m}$  into  $\eta$ .
Algorithm 2 Brute-Force Graph Learning
1: Input: Ordered skill set  $S = \{s_1, \dots, s_k\}$ . Number of training steps  $K$ , base model  $f_0$ .  
2: for  $j \in [k]$  do  
3: Train  $f_0$  on samples from  $\mathcal{X}_{s_j}$  for  $K$  steps and denote  $f_{K,j}$  to be the model after training.  
4: Observe change in loss,  $\delta_j^j = L_{\mathrm{eval},j}(f_0) - L_{\mathrm{eval},j}(f_{K,j})$ .  
5: end for  
6: for  $i,j \in [k]$  do  
7: Train  $f_0$  on samples from  $\mathcal{X}_{s_i} \cup \mathcal{X}_{s_j}$  for  $K$  steps and denote  $f_{K,i,j}$  to be the model after training.  
8: Observe change in loss,  $\delta_j^{i,j} = L_{\mathrm{eval},j}(f_0) - L_{\mathrm{eval},j}(f_{K,i,j})$ .  
9: if  $\delta_j^{ij} > \delta_j^j$  then  
10: Draw edge  $s_i \rightarrow s_j$  and set  $A_{ij} > 0$ .  
11: end if  
12: end for  
13: return Adjacency Matrix  $A \in \mathbb{R}^{k \times k}$
Algorithm 3 Approximate Graph Learning
1: Input: Ordered skill sets  $S_{train}$  and  $S_{eval}$ . Number of training steps  $K$ , base model  $f_0$ .  
2: for  $i \in [k]$  do  
3: Train  $f_0$  on samples from  $\mathcal{X}_{s_{train},i}$  for  $K$  steps and denote  $f_{K,i}$  to be the model after training.  
4: for  $j \in [m]$  do  
5: Observe change in loss,  $\delta_j^i = L_{eval,j}(f_0) - L_{eval,j}(f_{K,i})$ .  
6: If  $\delta_j^i > 0$ , draw edge  $s_{train,i} \rightarrow s_{train,j}$  and set  $A_{ij} > 0$ .  
7: end for  
8: end for  
9: return Bipartite Adjacency Matrix  $A \in \mathbb{R}^{k \times m}$
C.2 Graph Learning Method
We provide algorithms for learning the graph over an ordered skill set. In Algorithm 2, we discuss the brute-force approach for learning the adjacency matrix. This approach only works when  $S_{\mathrm{eval}} \subseteq S_{\mathrm{train}}$  (e.g. pre-training and fine-tuning cases), so we denote  $S = S_{\mathrm{train}}$  in the algorithm box. In Algorithm 3, we discuss the linear approach for learning the adjacency matrix. This approach works even in the out-of-domain case when  $S_{\mathrm{eval}}$  and  $S_{\mathrm{train}}$  are disjoint.
In both approaches, the exact value of  $A_{ij}$  can vary, but we can typically set it proportional to  $\delta_j^{i,j} - \delta_j^j$  in the brute-force case or  $\delta_j^i$  in the approximate case. The exact constructions and methods for learning each  $A$  in our experiments are in Appendix D.2.
We present details about each dataset used, including information on the skills and the validation dataset. A summary is presented in Table 2.
- Alpaca dataset [56]: the Alpaca dataset consists of 52K instruction examples that were generated from text-davinci-003. We applied the Berkeley Neural Parser [27, 28] to each instruction, keeping 40777 samples it was able to parse successfully. If the sample began with a question, we annotated it with the skill "question", and otherwise we annotated it with the verb identified from the parser. We grouped the data into a total of 38 skills, such as "list", "edit", "calculate", "describe" and "identify".
- Pile of Law [22]: the Pile of Law dataset consists of various sources of legal and administrative data, ranging from tax rulings to the world's constitutions. We evaluate on a subset of the Pile of Law validation dataset consisting of 13883 samples, where we selected max(645, source size) samples per source. We truncated each sample to be no more than 100K characters.
- LEGO [72]: for the LEGO synthetic, we set  $k = 5$  and sample 192000 points across the skills. Our validation dataset consisted of 100 samples per skill.
Table 2: We list each dataset used as well as its corresponding skill. We include the number of skills in the training dataset, as well as details on how the validation dataset is constructed.
Dataset	Skill	# skills	Validation data
Alpaca	Instruction type	38	50 samples per skill
Pile of Law	Legal data source	31	645 samples per skill
LEGO	Reasoning chain depth	5	100 samples per skill
Addition	Digit	3	100 samples per skill
NI (pre-training)	Task category	23	50 samples per task
NI (Spanish QG)	Task category × language	4	100 samples per task
NI (stance detection)	Task category	2	50 samples per task
NI (out-of-domain)	Task category	59,12	400 samples per task
RedPajama	Data source	7	LM eval harness
- Addition: for the 3-digit addition synthetic, we set  $k = 3$  and sample 192000 points across the skills. We use a validation dataset of 100 samples per skill.
- Natural Instructions [63, 41]: the Natural Instructions dataset is a large collection of tasks and their definitions in natural language. For the pre-training setting, we used a set of 23 task categories that had the largest degree (in-degree + out-degree) in the learned skills graph, for a total of 1,232,437 samples and 425 tasks to select from. We evaluated on 50 samples per task.
For the fine-tuning setting with Spanish question generation, we select data over 4 skills (Spanish question generation, Spanish question answering, English question generation, English question answering) for a total of 513210 samples and 212 tasks to select from. We evaluated on 100 samples per task.
For the fine-tuning setting with stance detection, we select data over 2 skills (stance detection, text matching) for a total of 50990 samples and 19 tasks to select from. We evaluated on 50 samples per task.
For the out-of-domain setting, we select data over all 59 task categories for a total of 2,417,867 samples and 753 tasks to select from. The test split consisted of 12 task categories and 119 tasks, and we evaluated on min(400, task size) samples per task.
- RedPajama [57]: the RedPajama dataset is a 1-trillion token dataset that aims to reproduce the LLaMA [59] training dataset. We select over the 7 data sources and evaluate using the LM evaluation harness [15].
D.2 Graph Learning Details
We describe how the skills graph was learned on each dataset.
- Alpaca (Figure 9): we use Algorithm 3 and train for  $K = 150$  steps per skill. Each edge  $i \rightarrow j$  has a weight of  $\delta_j^i$ , the difference in loss on skill  $j$  before and after training on  $i$ . Next, we compare the average validation loss of skill-stratified sampling versus random sampling when we train for  $K = 1000$  steps. We find that skill-stratified sampling only does 0.007 better than random sampling, confirming that Alpaca's dense skills graph suggests that random sampling is the best we can do.
- Pile of Law (Figure 10): we use Algorithm 3 and train for  $K = 150$  steps. Each edge  $i \to j$  has a weight of  $\delta_j^i$ , the difference in loss on skill  $j$  before and after training on  $i$ .
- LEGO (Figure 11): we use both Algorithm 2 and Algorithm 3 and train for  $K = 6000$  steps each. Each edge  $i \to j$  has a weight of 0.5 if the amount of data associated with skill  $j$  that is needed to reach 0.01 validation loss is less when training on  $(i,j)$  than on  $j$  (edges are set to 0 if 0.01 validation loss is not reached, even if loss is decreasing). Each edge  $i \to j$  is also set to 0.5 if training on  $i$  decreases loss directly on  $j$ . We set each diagonal entry of  $A$  to be 1.
- Addition (Figure 12): we use Algorithm 2 and train for  $K = 6000$  steps. Each edge  $i \rightarrow j$  has a weight of 0.5 if the amount of data associated with skill  $j$  that is needed to reach 0.01 validation loss is less when training on  $(i,j)$  than on  $j$  (edges are set to 0 if 0.01 validation loss is not reached, even if loss is decreasing). We set each diagonal entry of  $A$  to be 1.
- Natural Instructions (Figure 13, 14, 15): we use Algorithm 3. For the pre-training setting, we train for  $K = 600$  steps and assign each edge  $i \to j$  a weight  $\delta_j^i$  equal to the change in loss on  $j$  in the first 100 steps for all  $i, j \in [k]$ , including diagonal entries. For the fine-tuning setting, we train for  $K = 600$  steps and assign each edge  $i \to j$  a weight  $\delta_j^i$  equal to the change in loss before and after training. For the out-of-domain setting, we train for  $K = 600$  steps and assign each edge  $i \to j$  a weight  $\delta_j^i$  equal to the change in loss before and after training in the first 100 steps.
Figure 9: Alpaca heatmap where  $i, j$ th entry is  $\max(0, \delta_j^i)$  (the change in loss on  $s_j$  after training on  $s_i$  for 150 steps). Diagonal entries are set to 0 for clearer visualization.
- RedPajama (Figure 16): we use Algorithm 3 and train for 1 billion tokens per data source. We assign each edge  $i \rightarrow j$  a weight  $\delta_j^i$  equal to the change in perplexity on the validation data also before and after training.
D.3 Training Details
We describe the parameters used for SKILL-IT.
SKILL-IT pre-training
- LEGO:  $\eta = 0.5, T = 6, w = 3$ . We train for 6000 steps.
- Addition:  $\eta = 0.1, T = 5, w = 3$ . We train for 6000 steps.
- Natural Instructions (pre-training):  $\eta = 0.2, T = 1$ . We train for 5000 steps.
For the LEGO random baseline, when we selected points at random, we used an imbalanced training dataset with proportions 1:1:1:3:5. For the addition random baseline, we used an imbalanced dataset with randomly selected proportions: 13:14:18. For the curriculum learning baselines, the pacing function,  $g(i)$ , denotes the size of the subset of the highest scoring samples that we uniformly select from in the  $i$ th epoch. We define our pacing function as  $g(i) = \frac{iH}{M}$ , where  $H$  is the number of steps and  $M$  is 5 epochs for LEGO and NI, and 3 for addition.
SKILL-IT fine-tuning
- LEGO:  $\eta = 0.5, T = 10, w = 3$ . We train for 6000 steps.
- Addition:  $\eta = 0.1, T = 5, w = 3$ . We train for 6000 steps.
- Natural Instructions (Spanish QG):  $\eta = 0.8, T = 6, w = 3$ . We train for 600 steps.
- Natural Instructions (stance detection):  $\eta = 0.2, T = 6, w = 3$ . We train for 600 steps.
SKILL-IT out-of-domain
- Natural Instructions:  $\eta = 0.2, T = 10, w = 3$ . We train for 5000 steps.
- RedPajama:  $\eta = 100, T = 1$ . We train for 3 billion tokens.
All results are computed over 5 random seeds.
Figure 10: Pile of Law heatmap where  $i, j$ th entry is  $\max(0, \delta_j^i)$  (the change in loss on  $s_j$  after training on  $s_i$  for 150 steps). Diagonal entries are set to 0 for clearer visualization.
Figure 11: LEGO heatmap with  $k = 5$  where  $i, j$ th entry is set to 0.5 if the number of steps needed to reach 0.01 loss on skill  $j$  when training on a balanced mixture of skills  $i$  and  $j$  is less than when training on skill  $j$  only.
Figure 12: Addition heatmap with  $k = 3$  where  $i, j$ th entry is set to 0.5 if the number of steps needed to reach 0.01 loss on skill  $j$  when training on a balanced mixture of skills  $i$  and  $j$  is less than when training on skill  $j$  only.
Figure 13: Natural Instructions heatmap where  $i, j$ th entry is  $\max(0, \delta_j^i)$  (the change in loss on  $s_j$  after training on  $s_i$  for 100 steps). Diagonal entries are set to 0 for clearer visualization.
Figure 14: Spanish question generation and stance detection heatmaps where  $i, j$ th entry is  $\max(0, \delta_j^i)$  (the change in loss on  $s_j$  after training on  $s_i$  for 100 steps).
Figure 15: Natural Instructions heatmap for out-of-domain setting where rows are for the training skills and columns are for the evaluation skills. The  $i,j$ th entry is  $\max(0, \delta_j^i)$  (the change in loss on  $s_j$  after training on  $s_i$  for 100 steps).
Figure 16: RedPajama heatmap for out-of-domain setting where rows are for the training skills and columns are for the evaluation skills. The  $i,j$ th entry is  $\max(0, \delta_j^i)$  (the change in perplexity on  $s_j$  after training on  $s_i$  for 1B tokens).
Batch sizes of 32 and 64 were used for the LEGO and addition synthetic on the 125M and 1.3B parameter model, respectively. Batch sizes of 4 and 16 were used for the Natural Instructions experiments on the 125M and 1.3B parameter model.