How Far Can Foundation Models Be Distilled for Efficient Jet Tagging?

1 Introduction

• LHC experiments impose strict inference-cost requirements, while pairwise interactions, equivariant layers, and large pretrained models increase evaluation cost.Parameter count alone does not show how much a model’s inference costs.Source: §1 · p. 1
• Earlier studies could not separate teacher strength from pretraining and left open whether distillation could approach the strongest taggers.They used one or two teachers and simple students such as Deep Sets or an MLP.Source: §1 · pp. 1–2
• The study compares six teachers spanning strength, architecture, symmetry, and pretraining against five roughly 200-thousand-parameter students under a shared distillation setting.It then scales two strong students in width and depth, examines student-teacher agreement and shared mistakes, and benchmarks inference cost.Source: §1 · p. 2
• Every teacher improves every student, with the gain depending mainly on teacher strength rather than architecture or pretraining.A 20-thousand-parameter distilled L-GATr-slim outperforms every published tagger trained from scratch at a fraction of every teacher’s inference cost; at 2,000 parameters, it matches 2.1-million-parameter ParT.Source: §1 · p. 2

Jet tagging

• Jet taggers range from particle-set models to pair-interaction transformers and Lorentz-equivariant networks.ParT adds pair features to attention scores [35], while L-GATr and LLoCa use geometric algebra [37] and local reference frames [38]; L-GATr-slim retains scalar and vector channels at lower cost [33].Source: §2 · p. 2

Foundation models and inference cost

• Pretraining improves fine-tuned tagging, but the inference cost of foundation models has received less attention than their performance and data efficiency.Prior work also studies fixed-cost comparisons and shrinking or quantizing models [17] [33].Source: §2 · p. 2

Knowledge distillation

• Distillation transfers teacher predictions to a student through soft labels, but a much larger teacher can be difficult to follow and students can inherit teacher errors.Soft labels encode teacher confidence beyond the hard class label [24]; prior findings describe imperfect reproduction and inherited errors.Source: §2 · pp. 2–3

Distillation in particle physics

• Particle-physics distillation studies include LorentzNet, graph-to-particle, anomaly-detection, and OmniLearned compression, but the closest study confounds pretraining with teacher strength.The present study compares six teachers and five students and reports that strength accounts for most of the observed effect; its students also outperform Deep Sets + GNN at a fraction of its operations per jet [5].Source: §2 · p. 3

3 Experimental setup

• The top-tagging dataset contains 1.2 million training, 0.4 million validation, and 0.4 million test jets, with up to 200 particles per jet.Scratch-trained teachers and all students use seven features per particle, while the two pretrained teachers retain their pretraining inputs.Source: §3 · p. 3
• Performance is measured by AUC and background rejection at signal efficiencies of 0.5 and 0.3 on the full test set.Background rejection is 1/ϵB; the paper defines teacher strength using these metrics.Source: §3 · p. 3
• Six teachers span different strengths and include four trained from scratch and two pretrained foundation models.They include ParT, MIParT, L-GATr, L-GATr-slim, OmniLearned-S, and ParT fine-tuned after JetClass pretraining.Source: §3 · p. 3
• Five students have about 200 thousand parameters each, roughly one-tenth the size of a typical teacher.The set includes a standard Transformer, LorentzNet, and smaller ParT, MIParT, and L-GATr-slim models.Source: §3 · pp. 3–4

Appendix A.2 gives the configuration of every model.

• Distilled students learn from both hard labels and teacher probabilities softened by temperature, using fixed α = 0.7 and T = 2 for every teacher-student pair.Teacher logits are computed once and stored, so the teacher does not run during student training.Source: Appendix A.2 · p. 4
• Students keep the checkpoint with the lowest validation binary cross-entropy, and each Section 4 student is trained three times.Training recipes follow the literature and are documented with their sources.Source: Appendix A.2 · p. 4

Gain for every student

