Leverage score sampling improves language model training data diversity
Towards Scalable Data Diversification for Language Model Pretraining via Leverage Score Sampling
Artificial Intelligence
Summary
Choosing good and varied text data for training language models is tricky because focusing only on high-quality text can miss out on important topics. The authors introduce a new way called Leverage Score Sampling that picks data to keep it diverse without slowing things down. This method works faster and creates a more balanced dataset, improving performance on many language tasks and even helps when it's hard to decide what good data means, like for computer code. They also found that some common quality filters miss important kinds of data, showing why diversity matters in data selection.
What this means in practice
- •For language model engineers: Select more varied and representative datasets for pretraining to improve model accuracy across multiple tasks and domains.
- •For software development teams: Curate diverse code training data automatically to boost code generation models’ performance when quality criteria are unclear.
Authors
Zailin Ma, Quzhe Huang, Yujun Li, Congyuan Rao, Yaodong Yang
Abstract
Data selection for language model pretraining faces a fundamental tension between quality and diversity. While quality filtering is empirically effective, it often induces diversity collapse: by favoring texts similar to high-quality reference corpora (e.g., educational or QA-style data), it systematically excludes valuable data from underrepresented domains. In contrast, diversified selection preserves domain balance and encourages robust downstream performance, yet existing methods either focus on coverage-oriented objectives that indirectly enhance diversity, or directly optimize for diversity via costly covariance matrix recomputation that limits scalability. To address these issues, we introduce \textbf{Leverage Score Sampling (Lev)}, which iteratively selects samples that maximally expand the determinantal volume of the embedded data via leverage scores, a computationally efficient criterion that eliminates matrix recomputation and enables scalable selection. Empirically, Lev delivers up to $72\times$ speedup and improves dataset diversity, measured by the Vendi score, by $9.2\%$ over the strong diversification baseline \textbf{DiSF}. On CommonCrawl (CC) web data selection, Lev improves accuracy across seven downstream tasks by up to $1.31\%$ over existing baselines. For domains where robust quality criteria are inherently difficult to define (e.g., code), Lev serves as an effective unsupervised curation alternative: on StarCoderData, the selected subset reduces bits-per-byte by $3.08\%$ over DiSF. Notably, we uncover a cross-domain collapse of quality filtering: CC data filtered by DCLM-fastText fail to retain sufficient code-related content, yielding inferior code performance relative to Lev-selected data. These findings advocate for integrating diversity-aware practices into quality filtering for more effective data curation in language model pretraining.