Source-linked AI summary

Skill-it! A Data-Driven Skills Framework for Understanding and Training Language Models

Mayee F. Chen, Nicholas Roberts, Kush Bhatia, Jue Wang, Ce Zhang, Frederic Sala, Christopher Ré

arXiv:2307.14430v1cs.CLcs.LG

TL;DR

Given a fixed token budget, the paper asks how to select training data to improve LM capabilities when current selection methods lack a formal account of how data relates to skills. It defines prerequisite-linked skills and proposes Skill-It, reporting gains across training settings, including lower Natural Instructions target-skill loss than training on that skill alone.

  • Problem

    With a fixed token budget, it remains unclear how to select corpus data for desired LM capabilities, and existing methods mostly rely on heuristics rather than a formal account of how data shapes capabilities.

    Source: §1 · p. 1
  • Method

    The paper defines skills by associated data and prerequisite relations, then proposes Skill-It, an online algorithm that selects skill mixtures to reduce evaluation-skill loss.

    Source: §1 · p. 2
  • Results

    Skill-It reports better outcomes across continual pre-training, fine-tuning, and RedPajama; on Natural Instructions, it achieves up to 13.6% lower target-skill loss than training only on that skill.

    Source: §1 · pp. 2–3
  • Takeaways & Limitations

    Ordered skill sets can support learning skills with less data when training includes their prerequisites.

    Source: §1 · p. 2

    SKILL-IT assumes the ordered skill set is provided, while recovering it from metadata attributes or embedding clusters is challenging.

    Source: A Broader Impacts and Limitations · p. 16

Abstract

from arXiv · show

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 36.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 recent 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.

1 Introduction

The introduction frames data selection as a fixed-token-budget problem and proposes that prerequisite relationships among skills can guide more efficient LM training. It introduces a skill-based framework and Skill-It, then reports results across synthetic, instruction-following, and RedPajama settings.

1 Introduction

  • Partitioning Alpaca by instruction type produced performance similar to random sampling, suggesting that existing semantic groupings do not necessarily capture learnable skills.Source: §1 · p. 2
  • The paper identifies a gap in principled fixed-budget data selection: random sampling ignores uneven skill prevalence and prerequisite ordering, while curriculum learning orders samples rather than skills.For example, Spanish QG is 0.2% of Natural Instructions, compared with 5% Spanish and 4% QG.Source: §1 · pp. 1–2
  • The framework defines skills through associated data and ordered skill sets through prerequisite edges that reduce the data needed to learn a skill.The paper reports that such ordered skill sets exist in synthetic and real datasets.Source: §1 · p. 2
  • Skill-It selects mixtures online, weighting skills that remain unlearned or influential prerequisites to optimize evaluation-skill loss under a fixed data budget.The paper contrasts it with static skill-stratified sampling, which can oversample skills already learned early in training.Source: §1 · p. 2
  • On LEGO, Skill-It improves accuracy by 35.8 points over random sampling and curriculum learning; on Natural Instructions, it achieves up to 13.6% lower target-skill loss.The Natural Instructions comparison is against training solely on the target skill with the same overall training budget.Source: §1 · p. 2
  • On RedPajama, Skill-It achieves higher accuracy with 1B tokens than uniform sampling over data sources with 3B tokens.Source: §1 · p. 3

2 Skills framework

The framework defines skills through associated data and models their learning dependencies as ordered skill graphs. Examples across synthetic and real datasets show that prerequisite mixtures can improve learning efficiency, while loss-based clustering can recover LEGO skills better than embedding-based clustering.

2.1 Definitions

  • A skill is a data-associated behavior learned through improved performance on held-out skill data; an ordered skill set links prerequisite skills to skills whose data requirements can thereby be reduced.Learning is measured with cross-entropy validation loss on next-token predictions.Source: §2.1 · p. 3
  • Alpaca instruction types yield a 97.4%-dense graph and only 0.007 points average per-skill validation-loss improvement over random sampling; Pile of Law sources yield 3.9% density, while Natural Instructions task categories yield 42.7%.These cases illustrate that semantic groupings may produce complete-like, empty-like, or intermediate graphs; not every grouping induces an ordered skill set.Source: §2.1 · pp. 3–4