• All 30 distilled students improve AUC and background rejection at ϵS = 0.5; at ϵS = 0.3, all improve except MIParT distilled from MIParT.Figure 1 shows background rejection at both signal efficiencies for students trained from scratch and distilled from each teacher.Source: §4 · p. 4, Figure 1 · p. 5

Outperforming a larger model

• With a strong teacher, every student surpasses scratch-trained L-GATr-slim, and Transformer, LorentzNet, and L-GATr-slim distilled from ParT-f.t. also surpass the ten-times-larger L-GATr-slim teacher.Most students distilled from the strongest teachers also outperform the L-GATr teacher, which has five times their parameters.Source: §4 · p. 5

Effect of teacher strength

• Mean background-rejection gain at ϵS = 0.5 rises from 9 percent with ParT to 40 percent with ParT-f.t., as stronger teachers yield larger gains.The teacher-strength ordering holds across students, regardless of shared architecture or symmetry.Source: §4 · p. 5, Figure 2 · p. 6

Effect of pretraining

• Students distilled from a scratch-trained L-GATr-slim ensemble matched to OmniLearned-S perform about as well as OmniLearned-S students, indicating pretraining helps mainly by strengthening the teacher.The ensemble students outperform those distilled from a single L-GATr-slim run.Source: §4 · p. 6

Weak teachers

• Every student distilled from ParT exceeds ParT in AUC and both background rejections, while three of five students distilled from MIParT exceed MIParT in AUC.Students also learn from hard labels; the paper notes that soft labels regularize them much like label smoothing.Source: §4 · p. 6

Effect of width

• A distilled student beats the next wider scratch-trained student of its architecture, with five to twelve times as many parameters, down to 16 thousand for Transformer and two thousand for L-GATr-slim.The comparison holds in AUC and both background rejections.Source: §5 · p. 7

Effect of depth

• Distilled two-block Transformer beats scratch-trained ten-block Transformer; L-GATr-slim improves with depth on its own, while stronger OmniLearned-S and ParT-f.t. keep improving it at every depth.Figure 4 plots background rejection at ϵS = 0.3 across block counts, with scratch-trained and distilled students shown separately.Source: §5 · p. 8, Figure 4 · p. 8

Beyond the state of the art

• With 20 thousand parameters, L-GATr-slim distilled from ParT-f.t. outperforms every tagger trained from scratch, including published taggers with up to 150 times its parameters.L-GATr and OmniLearned-S also help at this size, raising background rejection of both approximately 20-thousand-parameter students by about half.Source: §5 · p. 8

Architecture of small students

• At two thousand parameters, scratch-trained L-GATr-slim performs about as well as 84-thousand-parameter Transformer, and architecture matters most for the smallest students.At 200 thousand parameters, distillation from a strong teacher closes the gap between the architectures; at the smallest size, only L-GATr-slim gains.Source: §5 · p. 8

Limits of distillation

• At one thousand parameters, Transformer gains nothing from distillation because capacity limits what it can learn; two-thousand-parameter L-GATr-slim still improves with every teacher.Distilled from ParT-f.t., the small L-GATr-slim matches ParT's AUC and exceeds its background rejections despite ParT having over a thousand times as many parameters.Source: §5 · p. 8

Agreement with the teacher

• Distillation lowers teacher-to-student KL divergence for all 30 pairs by more than half on average, but the reduction shrinks below 16 thousand parameters.At one thousand parameters Transformer moves no closer to its teacher, whereas two-thousand-parameter L-GATr-slim moves slightly closer.Source: §6 · p. 9, Figure 5 · p. 9

Mistakes inherited from the teacher

• For all 30 teacher-student pairs, distilled students repeat more teacher false positives than scratch-trained counterparts, especially on background jets accepted by their own teacher.When teachers disagree on a background jet, the student is more likely to accept it if its teacher accepts it.Source: §6 · p. 10

Mistakes removed by distillation

• Students remove more background mistakes on teacher-rejected jets than they add on teacher-accepted jets, accounting for their improved background rejection.The student learns from jets its teacher classifies correctly.Source: §6 · p. 10

