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

AI summary

Overview

Research area: Machine learning for particle physics, specifically knowledge distillation and efficient model design for jet tagging at the Large Hadron Collider.

Technical level: Intermediate. Readers should be comfortable with standard deep-learning concepts (transformers, distillation, parameter counts), but the paper explains the physics setup and the benchmark in accessible terms.

Scope: A controlled study of how six teacher models (including two pretrained foundation models) distil into five small students on the top-tagging benchmark, followed by width and depth scaling studies, a jet-by-jet analysis of what students inherit, and a direct benchmark of inference cost on GPU and CPU.

What This Paper Is About

Jet taggers at the LHC have grown more accurate but also more expensive to run, because experiments impose strict limits on inference cost and parameter count does not reveal how costly a model really is. The authors ask whether knowledge distillation can carry the performance of large, physics-aware and pretrained foundation models into tiny students that are cheap to evaluate, and how small those students can get while still competing with the state of the art. They also ask which property of a teacher—its raw strength, its symmetry, its pretraining, or its architecture—actually drives the benefit to the student.

Key Contributions

1. A systematic, controlled distillation study on the top-tagging benchmark that crosses six teachers (ParT, MIParT, L-GATr, L-GATr-slim, OmniLearned-S, and ParT fine-tuned after JetClass pretraining) with five students of roughly 200 thousand parameters each, using one fixed distillation setting (α = 0.7, T = 2) for all 30 pairs so that differences come only from the teacher–student pairing.
2. Two scaling studies that shrink the two strongest distilled students, Transformer and L-GATr-slim, in width (from about one thousand to 181 thousand parameters) and in depth, to locate where distillation stops helping and how small a competitive student can be.
3. A jet-level analysis of what the student learns: fidelity to the teacher measured by KL divergence, inherited false positives, mistakes removed versus added, and agreement as a function of teacher confidence.
4. A direct inference-cost benchmark reporting operations per jet, GPU throughput and peak memory at a batch of 512 jets, and single-jet CPU latency, showing that same-sized students can differ by almost an order of magnitude in speed; code and all trained weights are released.

Main Findings

• Every teacher helps every student: All 30 distilled students have higher AUC and higher background rejection at ε_S = 0.5 than the same student trained without a teacher. At ε_S = 0.3, every distilled student improves except MIParT distilled from MIParT.
• Teacher strength dominates: The mean gain in background rejection at ε_S = 0.5 rises from 9 percent with ParT as teacher to 40 percent with ParT-f.t. Neither shared architecture nor shared symmetry changes this ordering—both Lorentz-equivariant students learn more from ParT-f.t., which has no built-in symmetry, than from either Lorentz-equivariant teacher.
• Pretraining works mainly by making the teacher stronger: The two strongest teachers are both pretrained. As a control, the authors distil from an ensemble of three L-GATr-slim runs that matches OmniLearned-S in AUC but was trained from scratch; its students perform about as well as those of OmniLearned-S.
• Weak teachers still help, and can be overtaken: Every student distilled from ParT exceeds ParT in AUC and in both background rejections, even though ParT is weaker than the MIParT and L-GATr-slim students trained without a teacher. The authors attribute this to the student also learning from hard labels, with soft labels acting as a regularizer.
• A 20-thousand-parameter student beats all from-scratch taggers: Distilled from ParT-f.t., L-GATr-slim with 20 thousand parameters (AUC 0.98729, rejection 596 ± 38 at ε_S = 0.5 and 2,517 ± 122 at ε_S = 0.3) outperforms every published tagger trained from scratch, including the authors' own teachers, which have up to 150 times its parameters.
• A two-thousand-parameter student matches ParT: L-GATr-slim with 1.8 thousand parameters distilled from ParT-f.t. reaches AUC 0.98567 and rejections 416 ± 22 and 1,872 ± 153, matching the AUC of the 2.1M-parameter ParT teacher (0.98571) and exceeding its background rejections (394 ± 14 and 1,586 ± 19).
• Distillation has a floor for plain Transformers: With 1.2 thousand parameters, Transformer does not improve with any teacher, whereas L-GATr-slim at 1.8 thousand still gains from every teacher. Architecture choice matters most at the smallest sizes.
• Width and depth both compress well: A distilled student outperforms the next wider student of its architecture trained without a teacher, which has five to twelve times its parameters, down to 16 thousand parameters for Transformer and two thousand for L-GATr-slim. Distilled from any of the three teachers, a two-block Transformer outperforms a ten-block Transformer trained without a teacher.
• Students copy their teachers, mistakes included: Distillation lowers the KL divergence from teacher to student for all 30 pairs, by more than half on average. Distilled students repeat more of their own teacher's false positives than undistilled students, and when two teachers disagree the student tends to follow the one it was distilled from. Distillation nevertheless removes far more mistakes than it adds, which is where the gain in background rejection comes from.
• Cost is not predicted by parameter count: The five ~200-thousand-parameter students differ by almost an order of magnitude in GPU throughput. ParT and MIParT need about as many operations per jet as Transformer but run several times slower because they compute and store features for every pair of particles, a cost that grows with the square of the particle count.
• Pareto front: Distilled Transformer and L-GATr-slim students form almost the entire Pareto front of background rejection against inference cost. The strongest student, L-GATr-slim with 181 thousand parameters (28.6 MFLOP, 40,269 jets/s on GPU), matches the AUC of ParT-f.t. with a seventh of its operations and twice its throughput.

