Epiplexity Guided Data Selection and Generation for Out-of-Distribution Generalization

Epiplexity, a recently proposed measure of the structural information a compute-bounded learner can extract from data, provides a mechanism to reason about this relationship. In this paper, we show how to operationalize epiplexity as an online training signal for data selection and synthetic data generation. For selection, we fit scaling laws to the training loss curves of natural data domains to predict the expected epiplexity gain as a function of training tokens, and use this signal to adaptively determine the sampling weights over domains during training. For synthetic data generation, we define a generator’s reward as the change in the epiplexity of the learner’s training data and use REINFORCE policy gradients to guide the generator toward an epiplexity-maximizing distribution. In both cases, higher epiplexity predicts improved downstream performance on zero-shot and fine-tuning based tasks, supporting the hypothesis that data rich in structural information yield representations that transfer across domains.

We introduce EpiSelect, a new method that directly uses epiplexity to adaptively select batches of data through training. Section 3.1 shows that epiplexity of different data sources is a strong predictor of downstream performance. Section 3.2 introduces the EpiSelect procedure. In Section 3.3, we find the trivial procedure of only training on PileCC outperforms sophisticated curriculum learning procedures that had been benchmarked on The Pile. As a consequence, we propose instead evaluating curriculum procedures on the Common Pile, which has greater representation from a diversity of 30 sources. In Section 3.4, we show that EpiSelect on the Common Pile provides state-of-the-art performance across 10 downstream tasks.

Using The Pile [23], a popular benchmark for curriculum learning procedures, we evaluate whether higher epiplexity domains correspond to better downstream generalization. We train separate 124M LLaMA 2 family architecture [36] models for 15B tokens on a few dominant (by size) domains in the Pile: arXiv, GitHub, PileCC, PubMed, and Wikipedia. We then compute a post-hoc estimate of the epiplexity of each domain using the prequential estimator (Equation 1) and evaluate each model on a suite of 7 zero-shot tasks in the Language Model Evaluation Harness [37]. As seen in Figure 2 (Left), epiplexity strongly predicts average downstream performance across data domains. In the leftmost panel, we train on the 5 largest individual Pile domains (since smaller domains lack sufficient in-domain tokens to train without multiple epochs, which violates the prequential estimator’s assumption in Equation 1 that credits code-length reduction to unseen tokens). We find that, across domains, epiplexity correlates with OOD accuracy (Pearson r = 0.88, two-tailed p = 0.05), whereas the model’s weight norm shows no such relationship (r = 0.01, ρ = 0.20, not significant). These results suggest that the predictive signal is specific to epiplexity rather than being a generic function of the trained checkpoint (details in Appendix B). In addition, we run two existing data selection methods DoReMi [3] and ADO [4] and show in Figure 3 that they implicitly upsample data domains with higher epiplexity. Together, these findings further motivate the design of our data selection method, EpiSelect, which explicitly maximizes epiplexity through training.

In data selection, we wish to adaptively select batches of training data, amongst various data domains, towards better downstream generalization. In other words, while standard training uses random minibatches of data, we wish to pursue curriculum learning and to adaptively choose the domain for each batch as we proceed through training. We propose a data selection procedure, EpiSelect, which chooses the batches at each iteration that maximizes epiplexity. EpiSelect computes a distribution π ∈ ∆K−1 over some K data domains in a given dataset in order to maximize the expected epiplexity gain during training. Following the prequential estimator of epiplexity (Equation 1), we estimate the total epiplexity Ŝ at training step t with Equation 2, where n(s) k is the total number of domain k tokens seen at iteration s, and Lm(s) is the loss of domain m at iteration s. In order to make online predictions for epiplexity gain during training, we follow prior work [4] to propose a cross-domain scaling law in Equation 3, which predicts the loss on each domain m as a function of the current token counts seen from every domain. During training, we repeatedly fit the scaling law for each domain L̂m to the observed per-domain losses Lm and per-domain tokens counts nm collected from the training trajectory so far, using the form L̂m(n1,...,nK) = ϵm +βm(ΣK k=1 γm,k nk)−αm where αm and βm are positive constants, ϵm the irreducible error, and γm,k > 0 (with ΣK k=1 γm,k = 1) the normalized influence of all the domains on domain m (fitting details in Appendix C). While our model includes γm,k terms to capture cross-domain interactions, the single-domain scaling law can easily be recovered by setting γm,k = 0 for k ≠ m. Substituting the scaling laws into Equation 2 and taking the derivative with respect to the number of tokens from domain j (full derivation in Appendix D), Equation 4 arrives at a proxy for how much additional structure would be gained by training on additional tokens from domain j at time t. To maximize the overall epiplexity accumulated throughout training, we greedily maximize ∂Ŝ/∂nk to select more tokens from domains which are predicted to contribute a great increase in epiplexity. In Equation 5, we compute the distribution over domains, πk, as a softmax over ∂Ŝ/∂nk, scaled by temperature parameter τ, which allows for interpolation between the uniform distribution τ → ∞ and greedy maximization of marginal epiplexity τ → 0. In our experiments, we fixed τ = 1. πk serves as the domain sampling weights when selecting the next batch of training data. Algorithm 1 describes how we update the domain sampling weights with EpiSelect.

