Research map / Large language models / Language, safety and interpretability
Harnessing Diversity for Important Data Selection in Pretraining Large Language 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 Quad, a pretraining data selection method that clusters SlimPajama, uses a multi-armed bandit with UCB cluster scores to balance influence-based quality and diversity, and computes influence on attention and MLP layers via a Kronecker-factored Hessian approximation.
Problem
Selecting high-quality and diverse pretraining data for LLMs is important, but exact influence-based selection is computationally expensive over huge corpora and top-k influence selection yields homogeneous data that hurts generalization.
Approach
Cluster the candidate pool into groups, treat each cluster as an arm of a UCB-style multi-armed bandit where the reward is sampled influence, and extend K-FAC-style influence estimation to attention (QKV and output) layers so influence captures more semantics, enabling cheap and diverse selection.
Key results
- Quad achieves overall average 42.94 vs 41.55 for Random (+1.39) and 42.01-42.04 for the best baseline Qurating, using 9.15e19 FLOPs vs 13.66e19 for Qurating.
- Quad improves General Knowledge by 1.75 and Reading Comprehension by 1.98 points over Random.
- Selecting top-500 clusters by influence alone gives overall 42.48, while the bandit method gives 42.94, supporting the quality-diversity trade-off.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use MAB-selected clusters rather than selecting only top-scoring clusters (ablation) | Quad data selection · SlimPajama downstream tasks (Overall) | accuracy: 42.48 -> 42.94 | Helped |
| Set influence threshold τ to 0.0015, 0.0020, or 0.0030 instead of 0.0025 (ablation) | Quad · SlimPajama downstream tasks (Overall) | accuracy: 42.94 -> 42.16; 42.94 -> 42.45; 42.94 -> 42.10 | Hurt |
| Set MAB diversity coefficient α to 0.0015, 0.0025, or 0.0030 instead of 0.0020 (ablation) | Quad · SlimPajama downstream tasks (Overall) | accuracy: 42.94 -> 42.38; 42.94 -> 42.34; 42.94 -> 42.07 | Hurt |
| Omit the Hessian in influence calculation (ablation) | Quad influence calculation · SlimPajama downstream tasks (Overall) | accuracy: 42.06 vs 42.48 | Hurt |
| Calculate influence on MLP layers only, omitting attention layers (ablation) | Quad influence calculation · SlimPajama downstream tasks (Overall) | accuracy: 42.21 vs 42.48 | Hurt |
| Use MLP and attention layers for influence calculation rather than MLP-only influence (ablation) | Quad influence calculation · SlimPajama downstream tasks (Overall) | accuracy: 42.21 -> 42.48 | Helped |
| Select from top-150, top-500, or top-1000 influence-scored clusters instead of using MAB cluster selection (ablation) | Quad data selection · SlimPajama downstream tasks (Overall) | accuracy: 42.04; 42.48; 41.67 vs 42.94 | Mixed |
| Use Quad data selection instead of Random | 1.3B decoder-only Transformer pretrained on SlimPajama · SlimPajama downstream tasks (Overall) | accuracy: +1.39 | Helped |
| Use Quad data selection instead of Qurating | 1.3B decoder-only Transformer pretrained on SlimPajama · SlimPajama downstream tasks (Overall) | accuracy: 42.01 -> 42.94 | Helped |
| Use Quad data selection instead of Qurating | 1.3B decoder-only Transformer pretrained on SlimPajama · SlimPajama data selection | FLOPs: 13.66 -> 9.15 | Helped |
Surprising findings
- Top-k selection by influence does not always beat random: top-1000 clusters gives 41.67, below the 42.48 of top-500 and close to Random's 41.55.
Limitations
- Evaluation limited to a 1.3B model and 30B selected tokens, with LAMBADA as the sole reference set.
- Performance differences among methods are small (often under 1-2 points) and results come from a single run as reported.
- Hyperparameters (alpha, tau, sample ratio) were tuned on the same downstream evaluation tables.
Experimental setup
Datasets: SlimPajama LAMBADA
Benchmarks: ARC-C ARC-E SciQ HellaSwag SIQA WinoGrande LogiQA OpenbookQA BoolQ
Metrics: zero-shot accuracy overall average accuracy FLOPs
Baselines: Random DSIR PPL SemDedup Qurating MATES
Builds on: LESS MATES K-FAC DSIR Qurating
Closest papers by meaning
- MATES: Model-Aware Data Selection for Efficient Pretraining with Data Influence Models (NeurIPS 2024 · similarity 0.92)
- Group-Level Data Selection for Efficient Pretraining (NeurIPS 2025 · similarity 0.91)
- Towards Scalable Data Diversification for Language Model Pretraining via Leverage Score Sampling (NeurIPS 2026 · similarity 0.91)
- Predictive Data Selection: The Data That Predicts Is the Data That Teaches (ICML 2025 · similarity 0.91)
- The Best of Both Worlds: Bridging Quality and Diversity in Data Selection with Bipartite Graph (ICML 2025 · similarity 0.90)
- Diversity as a Reward: Fine-Tuning LLMs on a Mixture of Domain-Undetermined Data (NeurIPS 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.