2.2 Examples of skills and ordered skill sets

  • 35.9% fewer training steps reached validation loss 0.01 on LEGO when training on a balanced mixture of skills s1, s2, and s3 rather than s3 alone.The skills correspond to knowing variables at different positions in a reasoning chain.Source: §2.2 · p. 4
  • 32% fewer training steps reached validation loss 0.01 on the ones-digit addition skill when training on a balanced mixture of ones- and tens-digit skills rather than ones-digit data alone.The reported benefit occurs even though the ones digit should not require information from the tens digit.Source: §2.2 · p. 5
  • 4.1% lower validation loss on Spanish question generation came from training with equal parts English question generation, Spanish question answering, and English question answering, using 25% of the Spanish question-generation data.The examples suggest signal in real data even though task-category skill labels may be noisy.Source: §2.2 · p. 5
  • Training on Text Matching with Stance Detection decreased Stance Detection loss by 11%, suggesting an edge between categories that both involve relations between two texts.The passage presents this as a Natural Instructions example of learning a skill with another category.Source: §2.2 · p. 5

2.3 Skill recovery

  • On LEGO, loss-based clustering over ten random training runs recovered skills with 61% accuracy, while trained and pretrained embedding approaches did not exceed 39%.The method clusters data by its loss per timestep across runs, based on similar loss trajectories for points from the same skill.Source: §2.3 · p. 5

3 Skills-based data selection

The section frames data selection as choosing a mixture of training skills to meet different evaluation goals, using a skills graph to guide sampling. It introduces static skill-stratified sampling and SKILL-IT, which adapts sampling as learning progresses.

3.1 Problem statement

  • The sampling objective selects a mixture of training skills to learn all skills in continual pre-training, a target skill or subset in fine-tuning, or disjoint evaluation skills out of domain.The skills graph is a weighted adjacency matrix whose entries represent the strength of edges from training skills to evaluation skills.Source: §3.1 · p. 5

3.2 Skills graph learning

  • The skills graph can be learned by testing skill pairs, or approximated by checking whether training on each skill reduces evaluation loss.The pairwise approach has runtime O(Hk^2); the linear approximation has runtime O(hk) and is necessary when training and evaluation skills are disjoint.Source: §3.2 · pp. 5–6

3.3 Skills graph-aware sampling

  • The framework offers graph-aware sampling through uniform sampling over relevant skills or SKILL-IT’s online adaptation to how skills are being learned.The skills graph determines relevance; SKILL-IT additionally incorporates learning progress during training.Source: §3.3 · p. 6

3.3.1 Skill-stratified sampling

  • Skill-stratified sampling uniformly distributes samples across relevant skills: all training skills for continual pre-training, target and prerequisite skills for fine-tuning, and prerequisite skills for out-of-domain evaluation.This static approach discards training skills that do not benefit the evaluation skills.Source: §3.3.1 · p. 6

3.3.2 SKILL-IT online data selection algorithm

  • SKILL-IT addresses the static sampler’s continued allocation to already-learned skills by prioritizing prerequisite skills and skills not yet learned.It adjusts a skill’s weight based on the losses of skills it influences; the update rule is derived via online mirror descent.Source: §3.3.2 · pp. 6–7

4 Experimental results

The experiments evaluate SKILL-IT for continual pre-training, fine-tuning, out-of-domain learning, and RedPajama data mixing. Across these settings, the reported results favor skill-aware or graph-based sampling over the stated baselines.

4 Experimental results

  • The experiments measure per-skill validation loss after fixed training steps on imbalanced synthetic datasets and use 23 Natural Instructions task categories as skills.The comparisons include random, curriculum, and anticurriculum sampling, plus skill-stratified sampling that uses skill information without online adaptation.Source: §4.1 · pp. 7–8

4.1 Continual pre-training

  • In continual pre-training, SKILL-IT improves LEGO accuracy by 9.9–25.9 points halfway through training and reaches 99.4 final accuracy; on addition, it reaches near-zero loss faster on skills 1 and 2.On LEGO, skill-stratified sampling also reduces loss on all skills but does so more slowly; on Natural Instructions, SKILL-IT validation loss is 3.2% lower than random sampling.Source: §4.1 · p. 8, Figure 4 · p. 8

4.2 Fine-tuning

  • On Natural Instructions, SKILL-IT lowers validation loss by 5.3% for Spanish question generation and 13.6% for Stance Detection versus training on each target skill alone.Identifying prerequisite skills accounts for a significant portion of the gain, while dynamic prerequisite weighting improves performance further.Source: §4.2 · pp. 8–9, Figure 5 · p. 8

