Spokes: Optimizing for Diverse Pretraining Data Selection
Abstract Overview
This paper introduces Spokes, a scalable method for selecting diverse pretraining subsets by directly optimizing the G-Vendi diversity score in gradient space. The method relaxes fixed-size subset selection into a probabilistic weighting problem over examples and solves it with exponentiated gradient descent, while allowing an explicit trade-off between diversity and per-example quality. To make the approach practical at large scale, the authors use truncated gradients from the last model layers and Johnson–Lindenstrauss random projections to reduce representation cost. Experiments on DCLM and FineWeb evaluate both diversity-only and quality-weighted variants against random sampling, semantic deduplication, and quality filtering.
Novelty
The main novelty is treating gradient-space G-Vendi as a direct optimization objective for pretraining data selection rather than only as a post-hoc diversity metric or relying on clustering-based proxies. The paper also combines this objective with scalable approximations—layer truncation, low-dimensional random projection, and exponentiated-gradient optimization—to operate on large web-scale corpora.
Results
Spokes produces much more diverse subsets than random sampling, with a reported G-Vendi increase from 315 to 520 at the batch level and a +489 increase on a 500k-sample subset. In downstream pretraining, diversity-only Spokes improves average benchmark scores over random sampling on both DCLM (52.5 to 52.9) and FineWeb (45.1 to 45.6). Joint optimization of quality and diversity performs best overall, reaching 54.0 on DCLM and 46.5 on FineWeb, ahead of semantic deduplication and quality-only filtering under matched budgets.
Key Points
- Spokes directly optimizes subset diversity in gradient space using a relaxed weighted objective based on G-Vendi and exponentiated gradient descent.
- The method is designed for scale through practical approximations, including last-layer gradient truncation and Johnson–Lindenstrauss projection, while retaining stable subset selection.
- Empirically, diversity-only selection gives consistent gains over random sampling, and combining diversity with quality yields the strongest pretraining performance across DCLM and FineWeb.