In our first data selection experiment, following prior work, we trained 124M transformer models of the LLaMA 2 family [36] on the uncopyrighted subset of the Pile dataset [23], which contains around 210B tokens across 17 domains (excluding the original Books3, BookCorpus2, OpenSubtitles, YTSubtitles, and OWT2 domains due to copyright issues). However, due to the inherent domain size imbalance in the Pile uncopyrighted dataset, with one domain, PileCC, dominating at 41% of the dataset tokens, we found that the performance of data selection methods on Pile is saturated. Not only were the performances across selection methods highly similar, but the differences in downstream performance could be predicted by the differences in how each method selected from PileCC. We then trained our model exclusively on the PileCC domain, and Figure 4 shows that this model was able to outperform ADO [4], the SOTA selection method. While it is valid that perhaps PileCC (a dataset of common crawl web text) is the most useful training data for model generalization, this result means that the ideal selector would only ever select from the PileCC domain, which makes Pile uninformative as a benchmark for distinguishing between selection strategies. Therefore, in subsequent data selection experiments, we evaluate the selection methods on the Common Pile dataset, an 8TB collection of public-domain text. Common Pile spans 30 domains, and the largest domain (Stack V2) accounts for roughly 14% of the entire dataset. We additionally built a token-balanced version of the Common Pile dataset (details in Appendix A) and report subsequent data selection results on this construction.

We now show that by directly selecting for domains with high epiplexity, EpiSelect provides a practically compelling framework for curriculum learning. In these experiments, we train decoder-only transformer models (124M and 1.3B) of the LLaMA 2 family architecture [36, 38] on the Common Pile dataset, an 8TB collection of public-domain text comprising 30 sources. We compare our method (EpiSelect) against a no selection baseline (Natural), which samples batches proportionally to domain size (balanced according to the procedure in Appendix A), and the ADO selection method [4], a state-of-the-art approach for curriculum learning. We chose baselines which have comparable training overhead, omitting prior works, such as DoReMi [3], which require trained smaller proxy models, incur larger costs, and underperform ADO. We train all models for a fixed budget of 15B tokens and evaluate all models on 10 common zero-shot downstream tasks in the Language Model Evaluation Harness [37]. Further details in E.1. Code available at https://github.com/eysu35/EpiSelect. We evaluate our epiplexity maximizing selector on a set of 10 zero-shot language understanding tasks from the LM Evaluation Harness [37]: ARC, BBQ, BoolQ, CSQA, HellaSwag (HSwag), LAMBADA (LAM), OBQA, PIQA, SciQ, and WinoGrande (WinoG). In Table 1, we show that for both model sizes, EpiSelect achieves the highest average task accuracy (0.394 and 0.431 for 124M and 1.3B, respectively), outperforming both the Natural (no selection) baseline and ADO by margins comparable to those reported in prior work [4]. The improvements are consistent across task categories, with EpiSelect attaining the top score on 6 of 10 tasks for both model sizes. These results provide the initial support for our central hypothesis: by reweighting the batches toward domains predicted to contribute the largest marginal increase in epiplexity, the selector induces a training curriculum that produces representations which transfer more effectively to a diverse set of OOD tasks. A key component of estimating epiplexity correctly is to model the cross-domain effects as seen in Equation 3. Figure 5 plots the matrix of fitted values γm,k values at the end of training, which captures the weighting that domain k tokens have when computing the predicted loss in domain m. We see that, beyond improved downstream performance, our cross-domain scaling laws surface an interpretable map of inter-domain influence. By the end of training, the fitted γm,k matrix is diagonally dominant, which aligns with our intuition that training on domain k tokens should have the highest impact toward reducing loss on domain k. We also observe a sparse set of strong cross-domain effects, in the off-diagonal values, and note an asymmetry in the domain interactions.

