Spokes: Optimizing for Diverse Pretraining Data Selection

2026-06-13

AI Summary

This work addresses the critical challenge of enhancing pretraining data diversity under a fixed data budget to minimize redundancy and improve model performance. The authors propose the first probabilistic framework that directly optimizes set-level diversity, introducing the G-Vendi score as a principled diversity metric and leveraging exponentiated gradient descent for efficient subset selection. Their approach jointly optimizes both data quality and diversity. Experiments on FineWeb and DCLM demonstrate that optimizing diversity alone yields consistent downstream performance gains of 0.4–0.5 points over random sampling, while joint optimization achieves substantial improvements of 1.4–1.5 points, significantly outperforming existing baselines.

Abstract

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.

Problem

• diversity
• data selection
• pretraining
• redundancy
• G-Vendi score

Innovation

• diversity optimization
• G-Vendi score
• exponentiated gradient descent
• pretraining data selection
• SPOKES