Methodology in Plain English

The authors take a standard top-tagging dataset of 1.2 million training, 0.4 million validation and 0.4 million test jets, where each jet is a list of up to 200 particles described by seven features. They train a set of large "teachers" and a set of small "students" on it, then train each student a second time while also asking it to match the teacher's predicted probabilities (soft labels) through a temperature-softened Kullback–Leibler loss mixed with the ordinary label loss. The mixing weight and temperature are fixed in advance at 0.7 and 2 for every pair, and teacher predictions are precomputed once so no teacher runs during student training.

To separate the effect of teacher strength from architecture, symmetry and pretraining, they choose teachers that span a range of measured performance and students that include both plain architectures and shrunken copies of the teachers. They then take the two best-performing distilled students and systematically shrink them—first reducing width at fixed depth of ten blocks, then reducing depth at fixed width—retraining each variant with and without a teacher.

Beyond accuracy, they compare the predicted probabilities of students and teachers jet by jet, split test jets into ten equal groups by how confident the teacher is, and count shared false positives. Finally, they time every model directly on an NVIDIA RTX 6000 Ada GPU and an Intel Xeon Platinum 8352Y CPU in full 32-bit precision, reporting operations per jet, throughput, peak memory, and single-core latency.

Why This Matters

This work gives the particle-physics community a practical recipe: a small, cheap tagger can inherit most of the performance of a large foundation model, provided the teacher is strong and the student architecture remains capable at small size. It also corrects a common inference in prior distillation work—that pretraining itself confers an advantage—by showing the effect is largely explained by the resulting teacher strength. Because the authors release code and all trained weights, the result is directly reusable and comparable.

Real-world applications:

• Trigger and real-time event selection at collider experiments, where inference must fit inside strict latency and compute budgets.
• Deployment on constrained hardware, such as CPU cores or memory-limited accelerators, where the paper shows large speed differences between equally sized models.
• Reproducible benchmarking of model efficiency, since the released cost protocol can be reused to compare future taggers at fixed cost rather than at fixed parameter count.
• Transfer of large pretrained scientific foundation models into compact deployable surrogates, a pattern that extends beyond particle physics to other domains with expensive simulation-based pretraining.

Industry relevance: the finding that parameter count poorly predicts latency—pairwise feature computation can make an equally sized model several times slower—is directly relevant to anyone deploying transformers under latency constraints, and the distillation recipe requires no teacher at inference time.

Future Directions

• Tuning the distilled students specifically for deployment, where optimized settings may improve them further beyond the fixed training recipes used here.
• Testing the approach on larger datasets such as JetClass, rather than the single binary top-tagging benchmark used in this study.
• Distilling larger foundation models, such as OmniLearned-L, into students like L-GATr-slim.
• Using interpretability studies of jet taggers to reveal what students learn from their teachers beyond matching predictions.
• Exploring distillation losses and hyperparameters beyond the single fixed setting (α = 0.7, T = 2) applied uniformly here.

Target Audience

Machine-learning researchers working on efficient models and distillation, particle physicists deploying taggers under real inference constraints, and practitioners interested in compressing large pretrained models into deployable students. The paper is most useful to readers who already know basic deep learning and want a concrete, reproducible case study with measured hardware costs rather than parameter counts alone.

Authors’ abstract

Jet tagging is one of the most important applications of machine learning at the Large Hadron Collider. Recent models build in more physics knowledge or rely on large-scale pretraining, but the experiments place strict limits on their inference cost. Knowledge distillation can transfer their performance to small students that are cheap to run. We ask which properties of a teacher help the student, how small a student can get and how close it can come to the state of the art. On a jet tagging benchmark, we distil six teachers that differ in architecture, symmetry and pretraining, including two foundation models, into five small students. Every teacher improves every student, and the gain depends mainly on the strength of the teacher. We then shrink our two strongest distilled students, Transformer and L-GATr-slim, in width and in depth. With 20 thousand parameters, a distilled L-GATr-slim outperforms every published jet tagger trained from scratch at a fraction of the inference cost of every teacher. Even with two thousand parameters, it matches ParT, a model with a thousand times as many parameters. We also study what the students learn from their teachers. Our benchmark of inference cost shows that students of the same size can differ by almost an order of magnitude in speed. We release the code and the weights of all trained models.