We have shown that epiplexity can provide a practically compelling signal for data selection. Curriculum learning is often faced with the dilemma of whether we should learn on “hard” or “easy” problems first. Epiplexity provides a more nuanced resolution: train on batches that challenge the model, which incur sustained drops in loss, but are still learnable. What those batches are will depend on the state of the model throughout training. We now move to the question of data generation. Purely synthetic data generation, not in the form of simple transformations, such as in data augmentation, is almost uncharted territory. When a training dataset is not neatly divided into domains or we require additional training data, synthetic data is especially promising. It is our contention that, in the future, synthetic data will become an indispensable component of the pretraining pipeline. We show for the first time how epiplexity can be operationalized to select over generator model weights to effectively guide synthetic data generation. Intuitively, the generated data should challenge the learner, leading to sustained and substantial decreases in loss, while not being trivial (zero loss) or noise (high irreducible loss). As previously shown in panel (a) of Figure 1, the prequential coding estimate of epiplexity captures this notion and is well-suited to guide a generator model toward producing optimal synthetic data while mitigating reward hacking. We train a generator model, Pθg, to generate high-epiplexity synthetic data used to train a separate learner model, Pθ. To estimate the epiplexity for the learner model, we implement a generation buffer B, which accumulates the generated data throughout training, including the current batch Xt. Our use of a buffer differs in function from experience replay [39, 40], where a buffer stores past transitions for off-policy training updates. In our method, the buffer is not used to train either model; instead, it serves as a reference distribution for reward computation and provides a tractable estimate of epiplexity. For each batch of generated data, Equation 6 estimates the change in epiplexity as the difference between the learner loss on the samples in the buffer before (Pθt−1) and after (Pθt) training on the current batch. In order for the generator to learn to maximize the resulting epiplexity, we can set the reward rt for the model as this step-wise change in epiplexity and optimize the objective J(θg) = E[rt] using REINFORCE policy gradients [22]. In Equation 7, the gradient can be expressed as a product of an advantage term (where baseline b ← β·b+(1−β)·rt tracks the reward history with an exponential moving average) and the log-likelihood of the batch. Algorithm 2 describes this training procedure.

