# DRIFT: Data Selection for LLM Instruction Tuning via On-Policy Attribution

arXiv: https://arxiv.org/abs/2606.18307 · HTML: https://mlsanity.com/paper/2606.18307/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 48 of its 55 references that are in the map.
- First posted on arXiv 16 Jun 2026.

## Builds on (most cited first, 30 of 48)

- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [Measuring Mathematical Problem Solving With the MATH Dataset](https://mlsanity.com/paper/2103.03874/) (2021)
- [Finetuned Language Models Are Zero-Shot Learners](https://mlsanity.com/paper/2109.01652/) (2021)
- [GPQA: A Graduate-Level Google-Proof Q&A Benchmark](https://mlsanity.com/paper/2311.12022/) (2023)
- [LiveCodeBench: Holistic and Contamination Free Evaluation of Large Language Models for Code](https://mlsanity.com/paper/2403.07974/) (2024)
- [MMLU-Pro: A More Robust and Challenging Multi-Task Language Understanding Benchmark](https://mlsanity.com/paper/2406.01574/) (2024)
- [Challenging BIG-Bench Tasks and Whether Chain-of-Thought Can Solve Them](https://mlsanity.com/paper/2210.09261/) (2022)
- [Understanding Black-box Predictions via Influence Functions](https://mlsanity.com/paper/1703.04730/) (2017)
- [Is Your Code Generated by ChatGPT Really Correct? Rigorous Evaluation of Large Language Models for Code Generation](https://mlsanity.com/paper/2305.01210/) (2023)
- [OlympiadBench: A Challenging Benchmark for Promoting AGI with Olympiad-Level Bilingual Multimodal Scientific Problems](https://mlsanity.com/paper/2402.14008/) (2024)
- [On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes](https://mlsanity.com/paper/2306.13649/) (2023)
- [Estimating Training Data Influence by Tracing Gradient Descent](https://mlsanity.com/paper/2002.08484/) (2020)
- [CCNet: Extracting High Quality Monolingual Datasets from Web Crawl Data](https://mlsanity.com/paper/1911.00359/) (2019)
- [LESS: Selecting Influential Data for Targeted Instruction Tuning](https://mlsanity.com/paper/2402.04333/) (2024)
- [Policy Distillation](https://mlsanity.com/paper/1511.06295/) (2015)
- [DoReMi: Optimizing Data Mixtures Speeds Up Language Model Pretraining](https://mlsanity.com/paper/2305.10429/) (2023)
- [Beyond Human Data: Scaling Self-Training for Problem-Solving with Language Models](https://mlsanity.com/paper/2312.06585/) (2023)
- [Data Selection for Language Models via Importance Resampling](https://mlsanity.com/paper/2302.03169/) (2023)
- [Studying Large Language Model Generalization with Influence Functions](https://mlsanity.com/paper/2308.03296/) (2023)
- [TRAK: Attributing Model Behavior at Scale](https://mlsanity.com/paper/2303.14186/) (2023)
- [GRAD-MATCH: Gradient Matching based Data Subset Selection for Efficient Deep Model Training](https://mlsanity.com/paper/2103.00123/) (2021)
- [Prioritized Training on Points that are Learnable, Worth Learning, and Not Yet Learnt](https://mlsanity.com/paper/2206.07137/) (2022)
- [GLISTER: Generalization based Data Subset Selection for Efficient and Robust Learning](https://mlsanity.com/paper/2012.10630/) (2020)
- [Datamodels: Predicting Predictions from Training Data](https://mlsanity.com/paper/2202.00622/) (2022)
- [Data Mixing Laws: Optimizing Data Mixtures by Predicting Language Modeling Performance](https://mlsanity.com/paper/2403.16952/) (2024)
- [RegMix: Data Mixture as Regression for Language Model Pre-training](https://mlsanity.com/paper/2407.01492/) (2024)
- [QuRating: Selecting High-Quality Data for Training Language Models](https://mlsanity.com/paper/2402.09739/) (2024)
- [A Minimalist Approach to LLM Reasoning: from Rejection Sampling to Reinforce](https://mlsanity.com/paper/2504.11343/) (2025)
- [If Influence Functions are the Answer, Then What is the Question?](https://mlsanity.com/paper/2209.05364/) (2022)
- [Superfiltering: Weak-to-Strong Data Filtering for Fast Instruction-Tuning](https://mlsanity.com/paper/2402.00530/) (2024)

## Built on by (most cited first, 1 of 1)

- [Every Batch Is Its Own Validation Set: Leave-One-Out Gradient Matching for Online Data Selection in LLM Fine-Tuning](https://mlsanity.com/paper/2610.00436/) (2026)

## Abstract (arXiv)

Data selection for instruction tuning is usually studied for efficiency, reaching strong performance with a fraction of the training data. We ask whether data selection can instead raise the performance ceiling, producing a better model than training on all the available data. To test this, we take a large language model (LLM) that has already completed supervised fine-tuning (SFT) on its full corpus and train it further on examples selected from that same corpus. Any gain therefore lifts the model above the ceiling set by full-data training. Existing selection methods yield little or no gain in this setting. We propose DRIFT (Data Refinement via On-Policy Influence Functions for Supervised Fine-Tuning), a data attribution method built on influence functions. DRIFT samples the model's own responses to validation queries and weights them by correctness to form a validation objective. It ranks the examples of the original corpus by their estimated influence on this objective, and the model then continues SFT on the top-ranked subset. This design is motivated by the proximity gap in influence function theory: we hypothesize that these on-policy responses are better validation targets than external references in terms of locality. DRIFT also corrects influence scores for their dependence on gradient norm within each validation task. DRIFT outperforms all evaluated data selection baselines on Olmo3-7B-Instruct-SFT and OpenR1-Distill-7B, raising average accuracy across eight benchmarks from 39.17 to 40.14 and from 57.23 to 58.66, respectively, with gains on benchmarks not used for attribution.