4.3 Out-of-domain setting

  • On Natural Instructions test tasks, SKILL-IT achieves the lowest loss on 11 of 12 categories compared with random and skill-stratified sampling.The method uses the graph between training and evaluation skills to generate an online mixture on the training data.Source: §4.3 · p. 9, Figure 6 · p. 9
  • With 1B additional tokens, SKILL-IT achieves higher LM Evaluation Harness accuracy than uniform sampling with 3B tokens on a 3B-parameter RedPajama model.The mixture treats RedPajama sources as training skills and evaluation tasks as evaluation skills.Source: §4.3 · p. 9, Figure 7 · p. 9

5 Related work

Related work spans data selection, task selection, learning hierarchies, and theories of how language models learn. The paper distinguishes its focus on skill-level data selection and skill relationships from approaches targeting samples, data sources, or tasks.

Data selection for LMs

  • Large-scale LM data selection includes deduplication, quality filtering [7], importance reweighting, and influence-function selection.The paper contrasts these approaches with studying how data choices relate to the skills LMs learn.Source: §5 · pp. 9–10

Data selection methods

  • DoReMi selects data-source mixtures for worst-case performance across sources, while this work models skill-graph structure for varied target skills.The paper also contrasts skill-level mixture selection with supervised methods such as coresets and sample-level gradient or uncertainty selection.Source: §5 · p. 10
  • Domain-adaptation methods match target distributions using cross-entropy, source-target discrimination, or embedding divergence; the authors report that embedding heuristics do not fully identify skills.Data-attribution work explains individual samples or dataset counterfactuals [28, 16, 23, 20], while this paper studies groups of data associated with skills.Source: §5 · p. 10

Curriculum learning

  • Curriculum learning orders samples by difficulty, while education research studies prerequisite skill hierarchies and LM-learning accounts propose quanta, latent-variable clusters, or topic learning.The paper frames its contribution as studying how LMs learn skills and using that understanding for data selection.Source: §5 · pp. 10–11

Task selection

  • Task-selection work uses feature similarity, learned task hypergraphs in Taskonomy, or task relationships in TaskWeb, whereas this paper models skill graphs for dynamic LM data sampling.TaskWeb compares adding data for one task while this paper fixes the total data budget across skills.Source: §5 · p. 11

6 Conclusion

The conclusion presents a skills-based framework for understanding how LMs learn and selecting training data. It frames the work as an initial step toward aligning skills with data under fixed data budgets.

6 Conclusion

  • Selecting data to induce LM capabilities remains challenging when the data budget is fixed.Source: §6 · p. 11
  • The paper introduces a skills-based framework for understanding LM learning and selecting training data.Source: §6 · p. 11
  • The authors invite further work on the notion of skill and on aligning skills with data.Source: §6 · p. 11

A Broader Impacts and Limitations

The broader-impacts discussion notes that data selection can serve harmful or exclusionary goals and that pre-trained LMs have documented biases. The limitations include the assumption of an ordered skill set and the cost and difficulty of discovering one.

Broader Impacts

  • Data selection could be used to optimize for malicious skills or to target or exclude specific groups.The authors also note that pre-trained LMs have been found to exhibit biases.Source: A Broader Impacts and Limitations · p. 16

Limitations

  • SKILL-IT assumes an ordered skill set, which can be provided or learned; recovering it from metadata attributes or embedding clustering is challenging.Learning the graph requires initial training runs, and tradeoffs from using smaller models or fewer steps have not been thoroughly investigated.Source: A Broader Impacts and Limitations · p. 16
  • The authors identify end-to-end unsupervised skill discovery, graph learning, and data selection as future work.Source: A Broader Impacts and Limitations · p. 16

B.1 Derivation of SKILL-IT Update Rule

SKILL-IT’s update is derived by approximating the next-round average per-skill loss and regularizing the resulting sampling choice. The graph-learning methods estimate which training skills influence evaluation skills, with different approaches for in-domain and out-of-domain settings.

B.1 Derivation of SKILL-IT Update Rule

  • The objective approximates average per-skill evaluation loss after the model is trained on the mixture chosen for the current round.The derivation uses a linear approximation of that loss based on the model and mixtures from previous rounds.Source: §B.1 · p. 16
  • Bregman-divergence regularization stabilizes the sampling update after the unregularized approximation favors the skill with the largest gradient.The derivation drops terms independent of the sampling distribution; the step size does not affect the unregularized solution.Source: §B.1 · p. 16

