Research map / Large language models / Language, safety and interpretability
MATES: Model-Aware Data Selection for Efficient Pretraining with Data Influence Models
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 MATES, which continuously trains a small BERT-base data influence model on locally probed one-step oracle influence from the current pretraining model to select pretraining data on the fly, improving zero-shot downstream accuracy and cutting pretraining FLOPs by more than half.
Problem
Static data selection heuristics and reference-model methods do not track how a language model's data preferences evolve during pretraining, limiting pretraining efficiency and downstream performance.
Approach
A data influence model is periodically fine-tuned to predict oracle data influence, computed by locally probing the current pretraining model with one-step training on each hold-out example and measuring the change in loss on a reference task (LAMBADA); the model then selects data for the next pretraining stage via Gumbel-Top-k sampling.
Key results
- In the 410M setting, MATES averages 45.8 zero-shot accuracy versus 44.5 for random selection and 45.2 for QuRating (which uses GPT-3.5 ratings).
- In the 1B setting, MATES averages 47.5 versus 46.4 for random and 46.9 for QuRating, and reaches random-selection performance with only 43.3% of total FLOPs.
- Static data influence models reach Spearman correlations below 0.5, while MATES's dynamically updated model reaches around 0.7 at 40k steps.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use MATES instead of random data selection | Pythia-410M pretraining on C4 · Nine downstream tasks, average | zero-shot accuracy: 44.5 -> 45.8; 1.3% | Helped |
| Use MATES instead of random data selection | Pythia-1B pretraining on C4 · Nine downstream tasks, average | zero-shot accuracy: 46.4 -> 47.5; 1.1% | Helped |
| Use MATES instead of QuRating | Pythia-410M pretraining on C4 · Nine downstream tasks, average | zero-shot accuracy: 45.2 -> 45.8 | Helped |
| Use MATES instead of QuRating | Pythia-1B pretraining on C4 · Nine downstream tasks, average | zero-shot accuracy: 46.9 -> 47.5 | Helped |
| Select data using the locally probed oracle instead of DsDm scores | Pythia-410M · Nine downstream tasks, 40k-step decay | average zero-shot accuracy: 44.8 -> 45.6; 0.8 zero-shot accuracy gain | Helped |
| Reduce the update interval from 10k to 5k steps (ablation) | MATES on Pythia-410M · C4, 40k–50k steps | average accuracy: 45.8 -> 45.7 (zero-shot); 46.4 -> 46.6 (two-shot) | No clear effect |
| Reduce the update interval from 10k to 2.5k steps (ablation) | MATES on Pythia-410M · C4, 40k–50k steps | average accuracy: 45.8 -> 45.6 (zero-shot); 46.4 -> 46.0 (two-shot) | No clear effect |
| Increase sampling temperature from 1.0 to 2.0 (ablation) | MATES on Pythia-410M · C4, 40k–50k steps | average accuracy: 45.8 -> 45.0 (zero-shot); 46.4 -> 45.4 (two-shot) | Hurt |
| Set sampling temperature from 1.0 to 0.0 (ablation) | MATES on Pythia-410M · C4, 40k–50k steps | average accuracy: 45.8 -> 45.2 (zero-shot); 46.4 -> 45.5 (two-shot) | Hurt |
Surprising findings
- Static influence models based on a single random-pretrained checkpoint cause a notable performance drop, especially early in pretraining.
- Oracle influence is insensitive to one-step training batch size (1 to 16) as long as the model state is the same.
Limitations
- Assumes each data point contributes independently, ignoring combinational effects of batched data.
- Experiments limited to moderate scale (410M and 1B parameters, 25B tokens), so scaling to production models is untested.
- Oracle collection is still costly relative to a pure heuristic method; selection uses a single reference task (LAMBADA).
Experimental setup
Datasets: C4 LAMBADA SciQ ARC-E ARC-C LogiQA OBQA BoolQ HellaSwag PIQA WinoGrande FLAN
Benchmarks: SciQ ARC-E ARC-C LogiQA OBQA BoolQ HellaSwag PIQA WinoGrande
Metrics: zero-shot accuracy normalized accuracy total GPU FLOPs validation Spearman correlation
Baselines: Random DSIR LESS SemDeDup DsDm QuRating
Builds on: influence functions DsDm Pythia MiniCPM WSD scheduler LESS
Closest papers by meaning
- Group-Level Data Selection for Efficient Pretraining (NeurIPS 2025 · similarity 0.95)
- Harnessing Diversity for Important Data Selection in Pretraining Large Language Models (ICLR 2025 · similarity 0.92)
- Predictive Data Selection: The Data That Predicts Is the Data That Teaches (ICML 2025 · similarity 0.92)
- Efficient Data Selection at Scale via Influence Distillation (NeurIPS 2025 · similarity 0.89)
- Analyzing Similarity Metrics for Data Selection for Language Model Pretraining (NeurIPS 2025 · similarity 0.89)
- Improving Pretraining Data Using Perplexity Correlations (ICLR 2025 · similarity 0.89)
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.