How to scale your HEP ML models: A recipe for robust architecture comparisons at scale

Matthias Vigl; Nikita Pond; Jackson Barr; Alexander Froch; Dan Guest; Nicole Hartman; Michael Kagan; Lukas Heinrich

hep-exOct 5, 2026

How to scale your HEP ML models: A recipe for robust architecture comparisons at scale

Authors: Matthias Vigl, Nikita Pond, Jackson Barr, Alexander Froch, Dan Guest, Nicole Hartman, Michael Kagan, Lukas Heinrich

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 $\sqrt{C}$ 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 $L(\eta,b)$ of the flavor tagger trained with all auxiliary tasks included, at fixed model size $N\sim 6.7$ M, for training-token budgets $T=4\times 10^{7}$ , $4\times 10^{8}$ , and $4\times 10^{10}$ . The dashed line traces the optimal learning rate $\eta^{\star}(b;T)$ at each batch size for a given token budget, and the red marker the joint optimum $(\eta^{\star}(T),b^{\star}(T))$ . 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 $(\eta^{\star},b^{\star})$ and at the compute-optimal allocation $(N^{\star},D^{\star})$ at every budget $C$ .

Figure 3 : Sketch of how empirical scaling laws emerge. (a) loss against dataset size

D

for a sequence of model sizes

N

: every curve follows a data-driven descent and saturates at a model-capacity plateau, lower for larger

N

, with all sufficiently large models approaching the irreducible floor

L_{\infty}

(dashed). (b) the same curves against compute

C=6NT\propto ND

: larger models shift rightward and the lower envelope (red) traces the compute-optimal frontier

L^{\star}(C)

Figure 4 : Optimal validation loss

L^{\star}_{\mathrm{cell}}(N,T)

at fixed

N

as a function of token budget

T=42.6\,D

. Each loss curve is a separate run with

(b,\eta)

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

N

yields the envelope targeted by scaling laws.

Figure 5 : Compute-optimal envelope

L^{\star}(C)

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

L_{\infty}+A\,C^{\gamma}