B.2 Graph Learning Method

  • Brute-force graph learning applies when evaluation skills are a subset of training skills, while the linear method also supports disjoint out-of-domain sets.Edge weights can be based on differences in loss changes for the brute-force method or on loss change itself for the approximate method.Source: §B.2 · pp. 17–18

C.1 Datasets

These appendices describe the datasets, graph-learning choices, training settings, and additional LEGO examples used to evaluate SKILL-IT. The examples show that measured skill influence need not match the synthetic reasoning graph.

C.1 Datasets

  • LEGO and addition use synthetic skill sets with five and three skills, respectively, while Natural Instructions experiments select task categories and tasks for pre-training, fine-tuning, and out-of-domain evaluation.The Natural Instructions pre-training selection uses 23 categories and 425 tasks; its out-of-domain test uses 12 categories and 119 tasks.Source: §C.1 · p. 18
  • The datasets also include RedPajama’s seven data sources, Alpaca instructions grouped into 38 parsed skills, and a Pile of Law subset drawn from legal and administrative sources.The RedPajama dataset is described as a one-trillion-token dataset intended to reproduce the LLaMA training dataset.Source: §C.1 · p. 18

C.2 Graph Learning Details

  • Graph-learning configurations vary by dataset: synthetic graphs use thresholded training-step comparisons, while real-data graphs use changes in loss or perplexity.For RedPajama, the linear method trains for one billion tokens per source and uses validation perplexity change as edge weight.Source: §C.2 · p. 19

C.3 Training Details

SKILL-IT pre-training

  • For continual pre-training, the appendix gives separate configurations for LEGO, addition, and Natural Instructions, alongside imbalanced random baselines and curriculum-learning pacing details.The LEGO random baseline uses skill proportions 1:1:1:3:5.Source: SKILL-IT pre-training · p. 24

SKILL-IT fine-tuning

  • Fine-tuning configurations cover LEGO, addition, Spanish question generation, and stance detection, while out-of-domain configurations include Natural Instructions and RedPajama.The appendix also reports context-window, batch-size, and optimizer settings for the experiments.Source: SKILL-IT fine-tuning · p. 24, SKILL-IT out-of-domain · p. 24

D.1 Additional examples of LEGO ordered skill sets

  • Training skills 2 and 4 together lower skill 4 loss, unlike skills 3 and 4, despite their adjacent positions in the LEGO reasoning chain.In a tree-shaped variant, skills 2 and 3 improve each other despite lacking a direct edge, while adding either to skill 4 gives little benefit.Source: §D.1 · p. 25

D.2 Unsupervised skill recovery

The appendix reports loss-based clustering as the strongest tested approach for recovering LEGO skills and provides per-skill breakdowns for several experiments. Those breakdowns include an average Natural Instructions pre-training comparison against random sampling.

D.2 Unsupervised skill recovery

  • 38.4 points was the highest accuracy among the four embedding-based clustering methods, achieved with trained-model last-token embeddings.The comparisons include pre-trained-model embeddings and Sentence-BERT embeddings.Source: §D.2 · pp. 25–26
  • Across 10 runs of 6,000 steps with randomly sampled skill mixtures, k-means on validation-sample loss vectors from 120 checkpoints achieved 61.0 points accuracy, significantly above the 38.4-point second-best result.Source: §D.2 · p. 26

D.3.1 Per-skill performance

  • SKILL-IT outperforms random sampling by 3.2% on average across skills in Natural Instructions continual pre-training.The appendix also provides per-skill validation-loss breakdowns for LEGO, addition, Natural Instructions, and out-of-domain evaluations.Source: §D.3.1 · p. 26, Table 4 · p. 27, Table 6 · p. 27

D.3.2 Weight trajectories

The appendix traces how SKILL-IT’s sampling weights change as skill losses evolve across pre-training, fine-tuning, and out-of-domain experiments. Weights can shift toward harder or higher-loss skills and converge toward uniform or fixed values as losses saturate.

