What a Data Generation Budget Can Buy
How to split a data generation budget between breadth and depth, and which parameters determine whether any split is good enough.
Data generation increasingly dominates the costs of scientific ML. Millions of dollars go in over months or years, often before anyone knows whether the data generation strategy was sufficient. The choice between breadth and depth when constructing a dataset is not a matter of taste, but rather depends on the objective and a handful of parameters you can estimate. These same parameters also determine whether the given objective is possible at all on your budget.
DeepMind's Chinchilla scaling laws showed that the field had been splitting a fixed compute budget suboptimally across the tradeoff of model size and training tokens. Yet Chinchilla takes the data corpus as a given. Even for LLMs, that assumption is weakening, as the returns to scaling seem less predictable than previously believed, suggesting we need better datasets and not just more data. In scientific ML, data has always been scarce and expensive.
Most scientific datasets were built for purposes other than ML model training and then repurposed, like the Protein Data Bank being used to train AlphaFold. With more funding going into purpose-built scientific data generation, datasets increasingly become something you can design rather than inherit. The question is what a fixed budget can buy. Assuming we can choose sensible data points to measure, how do we decide when to gather more samples versus gather more accurate measurements per sample?
Coverage against accuracy
Assume we have a fixed data generation budget $G$, and we wish to develop an ML model with the objective of effectively learning the map $c \mapsto p(c)$, where $c$ is an arbitrary condition in the condition space $\mathcal{C}$. In practice we rarely care about all of $\mathcal{C}$, but rather the subset $\mathcal{C}' \subseteq \mathcal{C}$ where we expect our model to excel.
With a fixed total budget of $G$, we must choose how many samples $K$ we wish to collect data for, which samples $\{c_1, \ldots, c_K\}$ to choose that provide maximal coverage of the input domain $\mathcal{C}'$ we are interested in, and how much to spend ($b_1, \ldots, b_K$) on each sample, where increased spending can increase labeling fidelity and lower statistical noise. With a fixed budget, the number of samples in the set $c_1, \ldots, c_K$ and the spend on each sample $b_1, \ldots, b_K$ directly pull against each other.
To make the problem tractable in this essay, I make the five following assumptions:
- Homogeneous spend on each condition, so $\{b_i\} = b$ for all $i$ and $G = Kb$.
- Conditions are well-chosen to spread over the condition space $\mathcal{C}'$. One thousand conditions clustered in the same region can be less useful than one hundred well-chosen conditions.
- There is a known effective dimension for $\mathcal{C'}$, which is a $d$-dimensional manifold with a metric on it, and $d$ represents the dimensionality of the space where the conditions lie. Estimating $d$ for a given condition space is its own research problem, since it depends on both the metric chosen and the region of the space that you actually care about, and calculating this is out of scope here.
- We assume a model where the architecture and model size are held fixed, and improving the loss $L$ depends solely on the dataset.
- We assume additive separability between coverage error and label noise, with each independently contributing towards the loss. We can argue this is the worst case for $K$, since in practice a model can pool information from nearby conditions and average away per-sample noise, which would shift the optimum number for $K$ to having more conditions.
Under these assumptions, there are two error sources:
- How well do we cover the input domain $\mathcal{C}'$? This depends on the average fill distance $h$ of the space and the smoothness $\gamma$ of the map $c \mapsto p(c)$. Practically, this asks whether a small perturbation in $c$ produces a bounded change in $p(c)$. For $K$ points on a $d$-dimensional manifold, $h \propto K^{-1/d}$. Per the curse of dimensionality, as the dimensionality of the space we care about increases, the number of points needed to cover the space with same average fill distance increases nonlinearly. For a $\gamma$-smooth map, moving a distance of $h$ from condition $c$ changes $p(c)$ by at most a constant times $h^{\gamma}$. Thus, the error from predicting an unseen condition using its nearest sampled condition scales as $h^{\gamma}$, and so we write the squared coverage error as $h^{2 \gamma} \propto (K^{-1/d})^{2\gamma} = K^{-2\gamma/d}$.
- How noisy is each sample point? The statistical error at a condition $\varepsilon(b)$ can be reduced by spending more for that condition, so we can write this as $\varepsilon(b) \propto b^{-\rho}$, where $\rho$ represents how well increasing the budget for a sample reduces the error. Systematic error is handled separately by $\varepsilon_0$ since it does not respond to $b$.
Let $L$ be the model's expected squared error on the condition space $\mathcal{C}'$. Then,
$$L = \varepsilon_0 + A_1 K^{-2\gamma/d} + A_2\, \varepsilon(b)^2,$$where $A_1, A_2$ are empirically fitted coefficients and $\varepsilon_0$ is irreducible loss, set by the systematic error in the generating process, independent of $K$ and $b$. In practice, $\varepsilon_0$ is usually not known exactly, but can be estimated and refined as more data is gathered.
Given $Kb = G$ and $\varepsilon(b) \propto b^{-\rho}$, we minimize by solving $dL/dK = 0$ to get the optimal number of conditions $K^*$ and the corresponding per-condition spend $b^* = G/K^*$.1
$$K^* \propto G^{\rho d/(\rho d+\gamma)}, \quad b^* \propto G^{\gamma/(\rho d+\gamma)}.$$This is the framework. As $G$ grows, the optimal split changes as determined by the parameters $\gamma$, $\rho$, and $d$. Smooth maps favor higher $b$ whereas high-dimensional or jagged maps would favor more samples $K$.
Two further properties of data generators are worth keeping in mind when applying this in practice.
Feasibility floor on $b$. Generators often require a minimum budget spend $b$ for any particular condition, under which either an experiment cannot be run, or the result is dominated by noise. For example, in a biophysical assay this might be a requirement for a particular number of replicates to be run, or perhaps this is requiring a simulator to run for a minimum amount of time such that we learn the dynamics and not just strongly bias towards the initial seed. Consequently, this constraint requires all $b_i > b_{\text{min}}$, where $b_{\text{min}}$ is the minimum spend required to overcome this threshold. This implies $K \leq G/b_{\min}$.
Per-observable dependence. The observable we seek to learn for a condition determines how much spend a condition needs. If we simply seek a low-dimensional estimate of a distribution, the budget $b$ required for a sample may be much lower than if we want a precise, high-dimensional understanding with minimal error. Researchers with different objectives may look at the same condition space yet need different measurements; consequently, two labs targeting different observables need different $b$ to reach the same error on the same conditions.
What this looks like in molecular dynamics
The field already makes these tradeoffs, but they are often implicit and not following a formalized framework. We can see clear examples of these tradeoffs with dataset generation using molecular dynamics. Molecular dynamics is a physics-based simulation approach that integrates Newton's equations over time, giving a trajectory of continuous conformations that a protein takes. Choose $K$ proteins and simulate each for a time $b$. Given comparable system sizes and hardware, simulated time is proportional to compute cost, so we treat simulated time as the unit of spend and the budget is $G = Kb$.
Published datasets show a huge range in the ratios used between $K$ and $b$. mdCATH spends its budget on 5,398 domains at roughly 11.6 microseconds each (across 5 replicas and 5 temperatures for each system). By comparison, DE Shaw Research's 1.03-millisecond simulation of BPTI spends nearly 90 times the sampling of a single mdCATH system in one continuous run, roughly 2,300 times the length of any individual mdCATH trajectory. Note that these are not splits of the same budget, as mdCATH's total simulated time is roughly 60 times larger than BPTI's, and the two papers use different hardware for their simulations. Both datasets were expensive to generate and are at opposite ends of the $K$ vs $b$ ratio, yet nothing tells you whether, given your fixed objective, the next unit of budget should go towards continuing longer runs or starting new ones.
Note that the y-axis measures time per system, which stands in for cost as described earlier.
Molecular dynamics integrates continuous dynamics at tiny timesteps, and so there is often high correlation between contiguous frames. Consider $\tau$, the autocorrelation time of the observable beyond which two measurements are effectively independent. This means the number of independent samples in a trajectory of length $b$ is roughly proportional to $b/\tau$, and we can estimate the error $\varepsilon^2 = \sigma^2\tau/b$, where $\sigma^2$ is the variance (Flyvbjerg & Petersen). Since $\varepsilon \propto b^{-1/2}$, MD gives $\rho = 1/2$, and so the general result is $K^* \propto G^{d/(d+2\gamma)}$. This $\tau$ also sets $b_{\text{min}}$, as a trajectory shorter than $\tau$ has not decorrelated from its initial condition and falls below the feasibility floor described earlier.
In MD, $\varepsilon_0$ is set by the force field. The force field used in MD parametrizes how different atomic interactions should be treated, and is typically determined based on experimental measurements. However, different force fields have different parameterization, and any artifacts here set the upper bound for the useful signal we can expect to extract from molecular dynamics data generated with a single force field.
Finally, in the context of MD, how do we think about the variation in the costs required depending on the observable? Consider if we have an observable that converges quickly, like a scalar on the general size or compactness of a biomolecular system. In this context, the autocorrelation time $\tau$ may be relatively low, and we can gather lots of diverse data for relatively cheap. On the other hand, say we want to understand the rare states of a particular protein. This may require a large number of independent samples from a low-probability region of the protein's distribution of shapes. This suggests an observable with a high autocorrelation time $\tau$, consequently requiring more substantial budget resources per condition. This concretely shows how two different research groups with the same budget but interested in different observables should generate their datasets differently, showing how simply pulling existing public datasets may be insufficient.
What to go measure
How do we make use of this framework? I argue that this framework identifies which quantities determine the optimal allocation, and that these quantities are measurable. Hence, when constructing an effective dataset for your task, the key parameters to be determined are:
- $d$: the effective dimension of the condition space. No standard method, and open research explores dimensionality for different condition spaces.
- $\gamma$: the smoothness of the map. For MD, this can be estimated by simulating a few different conditions and measuring how much the resulting ensembles differ, or by fitting $2\gamma/d$ directly, which can be done by holding $G$ fixed, varying $K$, and fitting the two-term loss curve.
- $\tau$, $\sigma^2$: the relaxation time and variance per condition, which can be estimated from pilot runs before committing serious budget, and can determine whether the budget clears the feasibility floor. Block averaging is a standard method in MD.
This effectively turns dataset construction from an intuitive problem into an estimation problem. Each set of experiments we run can be used not only to feed a model but also to better estimate these parameters and gain better theoretical grounding for future dataset construction decision making.
Estimating these parameters is also what can help us reason through what objectives are even worth going after. Coverage of the condition space depends on $K^{-2\gamma/d}$, which shows that low $\gamma$ (a jagged map) and high $d$ can quickly make the condition space difficult to cover. Additionally, the feasibility floor caps the number of useful conditions at $G/b_{\text{min}}$. Taken together, these two constraints show some objectives are unreachable at any allocation of a given budget, not just poorly allocated ones. Fitting these parameters can help us estimate what objectives we may have a reasonable budget to tackle, rather than burning through half a budget just to find out our budgeting is orders of magnitude too small for the broad and complex condition space we wished to learn. If we can get better estimates a priori, we can readjust to model objectives the budget can actually reach.
While my examples here were concretely focused on molecular dynamics, this same framework applies to any generator where budget trades off between number of samples and per-sample accuracy. More ambitiously, we can put simulators, labeling pipelines, and physical experiments on the same footing, and reason about what a unit of budget buys from each. This would let us reason about how to spend a data budget across every data generator we could possibly use.
-
Dropping the constant $\varepsilon_0$ and substituting $b = G/K$ into $\varepsilon(b) \propto b^{-\rho}$ gives
$$L(K) = A_1 K^{-2\gamma/d} + A_2 G^{-2\rho} K^{2\rho}.$$Setting the derivative to zero,
$$\frac{dL}{dK} = -\frac{2\gamma}{d} A_1 K^{-2\gamma/d - 1} + 2\rho A_2 G^{-2\rho} K^{2\rho - 1} = 0,$$and collecting powers of $K$,
$$K^{2\rho + 2\gamma/d} \propto G^{2\rho}.$$Hence $K^* \propto G^{\rho d/(\rho d + \gamma)}$, and $b^* = G/K^* \propto G^{\gamma/(\rho d + \gamma)}$. ↩