How should HEP ML models be scaled and compared on budget axes?

Fixed hyperparameters mix real scaling with a moving optimum

A scaling law relates achievable loss to model size N, dataset size D, and compute C. The paper's central methodological point is that this relation only holds if each run is trained under appropriately tuned dynamics. At fixed model size and shape, the optimal learning rate and batch size change systematically with the training budget: as the token budget grows, the optimal learning rate decreases at fixed batch size, while at fixed learning rate the optimal batch size grows, and the joint optimum migrates toward larger batch sizes and learning rates.

The consequence is concrete. If hyperparameters are fixed at a reference budget, the measured loss curve mixes genuine scaling behavior with the movement of the training optimum. The paper's toy teacher–student experiment makes the size of that contamination visible: a learning rate tuned once at the small end and then held fixed does not sit at a constant offset from the per-cell tuned baseline. It departs from the tuned power law and turns upward at scale, so a grid trained this way biases the fitted exponents and can hide the scaling entirely.

The same effect appears in the compressed-representation system, where one-hidden-layer MLP probes are trained on latents from a frozen set-transformer autoencoder. Against the per-cell tuned baseline, a learning rate tuned at a single cell and then held fixed leaves a loss gap that grows toward both the large-D and large-width ends of the grid. Refitting the L(N, D) surface on that fixed-learning-rate grid returns a different law altogether, with a floor estimate above the best measured cell, so the mis-tuned grid misreports both the exponents and where scaling ends.

The joint optimum is predictable and transfers across scales

The recipe builds on theoretical descriptions of the joint (learning rate, batch size, token budget) optimum and on parameterizations that make hyperparameters transferable across width and depth. The authors fit how the joint optimal learning rate, batch size, and model shape move with the token budget, then check whether those predictions survive on configurations the fit never saw.

The validation is direct. On held-out combinations of batch size, token budget, width, and depth, the predicted optimal learning rate lands at or near the bottom of every loss-versus-learning-rate curve, so it transfers across batch size, token budget, and model shape. The joint optimum (learning rate, batch size) also transfers up to a possible shift in the optimal batch size as a function of model size at the optimal learning rate. For jet classification that shift is trivially absent because the loss is flat in batch size at the optimal learning rate. For the multi-task objective the optimal batch size grows as the square root of the token budget, leaving that shift as the only possible residual freedom as model size varies from the reference.

Using the predicted batch size instead of the lowest-loss batch size in a sweep costs at most a few times 10^-3 in loss, a per-mille-level penalty comparable to run-to-run noise. The authors note this transfer may be helped by the flatness of the loss around the optimum in batch size. The practical reading is that the joint optimum can be carried across scales without adding a further model-size dependence, which is what makes a compute-optimal frontier comparable across configurations.

Compute-optimal scaling splits budget near-equally between N and D

Applied to the GN family of transformer taggers on JetSet2, the extracted scaling laws show a near-equal square-root-of-C dependence of model and dataset size. In other words, as compute grows, model parameters and unique training jets grow at comparable rates rather than one dominating the allocation. The compute-optimal models are trained in a single pass over the data at the optimal model size, and the scaling-law projections assume the same.

The allocation is not uniform across settings. The paper reports compute-optimal allocation for five settings: two training objectives (all-aux multi-task and jet-classification-only single task) and three reduced input sets. The jet-classification-only setting is the clear outlier, preferring more data and a smaller model at every budget, whereas the input level does not appear to shift the allocation. Auxiliary objectives shift the compute-optimal split toward larger models and smaller datasets.

This distinction matters for how readers compare designs. A single-task tagger and a multi-task tagger do not sit on the same allocation line, so comparing them at one fixed model size or one fixed dataset size would confound the objective with the budget split. The paper's framework instead gives each configuration its own compute-optimal allocation before comparison.

Auxiliary tasks help at equal compute; richer inputs lower loss

Two design axes are compared along their own compute-optimal frontiers. The first is the training objective. Auxiliary objectives lower the primary jet-classification loss at equal compute budget, and they also shift the compute-optimal split toward larger models and smaller datasets. The paper notes that auxiliary objectives lower the jet-classification loss at a fixed number of training samples but also increase the computational cost of training, so their usefulness depends on whether the classification gain outweighs that added cost. The compute-budget comparison is what resolves that trade-off in their favor here.

The second axis is input representation. Expanding the inputs toward lower-level data systematically lowers the loss while leaving the scaling exponent nearly unchanged. The reduced input sets studied are the GN2 input set, a no-flows variant, and a four-momenta-only model. Lower-level inputs therefore improve the level of the loss curve without changing its slope, which is a different kind of gain from the one produced by auxiliary objectives.

Keeping these two effects separate is the point. An objective change moves both the loss level and the optimal allocation; an input change moves the loss level while leaving the allocation and the exponent approximately fixed. Readers comparing tagger designs at scale need to know which of the two they are looking at.

