Spokes: Optimizing for Diverse Pretraining Data Selection
This work addresses the challenge of selecting diverse subsets for large-scale pretraining, offering a principled optimization approach that outperforms existing heuristics and quality-based methods.
The authors introduce a probabilistic diversification framework (SPOKES) that directly optimizes for diversity in pretraining data selection, achieving a +489 increase in G-Vendi score on a 500k-sample subset and improving downstream performance by up to +1.5 points over random sampling on DCLM and FineWeb benchmarks.
Diversity plays a critical role in data selection, improving performance under fixed data budgets by reducing redundancy and repetition. However, optimizing for diversity is inherently challenging, as it is a set-level property that depends on interactions between data points rather than individual examples. As a result, existing approaches typically rely on proxies or approximations, which often fail to ensure sufficiently diverse subsets. In this work, we directly optimize diversity by introducing a probabilistic diversification framework based on the G-Vendi score, optimized via exponentiated gradient descent. Our method produces subsets that are substantially more diverse than those obtained via random sampling, achieving a +489 increase in G-Vendi score on a 500k-sample subset. We evaluate our approach on FineWeb and DCLM, where it consistently outperforms existing methods. Notably, SPOKES (diversity-only) improves average downstream performance by +0.4 and +0.5 points over random sampling on DCLM and FineWeb, respectively. More importantly, jointly optimizing for both quality and diversity yields the strongest results: SPOKES achieves gains of +1.5 and +1.4 points on DCLM and FineWeb, outperforming all baselines, including semantic deduplication and quality filtering.