Spokes: Optimizing for Diverse Pretraining Data Selection

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.

• The paper presents a novel selection method that directly maximizes gradient-space diversity using the G-Vendi score, achieving up to +489 improvement over random sampling.
• It introduces a probabilistic, continuous relaxation framework that balances quality and diversity via exponentiated gradient descent, enabling efficient control over data curation budgets.
• Empirical results on DCLM and FineWeb benchmarks show that jointly optimizing for diversity and quality leads to enhanced language model performance across various tasks.

Spokes: Probabilistic Optimization for Diverse Pretraining Subset Selection

Introduction and Motivation

The increasing scale of LLM pretraining has intensified the need for principled, efficient data selection paradigms. While diversity is empirically linked to improved generalization and data efficiency, most pipelines resort to either clustering heuristics or embedding-based similarity deduplication, which offer only proxies for true set-level diversity. "Spokes: Optimizing for Diverse Pretraining Data Selection" (2606.15216) presents a direct, scalable subset selection approach that jointly optimizes both gradient-space diversity and per-example data quality—a critical capability in large-scale corpus curation.

Theoretical Framework: Gradient-space Diversity via G-Vendi

Traditional diversity measures based on input or embedding similarity are limited by representation collapse and diminishing marginal utility in large pools. Spokes instead operationalizes diversity with the G-Vendi score, which quantifies the effective dimension of gradients induced by each data point on a proxy model. Formally, for gradients $g_i$ on data points $x_i$ and associated kernel $K$, the G-Vendi score is defined as

$$\mathrm{G\text{-}Vendi}(K) = \exp\left(-\sum_i \lambda_i \log \lambda_i\right)$$

where $\lambda_i$ are normalized eigenvalues of $K$. This metric rewards gradient directions with higher mutual orthogonality, mitigating redundancy of learning signals. Empirically, Spokes demonstrates that random sampling quickly saturates G-Vendi scores, whereas its data-driven selection can significantly increase this set-level diversity—by up to +489 points relative to random subsets of size 500k.

Algorithmic Approach: Probabilistic Subset Optimization

Spokes addresses the intractability of discrete subset maximization by relaxing selection into a continuous weight vector $w \in \Delta^n$ over the dataset, combining quality ($Q(w)$) and diversity ($\mathrm{Vendi}(K(w))$) objectives:

$$\max_{w \in \Delta^n}~\alpha \ln Q(w) + (1-\alpha) \ln \mathrm{Vendi}(K(w))$$

Exponentiated gradient descent optimizes this criterion over the simplex, enabling efficient convergence. The balance between diversity and quality is tunable via $x_i$0. Discrete subsets are recovered via top-$x_i$1 selection on optimized weights, providing deterministic control over data budgets.

Figure 1: Quality–diversity trade-off across $x_i$2: increasing emphasis on quality decreases G-Vendi diversity in a controlled manner.

This smooth trade-off allows practitioners to select operating points with tailored balance, as demonstrated by varying subset overlap with quality- and diversity-optimized extremes.

Scaling and Implementation: Gradient Projection and Computational Savings

Spokes targets the scalability required for web-scale curation. First, gradient extraction is restricted to the last $x_i$3 transformer layers, reducing costs while preserving high Spearman correlation with the full-model kernels (0.93 for the last 2 layers of Qwen3-0.6B-Base). Second, random projections via the Johnson–Lindenstrauss lemma compress gradient vectors from $x_i$4M dimensions to as little as 1024, retaining ≥95% subset overlap with reference selections even in this low-dimensional regime.

Figure 2: G-vendi scores scale as the number of top-k increases—a monotonic effect driven by Spokes’s optimization.

Time complexity is $x_i$5 for $x_i$6 iterations over subset size $x_i$7 and projected dimension $x_i$8, as opposed to the quadratic $x_i$9 scaling of clustering-deduplication techniques such as SemDeDup. Spokes can be further accelerated by data partitioning.

Empirical Results

Across two curated benchmarks—DCLM and FineWeb, representing filtered and noisy web-scale corpora respectively—Spokes consistently selects subsets with higher batch-level G-Vendi diversity compared to random sampling or semantic deduplication. Notably, gradient-based diversity metrics remain sensitive and discriminative under selection, whereas embedding-based metrics saturate.

Jointly optimizing for quality and diversity (with $K$0 adjusted per corpus) delivers further improvements. For instance, Spokes delivers average downstream benchmark gains of +1.5 (DCLM) and +1.4 (FineWeb) over baselines, and achieves higher OLMES suite scores than both quality-only and diversity-only selection. These effects persist across downstream task groups; diversity bias improves commonsense tasks (HellaSwag), while quality filtering dominates knowledge-intensive tasks (MMLU), yet joint optimization with Spokes offers Pareto improvements.

Figure 3: Subset overlap (%) with $K$1 projection dimension shows high stability of the selection when $K$2.

Practical and Theoretical Implications

Spokes demonstrates that gradient-based diversity is a robust, transferable signal for pretraining subset selection, complementing conventional quality scores. The method’s scalability and deterministic subset sizing address operational requirements for industrial data pipelines, and its framework generalizes to arbitrary per-example quality annotations.

On the theoretical side, the efficacy of gradient diversity supports the interpretation that set-level orthogonality among update directions enhances training signal orthogonality, resulting in more efficient parameter updates and better generalization per token.

Limitations and Future Work

While Spokes’s projection and layer truncation yield substantial efficiency gains, gradient extraction remains more costly relative to embedding methods. Further acceleration is possible via proxy gradient approximation or techniques such as Cut Cross Entropy. Extensions may also consider incorporating dynamic, context-aware proxy models for gradient extraction.

Future work may investigate integration with mixture-of-experts training or federated selection scenarios, and could explore more adaptive diversity-quality trade-off schedules.

Conclusion

Spokes provides a scalable, principled framework for pretraining data curation via direct optimization of gradient-space diversity, outperforming heuristic and embedding-based deduplication approaches. By enabling explicit joint optimization of quality and diversity, it reliably produces subsets that enhance downstream LLM performance. Spokes thus stands as a recommendation for practitioners seeking efficient, diversity-enriched pretraining pools in large-scale LLM training regimes.