SlowGuess/NeurIPS_2023_OCR huggingface dataset:
- Continual pre-training: when  $S_{\mathrm{eval}} = S_{\mathrm{train}}$ , our goal is to select a mixture of training skills to learn all of them.
- Fine-tuning: when  $S_{\text{eval}} \subset S_{\text{train}}$ , our goal is to select a mixture of training skills to learn an individual target skill or subset of these skills.
- Out-of-domain: when  $S_{\mathrm{eval}} \cap S_{\mathrm{train}} = \emptyset$ , our goal is to select a mixture of training skills to learn a disjoint set of evaluation skills we cannot train on. This can arise when we have a separate downstream validation dataset or the skills identified in the training dataset are noisy.
Furthermore, we have a skills graph  $G = (\mathcal{S}_{\mathrm{train}} \cup \mathcal{S}_{\mathrm{eval}}, E)$ , where  $E \subseteq \mathcal{S}_{\mathrm{train}} \times \mathcal{S}_{\mathrm{eval}}$  and  $A \in \mathbb{R}^{k \times m}$  is a weighted adjacency submatrix, where  $A_{ij}$  describes the strength of the edge from  $s_{\mathrm{train},i}$  to  $s_{\mathrm{eval},j}$ . In Table 1, we summarize how the three different settings are constructed and how  $A$  varies across them. Next, we discuss how  $A$  can be estimated from the data.
4.2 Skills graph learning
The skills graph is important for determining how to sample from the ordered skill set for training efficiently. We present two approaches for learning the skills graph—brute-force and linear approximation. Algorithms are provided in Appendix C.2. By definition 2, the brute-force way of identifying edges involves fixing an overall training budget of  $H$  steps and 1) training and evaluating the model on each  $s_{\text{train},i}$  and 2) training the model on each pair of  $(s_{\text{train},i}, s_{\text{eval},j})$  and evaluating on  $s_{\text{train},i}$  and  $s_{\text{eval},j}$ . If the loss on  $s_{\text{eval},j}$  when trained on both  $s_{\text{train},i}$  and  $s_{\text{eval},j}$  is lower, there exists an edge from  $s_{\text{train},i}$  to  $s_{\text{eval},j}$  with edge weight proportional to the difference in loss. This approach has runtime  $\mathcal{O}(Hkm)$  and is only feasible when  $k$  is small and when we have access to  $S_{\text{eval}}$  at training
Algorithm 1 SKILL-IT Online Data Selection Algorithm
1: Input: Ordered training skill set  $S_{\text{train}}$ , ordered evaluation skill set  $S_{\text{eval}}$ . Learning rate  $\eta$ ,  $T$  rounds,  $n$  samples,  $H$  training steps per run for graph learning, model  $f_1$ , window parameter  $w$ .
2:  $A\gets$  LEARNGRAPH  $(S_{train},S_{eval},H,f_1)$  (Alg. 2, 3).
3: Initialize  $p_1^i = \exp(\eta \sum_{j=1}^m A_{ij})$  for all  $i \in [k]$ , the softmax over  $A$ .
4: for  $t = 1,\dots,T - 1$  do
5: Observe losses  $L_{\mathrm{eval},j}(f_t)$  for all  $s_{\mathrm{eval},j} \in S_{\mathrm{eval}}$ .
6: Train model  $f_{t}$  with  $n / T$  samples from mixture  $p_t$  over  $S_{\mathrm{train}}$ . Update model  $f_{t + 1} = \Phi (f_t,p_t)$ .
7: Set  $p_{i + 1}^{t} = \exp (\eta \sum_{\tau = t - w + 1}^{t}\sum_{j = 1}^{m}A_{ij}L_{\mathrm{eval},j}(f_{\tau}))$ .
8: end for
time. Otherwise, we can approximate this approach in linear time by training on each  $s_i$  for  $h < H$  steps and setting  $A_{ij} > 0$  if the loss on  $s_j$  decreases over  $h$  steps for a runtime of  $\mathcal{O}(hk)$ . This linear approach is necessary in the out-of-domain setting, since it does not require training on  $S_{\mathrm{eval}}$ . In addition, both graph learning approaches can be performed on a smaller model, and the learned graph can be used for data selection for training a larger model (Appendix E.4).
4.3 Skills graph-aware sampling
We present two approaches for sampling over the mixture of training skills according to the skills graph: skill-stratified sampling, which samples uniformly over relevant training skills according to  $A$ , and SKILL-IT, an online generalization that incorporates feedback of how skills are being learned so far.
4.3.1 Skill-stratified sampling
A straightforward sampling approach is to discard training skills that do not benefit the evaluation skills and sample uniformly over the set of relevant training skills, which we call skill-stratified sampling. For continual pre-training, the relevant skills are the entire training skill set; for each  $s_{\text{train},i} \in S_{\text{train}}$ ,  $\operatorname*{Pr}(s_{\text{train},i}) = \frac{1}{k}$ . This enables each skill to have sufficient training data. For fine-tuning, the relevant skills are the target skills and prerequisite skills, which can be identified via positive entries of the  $i$ th column of  $A$  with  $S_{\text{prerequisite}} = \{s_{\text{train},i} : \exists s_{\text{eval},j} \text{ s.t. } A_{ij} > 0\}$ . We then set  $\operatorname*{Pr}(s) = \frac{1}{|S_{\text{prerequisite}} \cup S_{\text{eval}}|}$  for  $s \in S_{\text{prerequisite}} \cup S_{\text{eval}}$ . For the out-of-domain setting, skill-stratified sampling is over the set of prerequisite skills. For each  $s \in S_{\text{prerequisite}}$ , we set  $\operatorname*{Pr}(s) = \frac{1}{|S_{\text{prerequisite}}|}$ . Next, we propose our online algorithm that exploits the graph dynamically for more efficient training.
4.3.2 SKILL-IT online data selection algorithm
Despite accounting for prerequisite skills, one shortcoming of skill-stratified sampling is that even if a skill has already obtained sufficiently low validation loss early during training, we will continue to allocate the same weight to that skill throughout training. Therefore, we formulate our data selection problem as an online learning problem and propose SKILL-IT, which both prioritizes prerequisite skills and skills that are not yet learned.
We are given a budget of  $T$  rounds and  $n$  total samples to train on. At round  $t$ , we select a mixture  $p_t \in \Delta^{k-1}$  from the  $k$ -dimensional unit simplex, and for each training skill  $s_{\mathrm{train},i} \in S_{\mathrm{train}}$ , we sample from  $\mathcal{X}_{s_{\mathrm{train},i}}$  with proportion  $p_t^i$  for a total of  $\frac{n}{T}$  samples per round. Let  $f_t$  be the model at the start of round  $t$ . We can define  $f_t$  recursively as a function of the previous round's model  $f_{t-1}$  and mixture  $p_{t-1}$  via a dynamics function  $\Phi: \mathcal{F} \times \Delta^{k-1} \to \mathcal{F}$ ; that is,  $f_t = \Phi(f_{t-1}, p_{t-1})$ . Let  $L_{\mathrm{eval},j}(f_t)$  be the validation loss of  $f_t$  on  $s_{\mathrm{eval},j}$ . Our goal is to select  $p_1, \ldots, p_T$  to minimize loss per evaluation skill at the end of training:
$$
\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 {1}
$$
This optimization problem is challenging to solve without additional assumptions. In order to make the problem tractable, we impose an explicit dynamics rule for the each evaluation skill's loss  $L_{\mathrm{eval},j}$  in terms of the current loss and data mixture. Assuming for simplicity that  $S_{\mathrm{eval}} \subseteq S_{\mathrm{train}}$ , a simple rule would be  $L_{\mathrm{eval},j}(f_t) = L_{\mathrm{eval},j}(\Phi(f_{t-1}, p_{t-1})) := L_{\mathrm{eval},j}(f_{t-1})(1 - \alpha p_{t-1}^j)$  for  $\alpha \in [0,1]$ . That is, we expect that allocating more data to skill  $j$  should result in the validation loss on skill  $j$  decreasing. However, such
an expression assumes that only training on the  $j$ th skill will help learn the  $j$ th skill. Instead, Section 3.2 suggests that there are other skills that may help with the  $j$ th skill. We propose the following dynamics:
$$
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 - A _ {., j} ^ {\top} p _ {t - 1}\right), \tag {2}
$$
where  $A_{\cdot ,j}$  is the column with weights of all skills that influence  $s_{\mathrm{eval},j}$ , and we absorb the scalar  $\alpha$  into  $A$ . The optimization problem in (1) can thus be simplified as follows:
$$
\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 {3}
$$
$$
s. t f _ {t} = \Phi \left(f _ {t - 1}, p _ {t - 1}\right) \forall t = 1, \dots 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 - A _ {:, j} ^ {\top} p _ {t - 1}\right) \forall j \in [ m ]
$$
In Appendix C, we derive the following update rule via online mirror descent [46] for learning rate  $\eta > 0$ :
$$
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 {4}
$$
In addition, when equation 4 is expanded, we have that  $p_{t+1}^i = p_1^i \exp \left( \eta \sum_{\tau=1}^{t} \sum_{j=1}^{m} A_{ij} L_{\text{eval},j}(f_{\tau}) \right)$ . Since this summation over  $\tau$  results in diminishing strength of updates, we change it to a moving window of size  $w$ . Our full method is in Algorithm 1.
Intuitively, at each step we adjust the weight on skill  $i$  based on the losses of skills that  $i$  influences, with the assumption that more training data helps decrease loss. Note that when we use our algorithm with a complete graph or empty graph, we achieve expected behavior discussed in Section 3.1. For the complete graph, our algorithm reduces to stratified sampling. When we have a skill set with an empty graph, the update rule reduces to sampling proportional to each skill's validation loss.
Given an ordered skill set, we aim to validate SKILL-IT's ability to select data for efficiently learning skills in the continual pre-training, fine-tuning, and out-of-domain settings. We provide full tables of results in Appendix E.3.1 and results where we learn the skills graph on the 125M model and use it for the 1.3B parameter model in Appendix E.4. Skills graphs are in Appendix D.2, weight trajectories for SKILL-IT are in Appendix E.3.2, and ablations on the graph and online components of SKILL-IT are in Appendix E.5.
5.1 Continual pre-training
Setup We evaluate the ability of SKILL-IT to select data for efficiently learning over all skills. We measure average validation loss per skill after a fixed number of training steps. We construct the LEGO synthetic and addition synthetic with  $k = 5$  and 3, respectively, and an imbalanced dataset over the skills. On the Natural Instructions dataset, we use 23 of the task categories as skills.
Baselines We compare SKILL-IT against three baselines that do not account for skills: random sampling, curriculum learning, and anticurriculum learning. Random sampling is a standard procedure for selecting samples given no additional information. Curriculum learning [4] and anticurriculum learning [67] score the samples from easiest to hardest and vice versa, respectively, and sample over an expanding set of the lowest scored samples at every epoch; we use the pre-trained model's loss to rank points. We evaluate skill-stratified sampling, which uses knowledge of the skills but is not online, and include an additional skills curriculum and anticurriculum baseline from [60], which samples from skills in order of average loss per skill.
Analysis Across our experiments we find that SKILL-IT outperforms baselines that do not use skills as well as skill-stratified sampling and skill curriculum learning. Our results on the LEGO dataset are shown in Figure 4. SKILL-IT and skill-stratified sampling attain lower loss than other approaches on skills 2, 3, 4 and on average, and while curriculum and anticurriculum learning attain lower loss on skill 5, they fail to learn other skills, resulting in at most 68 points accuracy on skill 3. SKILL-IT reaches
Figure 4: Performance of SKILL-IT on each skill in the continual pre-training setting (learning over all skills in the ordered training skill set) on the LEGO synthetic.
Figure 5: Performance of SKILL-IT in the continual pre-training setting on the addition synthetic.
a high average accuracy earlier in training than other approaches; halfway through training, SKILL-IT obtains between 8.3 and 33.5 points higher average accuracy than other approaches, reaching a final accuracy of 99.3 (Figure 20). SKILL-IT initially allocates more weight to prerequisite skills such as skill 2 as suggested by Figure 11 and later on allocates more weights to the skills that are learned more slowly, such as skills 4 and 5 (Figure 22). On the addition synthetic with  $k = 3$ , SKILL-IT obtains lower validation loss than the baselines on skills 1 and 2 in Figure 5. While most approaches aside from skill curriculum learning eventually obtain  $100\%$  accuracy on all skills, SKILL-IT requires less training to reach sufficiently high accuracy; halfway through training, SKILL-IT has accuracy between 8.7 and 73.1 points higher than other approaches (Figure 21). Finally on Natural Instructions, the average validation loss from SKILL-IT is  $3.2\%$  lower than from random sampling (Table 7). Our results suggest that exploiting the construction and ordering of skills is critical to learning skills quickly.
5.2 Fine-tuning
Figure 6: Performance of SKILL-IT in the fine-tuning setting on LEGO, addition, and NI.
Setup We evaluate the ability of SKILL-IT to select data from an ordered training skill set for learning a target skill. Mirroring Figure 3, we evaluate on LEGO target skill 3 (third in reasoning chain), on the addition synthetic's skill 1 (ones place digit addition), and on NI's Spanish QG and Stance Detection.
Baselines We compare SKILL-IT against training on the target skill only and skill-stratified sampling over prerequisite skills and the target skill. The skill-stratified sampling approach uses the ordered skill set to identify prerequisite skills, but does not exploit them dynamically.
Analysis Our results are shown in Figure 6. On LEGO, SKILL-IT results in the same validation loss of 0.01 as training only on the target skill in  $38.1\%$  fewer steps. We observe a similar trend on addition, with SKILL-IT converging to a validation loss of 0.01 in  $59\%$  fewer steps required to do so when training only on the target skill. Finally, on NI, SKILL-IT improves validation loss on Spanish question generation by  $5.3\%$  and Stance Detection by  $13.6\%$  over just training on the respective target skill only. In this setting,
Figure 7: Performance of SKILL-IT in the out-of-domain setting for the NI test task split. SKILL-IT uses the graph between the train and evaluation skills to produce an online mixture on the training dataset.
Number of tokens (billion)
Figure 8: Left: Accuracy on LM Evaluation Harness for continual pre-training of a 3B parameter model using SKILL-IT on the RedPajama dataset. We achieve higher accuracy at 1B additional tokens than uniform at 3B tokens. Right: SKILL-IT mixture over RedPajama sources.
RedPajama source	SKILL-IT mixture
ArXiv	0.1370
Books	0.0437
C4	0.4195
CommonCrawl	0.0732
GitHub	0.1890
StackExchange	0.0892
Wikipedia	0.0484
a significant portion of the improvement over training only on the target skill comes from identification of prerequisite skills through the learned graph in the skill-stratified sampling method. SKILL-IT is further able to improve performance with finer-grained dynamic weighting on prerequisite skills.
5.3 Out-of-domain setting
Natural Instructions We evaluate the ability of SKILL-IT to select data from a set of training skills for learning a disjoint set of evaluation skills that we cannot train on. We use all 59 task categories in the NI train tasks split as the training skills and the 12 task categories in the test tasks split as our evaluation skills. We compare SKILL-IT against random and skill-stratified sampling, both of which do not exploit the relationships between training skills and evaluation skills. SKILL-IT achieves the lowest loss on 11 out of 12 task categories over random and skill-stratified sampling (Figure 7, tables in Appendix).
RedPajama We use SKILL-IT to produce a data mixture on the RedPajama dataset. The training skills are the data sources comprising the dataset, and the evaluation skills are several tasks from the Language Model Evaluation Harness [15]. SKILL-IT with  $T = 1$  yields the mixture in Figure 8 (right). We continually pre-train a 3B parameter model trained on 1T tokens for 3B additional tokens using this mixture, and see that it outperforms uniform sampling over the data sources (Figure 8 left). In particular, SKILL-IT achieves higher accuracy with 1B additional tokens than uniform with 3B additional tokens.
Given a fixed budget of data, knowing what data to train on to induce various capabilities in an LM is challenging. As LMs continue to improve, it will become increasingly important to extract as much signal as possible from the data and to direct that signal towards acquiring a broad variety of capabilities. In this paper, we introduce a skills-based framework for understanding how LMs learn and for selecting training data. We hope our study invites others to build on such a notion of skill and further explore how to align skills with data.
We thank Together AI (https://together.ai/) and Stanford HAI for providing portions of the compute used to train models in this paper. We thank Sabri Eyuboglu, Karan Goel, Arjun Desai, Neel Guha, Michael Zhang, Vishnu Sarrukai, Simran Arora, Ben Spector, Brandon Yang, Gautam Machiraju, and Sang Michael Xie for their helpful feedback and discussion.
We gratefully acknowledge the support of NIH under No. U54EB020405 (Mobilize), NSF under Nos. CCF1763315 (Beyond Sparsity), CCF1563078 (Volume to Velocity), and 1937301 (RTML); US DEVCOM ARL under No. W911NF-21-2-0251 (Interactive Human-AI Teaming); ONR under No. N000141712266 (Unifying Weak Supervision); ONR N00014-20-1-2480: Understanding and Applying Non-Euclidean Geometry in Machine Learning; N000142012275 (NEPTUNE); NXP, Xilinx, LETI-CEA, Intel, IBM, Microsoft, NEC, Toshiba, TSMC, ARM, Hitachi, BASF, Accenture, Ericsson, Qualcomm, Analog Devices, Google Cloud, Salesforce, Total, the HAI-GCP Cloud Credits for Research program, the Stanford Data Science Initiative (SDSI), and members of the Stanford DAWN project: Facebook, Google, and VMWare. FS is supported by NSF CCF2106707 and the Wisconsin Alumni Research Foundation (WARF). The U.S. Government is authorized to reproduce and distribute reprints for Governmental purposes notwithstanding any copyright notation thereon. Any opinions, findings, and conclusions or recommendations expressed in this material are those of the authors and do not necessarily reflect the views, policies, or endorsements, either expressed or implied, of NIH, ONR, or the U.S. Government.
[1] Amro Abbas, Kushal Tirumala, Daniel Simig, Surya Ganguli, and Ari S. Morcos. Semdedup: Data-efficient learning at web-scale through semantic dedduplication, 2023.
[2] Sanjeev Arora and Anirudh Goyal. A theory for emergence of complex skills in language models, 2023.
[3] Yuntao Bai, Andy Jones, et al. Training a helpful and harmless assistant with reinforcement learning from human feedback, 2022.
[4] Yoshua Bengio, Jérôme Louradour, Ronan Collobert, and Jason Weston. Curriculum learning. In Proceedings of the 26th annual international conference on machine learning, pages 41-48, 2009.
[5] Stella Biderman, Hailey Schoelkopf, Quentin Anthony, Herbie Bradley, Kyle O'Brien, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, Aviya Skowron, Lintang Sutawika, and Oskar van der Wal. Pythia: A suite for analyzing large language models across training and scaling, 2023.
[6] Sid Black, Gao Leo, Phil Wang, Connor Leahy, and Stella Biderman. GPT-Neo: Large Scale Autoregressive Language Modeling with Mesh-Tensorflow, March 2021. If you use this software, please cite it using these metadata.
[7] Rishi Bommasani, Percy Liang, et al. On the opportunities and risks of foundation models, 2021.
[8] Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. arXiv preprint arXiv:2005.14165, 2020.
[9] Mark Chen, Jerry Tworek, et al. Evaluating large language models trained on code, 2021.
[10] Wei-Lin Chiang, Zhuohan Li, Zi Lin, Ying Sheng, Zhanghao Wu, Hao Zhang, Lianmin Zheng, Siyuan Zhuang, Yonghao Zhuang, Joseph E. Gonzalez, Ion Stoica, and Eric P. Xing. Vicuna: An open-source chatbot impressing gpt-4 with  $90\%$  * chatgpt quality, March 2023.
[11] Cody Coleman, Christopher Yeh, Stephen Mussmann, Baharan Mirzasoleiman, Peter Bailis, Percy Liang, Jure Leskovec, and Matei Zaharia. Selection via proxy: Efficient data selection for deep learning, 2019.
[12] Robert M Gagne. The acquisition of knowledge. Psychological review, 69(4):355, 1962.
[68] Sang Michael Xie, Hieu Pham, Xuanyi Dong, Nan Du, Hanxiao Liu, Yifeng Lu, Percy Liang, Quoc V. Le, Tengyu Ma, and Adams Wei Yu. Doremi: Optimizing data mixtures speeds up language model pretraining, 2023.
[69] Sang Michael Xie, Shibani Santurkar, Tengyu Ma, and Percy Liang. Data selection for language models via importance resampling, 2023.
[70] Amir R. Zamir, Alexander Sax, William Shen, Leonidas Guibas, Jitendra Malik, and Silvio Savarese. Taskonomy: Disentangling task transfer learning. 2018 IEEE/CVF Conference on Computer Vision and Pattern Recognition, Jun 2018.
[71] Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, Todor Mihaylov, Myle Ott, Sam Shleifer, Kurt Shuster, Daniel Simig, Punit Singh Koura, Anjali Sridhar, Tianlu Wang, and Luke Zettlemoyer. Opt: Open pre-trained transformer language models, 2022.
[72] Yi Zhang, Arturs Backurs, Sébastien Bubeck, Ronen Eldan, Suriya Gunasekar, and Tal Wagner. Unveiling transformers with lego: a synthetic reasoning task, 2022.
A Broader Impacts and Limitations
Broader Impacts As more LMs are developed, a key criteria for their adoption and utility is if they exhibit a wide array of useful capabilities, such as generating harmless content, summarizing essays, and being conversational with the user. While improvements in other parts of the LM development pipeline such as training and architecture are important, many recent advances in building LMs with a wide array of useful capabilities have come from the data itself [18, 56, 10, 16, 43]. Our work is fundamental in investigating how LMs learn and how to select data to learn skills more efficiently. However, we recognize that data selection methods can always be utilized to optimize for particular skills that may be considered malicious or negatively target or exclude specific groups [3]. Furthermore, pre-trained LMs have been found to have various biases [26, 44, 36, 7].
**Limitations** The skills graph can either be provided (e.g., using a knowledge graph) or learned. Our work learns the skills graph using Algorithm 2 or Algorithm 3, which requires initial training runs on pairs of skills or each skill, respectively. This can be made more efficient by performing these training runs on a smaller model and for fewer number of steps, but tradeoffs here have yet to be thoroughly investigated. SKILL-IT also assumes that the ordered skill set is provided; as discussed in sections 3.1 and 3.3, it is challenging to recover ordered skill sets simply via metadata attributes or embedding clustering. Otherwise, the best way to sample over collections of skills that form a complete or empty graph is random or stratified sampling with no ordering to exploit. Our loss-based clustering approach presented in section 3.3 demonstrates that grouping by losses can provide an explanation for how skills are defined over data. An important direction for future work is to use such a clustering approach or other unsupervised algorithms in an end-to-end pipeline for skill discovery, skill graph learning, and data selection based on such skills.
Code release The code for SKILL-IT is available at https://github.com/HazyResearch/skill-it.