D.3.2 Weight trajectories

  • LEGO pre-training weights shift from initially favoring skill 2 toward harder skills 3, 4, and 5, then become roughly uniform as losses approach zero.The weight on quickly learned skill 1 falls below 0.1 at 1,000 steps.Source: §D.3.2 · p. 29, Figure 20 · p. 30
  • Addition pre-training initially favors skill 2, assigns little weight to quickly learned skill 3, then shifts toward hardest skill 1 before approaching uniform weights.The described skill relationships include an edge from skill 2 to skill 1.Source: §D.3.2 · p. 29
  • In the fine-tuning experiments, the mixture weights approach uniform as the single target skill’s loss approaches zero.The authors report that preliminary attempts to decay edge weights and regularization did not appear better.Source: §D.3.2 · pp. 29–30
  • For Natural Instructions out-of-domain evaluation, initial graph-based weights roughly order the categories, later weight differences widen, and weights settle as validation losses saturate.The figure legend shows the ten task categories with the largest weights.Source: §D.3.2 · p. 30

D.4 Experiments on 1.3B parameter model

The 125M-learned skills graph transfers to 1.3B-parameter continual pre-training, where Skill-It outperforms random and skill-stratified sampling on average for LEGO and Natural Instructions. On LEGO, its weights shift from prerequisite skill 2 toward harder skills 4 and 5 before converging to uniform as losses approach zero.

D.4 Experiments on 1.3B parameter model

  • The experiments use the GPT-Neo-1.3B model for continual pre-training on LEGO synthetic and Natural Instructions, with skill graphs learned on the 125M model.Results average three random seeds; LEGO training runs 1,500 steps and Natural Instructions runs 5,000 steps.Source: §D.4 · p. 30
  • With the 1.3B model, Skill-It outperforms random and skill-stratified sampling on average for LEGO and Natural Instructions using the 125M-learned graph.For LEGO, performance differences are larger on skill 5 than on early skills.Source: §D.4 · p. 30
  • On LEGO, Skill-It initially weights skill 2, then skills 4 and 5, and converges to uniform sampling as losses approach zero.The weight trajectories resemble those observed with the 125M-parameter model.Source: §D.4 · p. 30

D.5 Ablations

Ablations indicate that both the skills graph and online loss-based weighting matter most for LEGO continual pre-training, while effects vary across Addition and fine-tuning tasks. Skill-It also achieves the lowest average loss among the compared approaches in the Natural Instructions out-of-domain experiment.

D.5 Ablations

  • “No graph” replaces the skills adjacency matrix with the identity, while “Static” sets T = 1 and uses a softmax over the graph for fixed weights.The complete-graph case corresponds to skill-stratified sampling; Natural Instructions continual pre-training is omitted from these ablations because T = 1 is already used there.Source: §D.5 · pp. 30, 33
  • For LEGO pre-training, Skill-It shifts weight toward skill 2 initially and later toward harder skills; for Addition it favors skill 2 before weighting harder skill 1.In both experiments, the weights eventually converge to uniform as losses approach zero.Source: Figure 21 · p. 31, Figure 25 · p. 32
  • For LEGO continual pre-training, removing the graph or online dynamics raises loss across skills relative to Skill-It.Without the graph, harder skills receive weight too early; static selection continues weighting prerequisites after their losses approach zero.Source: §D.5 · p. 33, Figure 26 · p. 34, Figure 27 · p. 34
  • On Addition continual pre-training, graph-free and static baselines perform similarly to Skill-It on average, though Skill-It is best on skill 2 versus vanilla multiplicative weights and skill 1 versus static selection.The authors describe Addition as somewhat easier than their other datasets.Source: §D.5 · p. 33, Figure 28 · p. 35, Figure 29 · p. 35
  • In LEGO fine-tuning, Skill-It reaches validation loss 0.05 in 2,000 steps, versus 2,050–2,200 steps with no graph.Skill-It and static selection behave similarly; Skill-It is slightly better than no graph.Source: §D.5 · p. 33, Figure 30 · p. 35
  • For Natural Instructions Spanish question generation and stance detection fine-tuning, Skill-It attains lower loss than no-graph and static-selection alternatives.These comparisons use validation loss.Source: §D.5 · p. 33, Figure 31 · p. 36, Figure 32 · p. 36
  • In Natural Instructions out-of-domain evaluation, Skill-It has the lowest validation loss on 7 of 12 skills and average loss 2.540 versus 2.541–2.551 for static selection.The comparison is between Skill-It and static data selection.Source: §D.5 · pp. 33, 37, Figure 33 · p. 37
Loading 2307.14430v1…