that recovers the known floor (

L_{\infty}=0.0101

against the true

\sigma^{2}=0.01

, 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 ( $\ell=128$ ) against training compute, colored by probe width $w$ (left) and by dataset size $D$ (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 $w=8$ 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 ( $\sim 10^{5}$ jets here), after which the allocation $(N^{\star},D^{\star})$ starts to matter.

Figure 7 : Compressed-representation scaling across latent dimensions

\ell=8

–

128

: test cross-entropy on the

10

-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 $\epsilon$	weight decay
Embedding (input)	$\eta_{\mathrm{base}}$	$\epsilon_{\mathrm{base}}$	$\lambda_{\mathrm{base}}$
Hidden weights (attn, FFN)	$\eta_{\mathrm{base}}\cdot m_{N}^{-1}$	$\epsilon_{\mathrm{base}}\cdot m_{l}^{-1}$	$\lambda_{\mathrm{base}}\cdot m_{N}$
Unembedding / readout	$\eta_{\mathrm{base}}$	$\epsilon_{\mathrm{base}}$	$\lambda_{\mathrm{base}}$
Biases, LN/RMSNorm gains	$\eta_{\mathrm{base}}$	$\epsilon_{\mathrm{base}}$	$\lambda_{\mathrm{base}}$
QK-norm gains	$\eta_{\mathrm{base}}\cdot m_{l}^{-1}$	$\epsilon_{\mathrm{base}}\cdot m_{l}^{-1}$	$\lambda_{\mathrm{base}}$

Table 1 : Per-parameter-type learning-rate (LR), AdamW-

\epsilon

and weight-decay multipliers of the Complete(d)P recipe used for every run in this work, relative to the base values

\eta_{\mathrm{base}}

\epsilon_{\mathrm{base}}

\lambda_{\mathrm{base}}

;

m_{N}

and

m_{l}

are the width and depth ratios to the base model. The learning-rate and weight-decay columns follow Table 1 of [ 61 ] at

\alpha=1

and the QK-norm row in Table 1 of [ 63 ] ; weight decay is the PyTorch AdamW convention, so

\lambda_{\mathrm{base}}\cdot m_{N}

keeps

\eta\lambda

width-independent.

Figure 8 : Empirical validation of the Complete(d)P recipe on the

(w,l,\eta)

coordinate sweep at

(b,T)=(500,2{\times}10^{7})

. (a) Final validation loss against the effective peak learning rate

\eta\,(l/4)^{0.46}

for seven widths at

l=4

and

l=8

(top) and five depths at

w=256

and

w=768

(bottom). The optimum lies at the same

\eta

for all widths, with the residual drift in

l

factored out. (b) Training loss against the training step for the

2{\times}2

slice

w\in\{128,1024\}

l\in\{4,8\}

, one curve per peak learning rate

\eta

optimization target	$\eta_{\mathrm{base}}\,[10^{-4}]$	$\kappa$	$\lambda$
Total (multi-task)	$6.05\pm 0.18$	$+0.429\pm 0.014$	$+0.098\pm 0.009$
Jet-classification-only	$2.91\pm 0.16$	$+0.548\pm 0.025$	$+0.174\pm 0.019$

Table 2 : Multiplicative calibration of the optimal learning rate

\eta^{\star}=\eta_{\mathrm{base}}\,(b/b_{\mathrm{ref}})^{\kappa}\,(T_{\mathrm{ref}}/T)^{\lambda}

( Eq. 11 ), with anchors

b_{\mathrm{ref}}=500

and

T_{\mathrm{ref}}=2\times 10^{7}

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

N_{\mathrm{ref}}

and

b=100

, one curve per dataset size

D

. The minimum of each curve is the extracted

\eta^{\star}(b,T)

Panel of Figure 10 — Figure 10 : The $(b,T)$ landscape at the reference architecture $N_{\mathrm{ref}}$ for the jet-classification-only variant. (a) the measured optimal learning rate $\eta^{\star}(b,T)$ , colored by $\log_{10}\eta^{\star}$ , with iso- $\eta^{\star}$ contours from the multiplicative fit of Eq. 11 . (b) the validation loss at that optimal learning rate, $\mathrm{loss}^{\star}(b,T)=\min_{\eta}L$ , with iso-loss contours. The loss is essentially flat in $b$ and steepens only in the few-step corner (large $b$ , small $T$ ), so $b^{\star}$ is left largely unconstrained. White circles mark the swept cells.

Panel of Figure 11 — Figure 11 : The $(b,T)$ landscape at the reference architecture $N_{\mathrm{ref}}$ for the multi-task variant. (a) the measured optimal learning rate $\eta^{\star}(b,T)$ , colored by $\log_{10}\eta^{\star}$ , with iso- $\eta^{\star}$ contours from the multiplicative fit of Eq. 11 . (b) the validation loss at that optimal learning rate, $\mathrm{loss}^{\star}(b,T)=\min_{\eta}L$ , with iso-loss contours. The loss now forms a shallow minimum in $b$ whose location grows as $b^{\star}\propto\sqrt{T}$ . White circles mark the swept cells.

	Batch Size	Learning Rate	Optimizer
This work	$b^{\star}\propto T^{0.45\pm 0.05}$	$\eta^{\star}\propto T^{+0.10\pm 0.02}$	AdamW [ 70 ]
DeepSeek [ 58 ]	$b^{\star}\propto T^{0.3271}$	$\eta^{\star}\propto T^{-0.1250}$	AdamW [ 70 ]
StepLaw [ 56 ]	$b^{\star}\propto T^{0.571}$	$\eta^{\star}\propto N^{-0.713}\,T^{0.307}$	AdamW [ 70 ]
von Rütte et al. [ 84 ]	$b^{\star}\propto T^{0.8225}$	$\eta^{\star}\propto T^{0.2806}$	LaProp [ 85 ]
Filatov et al. [ 86 ]	$b^{\star}\propto T^{0.45\pm 0.07}$	$\eta^{\star}\propto T^{-0.28\pm 0.07}$	Scion [ 87 ]

Table 3 : Fitted power-law exponents for optimal batch size and learning rate as a function of token budget

T

, as reported by studies with differing optimizers and tuning protocols. Adapted from [ 59 ] .

Figure 12 : Final validation loss against aspect ratio

\mathrm{AR}=w/l

for four model sizes,

N\approx 10

M, 30M, 50M and 100M, with each point trained at its own

(\eta^{\star},b^{\star})

. Large AR is a wide, shallow network and small AR a narrow, deep one. The minimum

\mathrm{AR}^{\star}(N)

drifts to larger AR as

N

grows, so width should grow faster than depth.

Figure 13 : Measured throughput on an H200 against the parameter count

N

, one color per depth

l

and one column per batch size

b

. Points are measurements of a full training step, and lines are linear fits in

N

at fixed

b

and

l

. (a) Time per jet. Throughput improves as

b

grows and amortizes the per-step overhead, leveling off once the device is compute-bound. (b) The same runs as GPU-hours per

10^{18}

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

1.7\times

more optimizer steps per GPU-hour.

Figure 15 : Onset of scaling on the multi-task loss variant. Jet classification loss against compute

C=6NT

. The portion of the compute-optimal frontier a few-hundred-million-jet dataset can probe (up to

C\sim 10^{17}

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

N

(left) and by dataset size

D

(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 $C=6N\,T$ , one lightly smoothed curve per run colored by $N$ . The black line is the lower envelope, the empirical compute-optimal frontier $L^{\star}(C)$ . (b,c) The model size $N^{\star}$ and dataset size $D^{\star}$ 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 $N$ 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 $L(N,D)$ , 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 $(N^{\star},D^{\star})$ that minimize the loss at each compute budget.

Method	$a$	$b$
This work (Approach 1)	$0.52\pm 0.01$	$0.48\pm 0.01$
This work (Approach 2)	$0.55\pm 0.02$	$0.45\pm 0.02$
This work (Approach 3)	$0.49\pm 0.03$	$0.51\pm 0.03$
Carpe Datum [ 42 ]	$0.10$	$0.90$
Kaplan et al. [ 24 ]	$0.73$	$0.27$
Hoffmann et al. [ 25 ]	$0.49$	$0.51$

Table 4: Compute-optimal allocation exponents, where

N^{\star}(C)\propto C^{a}

and

D^{\star}(C)\propto C^{b}

. All three approaches agree on

a\simeq b\simeq 0.5

, 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

\sqrt{C}

dependence of model and dataset size.

Figure 19 : (a) Best early-stopped validation loss of runs targeting $n$ passes for a $31.5$ M model on $U_{D}=10^{7}$ 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 $n$ , and the single-pass law $L(N,nU_{D})$ on fresh jets (dashed). (c) The same model on $10^{7}$ , $10^{8}$ and $10^{9}$ 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

n=1

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

10^{6}

10^{7}

and

10^{8}

unique jets, with an independent fit per unique set: a power-law descent plus a bump in

\ln N

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) $U_{D}=10^{7}$ and (b) $U_{D}=10^{6}$ unique jets, with the compute-optimal allocations of Eq. 15 (solid blue) and counting repeated jets as fresh (dashed), with their optima at $C=3\times 10^{17}$ and $3\times 10^{15}$ 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

N\sim 6.7

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

L^{\star}(C)

with

1\sigma

bands. The multi-task frontier lies below the single-task one across the studied range. (b) The corresponding compute-optimal allocation

N^{\star}(D^{\star})

with

1\sigma

bands. Dotted diagonals are lines of constant compute at decades from

C=10^{17}

10^{24}

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

L^{\star}(C)

with its

1\sigma

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 $L^{\star}(C)$ . The data-repeated TN25-86M run lies above the single-pass frontier.

Figure 27 : Light-jet (blue) and

c

-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

70\%

global

b

-jet efficiency,

20<p_{\text{T}}<250

GeV,

|\eta|<2.5

on a

t\bar{t}

sample.

Study	Runs	Jets seen	Training FLOPs	GPU-hours (A100)
$\mu$ P transfer	324	$6.5\times 10^{9}$	$2.8\times 10^{19}$	0.34k
$\eta^{\star}$ / $b^{\star}$ + aspect ratio	1395	$1.1\times 10^{11}$	$1.8\times 10^{20}$	4.5k
Compute-optimal grid (all variants)	1219	$1.5\times 10^{12}$	$8.4\times 10^{21}$	92k
Total	2938	$\mathbf{1.6\times 10^{12}}$	$\mathbf{8.6\times 10^{21}}$	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	$p_{\text{T}}$	Transverse momentum
	$\eta$	Pseudo-rapidity
Track	$q/p$	Track’s charge over momentum (curvature)
	$\Delta\eta$	Pseudorapidity relative to the jet $\eta$
	$\Delta\phi$	Azimuthal angle relative to the jet $\phi$
	$d_{0}$	Closest distance from the track to PV in the transverse plane

Appendix

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

b

(markers) and

\eta

(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

10^{6}

and the second on

10^{7}

unique jets.

Figure 29 : Allocation under repetition. (a) Contours of the loss of Eq. 15 at $U_{D}=10^{7}$ 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 $C=3\times 10^{17}$ 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 $U_{D}=10^{5}$ to $10^{9}$ , 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

U_{D}=10^{7}

unique jets with the onset exponent held at 0.5, 1 and 1.5 and the other parameters refitted. (a)

n_{\mathrm{crit}}

; (b) the terms of the law at its best target (the parameter term does not depend on

\gamma

); (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 $U_{D}=10^{7}$ unique jets with the onset exponent $\gamma$ 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

c

-, light- and

\tau

-jet rejection at the

70\%

b

-jet efficiency working point against the jet-classification loss. The fourth shows the median relative

p_{\text{T}}

resolution against the

p_{\text{T}}

-regression loss. Curves are fits of

R(L)=A\,e^{BL}+C

, with the shaded

68\%

interval from bootstrapping. The top axis converts loss into training compute through the compute-optimal envelope

L(C)=E+A\,C^{-\gamma}

Figure 33 : Light-jet rejection (blue) and

c

-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

L^{\star}(C)

uncertainties. Evaluated at

70\%

global

b

-jet efficiency,

20<p_{\text{T}}<250

GeV,

|\eta|<2.5

on a

t\bar{t}

sample.

Figure 34 : All trained multi-task runs: training loss against cumulative compute $C$ . The same curves are colored by model size $N$ (left) and by dataset size $D$ (right) ; circles mark the final per-cell validation loss. The shared lower envelope is the empirical compute-optimal frontier $L^{\star}(C)$ that Approaches 1–3 estimate.

Figure 35 : Jet- $p_{\text{T}}$ regression loss against compute. The same curves are colored by model size $N$ (left) and by dataset size $D$ (right) ; circles mark the final per-cell validation loss.

Figure 36 : Track-origin loss against compute. The same curves are colored by model size

N

(left) and by dataset size

D

(right) ; circles mark the final per-cell validation loss.

Figure 37 : Track-type loss against compute. The same curves are colored by model size

N

(left) and by dataset size

D

(right) ; circles mark the final per-cell validation loss.

Figure 38 : Track-vertexing loss against compute. The same curves are colored by model size

N

(left) and by dataset size

D

(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

C=6NT

(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-

p_{\text{T}}

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 $(w,l)$ cell of the $\mu$ P sweep (rows: depth $l$ ; columns: width $w$ ), with one curve per peak learning rate $\eta$ (color). Anchors $(b,T)=(500,2{\times}10^{7})$ .

Figure 41 : Smoothed best training loss across the $\mu$ P sweep. Top: over peak $\eta$ and width $w$ at each depth $l$ ; bottom: over peak $\eta$ and depth $l$ at each width $w$ . Anchors $(b,T)=(500,2{\times}10^{7})$ .

Figure 42 : Fit of the depth exponent

\gamma_{l}=0.464\pm 0.049

\eta^{\star}(l)=(l_{\mathrm{ref}}/l)^{\gamma_{l}}\,\eta^{\star}(l_{\mathrm{ref}})

; after applying this multiplier the effective

\eta^{\star}

reduces to a single value across the studied

l

range.

optimization target	$\eta_{\mathrm{base}}\,[10^{-4}]$	$\kappa_{\mathrm{eff}}$	$\lambda_{\mathrm{eff}}$
Total (multi-task)	$6.05\pm 0.18$	$+0.429\pm 0.014$	$+0.098\pm 0.009$
jet classification	$5.68\pm 0.20$	$+0.453\pm 0.014$	$+0.125\pm 0.010$
jet $p_{\text{T}}$ regression	$5.54\pm 0.26$	$+0.450\pm 0.021$	$+0.065\pm 0.014$
track origin	$5.91\pm 0.21$	$+0.430\pm 0.014$	$+0.090\pm 0.010$
track type	$5.81\pm 0.20$	$+0.438\pm 0.014$	$+0.125\pm 0.009$
track vertexing	$5.85\pm 0.23$	$+0.422\pm 0.017$	$+0.108\pm 0.011$

Appendix

Table 7 : Full per-head multiplicative

\eta^{\star}

calibration, extending Table 2 with the five individual multi-task heads. Same functional form and anchors (

b_{\mathrm{ref}}=500

T_{\mathrm{ref}}=2\times 10^{7}

jets) as Table 2 .

Figure 43 : Per-task $(b,T)$ landscape for the jet-classification head: (left) the optimal learning rate $\eta^{\star}(b,T)$ with iso- $\eta^{\star}$ contours from the multiplicative fit, and (right) the loss at that optimum, $\min_{\eta}L$ , with iso-loss contours; white circles mark the swept cells.

Figure 44 : Per-task $(b,T)$ landscape for the $p_{\text{T}}$ regression head: (left) the optimal learning rate $\eta^{\star}(b,T)$ with iso- $\eta^{\star}$ contours from the multiplicative fit, and (right) the loss at that optimum, $\min_{\eta}L$ , with iso-loss contours; white circles mark the swept cells.

Figure 45 : Per-task $(b,T)$ landscape for the track-origin head: (left) the optimal learning rate $\eta^{\star}(b,T)$ with iso- $\eta^{\star}$ contours from the multiplicative fit, and (right) the loss at that optimum, $\min_{\eta}L$ , with iso-loss contours; white circles mark the swept cells.

Figure 46 : Per-task $(b,T)$ landscape for the track-type head: (left) the optimal learning rate $\eta^{\star}(b,T)$ with iso- $\eta^{\star}$ contours from the multiplicative fit, and (right) the loss at that optimum, $\min_{\eta}L$ , with iso-loss contours; white circles mark the swept cells.

Figure 47 : Per-task $(b,T)$ landscape for the track-vertexing head: (left) the optimal learning rate $\eta^{\star}(b,T)$ with iso- $\eta^{\star}$ contours from the multiplicative fit, and (right) the loss at that optimum, $\min_{\eta}L$ , with iso-loss contours; white circles mark the swept cells.

Figure 48 : Forward-plus-backward FLOPs per jet measured with FlopCounterMode (solid) against width

w

at fixed depth (left) and against depth

l

at fixed width (right) , compared with the per-jet estimate

6\,N\,T_{\mathrm{jet}}\approx 256\,N

of Eq. 1 (dashed). The measurement follows the estimate across the model family, confirming that

C\approx 6\,N\,T

is an adequate compute proxy for the architectures studied here.

Figure 49 : Background rejection at fixed $b$ -jet efficiency as a function of the discriminant fractions $(f_{c},f_{\tau})$ of Eq. 16 , for $c$ -jets (left) , light-jets (middle) , and $\tau$ -jets (right) , at global $\epsilon_{b}=70\%$ (top) and $90\%$ (bottom) , in $p_{\text{T}}\in[25,250]$ 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 $(f_{c},f_{\tau})=(0.20,0.10)$ (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 $1/\epsilon$ against jet $p_{\text{T}}$ and in-bin $b$ -jet efficiency $\epsilon_{b}$ , for $c$ -jets (top) and light-jets (bottom) , for a $w{=}768$ , $l{=}4$ classification-only model trained at four settings: $10^{9}$ and $5\times 10^{8}$ jets with ( $+$ sw) and without ( $-$ sw) sample weights. The discriminant uses the default GN3 fractions $(f_{c},f_{\tau})=(0.20,0.10)$ . At fixed in-bin efficiency the rejection landscape is nearly identical with and without sample weights; the dashed line, a global $\epsilon_{b}=70\%$ over $p_{\text{T}}\in[25,250]$ GeV, cuts across very different in-bin efficiencies and rejections.

Figure 51 : Background composition among non-

b

jets of the

t\bar{t}

test sample on the

(p_{\text{T}},|\eta|)

grid:

f_{c}

(left) ,

f_{\tau}

(middle) , and

f_{l}

(right) , with per-bin values annotated and the sample-global values in the titles. The composition is far from the fixed

(0.20,0.10,0.70)

convention and varies strongly with kinematics, so the Neyman–Pearson-optimal fractions are per-jet, per-analysis quantities.

Figure 52 : Compute-optimal allocation

(N^{\star},D^{\star})

for five settings: two training objectives (all-aux, jcls-only) and three reduced input sets, with

1\sigma

bands from the fitted allocation exponents. Dotted diagonals are lines of constant compute at decades from

C=10^{17}

10^{24}

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.

Panel of Figure 56 — Figure 56 : Validation of joint-optimum transfer. (a) final validation loss against $\eta$ for $(b,T,w,l)$ combinations not used in the multiplicative fit. The star marks the predicted $\eta^{\star}$ , which sits at or near every minimum, so $\eta^{\star}$ transfers across batch size, token budget, and model shape. (b) loss penalty $\Delta L=L(b^{\star}_{\mathrm{pred}})-L_{\min}$ from using the predicted batch size rather than the lowest-loss $b$ in the sweep ( $L_{\min}$ ), as a function of width (left sub-panel) and depth (right sub-panel) at three token budgets $T$ . The gray band marks the run-to-run spread, showing that the possible residual $b^{\star}$ shift from the prediction can be neglected once $\eta^{\star}$ is tuned.

Figure 57 : Tutorial scaling laws on the teacher–student toy problem: validation loss against data $D$ (left) and compute $C$ (right) , one curve per student size $N$ . Smaller students win at low compute and larger ones take over, with the envelope approaching the known floor $L_{\infty}=\sigma^{2}=0.01$ (dotted).

Figure 58 : Cost of not re-tuning the learning rate on the teacher–student toy ( Appendix L ), whose floor

E

is known exactly. Excess loss

L-E

for the per-cell tuned

\eta^{\star}(w,T)

(green) against a single learning rate tuned once at the small end and then held fixed (red), growing the data

D

at fixed width (left) and the model size

N

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 $D$ . The test-loss ridge follows the interpolation threshold $N\sim D$ (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 (

10

-way classification). (a) growing

D

at fixed width

w=64

: per-cell tuned

\eta^{\star}(w,T)

(blue) against a learning rate tuned at

D=1

M and then held fixed (red). (b) growing width at fixed

D

. A grid trained this way refits to visibly different scaling-law parameters.

Explore similar work

Jul 25, 2026hep-ex

Predict before you train: Scaling Laws for particle physics foundation models

The largest machine learning models in particle physics are also the most expensive to train, yet the return on scaling a given architecture cannot be estimated before that compute is spent. Scaling laws have been fit for jets, but none has yet been shown to predict the performance of models it was not fit on. We show that, for a generic transformer pretrained on collider jets, it can be forecast. Fitting a joint model-and-data scaling law on small models alone, spanning three orders of magnitude of training compute, we predict the loss of models trained afterward with more than one hundred times more compute to within one percent. We then connect the forecast to downstream physics performance: across two standard tagging benchmarks, lower pretraining loss yields systematically lower fine-tuning loss and higher background rejection after fine-tuning. Within this model family and these tasks, a compute budget can therefore be translated into expected physics performance before any large model is trained. The final frontier model is consistent with the published numbers for current state-of-the-art physics-aware foundation models trained on the same corpus, on accuracy, AUC, and quark/gluon rejection, with a residual edge for the physics-aware model only in the high-purity tail of top tagging. We release five pretrained models spanning multiple sizes, together with the complete training recipe and code.

Jan-Lucas Uslu, Benjamin Nachman, Christopher Re

Department of Applied Physics, Stanford University, Stanford, CA 94305, USA · Fundamental Physics Directorate, SLAC National Accelerator Laboratory, Menlo Park, CA 94025, USA · Department of Physics, Stanford University, Stanford, CA 94305, USA +1

Model Size Pretraining Scaling Laws Physics-Aware Models

Ask PDF

Jun 18, 2026hep-ex

Towards Engineering Scaling Laws with Pretraining Data Composition

Neural scaling laws describe how model performance improves as a power law in compute, model size, and dataset size. While well-established for large language models, these relationships are emerging for large models in particle physics. As with language, empirical studies show that the performance scales as a power law. However, unlike natural language or image domains, fundamental physics has high-fidelity simulators that produce synthetic data cheaply. This favors scaling regimes where additional data is cheaper than additional parameters, and allows the pretraining dataset itself to be engineered to influence the scaling. For the task of classifying hadronic jets produced in collisions of high-energy particle beams, we show that the scaling behavior can be engineered towards requiring more data rather than larger models by inclusion of pretraining data which is more diverse and better aligned with the downstream classification task.

Jan-Lucas Uslu, Kevin Greif, Daniel Whiteson +1

Department of Applied Physics Stanford University Stanford, CA 94305, USA · Department of Physics and Astronomy University of California, Irvine Irvine, CA 92697, USA · Department of Particle Physics and Astrophysics SLAC National Accelerator Laboratory Stanford University Menlo Park, CA 94025, USA

Scaling Laws Pretraining Scaling Synthetic Data

Ask PDF

May 27, 2026hep-ph

Neural Scaling Laws for Jet Generation

Recently observed empirical scaling laws describe the performance of foundation-type models as three independent key quantities -- dataset size, compute, and model parameters -- are modified. Extracting these scaling laws informs the training of large complex models for which the tuning of hyperparameters in traditional ways is not feasible. This work for the first time explores if scaling laws can also be observed for the task of particle jet generation -- both relevant as a pre-training objective for foundation models and as in-situ simulation by itself. We indeed replicate the key logarithmic scaling law behavior for model-size scaling. Beyond studying the next token prediction validation loss of the generative model, we also study the sliced Wasserstein distance of five physical quantities that are not immediately available to the model during training. Our study shows that this quantity is monotonically related to the next token prediction validation loss, meaning that this loss is indeed a good proxy for the physics performance. For the scaling with dataset size and compute, we observe substantially weaker scaling behavior of both the loss and the sliced Wasserstein distance. We analyze this behavior by introducing the concept of a learnable window, and argue that autoregressive next token prediction on jet constituents exhibits comparatively rapid saturation relative to language-model studies. We discuss possible origins of this behavior, including the stochastic nature of QCD radiation and differences between generative and supervised learning tasks in collider physics.

Oz Amram, Darius A. Faroughy, Tjarko Gerdes +5

1Fermi National Accelerator Laboratory, Batavia, IL 60510, USA · 2NHETC, Dept. of Physics and Astronomy, Rutgers University, Piscataway, NJ 08854, USA · Institute for Experimental Physics, Universität Hamburg, Luruper Chaussee 149, 22761 Hamburg, Germany +1

Scaling Laws Scaling Autoregressive Token Prediction Downstream Fine-Tuning

Ask PDF