Agreement by teacher confidence

• Distillation moves students toward teachers across confidence groups, but pulls them toward errors in the least-confident top-jet group, where about half are classified as background by the teacher.In the most-confident groups, every student already agrees with the teacher.Source: Figure 5 · p. 9, §6 · p. 10

Protocol

• The benchmark measures GPU throughput and peak memory at batch size 512, CPU single-jet latency on one core, and operations per jet in the fastest 32-bit setting.Timing includes feature construction and the network, with measurements repeated across three runs.Source: §7 · p. 10, Table 5 · p. 11

Operations per jet

• A 181-thousand-parameter L-GATr-slim performs as well as ParT-f.t. with a seventh of its operations, while similarly sized Transformer is slightly weaker and uses fewer operations.Both a 20-thousand-parameter L-GATr-slim and a 188-thousand-parameter Transformer outperform every scratch-trained teacher with a fraction of their operations.Source: §7 · p. 10, Table 5 · p. 11

Cost of pair features

• Five students with about 200 thousand parameters each differ by almost an order of magnitude in GPU throughput, largely because pair-feature models compute and store features for every particle pair.That pairwise cost grows quadratically with particle count; ParT and MIParT run several times slower than Transformer despite similar operations per jet, while LorentzNet and Deep Sets + GNN need more memory.Source: §7 · p. 10, Table 5 · p. 11

Performance at fixed cost

• Distilled Transformer and L-GATr-slim form almost all of the performance-versus-cost Pareto front, and distilled L-GATr-slim outperforms scratch-trained state-of-the-art taggers at a fraction of their cost.The strongest student matches ParT-f.t.'s AUC at twice its throughput; 20-thousand-parameter L-GATr-slim distilled from OmniLearned-S rejects more background than its teacher and processes about twenty times as many jets per second.Source: §7 · p. 11, Figure 6 · p. 12

Latency on the CPU

• On a one-jet-at-a-time CPU, student latency hardly depends on model size, so distillation yields less speed gain than on the GPU.Deep Sets + GNN is fastest on CPU, followed by Transformer with one and 16 thousand parameters.Source: §7 · p. 11, Figure 6 · p. 12

8 Discussion

• Distillation improves students down to 2,000 parameters, and its gain follows teacher strength; scratch-trained teachers teach about as well as similarly strong pretrained teachers.At 20,000 parameters, distilled L-GATr-slim outperforms every published scratch-trained tagger at a fraction of every teacher's inference cost; at 2,000 parameters it matches ParT's AUC with one-thousandth the parameters and exceeds its background rejection.Source: §8 · p. 13
• The study is limited to one binary benchmark and one distillation loss, using the same settings for every teacher-student pair.The authors identify deployment tuning, larger datasets, larger foundation-model teachers, and interpretability studies as next steps.Source: §8 · p. 13

AI use statement

• The authors used AI tools for code writing and debugging, experiment runs, and minor grammar and sentence edits, and reviewed the assisted work.They state that they take responsibility for the final content.Source: AI use statement · p. 13

A.1 Inputs

• Scratch-trained models use seven per-particle features relative to the jet axis, while the two pretrained teachers retain their pretraining inputs and use different particle and feature selections.ParT-f.t. uses the seven features and pair features for its 128 highest-transverse-momentum particles; OmniLearned-S uses four features for up to 150 particles. Lorentz-equivariant models also receive reference directions.Source: §A.1 · p. 17

Architectures

• All students are smaller versions of published architectures, and models read out jets by particle mean, learned class tokens, or an L-GATr global token.ParT-f.t. uses ParT's blocks after training on JetClass; MIParT attends over pair features in initial blocks, while OmniLearned-S has eight transformer blocks and two classification-head blocks.Source: §A.2 · p. 17, Table 6 · p. 18

ParT

