Technical Summary: SPOKES – Optimizing for Diverse Pretraining Data Selection

1. Problem Statement

Data diversity is a critical factor in constructing pretraining corpora, as it reduces redundancy and improves downstream performance under fixed data budgets. However, optimizing for diversity is inherently difficult because it is a set-level property dependent on interactions between data points rather than individual examples. Directly optimizing diversity is computationally intractable due to the combinatorial nature of subset selection, particularly at pretraining scale.

Existing approaches typically rely on proxies or approximations, such as coarse-grained clustering (by topics, skills, or unsupervised clusters) or semantic deduplication. While these methods ensure cluster representation, they rely on low-resolution partitions and often fail to capture information in long-tail or underrepresented knowledge. Furthermore, current methods often treat diversity metrics (like the G-Vendi score) as post-hoc evaluation tools rather than optimization objectives. The central challenge addressed by this work is how to reliably extract diverse subsets from large-scale pretraining corpora to improve downstream outcomes without relying on heuristics that fail to scale or capture fine-grained diversity.

2. Methodology: SPOKES

The authors introduce SPOKES (Spokes: Optimizing for Diverse Pretraining Data Selection), a scalable framework that directly optimizes for diversity using the G-Vendi score within a probabilistic diversification framework.

2.1 Core Objective

The method formulates data selection as a global optimization problem that jointly maximizes quality and diversity. Given a dataset of $n$ examples, the goal is to select a subset $S$ of size $k$ that maximizes:
$$\max_{S \subseteq [n], |S|=k} \alpha \ln \left( \frac{1}{k} \sum_{i \in S} q_i \right) + (1 - \alpha) \ln \text{Vendi}(K_S)$$
Where:

• $q_i$ represents the quality score of example $i$ (e.g., from a FineWeb-Edu classifier).
• $K_S$ is the similarity kernel restricted to the selected subset.
• $\alpha \in [0, 1]$ is a tunable parameter controlling the trade-off between quality and diversity.

2.2 Gradient-Based Diversity (G-Vendi)

Unlike traditional methods that use static embeddings, SPOKES operates in gradient space. It computes gradients $\nabla_\theta \ell(x_i)$ with respect to model parameters $\theta$ for each data point using a proxy model. The G-Vendi score measures diversity by examining the entropy of the spectrum of the similarity matrix constructed from these gradients. This approach encourages orthogonality among representation vectors, reducing redundancy in update directions and promoting more independent, informative optimization steps.

2.3 Optimization via Exponentiated Gradient Descent

To solve the combinatorial subset selection problem, SPOKES employs a relaxed optimization strategy:

1. Continuous Relaxation: Instead of discrete selection, a continuous non-negative weight vector $w$ is assigned to every data point, constrained to the simplex ($w \in \Delta_n$).
2. Kernel Construction: A weight-dependent kernel $K(w)$ is defined where entries are scaled by $\sqrt{w_i w_j}$.
3. Optimization: The objective is optimized using exponentiated gradient descent. At each iteration, weights are updated via $w_i^{(t+1)} \propto w_i^{(t)} \exp(-\eta g_i^{(t)})$ and normalized to prevent collapse onto a few high-weight examples.
4. Discrete Recovery: After convergence, the final subset is recovered by selecting the top-$k$ entries of the learned weight vector.

2.4 Scalability Enhancements

To make this approach feasible at pretraining scale, the authors implement several efficiency strategies:

• Truncated Gradients: Instead of computing full-model gradients, the method approximates them using only the last $n$ transformer layers (typically 2 layers). This significantly reduces computational cost while preserving the structure of the similarity kernel.
• Johnson-Lindenstrauss (JL) Projection: To handle high-dimensional gradient vectors (e.g., 200M parameters), the authors apply random projections (Rademacher matrices) to reduce dimensionality. Empirical results show that even low-dimensional projections (e.g., $k=1024$) maintain high subset stability (95.5% overlap with high-dimensional references).
• Parallelization: The dataset can be partitioned into disjoint subsets optimized independently, yielding linear speedup.

3. Key Contributions

1. SPOKES Framework: Introduction of a scalable, principled technique for diversity optimization that operates efficiently at the scale of modern pretraining datasets, moving beyond heuristic clustering.
2. Direct G-Vendi Optimization: Demonstration that the G-Vendi score can be directly optimized via exponentiated gradient descent to extract highly diverse subsets, rather than merely serving as a post-hoc metric.
3. Joint Quality-Diversity Optimization: Proof that jointly optimizing for quality and diversity yields superior results compared to optimizing for either objective alone or using standard baselines.
4. Efficiency Improvements: Validation that truncated gradient computation and dimensionality reduction via JL projection allow the method to scale to massive corpora without prohibitive computational costs.

4. Experimental Results

The method was evaluated on two large-scale pretraining corpora: DCLM (Dolmino subset, aggressively filtered) and FineWeb (less aggressively filtered). Models were pre-trained from scratch on a LLaMA-1B architecture using 100 billion tokens.

4.1 Diversity Metrics

• G-Vendi Score Increase: On a 500k-sample subset, SPOKES achieved a +489 point increase in G-Vendi score compared to random sampling. In contrast, semantic deduplication (SemDeDup) only achieved a +7 point increase.
• Batch Diversity: Throughout training, batches selected by SPOKES maintained consistently higher average diversity scores (increasing from 315 to 520 on DCLM) compared to random sampling.

4.2 Downstream Performance

Evaluation was conducted on 10 English benchmarks from the OLMES suite.

• Diversity-Only: SPOKES (diversity-only) improved average downstream performance by +0.4 (DCLM) and +0.5 (FineWeb) points over random sampling.
• Joint Optimization: The strongest results were achieved by jointly optimizing quality and diversity. SPOKES achieved gains of +1.5 (DCLM) and +1.4 (FineWeb) points over random sampling, outperforming all baselines, including semantic deduplication and quality-only filtering.
• Complementary Strengths: The results indicate that quality and diversity capture distinct but complementary properties. While quality filtering excels on knowledge-intensive benchmarks (e.g., MMLU), diversity optimization improves commonsense reasoning (e.g., HellaSwag). Joint optimization bridges this gap, performing competitively across all tasks.

5. Significance and Claims

The paper claims that direct optimization of diversity is both tractable at scale and practically beneficial. By moving away from coarse-grained approximations and heuristics, SPOKES demonstrates that:

• Diversity is a critical axis: Even without explicit quality supervision, optimizing for diversity alone yields consistent gains, suggesting that a diverse gradient representation of the data distribution improves generalization.
• Quality metrics are insufficient: Standard quality metrics do not fully capture data utility. Jointly optimizing for diversity can provide complementary training signals that lead to better downstream performance, even when the selected subset has a lower average quality score than a quality-only baseline.
• Scalability: The proposed method successfully extracts dense, highly diverse subsets from massive pretraining corpora, offering a viable alternative to existing selection pipelines that rely on semantic deduplication or random sampling.

The authors position SPOKES as a tool for constructing high-quality and diverse subsets, emphasizing that the benefits of diversity reliably translate across datasets with different filtering regimes and noise profiles.