Learning from the Best, Differently: A Diversity-Driven Rethinking on Data Selection

TL;DR

The paper tackles the problem that top-scoring data samples, when evaluated on a single or few dimensions, can bias pretraining data and harm downstream performance. It introduces Orthogonal Diversity-Aware Selection (ODiS), which uses 11 evaluation dimensions, applies PCA to obtain orthogonal principal components, and trains Roberta-based scorers to predict PC scores for scalable selection. Data are drawn by allocating budgets across PC dimensions with joint thresholds, ensuring both high quality and broad diversity. Empirical results on Nemotron-CC show that ODiS consistently improves downstream benchmarks over strong baselines, demonstrating that diversity-aware, orthogonally decomposed data selection yields more effective pretraining data than traditional score-based approaches.

Abstract

High-quality pre-training data is crutial for large language models, where quality captures factual reliability and semantic value, and diversity ensures broad coverage and distributional heterogeneity. Existing approaches typically rely on single or multiple-dimensional score-based selection. However, directly selecting top-scored data often degrades performance, and sampling from a broader range is required to recover results. The above non-monotonicity between dataset scores and downstream benchmark results reveals a fundamental bias: score-based methods collapse correlated dimensions, causing top-scored data to appear high-quality while systematically overlooking diversity. We argue that ensuring diversity requires decomposing correlated metrics into orthogonal feature dimensions, from which the top-scored data can be directly selected. Therefore, we proposed the Orthogonal Diversity-Aware Selection (ODiS) algorithm, which preserves both quality and diversity during data selection. First, ODiS evaluates data from multiple dimensions, covering language quality, knowledge quality, and comprehension difficulty. The multi-dimensional scores are then decorrelated via Principal Component Analysis (PCA), yielding orthogonal evaluation dimensions. For each dimension, a Roberta-based scorer is trained to regress the data onto PCA-projected scores, enabling scalable inference on large corpora. Finally, ODiS constructs the training dataset by selecting top-scored data within each orthogonal dimension, thereby ensuring both quality and diversity. Empirical results show that ODiS-selected data exhibit less than 2\% inter-dimension overlap, confirming orthogonality between dimensions. More importantly, models trained with ODiS-selected data significantly outperform other baselines on downstream benchmarks, highlighting the necessity of orthogonal, diversity-aware data selection for LLMs.

Figures (8)

• Figure 1: Performances across downstream tasks.
• Figure 2: Visualization of data distribution from different data pool sizes. The embeddings are generated using 1000 randomly sampled texts with the m3e-base model, followed by L2 normalization and removal of the top 3 principal components to suppress dominant common directions.
• Figure 3: Model performance under different selection methods and ranges.
• Figure 4: Data diversity visualization from different selection methods.
• Figure 5: Correlation analysis of evaluation dimensions. The correlations are calculated using 460k samples from the Fineweb-Edu dataset. Pairwise correlation coefficients are visualized as a heatmap.