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

What it claims

The paper's central methodological point is that scaling laws are only trustworthy if learning rate η and batch size b are re-tuned at every budget. At fixed model size, the optimal (η, b) moves systematically as the training budget grows. If you fix hyperparameters at one reference budget, that drift gets folded into the apparent scaling. The authors show the joint optimum (η⋆, b⋆) follows its own scaling law in the token budget T. When every run is trained at its predicted optimum and at the compute-optimal (N⋆, D⋆), loss against compute becomes a clean power law. The resulting frontier is then a property of the architecture, not of tuning luck, and only that frontier is used to compare design choices.

Two findings are absent from the summary:

• Lower-level inputs shift the curve down but not its slope. Feeding the tagger lower-level inputs lowers the loss systematically while leaving the scaling exponent nearly unchanged.
• Small datasets do not predict large-scale behaviour. The power-law regime only begins after a few hundred million jets. Below that threshold, loss carries little information about high-compute scaling, so small-data pilot studies can mislead.

How it works

• Model. The tagger is an ATLAS GN-family transformer, specifically GN3EPCLV01 without the charge task, trained in the salt framework.
• Inputs are tracks, particle-flow objects, and soft electron/muon information.
• Jet features are broadcast to each object and passed through per-modality init networks into a shared encoder.
• Attention pooling produces the jet representation, which feeds the task heads.
• Training objective. Five losses are summed with equal weight:
• jet-flavour cross-entropy, with class weights b=1.00, c=3.64, s=5.00, ud=1.82, g=1.54, τ=9.09;
• pairwise track-vertex binary cross-entropy;
• track origin;
• track type;
• an L1 loss on the true/reconstructed jet pT ratio.
• Variants. Single-task, multi-task, and pairwise variants are each hyperparameter-optimised independently.
• Tuned axes. Four axes are tuned per scale: width, depth, b and η. Width and depth use a transferable parameterisation. The head dimension is fixed at 32, and the optimal aspect ratio is fitted.
• Fixed settings.
• Optimiser: AdamW.
• Learning-rate schedule: linear warmup from η/100 over 8% of steps, then cosine decay to η/10.
• The authors note that a WSD-style schedule would make budget sweeps cheaper.
• Compute accounting. The paper uses C = 6NT ≈ 256·N·D, with T ≈ 42.6·D because the mean is 42.6 ± 12.7 constituents per jet. The paper reports this matches measured FLOPs across the model family.
• Fitting. The frontier is fitted as L⋆(C) = L∞ + (Cc/C)^γ, with N⋆ ∝ C^a and D⋆ ∝ C^(1−a). The authors also discuss the full Chinchilla form L(N,D).
• Validation. The procedure is first checked on a teacher–student toy problem with a known loss floor and on MLP probes over compressed jets.

Evidence

• Optimum drift (Fig. 1). At N ≈ 6.7M with all auxiliary tasks, the minimum total loss is:
• 2.471 at T = 4×10⁷;
• 2.191 at T = 4×10⁸;
• 1.910 at T = 4×10¹⁰.

As the budget grows, the joint optimum moves toward larger b and η.

• Frontier (Fig. 2). The well-behaved frontier spans roughly 10¹⁶–10²² FLOPs and models of about 10⁵–10⁷ parameters.
• Comparison setup. Training uses the same ~7.7B-jet mixture as the earlier TN25-86M model, so results can be compared against it.
• Missing numbers. The exponent values, auxiliary-loss gains and background-rejection projections are not in the fetched text.

How to use it

• Training code (salt): https://gitlab.cern.ch/aft/algorithms/salt
• JetSet2 open data (~11B jets): https://gitlab.cern.ch/atlas/open-data/jetset2-open-data

The practical recipe for your own scaling study has four steps:

1. Fit (η⋆(T), b⋆(T)) on small runs.
2. Transfer it across width and depth.
3. Train every frontier point at its predicted optimum.
4. Compare architectures, objectives or inputs only along their own compute-optimal frontiers.

Limits

• Auxiliary losses cost compute. They add training cost, so their benefit has to be judged at equal compute, not at an equal number of samples.
• Repeated data needs separate treatment. The main analysis assumes non-repeated data. Data repetition is handled separately under early stopping in Section 7.
• The loss floor is not the physics limit. L∞ depends on architecture, inputs and optimiser, and is bounded below by the Bayes risk. It is not the physics limit.
• Truncated page. The fetched text stops in Section 3. The fitted exponents, the data-constrained predictions, the design-choice comparisons and the rejection projections are not shown here and should be checked in the full paper.