To empirically validate EpiGen, we initialize a generator and a learner model from the pretrained GPT-2 weights of Radford et al. [41], and train the learner on batches of synthetic data produced by the generator. We train the generator to optimize for an epiplexity-maximizing distribution via REINFORCE policy gradients. Fortunately, the variance of our REINFORCE gradients is not high, leading to stable generation dynamics as seen in Appendix F. We compare the final EpiGen learner model against the pretrained GPT-2 checkpoint (Pretrained) and a set of baseline learner models trained by the following generators: a frozen generator to match the total token count (FrozenGen), a generator trained with a heuristic-based reward (PPL, negative log perplexity), and a generator trained with a learning-based reward that only considers the current batch of data (NoBuffer). To verify that the learner is shifting from the pretrained distribution, we track its performance on OpenWebText [42] (OWT), which is an open source reproduction of WebText, the pretraining dataset of GPT-2 [41]. Finally, we evaluate the learner model’s generalization capacity on subtasks from the General Language Understanding Evaluation (GLUE) benchmark [43]. Additional details are provided in Appendix E.2. Code available at https://github.com/eysu35/EpiGen. After training on synthetic data, we evaluate the learner on fine-tuning-based GLUE tasks [43] in order to assess the representation quality of the new learner model. For each model, we follow the standard procedure of fine-tuning all weights along with a linear classification head on seven GLUE natural language understanding tasks (CoLA, SST-2, MRPC, QQP, MNLI, QNLI, RTE) for 3 epochs using a learning rate of 2 × 10−5 and batch size 32, then evaluating the model on the development set. We report per-task scores using each task’s standard metric as given in the benchmark reference (Matthews correlation coefficient for CoLA, accuracy for others) and the average score. Table 2 shows that training the generator with the epiplexity-maximizing reward yields a learner model which outperforms the pretrained baseline on most GLUE tasks. Critically, none of the learner models evaluated in Table 2 have seen any more tokens of real data than the pretrained model. Yet we see the epiplexity-maximizing generator is able to produce synthetic data which yields a learner model that achieves the highest overall score across GLUE tasks (Table 2) while maintaining low perplexity on in-distribution OWT prediction (Appendix G). We note an uptick in performance even when the learner is trained on additional tokens sampled from the FrozenGen pretrained generator weights, which suggests that improvement is partially accredited to continued training on additional in-distribution tokens. However, the gap between the FrozenGen, PPL, and NoBuffer and EpiGen methods isolates the contribution of epiplexity-based guidance. One hypothesis we have for the source of the improvement is that computing the reward over the data buffer may partially regularize the generator model against mode collapse. Memorizing a particular batch Xt of noise data would not reduce learner loss on the remainder of B, meaning generating a non-generalizable sample will not guarantee a high reward. This hypothesis is partially supported by the fact that synthetic data quality seems to be maintained throughout the course of training, despite not having any external verifiers or grounding constraints in our setup. We provide a set of randomly selected synthetic text samples generated at steps 1000 and 60,000 (last) in Appendix H. To better understand the efficacy of our method, we first investigate the effect of model initialization to determine where the additional learning signal sources from. We ran our data generation pipeline with generator and learner model initialized with random weights and with pretrained GPT-2 weights. Figure 6 shows that, while epiplexity-guided training improves average GLUE performance from 0.743 to 0.770 with the pretrained GPT-2 weights, it provides essentially no change (-0.003) with the random initialization. This finding is consistent with the observation that while the ceiling of reinforcement learning is in principle not bound to the pretrained model, empirically the quality of the pretrained initialization still constrains the performance after RL [44]. In our case, having a pre-trained generator proves necessary for RL to discover useful synthetic data. Finally, although self-play paradigms are inherently useful in not requiring additional data to improve model performance, we conduct a data mixing experiment to assess how the epiplexity-guided synthetic data stack up against novel real data. Here, we initialize the generator model as the trained, epiplexity-maximizing generator from Section 4.1. Then, we run the data generation pipeline with another learner model (initialized with the pretrained GPT-2 weights) and keep the generator weights frozen. In this way, we train the learner on the epiplexity-maximizing synthetic dataset for the entire duration of training, and can incrementally replace fractions of the batches of synthetic dataset with samples from the OWT dataset, all of the same sequence length, to further ablate the impact of synthetic data training on downstream performance. We make note of an additional analysis in Appendix G.1, in which we ablate the fraction of OWT data in the evaluation buffer to more realistically approximate epiplexity over all prior training data, real and synthetic. Figure 7 shows that learner models trained on 50% and 25% of synthetic data (and, subsequently, 50% and 75% of OWT data) achieve higher average performance on the OOD GLUE benchmark, indicating that training on synthetic and real data outperforms training on either one individually. This result both further validates the performance improvement we observed in Table 2 and substantiates the need for principled synthetic data generation methods.

Machine learning research has historically operated under the assumption that we wish to maximize generalization performance on test points drawn from the same distribution as a set of training data. From the model selection side, this setting has led to principles such as Occam’s razor, where we wish to have the simplest explanation that is consistent with our observations. However, we are now entering uncharted territory, where we train models on heterogeneous data and deploy them in new unanticipated domains. This setting forces a shift in perspective from model selection to data selection. The principles here may be quite different: while we want the model for a given dataset to be as compressible as possible, perhaps we want to expose the model to data that makes the model as incompressible as possible. The rationale for this prescription is that we want our models to learn to the fullest extent, with the hope that what it is learning may be applied in a wide array of settings. This contention is at the heart of epiplexity. We showed that epiplexity not only predicts better downstream generalization, but can be used to develop novel selection procedures that exceed state-of-the-art. Perhaps surprisingly, we also showed that epiplexity has promise as a learning signal for generating purely synthetic data that improves out-of-distribution generalization. The idea of synthesizing data for OOD generalization is entirely new, with many exciting future directions to explore. Multi-modal and multi-objective generation, as well as efficiency and scale, will be key considerations. We note some limitations of our work. First, we conducted our experiments on small models, which are inherently constrained in their performance on downstream tasks, and additional experiments should be conducted on the scalability and generalizability of these methods across model architectures. In addition, in both data selection and generation, we relied on a tractable proxy for epiplexity rather than the quantity itself, and we leave to future work the theoretical guarantees that maximizing for this proxy during training maximizes the true epiplexity. Further ahead, we will in time have a better understanding of what abstract properties of data are shared across domains beyond epiplexity, so that we can generate increasingly universal data samples, leading to broadly improved generalization.

