hep-exOct 5, 2026
SaveHow to scale your HEP ML models: A recipe for robust architecture comparisons at scale
Organizations: Technical University of Munich · SLAC National Accelerator Laboratory · University College London · University of Geneva · Humboldt-Universität zu Berlin · NSF AI Institute for Artificial Intelligence and Fundamental Interactions · Munich Center for Machine Learning
Abstract
Much of the recent progress in machine learning domains such as language models has come from scaling laws that predict performance as a function of training effort. In high-energy physics (HEP) similar behavior has now been observed. To aid further study, we present a systematic procedure to derive robust scaling laws and compare design choices on the relevant budget axes for HEP tasks. We first validate the full scaling trajectory on toy problems and then apply the procedure to multi-task transformers on the ~11 billion-jet ATLAS JetSet2 dataset, in both the compute- and data-constrained regimes. For the latter, we predict, to the best of our knowledge for the first time, the jointly optimal model size, training horizon, learning rate and batch size under early stopping. At compute-optimal scaling, we recover a near-equal dependence of model and dataset size, and find that auxiliary objectives lower the primary jet-classification loss at equal compute budget. Expanding the inputs toward lower-level data systematically lowers the loss while leaving the scaling exponent nearly unchanged. 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, underscoring the value of large, high-quality full-simulation datasets as a foundation for scaling studies and the development of foundation models in HEP.
Figures & tables
Figure 1 : Validation-loss landscape of the flavor tagger trained with all auxiliary tasks included, at fixed model size M, for training-token budgets , , and . The dashed line traces the optimal learning rate at each batch size for a given token budget, and the red marker the joint optimum . As the token budget grows, the optimal learning rate decreases at fixed batch size, while at fixed learning rate the optimal batch size grows. The joint optimum migrates toward larger batch sizes and learning rates. Scaling these quantities optimally and transferring them across model sizes and shapes is a prerequisite for constructing robust scaling laws.
Figure 2 : Jet-classification loss against training compute, with every run trained at its predicted optimal and at the compute-optimal allocation at every budget .
Figure 3 : Sketch of how empirical scaling laws emerge. (a) loss against dataset size for a sequence of model sizes : every curve follows a data-driven descent and saturates at a model-capacity plateau, lower for larger , with all sufficiently large models approaching the irreducible floor (dashed). (b) the same curves against compute : larger models shift rightward and the lower envelope (red) traces the compute-optimal frontier .
Figure 4 : Optimal validation loss at fixed as a function of token budget . Each loss curve is a separate run with tuned for the budget at which it ends, and the envelope passes through their endpoints. This is the inner minimization of Eq. 7 . The outer minimization over yields the envelope targeted by scaling laws.
Figure 5 : Compute-optimal envelope of the teacher–student problem (Appendix L ), with three power-law fits: an initial data-driven descent, a scaling regime with a pure power law, and a tail asymptoting to that recovers the known floor ( against the true , dotted). The full-surface fit (Approach 3, green dashed) mis-fits the envelope across the regimes and flattens onto a biased floor. Inside the scaling regime nothing signals where the power law ends.
Figure 6 : Test cross-entropy of MLP probes on a frozen jet latent space ( ) against training compute, colored by probe width (left) and by dataset size (right). The dashed red line is the compute-optimal envelope. In the shaded data-hungry region the envelope is traced by the training curve of a single small model (the highlighted run): every larger model lies on or above it, and nothing yet distinguishes model sizes. Scaling in the compute-optimal sense only sets in past a task-dependent dataset budget ( jets here), after which the allocation starts to matter.
Figure 7 : Compressed-representation scaling across latent dimensions – : test cross-entropy on the -way classification target against compute. Gray: test cross-entropy per probe width. Blue: latent’s compute-optimal envelope. Larger latents lower the reachable loss while the envelope shape repeats.
| Parameter type | LR | AdamW | weight decay |
|---|---|---|---|
| Embedding (input) | |||
| Hidden weights (attn, FFN) | |||
| Unembedding / readout | |||
| Biases, LN/RMSNorm gains | |||
| QK-norm gains |
Table 1 : Per-parameter-type learning-rate (LR), AdamW- and weight-decay multipliers of the Complete(d)P recipe used for every run in this work, relative to the base values , , ; and are the width and depth ratios to the base model. The learning-rate and weight-decay columns follow Table 1 of [ 61 ] at and the QK-norm row in Table 1 of [ 63 ] ; weight decay is the PyTorch AdamW convention, so keeps width-independent.
Figure 8 : Empirical validation of the Complete(d)P recipe on the coordinate sweep at . (a) Final validation loss against the effective peak learning rate for seven widths at and (top) and five depths at and (bottom). The optimum lies at the same for all widths, with the residual drift in factored out. (b) Training loss against the training step for the slice , , one curve per peak learning rate .
| optimization target | |||
|---|---|---|---|
| Total (multi-task) | |||
| Jet-classification-only |
Table 2 : Multiplicative calibration of the optimal learning rate ( Eq. 11 ), with anchors and jets. Shown for the total (multi-task) and jet-classification-only objectives. The per-head multi-task breakdown is given in Table 7 .
Figure 9 : Final validation loss against peak learning rate at and , one curve per dataset size . The minimum of each curve is the extracted .
Figure 10 : The landscape at the reference architecture for the jet-classification-only variant. (a) the measured optimal learning rate , colored by , with iso- contours from the multiplicative fit of Eq. 11 . (b) the validation loss at that optimal learning rate, , with iso-loss contours. The loss is essentially flat in and steepens only in the few-step corner (large , small ), so is left largely unconstrained. White circles mark the swept cells.
Figure 11 : The landscape at the reference architecture for the multi-task variant. (a) the measured optimal learning rate , colored by , with iso- contours from the multiplicative fit of Eq. 11 . (b) the validation loss at that optimal learning rate, , with iso-loss contours. The loss now forms a shallow minimum in whose location grows as . White circles mark the swept cells.
| Batch Size | Learning Rate | Optimizer | |
|---|---|---|---|
| This work | AdamW [ 70 ] | ||
| DeepSeek [ 58 ] | AdamW [ 70 ] | ||
| StepLaw [ 56 ] | AdamW [ 70 ] | ||
| von Rütte et al. [ 84 ] | LaProp [ 85 ] | ||
| Filatov et al. [ 86 ] | Scion [ 87 ] |
Table 3 : Fitted power-law exponents for optimal batch size and learning rate as a function of token budget , as reported by studies with differing optimizers and tuning protocols. Adapted from [ 59 ] .
Figure 12 : Final validation loss against aspect ratio for four model sizes, M, 30M, 50M and 100M, with each point trained at its own . Large AR is a wide, shallow network and small AR a narrow, deep one. The minimum drifts to larger AR as grows, so width should grow faster than depth.
Figure 13 : Measured throughput on an H200 against the parameter count , one color per depth and one column per batch size . Points are measurements of a full training step, and lines are linear fits in at fixed and . (a) Time per jet. Throughput improves as grows and amortizes the per-step overhead, leveling off once the device is compute-bound. (b) The same runs as GPU-hours per FLOPs of theoretical compute. Equal theoretical compute costs more wall-clock time for deeper networks, and small models at small batch size reach only a fraction of the device peak.
Figure 14 : Training loss against optimizer steps (left) and A100 GPU-hours (right), for a model trained eagerly and with the training graph compiled by torch.compile . Both runs spent the same amount of compute. Model and hyperparameters are identical, so the loss per step coincides while the compiled run performs more optimizer steps per GPU-hour.
Figure 15 : Onset of scaling on the multi-task loss variant. Jet classification loss against compute . The portion of the compute-optimal frontier a few-hundred-million-jet dataset can probe (up to FLOP) is still the data-driven transient, and the asymptotic scaling slope sets in only beyond it. The same curves are colored by model size (left) and by dataset size (right) . Circles mark the final per-cell validation loss.
Figure 16 : Approach 1 for the multi-task variant. (a) Total training loss against cumulative compute , one lightly smoothed curve per run colored by . The black line is the lower envelope, the empirical compute-optimal frontier . (b,c) The model size and dataset size on the envelope at each compute budget, with the fitted power laws of Eq. 14 .
Figure 17 : Approach 2 (IsoFLOP profiles) for the multi-task variant. a) Total final validation loss against model size along iso-compute slices. The dashed line is the compute-optimal trajectory. b) the resulting model-size–dataset allocation.
Figure 18 : Approach 3 for the multi-task variant: iso-loss contours of the fitted total validation loss surface , with the measured cells on the same color scale. Dotted lines are iso-compute slices. The dashed line is the compute-optimal locus implied by the fit, the that minimize the loss at each compute budget.
| Method | ||
|---|---|---|
| This work (Approach 1) | ||
| This work (Approach 2) | ||
| This work (Approach 3) | ||
| Carpe Datum [ 42 ] | ||
| Kaplan et al. [ 24 ] | ||
| Hoffmann et al. [ 25 ] |
Table 4: Compute-optimal allocation exponents, where and . All three approaches agree on , matching the Hoffmann et al. result, where the tension with the earlier Kaplan et al. coefficients was attributed to non-embedding parameter count and hyperparameter tuning at small scale [ 77 , 53 , 74 ] . We observe a similar pattern relative to the Carpe Datum results: with tuned hyperparameter scaling and higher dataset sizes, the compute-optimal allocation recovers the equal dependence of model and dataset size.
Figure 19 : (a) Best early-stopped validation loss of runs targeting passes for a M model on unique jets, with Eq. 15 and the single-pass law counting repeated jets as fresh. (b) Validation (top) and training loss (bottom) of six of these runs against the number of passes , and the single-pass law on fresh jets (dashed). (c) The same model on , and unique jets against the training horizon.
Figure 20 : Top: early-stopped validation loss against the targeted number of passes for four model sizes, with Eq. 15 (solid) and the single-pass law with repeated jets counted as fresh (dotted). The open circles at show the single-pass scaling law. Bottom: excess of the early-stopped loss over the single pass of the same model on as many fresh jets, for three unique sets.
Figure 21 : Best early-stopped validation loss over all targeted numbers of passes against model size on , and unique jets, with an independent fit per unique set: a power-law descent plus a bump in that is only as large as the data require.
Figure 22 : Loss of Eq. 15 in the plane of training horizon and model size (contours) at (a) and (b) unique jets, with the compute-optimal allocations of Eq. 15 (solid blue) and counting repeated jets as fresh (dashed), with their optima at and FLOP (gray dashed lines). Dotted lines mark the number of passes. Plot style inspired by [ 76 ] .
Figure 23 : Jet-classification loss (jcls) for the classification-only baseline, for classification plus each single auxiliary task, and for all five tasks, against jets seen (left) and A100 GPU-hours (right) , at fixed model size M. Per jet seen, every auxiliary task lowers the jet-classification loss. Per GPU-hour the gain is partly offset by the higher per-jet cost of the extra heads.
Figure 24 : Multi-task (blue) versus single-task (brown) compute-optimal scaling, both with the full input set. (a) Final jet-classification validation loss on the compute-optimal frontier with bands. The multi-task frontier lies below the single-task one across the studied range. (b) The corresponding compute-optimal allocation with bands. Dotted diagonals are lines of constant compute at decades from to FLOP. At a given compute budget the multi-task objective prefers a larger model and a smaller dataset, so the auxiliary heads shift the compute-optimal split toward model size.
Figure 25 : Compute-optimal jet-classification loss for each input level: the runs on the compute-optimal frontier and the fitted frontier with its band, dashed beyond the measured range. Each added modality lowers the loss at fixed compute, and the four-momenta-only model (upper axis range) lies far above the rest. Apart from the 4-momenta-only variant, the frontiers are close to parallel, so the input level mainly shifts the loss rather than changing the scaling exponent.
Figure 26 : Jet-classification loss against compute, colored by model size (left) and dataset size (right) . The dashed line is the compute optimal loss . The data-repeated TN25-86M run lies above the single-pass frontier.
Figure 27 : Light-jet (blue) and -jet (orange) rejection of GN3 [ 6 ] , TN25-86M [ 42 ] , and the compute-optimal models (model sizes are annotated for a few decades in compute). For the scaling laws projections, all three fitting approaches are shown. Evaluated at global -jet efficiency, GeV, on a sample.
| Study | Runs | Jets seen | Training FLOPs | GPU-hours (A100) |
|---|---|---|---|---|
| P transfer | 324 | 0.34k | ||
| / + aspect ratio | 1395 | 4.5k | ||
| Compute-optimal grid (all variants) | 1219 | 92k | ||
| Total | 2938 | 97k |
Table 5: The number of runs, jets seen, training FLOPs, and A100 GPU-hours per study, with the total.
Appendix figures & tables35 assets
Supplementary material from the paper’s appendix.
Appendix
| Input Modality | Variable | Description | Full set (GN3) | Tracks electrons muons | Tracks electrons | Tracks (GN2) | 4-momenta |
| Jet | Transverse momentum | ||||||
| Pseudo-rapidity | |||||||
| Track | Track’s charge over momentum (curvature) | ||||||
| Pseudorapidity relative to the jet | |||||||
| Azimuthal angle relative to the jet | |||||||
| Closest distance from the track to PV in the transverse plane |
Table 6 : Input variables considered in this study. Muon variables are treated as additional track features, while particle-flow and electron inputs are separate input modalities. The five right-most columns indicate, with a colored cell, which variables enter each training variant: full set (GN3), Tracks electrons muons (GN3 minus the pflow modality), Tracks electrons (GN3 minus the the pflow and muon modality / GN2 plus the electron modality), tracks (GN2), 4-momenta only (all modalities).
Figure 28 : Best early-stopped loss of each point of a grid of (markers) and (colors), over the targeted pass counts run for it, relative to the best loss of the recipe (star), against the pass at which it is reached. The gray line connects the recipe’s own runs, one per targeted pass count. A setting better than the recipe on both counts would sit in the shaded corner. Triangles: points above the range shown. Of the two batch sizes of each marker, the first is used on and the second on unique jets.
Figure 29 : Allocation under repetition. (a) Contours of the loss of Eq. 15 at in the plane of training horizon and parameters, the measured early-stopped losses on the same color scale, the compute-optimal allocations of Eq. 15 and of the single-pass law with repeated jets counted as fresh, and their optima at FLOP; dotted lines mark the number of passes. (b) Constant compute when runs stop early (solid, labeled) and when the full schedule is charged (dashed), with the runs colored by their loss. (c) Loss on the compute-optimal frontier for to , the frontier with unlimited fresh jets (dashed), the measured runs on their lower envelope (circles) and compute-optimal single passes (triangles).
Figure 30 : Eq. 15 at unique jets with the onset exponent held at 0.5, 1 and 1.5 and the other parameters refitted. (a) ; (b) the terms of the law at its best target (the parameter term does not depend on ); (c) the law at its best target against model size; (d) the loss on the compute-optimal frontier, spending at most (dashed) and exactly (solid) the budget.
Figure 31 : Allocation implied by Eq. 15 at unique jets with the onset exponent held at 0.5, 1, 1.5 and 2: contours of the loss in the plane of training horizon and parameters, and the compute-optimal allocation of Eq. 15 (solid) and of the single-pass law with repeated jets counted as fresh (dashed). Where the law has a minimum at small models, the allocation stops there (lower diamond), and then restarts once the budget exceeds the band. Dotted lines mark the number of target passes.
Figure 32 : Downstream performance against training loss for the multi-task models. The first three panels show the -, light- and -jet rejection at the -jet efficiency working point against the jet-classification loss. The fourth shows the median relative resolution against the -regression loss. Curves are fits of , with the shaded interval from bootstrapping. The top axis converts loss into training compute through the compute-optimal envelope .
Figure 33 : Light-jet rejection (blue) and -jet rejection (orange) of the deployed ATLAS flavor taggers (DL1 to GN3), the TN25-86M model, and the compute optimal models. Error bars are the rejection-fit and uncertainties. Evaluated at global -jet efficiency, GeV, on a sample.
Figure 34 : All trained multi-task runs: training loss against cumulative compute . The same curves are colored by model size (left) and by dataset size (right) ; circles mark the final per-cell validation loss. The shared lower envelope is the empirical compute-optimal frontier that Approaches 1–3 estimate.
Figure 35 : Jet- regression loss against compute. The same curves are colored by model size (left) and by dataset size (right) ; circles mark the final per-cell validation loss.
Figure 36 : Track-origin loss against compute. The same curves are colored by model size (left) and by dataset size (right) ; circles mark the final per-cell validation loss.
Figure 37 : Track-type loss against compute. The same curves are colored by model size (left) and by dataset size (right) ; circles mark the final per-cell validation loss.
Figure 38 : Track-vertexing loss against compute. The same curves are colored by model size (left) and by dataset size (right) ; circles mark the final per-cell validation loss.
Figure 39 : Single-auxiliary-task ablation: compute-optimal jet-classification loss for the classification-only baseline, for classification plus each single auxiliary head, and for all five together, against compute (left) and A100 GPU-hours (right) . Bands are the seed-to-seed spread. The three track-level heads (origin, type, vertexing) each recover most of the full multi-task gain on their own, whereas jet- regression contributes least; per GPU-hour the ordering tightens as the throughput cost of the extra head is folded in.
Figure 40 : Training loss against optimizer step for every cell of the P sweep (rows: depth ; columns: width ), with one curve per peak learning rate (color). Anchors .
Figure 41 : Smoothed best training loss across the P sweep. Top: over peak and width at each depth ; bottom: over peak and depth at each width . Anchors .
Figure 42 : Fit of the depth exponent in ; after applying this multiplier the effective reduces to a single value across the studied range.
| optimization target | |||
|---|---|---|---|
| Total (multi-task) | |||
| jet classification | |||
| jet regression | |||
| track origin | |||
| track type | |||
| track vertexing |
Table 7 : Full per-head multiplicative calibration, extending Table 2 with the five individual multi-task heads. Same functional form and anchors ( , jets) as Table 2 .
Figure 43 : Per-task landscape for the jet-classification head: (left) the optimal learning rate with iso- contours from the multiplicative fit, and (right) the loss at that optimum, , with iso-loss contours; white circles mark the swept cells.
Figure 44 : Per-task landscape for the regression head: (left) the optimal learning rate with iso- contours from the multiplicative fit, and (right) the loss at that optimum, , with iso-loss contours; white circles mark the swept cells.
Figure 45 : Per-task landscape for the track-origin head: (left) the optimal learning rate with iso- contours from the multiplicative fit, and (right) the loss at that optimum, , with iso-loss contours; white circles mark the swept cells.
Figure 46 : Per-task landscape for the track-type head: (left) the optimal learning rate with iso- contours from the multiplicative fit, and (right) the loss at that optimum, , with iso-loss contours; white circles mark the swept cells.
Figure 47 : Per-task landscape for the track-vertexing head: (left) the optimal learning rate with iso- contours from the multiplicative fit, and (right) the loss at that optimum, , with iso-loss contours; white circles mark the swept cells.
Figure 48 : Forward-plus-backward FLOPs per jet measured with FlopCounterMode (solid) against width at fixed depth (left) and against depth at fixed width (right) , compared with the per-jet estimate of Eq. 1 (dashed). The measurement follows the estimate across the model family, confirming that is an adequate compute proxy for the architectures studied here.
Figure 49 : Background rejection at fixed -jet efficiency as a function of the discriminant fractions of Eq. 16 , for -jets (left) , light-jets (middle) , and -jets (right) , at global (top) and (bottom) , in GeV. The rejection varies by an order of magnitude or more across the plane and the three flavors trade off against one another. GN3 uses (red circle); the fractions maximizing a given rejection differ per flavor, and the optimal choice for an analysis is set by its background composition through the Neyman–Pearson lemma. The reported rejection is thus a property of the chosen projection.
Figure 50 : Background rejection against jet and in-bin -jet efficiency , for -jets (top) and light-jets (bottom) , for a , classification-only model trained at four settings: and jets with ( sw) and without ( sw) sample weights. The discriminant uses the default GN3 fractions . At fixed in-bin efficiency the rejection landscape is nearly identical with and without sample weights; the dashed line, a global over GeV, cuts across very different in-bin efficiencies and rejections.
Figure 51 : Background composition among non- jets of the test sample on the grid: (left) , (middle) , and (right) , with per-bin values annotated and the sample-global values in the titles. The composition is far from the fixed convention and varies strongly with kinematics, so the Neyman–Pearson-optimal fractions are per-jet, per-analysis quantities.
Figure 52 : Compute-optimal allocation for five settings: two training objectives (all-aux, jcls-only) and three reduced input sets, with bands from the fitted allocation exponents. Dotted diagonals are lines of constant compute at decades from to FLOP. The jet-classification-only setting is the clear outlier, preferring more data and a smaller model at every budget; the input level does not appear to shift the allocation.
Figure 53 : Jet-classification loss against compute for the GN2 input set.
Figure 54 : Jet-classification loss against compute for the no-flows input set.
Figure 55 : Jet-classification loss against compute for the four-momenta-only input set.
Figure 56 : Validation of joint-optimum transfer. (a) final validation loss against for combinations not used in the multiplicative fit. The star marks the predicted , which sits at or near every minimum, so transfers across batch size, token budget, and model shape. (b) loss penalty from using the predicted batch size rather than the lowest-loss in the sweep ( ), as a function of width (left sub-panel) and depth (right sub-panel) at three token budgets . The gray band marks the run-to-run spread, showing that the possible residual shift from the prediction can be neglected once is tuned.
Figure 57 : Tutorial scaling laws on the teacher–student toy problem: validation loss against data (left) and compute (right) , one curve per student size . Smaller students win at low compute and larger ones take over, with the envelope approaching the known floor (dotted).
Figure 58 : Cost of not re-tuning the learning rate on the teacher–student toy ( Appendix L ), whose floor is known exactly. Excess loss for the per-cell tuned (green) against a single learning rate tuned once at the small end and then held fixed (red), growing the data at fixed width (left) and the model size at fixed data (right) . The mis-tuned arm does not sit at a constant offset: 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.
Figure 59 : Double descent on the teacher–student toy problem: test loss (left) and train MSE (right) over model width and number of training samples . The test-loss ridge follows the interpolation threshold (dashed), where the train MSE drops to zero; beyond it the test loss second-descends.
Figure 60 : Cost of not re-tuning the learning rate on the compressed-representation system ( -way classification). (a) growing at fixed width : per-cell tuned (blue) against a learning rate tuned at M and then held fixed (red). (b) growing width at fixed . A grid trained this way refits to visibly different scaling-law parameters.