Data Selection

By Arya, 8/1/2026

Overview

In this project I explore three methods to learn a data selection strategy for pre-training. For the first method, I attempted the ClassAct algorithm¹ that uses two proxy models to compute learnability scores on each sequence in the data pool. For the second method, I used one proxy model instead of two, and computed the loss reduction in each sequence. And for the third method, I designed a learned selection policy algorithm to optimize selection.

The goal is to learn a data selection strategy to choose 100M tokens out of a pool of 300M, that would perform better than random selection, and beat the baseline val loss of 3.7343. The three methods I explored produce a lower train loss but a slightly negative val loss compared to the baseline, but train stably.

Experiments

Fig 1. validation loss

Fig 2. train loss

1. ClassAct

I wanted to implement an existing data selection method to see how it would perform on the nanoGPT slowrun setup. The original idea of class act algorithm was conducted on vision transformers by Evans et al. I adapted the online class act for language models.

For selecting the proxies for this method, I went with a 4 layer transformer with 256dim since the paper mentions smaller proxies are effective for models 1000x larger.

The 300M data pool is separated into ~146,412 sequences:

$$\frac{300\mathrm{M\ tokens}}{2048\ \mathrm{tokens/seq}} = 146{,}412\ \mathrm{sequences}$$

The method uses two proxy models - an online proxy that represents the learner (GPT) model, and the reference proxy that learns faster at 2x the learning rate. Both models are trained on the superbatch and returns the per-sequence loss. We then compute the score by measuring the difference between its losses:

$$s(x)=\ell_{\text{online}}(x)-\ell_{\text{reference}}(x)$$

A high score means that the reference model improves a lot more than the online model, suggesting that the sequence is learnable. At each step we choose the top 64 sequences with highest score, and 763 steps * 64 sequences = 48,832 which is ~100M tokens.

Algorithm

Results

Val loss: 3.74618

Train loss: 3.57908

The candidate pool for selecting data was smaller than in the paper's implementation. Here, our selection ratio rho is 3, which means we select 1 out of every 3 sequences. However the authors set rho as 10. This could be why it didn't show improvement. The results also show lower train loss than baseline, but higher val loss which suggests the model fit the selected data better but could not generalize.

Cost

Online model - 3 fwd passes + 1 bwd pass per step; x 762 steps

Reference model - 3 fwd passes + 3 bwd passes per step; x 762 steps. (so 300M tokens fwd and bwd)

Final training on the learner model is the same as baseline - 100M tokens fwd and bwd

2. Warmup delta

While thinking of the problem, the simplest/immediate approach I thought of was to have a smaller indexer or proxy model score each data sequence by measuring how much it decreases the proxy's loss during a brief warmup run.

This is different from the previous Class Act method because this uses only one proxy model.

The idea is to have a warmup phase before pretraining, to train the proxy transformer model over the entire data pool and save an early checkpoint at 10% warmup and final checkpoint at 100% warmup. We can then compute the per-sequence delta L_early - L_final and take the topK sequences to be the training data

Algorithm

Results

Val loss: 3.76578

Train loss: 3.60237

The val loss was higher than the baseline despite the lower train loss and optimized proxy learnability rather than generalization. Did some testing and noticed Δ = ℓ_early − ℓ_final disproportionately selected multilingual and corrupted text. The proxy began with very high loss, then learned & memorized local patterns, producing a large Δ.

You can see in the plot below the highest delta occurred at the top and bottom percentiles. With the top percentile being foreign language sequences and the bottom percentile being unwanted data. In a followup experiment I tried trimming the top and bottom percentiles and selected data from the remaining stable delta region and the validation performance for this improved (Val loss: 3.74097). However this method was more of a learned data filtering approach.

Fig 3. Per-sequence proxy learnability (warmup delta).

Cost

• The proxy does 1 fwd pass and 1 bwd pass to save early and final checkpoints
• loads the early checkpoint and does 1 fwd pass on the entire data pool
• loads the final checkpoint and does 1 fwd pass on the entire data pool
• = 3 fwd passes + 1 bwd pass

3. Learned Selection policy

This method uses policy gradients to learn a policy to select data. So far the methods use a hand designed score by computing a single proxy metric. So instead I wanted to RL a proxy to learn the scoring and selection mechanism itself. More importantly, I wanted to train on samples that move the policy toward lower target loss². The idea of using policy gradients and relative advantages was borrowed from Fan et al³.

A proxy LM is first used to run a forward pass on the entire data pool to compute the scores. Every sequence in the data pool is then featurized into 8 base features: final loss, delta, unique tokens, non-ascii tokens, log frequency of tokens, no. of BOS tokens, entropy and gradient alignment with reward set. The policy then samples a subset of sequences, computes reward, and uses REINFORCE to update the feature weights towards subsets that showed lower loss on the reward set.

Algorithm

Results

Val loss: 3.74523

Train loss: 3.78715

The reward signal was noisy likely due to several reasons - 1. I used only G=8 policy rollouts per round. Small group sizes can destabilize policy gradient training³. Setting G=128 or 256 would reduce gradient variance. Similarly, increasing rounds to R>=5000 is preferable under more compute. 2. The proxy loss rewards learnable data, but does not optimize for diversity, causing dead subsets. 3. The proxy was much smaller, in the next run I would try a proxy that is closer in size and architecture to the learner model.

Cost

• Building the features with a proxy: 3 fwd passes + 1 bwd pass over the 300M pool
• Each rollout trains a fresh proxy: 1 fwd + 1 bwd
• Each rollout evaluates on a reward sequence set: 1 fwd pass
• There are 60 rounds and 8 rollouts per round, 492 rollouts in total
• The main learner is the same — it still does 1 fwd and 1 bwd on the final train set

Future Direction

For this project I only tried sequence-level selection, since choosing contiguous sequences preserves context better. However, with more time and compute I would explore token-level selection. Another direction to explore is using metagradients⁴ to make sequence or token weight parameters differentiable. This way we can measure how much nudging each weight affects val loss by backpropagation during training.

References

• 1. Bad Students Make Great Teachers: Active Learning Accelerates Large-Scale Visual Understanding Evans, T., Pathak, S., Merzić, H., Schwarz, J., Tanno, R., Hénaff, O. J. (2023).
• 2. OPUS: Towards Efficient and Principled Data Selection in Large Language Model Pre-training in Every Iteration Wang, S., Ouyang, X., Xu, T., Hu, Y., Liu, J., Chen, G., Zhang, T., Zheng, J., Yang, K., Ren, X., Liu, D., Zhang, L. (2026).
• 3. Joint Selection for Large-Scale Pre-Training Data via Policy Gradient-based Mask Learning Fan, Z., Xian, Y., Sun, Y., Shen, L. (2025).
• 4. Optimizing ML Training with Metagradient Descent Engstrom, L., Ilyas, A., Chen, B., Feldmann, A., Moses, W., Mądry, A. (2025).
• 5. Predictive Data Selection: The Data That Predicts Is the Data That Teaches Shum, K., Huang, Y., Zou, H., Ding, Q., Liao, Y., Chen, X., Liu, Q., He, J. (2025).
• 6. Data Efficacy for Language Model Training Dai, Y., Huang, Y., Zhang, X., Wu, W., Li, C., Lu, W., Cao, S., Dong, L., Li, S. (2025).