Research map / Learning theory and algorithms / Prediction and decision losses
TSDS: Data Selection for Task-Specific Model Finetuning
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes TSDS, an optimal-transport-based regularized probability assignment framework with KDE-based near-duplicate penalty and nearest-neighbor algorithms for task-specific data selection, evaluated on instruction tuning and continued pretraining.
Problem
Selecting training data from a massive candidate repository to finetune foundation models for a specific target task, guided by a small set of representative query examples, while ensuring distribution alignment, diversity, and robustness to near-duplicates, at scale.
Approach
Formulate data selection as a regularized optimal transport problem from query examples to candidates, with a diversity regularizer (uniform or KDE-weighted by inverse density) so near-duplicates are down-weighted; show closed-form solutions reduce to k-nearest-neighbor assignments that can use approximate nearest neighbor search.
Key results
- Instruction tuning with 1% selection by TSDS (KNN-KDE) beats LESS by about 1.5 F1 points on average across three tasks and two 7B models, and outperforms the full 270K repository on TydiQA and BBH.
- Continued pretraining with TSDS outperforms DSIR and random selection in most of the 10 settings on four classification tasks, with gains of up to about 3 F1 points.
- Preprocessing a 150M-example corpus takes 28 hours and each task-specific selection takes under 1 hour.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use TSDS-selected instruction-tuning data instead of LESS-selected data | Llama-2-7B and Mistral-7B · TydiQA and BBH, 1.0% selection | F1 score on TydiQA; accuracy on BBH: beats the state-of-the-art method (LESS) by up to 6 points | Helped |
| Use TSDS-selected instruction-tuning data instead of LESS-selected data | Llama-2-7B and Mistral-7B · MMLU | accuracy | No clear effect |
| Use TSDS-selected instruction-tuning data instead of the full data repository | Llama-2-7B and Mistral-7B · TydiQA and BBH, 1% selection | F1 score on TydiQA; accuracy on BBH | Helped |
| Use TSDS-selected instruction-tuning data instead of LESS-selected data | Llama-2-7B and Mistral-7B · Instruction tuning across three tasks | F1 score: 1.5 points in F1 score on average | Helped |
| Increase TSDS instruction-tuning selection ratio from 1% to 5% | TSDS · TydiQA | F1 score | Hurt |
| Finetune on the full repository instead of using the base model | Mistral-7B · Instruction tuning | accuracy | Hurt |
| Finetune on any selected set | Mistral-7B · MMLU | accuracy | No clear effect |
| Use TSDS-selected data instead of DSIR-selected data | ALBERT · Domain-specific continued pretraining with 1K annotated examples | F1 score: 1.92 points over DSIR | Helped |
| Use TSDS-selected data instead of DSIR-selected data | ALBERT · Domain-specific continued pretraining with 3K annotated examples | F1 score: 0.38 points with an annotated set size of 3K | Mixed |
| Use TSDS-selected data instead of DSIR-selected data | ALBERT · ChemProt (3K) and AGNews (1K) | F1 score | No clear effect |
| Use TSDS KDE-based selection with near-duplicate candidates | TSDS · Candidate repository with 1% of examples duplicated up to 1,000 times | F1 score: baseline methods show a drop of 2 points in F1 score | Helped |
Surprising findings
- Finetuning Mistral-7B on the full data repository performs worse than the base model on TydiQA and BBH.
- Finetuning on any selected set does not increase Mistral-7B MMLU accuracy, possibly because the base model is already tuned for multiple-choice questions.
Limitations
- Relies on a set of representative query examples that may be hard to craft and may carry biases amplified by selection.
- Performance on MMLU is comparable to LESS rather than clearly better; TydiQA drops when selection ratio rises from 1% to 5%, possibly from overfitting.
- Hyperparameters (C, alpha, h) need tuning via validation set or surrogate model.
Experimental setup
Datasets: Flan V2 CoT Dolly Open Assistant The Pile TydiQA MMLU BBH ChemProt IMDB SCIERC AGNews
Benchmarks: TydiQA MMLU BBH ChemProt IMDB SCIERC AGNews
Metrics: F1 score accuracy micro-F1 macro-F1
Baselines: Rand LESS DSIR Full dataset
Builds on: LESS DSIR optimal transport HNSW
Closest papers by meaning
- Get more for less: Principled Data Selection for Warming Up Fine-Tuning in LLMs (ICLR 2024 · similarity 0.88)
- Compute-Constrained Data Selection (ICLR 2025 · similarity 0.87)
- Data Mixing Optimization for Supervised Fine-Tuning of Large Language Models (ICML 2025 · similarity 0.87)
- TAIA: Large Language Models are Out-of-Distribution Data Learners (NeurIPS 2024 · similarity 0.86)
- SmallToLarge (S2L): Scalable Data Selection for Fine-tuning Large Language Models by Summarizing Training Trajectories of Small Models (NeurIPS 2024 · similarity 0.86)
- Efficient Data Selection at Scale via Influence Distillation (NeurIPS 2025 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.