In the data-constrained regime four quantities are optimized jointly

The paper extends the procedure to the data-constrained regime, where the run is trained under early stopping rather than a single pass. Here the authors predict, to the best of their knowledge for the first time, the jointly optimal model size, training horizon, learning rate, and batch size under early stopping. This is a different optimization problem from the compute-constrained case: the training horizon is itself a free variable rather than being fixed by the single-pass assumption.

The distinction between the two regimes is not cosmetic. In the compute-optimal setting the models are trained in a single pass over the data at the optimal model size, and the projections assume the same. In the data-constrained setting, early stopping determines when training ends, so the optimal horizon has to be predicted alongside model size and the two training-dynamics hyperparameters.

The paper also reports a validation of the joint-optimum transfer on held-out combinations of batch size, token budget, width, and depth, where the predicted optimal learning rate sits at or near every minimum. The residual freedom is the possible batch-size shift as a function of model size at the optimal learning rate, and the measured penalty for using the predicted rather than the swept-best batch size stays at the per-mille level. That is what allows the joint optimum to be treated as transferable in practice.

The power-law regime only starts after enough data is seen

The onset of the power-law regime is itself set by scale. Below a threshold in dataset size, the loss carries little information about high-compute scaling. Fits based on smaller datasets can therefore return biased scaling exponents even when the resulting curves appear well described by a power law. This is a warning about apparent fits, not just about noisy ones: a curve can look like a clean power law and still yield the wrong exponent.

The authors connect this directly to data infrastructure. Large full-simulation datasets such as JetSet2 are a prerequisite for establishing how HEP models behave at scale, and the ability of LHC experiments to produce and retain such samples is described as a major asset for machine-learning research in HEP. The ~11 billion-jet JetSet2 dataset, built by combining ttbar and Z' to qqbar samples with VH to bbarb, ccbar, tau+tau-, Z' to ttbar, multi-jet, and gamma* to tau+tau- samples reweighted to the mean pT and eta spectrum across flavors, is the resource that makes the study possible.

The scale of the effort is also reported. The full study required approximately 97k A100 GPU-hours, with the compute-optimal grid across all variants accounting for 92k of those hours over 1219 runs and 1.5 x 10^12 jets seen. The authors describe this as substantial by HEP standards but modest compared with industrial-scale models.

Loss projections translate into rejection, not into a tagger ranking

The paper translates loss into background rejection on a ttbar sample at a fixed global b-jet efficiency of 70%, with 20 < pT < 250 GeV and |eta| < 2.5, using the GN3 discriminant fractions (fc, ftau) = (0.2, 0.1). It reports light-jet and c-jet rejection for GN3, for the TN25-86M model, and for projections of compute-optimal models trained on the full set of inputs and tasks, obtained with three fitting approaches.

The comparison needs a careful reading. GN3, TN25-86M, and the compute-optimal models share the same architecture and therefore belong to the same model class; they differ mainly in model size, amount of unique training data, and compute. The paper states explicitly that this is not a comparison of tagger designs. GN3 has about 12M parameters and is trained for several passes over 350M unique jets, while TN25-86M has 86M parameters and is trained for several passes over about 7.7B unique jets, and both are trained until the validation loss stops improving. The compute-optimal models, by contrast, are trained in a single pass at the optimal model size.

The authors also note that a deployed tagger is a frozen predictor that must be calibrated against data at a few working points, and that the optimal background fractions are in general functions of (pT, eta) and differ from one analysis to the next, so no single fixed choice is universally optimal. The rejection projections follow from a fitted relation between rejection and loss, and the spread between the three fitting approaches is shown alongside them.

What a scaling claim must show before it is trusted

The paper's own summary of the progression is a usable checklist. A method should demonstrate that its advantage persists as model size, dataset size, and compute increase; that it remains competitive when given the full input information available to the experiment; and ultimately that it holds up on real collision data. Each of these corresponds to a specific failure mode the study documents: an advantage that disappears once the training dynamics are scaled, an advantage that depends on withholding inputs, and an advantage that has only been shown on simulation.

The most decision-relevant threshold in the paper is the dataset-size threshold below which the loss carries little information about high-compute scaling. That threshold determines whether a scaling study can say anything about large budgets at all, and it is the reason the authors treat large full-simulation datasets as a prerequisite rather than a convenience. A reader evaluating a scaling claim should first ask how much unique data was seen before the fit was made, and whether the hyperparameters were re-tuned at each budget or held fixed from a reference run.

The forward-looking suggestion is a shift from separate models per application toward pretraining a single model at scale and adapting the resulting representation to different experiments, detector configurations, and physics tasks through fine-tuning. The observed scaling behavior is described as suggesting that considerably larger budgets could be productively invested, particularly under that foundation-model direction. The paper's framework is offered as a practical blueprint for such future scaling studies.

Sources used in this answer