• The ParT student uses half the width and heads and four of ParT's ten blocks, leaving about a tenth of its parameters while reaching ParT's AUC within run-to-run spread.It retains ParT's per-particle and pair features and reads out with a particle mean.Source: §A.2 · p. 17

Transformer

• Transformer in the width study has ten blocks and 48 channels, while its depth-study zero-block variant embeds each particle linearly and averages across particles.The width-study configuration uses six heads, a twice-width MLP, and no dropout.Source: §A.2 · p. 17

B.1 Recipes

• Models use their architectures' published training recipes, with gradient clipping at norm 1 and 32-bit precision for all students.The width study uses each size's published recipe; depth variants use their architecture's recipe.Source: §B.1 · p. 18, Table 7 · p. 18

B.2 Runs and seeds

• Section 4 and width-study students use three seeds, whereas each depth-study cell has one run; the 248 student runs took about 1,500 GPU hours.Some students are distilled from teacher run one, while ParT, L-GATr, and L-GATr-slim student run k is paired with teacher run k.Source: §B.2 · p. 19

B.3 Teacher ensembles

• Averaging three teacher runs' logits raises every teacher's AUC, with larger gains for scratch-trained teachers than for the two pretrained teachers.The L-GATr-slim ensemble matches OmniLearned-S in AUC and rejects more background; the table compares their students.Source: §B.3 · p. 19, Table 9 · p. 19, Table 10 · p. 20

C Background rejection of every student

• The student comparison table organizes background rejection by student columns and teacher-training condition rows, with results at signal efficiencies 0.5 and 0.3.It reports means and standard deviations over three runs.Source: Table 11 · p. 20

D Scaling studies

• The scaling studies examine background rejection as student width decreases and as block count varies, comparing undistilled students with students distilled from teachers.The width study uses ten-block students; the depth study fixes width. The displayed captions specify three-run width means and one-run depth results.Source: D Scaling studies · p. 20, Figure 7 · p. 21, Table 12 · p. 21

E What every student learns from its teacher

• The student-teacher analyses compare accepted background jets and student probabilities against teacher confidence, separating distilled from undistilled students.The captions specify comparisons involving teacher-rejected or teacher-accepted jets, and probability differences across confidence groups; they do not state the plotted outcomes.Source: Table 15 · p. 24, Figure 10 · p. 25, Figure 11 · p. 25

F Inference benchmark

• Figure 16 compares operations per jet, GPU throughput and peak memory at batch 512, and single-jet CPU latency against parameter count on real test jets.Measurements use each model's fastest 32-bit setting.Source: Figure 16 · p. 30

F.1 Execution settings

• Models are timed in eager and compiled modes with supported particle-list formats and attention kernels; the fastest setting is reported per device.The smallest models are fastest with padded lists, which can use more memory than packed lists and exceed larger models' peak memory.Source: §F.1 · p. 30, Table 16 · p. 31

F.2 Precision

• Full 32-bit precision is used; TF32 leaves Transformer outputs unchanged but reduces L-GATr and L-GATr-slim performance, especially for small L-GATr-slim students.The networks' Minkowski products can involve small differences of large four-momenta, making rounding consequential; Transformer uses order-one features and is unaffected.Source: §F.2 · pp. 30–31

F.3 Jets with 50 particles

• Table 17 measures inference cost on jets with exactly 50 particles and random momenta, so cost no longer depends on the number of particles in test jets.GPU throughput and peak memory use batches of 512; CPU latency is measured for one jet on one core.Source: §F.3 · p. 31, Table 17 · p. 32

• A student's inference cost is the same with every teacher, while its performance changes with the teacher.Source: G Performance and inference cost with… · p. 32
• Table 18 reports performance and inference cost for students trained without a teacher.Source: Table 18 · p. 32
• Tables 19–23 report performance and inference cost for students distilled from ParT, MIParT, L-GATr, L-GATr-slim, and OmniLearned-S, alongside each teacher.Source: Table 19 · p. 33, Table 20 · p. 33, Table 21 · p. 33