Based on this paper, here are the specific improvements I can make to AI systems and what the improved systems can do:

1. Adaptive Domain Sampling for Pretraining

• Implement EpiSelect’s online scaling-law fitting to dynamically reweight training domains based on predicted marginal epiplexity gain.

• The improved system can automatically shift sampling toward domains that maximize structural information gain at each training step, leading to better zero-shot and fine-tuning performance on out-of-distribution tasks (e.g., +0.02 average accuracy on 10 tasks for 124M and 1.3B models).

2. Epiplexity-Guided Synthetic Data Generation

• Use the REINFORCE-based generator (EpiGen) with a buffer-based reward that measures learner loss reduction on prior generated data.

• The improved system can generate synthetic text that improves downstream GLUE performance (e.g., from 0.743 to 0.770 average) without any additional real data, while avoiding mode collapse and reward hacking.

3. Cross-Domain Scaling Law for Loss Prediction

• Fit per-domain scaling laws with cross-domain influence parameters (γ m,k) to predict loss reduction from any domain’s tokens.

• The improved system can accurately forecast which training data (from any source) will most reduce loss on a target domain, enabling efficient data mixing and curriculum design.

4. Epiplexity-Based Data Quality Scoring

• Use the prequential estimator to score any dataset or domain’s structural information content.

• The improved system can rank data sources by their expected contribution to representation learning, allowing users to discard low-epiplexity data and prioritize high-value domains before training.

5. Dynamic Curriculum Learning Without Manual Tuning

• Replace heuristic curricula (e.g., easy-to-hard) with epiplexity-maximizing batch selection.

• The improved system can automatically determine the optimal order and mixture of training data based on the learner’s current state, outperforming prior curriculum methods (e.g., ADO) without requiring a separate proxy model.

6. Robust OOD Generalization via Epiplexity Maximization

• Train models to maximize incompressibility of learned representations (as opposed to minimizing loss on fixed data).

• The improved system can produce representations that transfer more effectively to unseen domains, as evidenced by higher accuracy on diverse zero-shot tasks (ARC, HellaSwag, etc.) and fine-tuning benchmarks.

7. Real-Synthetic Data Mixing Optimization

• Use epiplexity to determine the optimal ratio of real to synthetic data during training.

• The improved system can combine synthetic and real data (e.g., 50% synthetic + 50% real) to achieve higher OOD performance than either alone, enabling data-efficient scaling when real data is scarce.

8. Interpretable Inter-Domain Influence Mapping

• Extract the fitted γ matrix to visualize which domains help or hinder learning in others.

• The improved system can provide actionable insights into data dependencies, helping researchers identify complementary or redundant data sources for more efficient pretraining.

Modern systems are increasingly expected to transfer across tasks not specified during training. What data facilitates generalization in these new, unanticipated settings? One hypothesis is that data with more structural information could contain shared circuits and subprograms that could be recycled in a wider array of downstream settings. Epiplexity, a recently proposed measure of the structural information a compute-bounded learner can extract from data, provides a mechanism to reason about this relationship. In this paper, we show how to operationalize epiplexity as an online training signal for data selection and synthetic data generation. For selection, we fit scaling laws to the training loss curves of natural data domains to predict the expected epiplexity gain as a function of training tokens, and use this signal to adaptively determine the sampling weights over domains during training. For synthetic data generation, we define a generator's reward as the change in learner epiplexity over a buffer of previously generated data and use REINFORCE policy gradients to guide the generator toward an epiplexity-maximizing distribution. In both cases, higher epiplexity predicts improved downstream performance on zero-shot and fine-tuning based tasks, supporting the hypothesis that data rich in structural information yield representations that transfer across domains.