Pith. sign in

REVIEW 2 major objections 5 minor 2 cited by

This paper claims that feeding jets to a transformer as continuous vectors while keeping token IDs as prediction targets, and pre-training jointly with next-token and masked-token objectives, greatly improves downstream jet classification w

Reviewed by Pith at T0; open to challenge. T0 means a machine referee read the full paper against a public rubric. the ladder, T0–T4 →

Challenge this review Re-run · record.json Download PDF Read on arXiv ↗

T0 review · deepseek-v4-flash

2026-08-03 18:37 UTC pith:DLENBWWV

load-bearing objection Solid extension of OmniJet-α with real classification gains, but the 'no loss in generation' claim needs a quantitative check before I'd cite it. the 2 major comments →

arxiv 2512.04149 v2 pith:DLENBWWV submitted 2025-12-03 hep-ph cs.LGhep-exphysics.data-an

Enhancing next token prediction based pre-training for jet foundation models

Joschka Birk , Anna Hallin , Gregor Kasieczka , Nikol Madzharova , Ian Pang , David Shih This is my paper
classification hep-ph cs.LGhep-exphysics.data-an
keywords jet foundation modelsnext token predictionmasked particle modelingVQ-VAE tokenizationhybrid inputtransfer learningjet classificationgenerative model
verification ladder T0 review T1 audit T2 compute T3 formal T4 reserved

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

This paper addresses a trade-off in simulation-free jet foundation models: the original next-token-prediction design generates realistic jets and transfers across datasets, but classifies worse than other pretraining methods. The authors argue that two design choices cause the gap: feeding the model only discrete token IDs, and training with a purely causal next-token objective. They propose a hybrid setup in which the model receives continuous feature vectors as input while keeping token IDs only as the prediction target, and a joint pretraining objective that adds masked particle modeling to next token prediction. They report that these changes substantially improve downstream jet classification and background rejection, on both in-distribution and transfer tasks, without visible degradation of the generated jets. A controlled causal variant indicates that the classification gain comes largely from bidirectional attention rather than the masked objective alone.

Core claim

On the paper's own terms, the central discovery is that two modifications to autoregressive next-token prediction for jets yield the best of both worlds. First, using continuous (or pseudo-continuous) feature vectors as the transformer input while keeping token IDs only as the next-token prediction target removes the tokenization bottleneck from downstream classification without sacrificing the generative pre-training task. Second, combining next token prediction (NTP) with masked particle modeling (MPM) during pre-training produces a representation that classifies almost as well as pure MPM while remaining generative. A causal-attention version of MPM performs like NTP, not like bidirection

What carries the argument

The load-bearing mechanism is the hybrid input-target split: a VQ-VAE compresses particle features into token IDs, but the transformer receives the decoded continuous feature vectors while the token IDs remain only the prediction targets for the generative head. This preserves the generative objective while letting the classifier operate on full-resolution features. The second piece is a joint pre-training objective with two heads — one for next-token prediction with causal attention and one for masked-token prediction with bidirectional attention — whose losses are simply added. A causal-attention MPM variant (MPM-Causal) is used as a control to separate the effect of the masking objective

Load-bearing premise

The 'no loss in generative performance' claim holds only if comparing generated jets with the VQ-VAE's tokenized-and-decoded target is a sufficient fidelity test; if the VQ-VAE's known reconstruction errors (for example on trajectory displacement) hide a real drop from adding the masked objective, the trade-off claim weakens.

What would settle it

Generate jets from the NTP-only and joint NTP+MPM checkpoints and compare them directly to raw, never-tokenized JetClass jets using a quantitative distributional metric (e.g., Wasserstein distance or a trained classifier) on particle pT, eta, phi, mass, and d0; if the joint model is significantly worse on these raw features while both match the tokenized-decoded target, the claim of no generative loss is refuted.

Watch this falsifier — get emailed when new claim-graph text bears on it.

If this is right

  • Continuous inputs instead of token-ID inputs improve classification at every training-set size tested, and pre-training still helps on small datasets, though gains vanish at 100 million jets for in-distribution classification.
  • Joint NTP+MPM pre-training nearly matches pure MPM on classification accuracy and background rejection while retaining the generative fidelity of NTP-only training.
  • Because MPM-Causal performs like NTP rather than bidirectional MPM, the classification advantage of MPM is attributed mainly to bidirectional attention, not to the masking objective itself.
  • The pre-training pipeline remains simulation-free and label-free, so the improved classifier can still be trained on real LHC data and transferred to different datasets with minimal fine-tuning.
  • Extending the feature set with mass, charge, particle ID, and trajectory displacement further improves classification accuracy without visibly harming generated distributions.

Where Pith is reading between the lines

These are editorial extensions of the paper, not claims the author makes directly.

  • This result suggests a general principle for tokenized point-cloud models: put the compression on the target side of the loss, not on the input side, so that downstream tasks keep full-resolution information.
  • A natural next experiment is a single bidirectional transformer trained with MPM plus a generative decoding head, which might achieve both capabilities without running two backbone forward passes during pre-training.
  • Since the paper assesses generative fidelity only by visual histogram agreement against a tokenized-and-decoded target, a quantitative comparison of raw, never-tokenized events would be the decisive test of the 'no loss' half of the claim.
  • The MPM-Causal control points toward attention directionality as the key design axis for future jet foundation models, more than the specific choice of prediction target.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit.

Referee Report

2 major / 5 minor

Summary. The paper proposes two modifications to the OmniJet-α jet foundation model: (i) a hybrid input representation in which continuous particle feature vectors are fed to the transformer while token-IDs are retained only as next-token prediction targets, and (ii) a joint pre-training objective combining next token prediction (NTP) with masked particle modeling (MPM). The authors pre-train on 100M JetClass jets and evaluate by fine-tuning on JetClass 10-class classification and on the top-tagging dataset, including out-of-distribution transfer. They report large classification gains from the hybrid input and from adding MPM, while claiming no degradation in generative performance, and they decompose the MPM benefit through a causal-MPM ablation. The classifier experiments use five seeds and controlled comparisons; the generative evaluation is visual only, against the VQ-VAE-decoded target.

Significance. The contribution is potentially significant for simulation-free jet foundation models: if a single pre-trained model can both generate and provide a representation that fine-tunes to near-MPM classification accuracy, that would address a central limitation of the original OmniJet-α. The classification comparisons are well designed, and the use of five seeds, OOD transfer to top tagging, and a controlled MPM-Causal variant to separate objective from attention directionality are strengths. The paper also makes code available. However, the generative half of the central claim is currently supported only by side-by-side histogram inspection against a lossy tokenized-decoded reference, which is not sufficient for a quantitative 'no loss' claim. Because the joint NTP+MPM model is the main proposal, this gap is load-bearing; the paper needs a quantitative generative metric or a more qualified claim.

major comments (2)
  1. [III.B, Fig. 6; also Abstract and Conclusion] The claim 'without any loss in generative performance' is not supported quantitatively. The only evidence is normalized histograms of generated samples against the VQ-VAE-decoded JetClass target, with no error bars, no divergence measure, no held-out next-token loss/perplexity, and no multiple seeds for generation. Since the joint model's value proposition is exactly that it retains NTP-level generation while improving classification, a small but real degradation (e.g. in tail behavior or next-token perplexity) would be invisible. Please add quantitative generative metrics, such as held-out token-level negative log-likelihood/perplexity, JSD/Wasserstein distances on the plotted observables with bootstrap uncertainties, or a two-sample test between NTP and NTP+MPM samples, and report them for both models. At minimum, the abstract should be qualified until such evidence is provided.
  2. [Appendix C, Fig. 8] The no-loss generative claim is only tested for the kinematics-only feature set. The extended-feature generative model uses NTP only; no joint NTP+MPM model is trained or evaluated with the extended feature set. If the claim is meant to cover the extended-feature setup (as the abstract suggests), this is a gap. Either provide a joint extended-feature generative evaluation or restrict the claim to the kinematic-only configuration.
minor comments (5)
  1. [III.B, Fig. 5] The precise accuracy and R30 values corresponding to the curves in Fig. 5 are not given in the text; the inset makes differences difficult to read. A small table for the largest training set sizes would help quantify the gains of NTP+MPM over NTP and MPM.
  2. [II.C.3, Fig. 2] The description of the joint forward passes states that 'two backbone forward passes are performed,' but it does not specify that one pass uses the unmasked sequence with causal attention and the other uses the masked sequence with bidirectional attention. Please make this explicit, and state how the two losses are weighted (the text says 'simply added,' but a formula would be clearer).
  3. [III.B, Fig. 5 and Appendix D] The comparison between MPM-Causal and NTP in Fig. 5 is confounded by the different prediction head architectures (transformer blocks for MPM(-Causal) and a single linear layer for NTP). Although Appendix D provides a control showing no effect on final fine-tuning performance, the main text does not mention this difference. Please refer to the control when interpreting the MPM-Causal result.
  4. [Appendix A.1, Table I] The LayerScale initialization value of 1.0 and the use of pre-norm transformer blocks differ from OmniJet-α, but no ablation is given. A sentence explaining the choice of 1.0 and its stability would be useful.
  5. [Figs. 3, 6, 8] The caption does not state how many jets were generated or whether the histograms are from a single fixed seed. Adding this information would allow readers to assess the statistical precision of the visual comparisons.

Circularity Check

0 steps flagged

No significant circularity: classification gains are measured on external labels and generative comparisons use a fixed target distribution.

full rationale

The central classification claim is supported by fine-tuning on external supervised tasks (JetClass 10-class, top-tagging binary) with held-out labels, so the pre-training improvements are not fitted to the evaluation metric. The generative claim is checked by comparing generated, decoded jets against the fixed 'JetClass (tokenized & decoded)' reference — the correct target for an autoregressive token model — with both the NTP-only and NTP+MPM generators evaluated against the same reference; no parameter is adjusted to make the histograms agree. The MPM masking rate is imported from an external reference [5], and the joint objective simply adds the two loss terms without fitting weights to downstream performance. OmniJet-α [9] is used as a baseline and architectural starting point, but the paper's improvements are demonstrated by direct comparison to this baseline and to from-scratch models, so the self-citation is not load-bearing in the derivation. The acknowledged VQ-VAE reconstruction limitation (Appendix B, notably d0) means the generative evidence is less strong than a quantitative raw-data metric would be, but that is an evidence-quality concern, not circularity. No equation or fitted parameter reduces to the claimed result by construction.

Axiom & Free-Parameter Ledger

5 free parameters · 4 axioms · 0 invented entities

No physics derivation is attempted; this is an empirical ML study. The central results depend on the chosen training objectives, hyperparameters, data-transfer assumptions, and the VQ-VAE tokenizer. No new physical entities, forces, or conserved quantities are introduced; architectural constructs (mask embeddings, registers, trainable start token) are learned model parameters, not physical inventions.

free parameters (5)
  • MPM masking rate = 40%
    Adopted from [5] without re-tuning; directly controls MPM and joint pre-training difficulty and thus the magnitude of classification improvements.
  • NTP:MPM loss weighting = 1:1 (losses added without modification)
    The two losses are added without any modification; no sweep is reported, so the joint-training result is conditional on this hand choice.
  • VQ-VAE codebook size = 8,192 (kin.); 32,768 (ext.)
    Hand-selected; codebook size determines tokenization granularity, hence the generative target and information loss.
  • Backbone size = 8 transformer blocks, embedding dim 128 (≈1.6M params for continuous input)
    Chosen from 'a few configurations' with no dedicated HPO; all pre-training strategy comparisons use this fixed size.
  • LayerScale initialization = 1.0
    Hand-selected after informal experiments; affects training stability and representation quality.
axioms (4)
  • domain assumption VQ-VAE tokenization/reconstruction is a sufficient representation for evaluating generative fidelity (target = tokenized-decoded JetClass distributions)
    All generative comparisons (Figs. 3, 6, 8) use the VQ-VAE-decoded data as reference, not raw JetClass data; acknowledged d0 reconstruction problems (App. B) make this assumption non-trivial.
  • domain assumption Pre-training on JetClass (CMS-like Delphes simulation) transfers to the top-tagging dataset (ATLAS-like Delphes simulation) despite different detector cards and generator settings
    The out-of-distribution transfer results (Figs. 4b, 5) are the main evidence for representation quality; this assumes the two simulations share enough structure for learned features to transfer.
  • ad hoc to paper MPM masking rate 40% from [5] and 1:1 loss weighting are appropriate for this smaller backbone
    No re-tuning was performed (§II C 2, §A 5); if these choices are suboptimal for the 128-dim/8-block backbone, the reported NTP+MPM gains could change.
  • standard math Standard transformer/VQ-VAE training procedures in PyTorch/Lightning are implemented correctly
    No bespoke mathematics is introduced; the paper's correctness depends on the standard neural-network training stack and the libraries cited.

pith-pipeline@v1.3.0-alltime-deepseek · 18213 in / 13628 out tokens · 116596 ms · 2026-08-03T18:37:29.824102+00:00 · methodology

0 comments
read the original abstract

Next token prediction is an attractive pre-training task for jet foundation models, in that it is simulation free and enables excellent generative capabilities that can transfer across datasets. Here we study multiple improvements to next token prediction, building on the initial work of OmniJet-$\alpha$. Instead of tokenizing particles and subsequently only using the token-ID as the model input for both the generative and the classification task, we adopt a hybrid setup, which allows us to use continuous feature vectors as model input while only using token-IDs in the next token prediction target. Secondly, we explore a combined pre-training strategy that combines masked particle modeling and generative learning objectives. Taken together, these changes greatly improve the performance in downstream classification tasks without any loss in generative performance.

Figures

Figures reproduced from arXiv: 2512.04149 by Anna Hallin, David Shih, Gregor Kasieczka, Ian Pang, Joschka Birk, Nikol Madzharova.

Figure 1
Figure 1. Figure 1: FIG. 1. The high-level architecture of the OmniJet workflow: (a) the pre-training based on next token prediction, (b) the generative model [PITH_FULL_IMAGE:figures/full_fig_p003_1.png] view at source ↗
Figure 2
Figure 2. Figure 2: FIG. 2. Schematic overview of the di [PITH_FULL_IMAGE:figures/full_fig_p003_2.png] view at source ↗
Figure 3
Figure 3. Figure 3: FIG. 3. Comparison of the jets that are generated by the token-ID-input and the continuous-input model, as well as the jets from the JetClass [PITH_FULL_IMAGE:figures/full_fig_p005_3.png] view at source ↗
Figure 4
Figure 4. Figure 4: FIG. 4. Comparison of token-ID input vs. continuous input classification performance: (a) for multi-class classification performance on the [PITH_FULL_IMAGE:figures/full_fig_p005_4.png] view at source ↗
Figure 5
Figure 5. Figure 5: FIG. 5. Classification performance on the top tagging dataset with di [PITH_FULL_IMAGE:figures/full_fig_p006_5.png] view at source ↗
Figure 6
Figure 6. Figure 6: FIG. 6. Comparison of the generative models trained with NTP only vs. joint NTP and MPM. The jets shown for the JetClass dataset are [PITH_FULL_IMAGE:figures/full_fig_p007_6.png] view at source ↗
Figure 7
Figure 7. Figure 7: FIG. 7. Comparison of the original JetClass dataset and the corresponding distributions obtained after tokenization. The tokenized datasets [PITH_FULL_IMAGE:figures/full_fig_p012_7.png] view at source ↗
Figure 8
Figure 8. Figure 8: FIG. 8. Comparison of the jets that are generated by the continuous-input model using the kinematics-only (kin.) and the extended feature [PITH_FULL_IMAGE:figures/full_fig_p013_8.png] view at source ↗
Figure 9
Figure 9. Figure 9: FIG. 9. Multi-class classification performance on the JetClass dataset (all 10 jet types) as a function of the number of training jets comparing [PITH_FULL_IMAGE:figures/full_fig_p013_9.png] view at source ↗
Figure 10
Figure 10. Figure 10: FIG. 10. Classification performance on the top tagging dataset when [PITH_FULL_IMAGE:figures/full_fig_p014_10.png] view at source ↗
Figure 11
Figure 11. Figure 11: FIG. 11. Evolution of the backbone’s gradient norm, as well as the validation loss and [PITH_FULL_IMAGE:figures/full_fig_p015_11.png] view at source ↗
Figure 12
Figure 12. Figure 12: FIG. 12. Classification performance on the top tagging dataset for [PITH_FULL_IMAGE:figures/full_fig_p015_12.png] view at source ↗

discussion (0)

Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.

Forward citations

Cited by 2 Pith papers

Reviewed papers in the Pith corpus that reference this work. Sorted by Pith novelty score.

  1. Neural Scaling Laws for Jet Generation

    hep-ph 2026-05 unverdicted novelty 7.0

    Scaling laws hold logarithmically for model size in autoregressive jet generation, with next-token loss correlating to physical metrics via sliced Wasserstein distance, but show weaker scaling for dataset size and com...

  2. Articulating Assumptions in AI-Generated Scientific Analyses through Task Decomposition

    cs.SE 2026-07 conditional novelty 6.0

    Quantity-grounded multi-agent decomposition makes LLM-generated collider analysis code inspectable and reliable with 14B-scale models, outperforming prior single-prompt approaches.

Reference graph

Works this paper leans on

76 extracted references · 27 linked inside Pith · cited by 2 Pith papers

  1. [1]

    Continuous-input next-token prediction The main advantage of using continuous input vectors is that the downstream classification task can operate directly on the full resolution feature vectors ⃗ci, ensuring that the model does not suffer from tokenization-induced loss of information in this task. With this goal of improving downstream clas- sification p...

  2. [2]

    Although MPM does not produce a generative model out of the box, it has demonstrated competitive classification performance

    Masked token prediction Following masked modeling approaches in natural lan- guage processing and computer vision [51–56], masked mod- eling was introduced for jet physics in the form of Masked Particle Modeling (MPM) [4, 5]. Although MPM does not produce a generative model out of the box, it has demonstrated competitive classification performance. Since ...

  3. [3]

    An illustration of the joint setup is shown in Figure 2c

    Joint next and masked token prediction To investigate if the advantages of both next and masked token prediction can be combined, we implement a joint pre-training. An illustration of the joint setup is shown in Figure 2c. This joint pre-training is realized by having two model heads, predicting the next token logits and the masked token logits, respectiv...

  4. [4]

    Bommasaniet al., On the opportunities and risks of founda- tion models, arXiv (2021), 2108.07258

    R. Bommasaniet al., On the opportunities and risks of founda- tion models, arXiv (2021), 2108.07258

    Pith/arXiv arXiv 2021

  5. [5]

    Hallin, Foundation models for high-energy physics (2025) arXiv:2509.21434 [hep-ph]

    A. Hallin, Foundation models for high-energy physics (2025) arXiv:2509.21434 [hep-ph]

    arXiv 2025

  6. [6]

    Butteret al., The Machine Learning landscape of top tag- gers, SciPost Phys.7, 014 (2019), arXiv:1902.09914 [hep-ph]

    A. Butteret al., The Machine Learning landscape of top tag- gers, SciPost Phys.7, 014 (2019), arXiv:1902.09914 [hep-ph]

    Pith/arXiv arXiv 2019

  7. [7]

    Golling, L

    T. Golling, L. Heinrich, M. Kagan, S. Klein, M. Leigh, M. Os- adchy, and J. A. Raine, Masked particle modeling on sets: to- wards self-supervised high energy physics foundation models, Mach. Learn. Sci. Tech.5, 035074 (2024), arXiv:2401.13537 [hep-ph]

    Pith/arXiv arXiv 2024

  8. [8]

    Leigh, S

    M. Leigh, S. Klein, F. Charton, T. Golling, L. Heinrich, M. Ka- gan, I. Ochoa, and M. Osadchy, Is Tokenization Needed for Masked Particle Modelling?, Machine Learning: Science and Technology6, 025075 (2025), arXiv:2409.12589 [hep-ph]

    Pith/arXiv arXiv 2025

  9. [9]

    Katel, H

    S. Katel, H. Li, Z. Zhao, F. Mokhtar, J. Duarte, and R. Kansal, Learning Symmetry-Independent Jet Representations via Jet- Based Joint Embedding Predictive Architecture, inMachine Learning and the Physical Sciences: Workshop at NeurIPS 2024(2024) arXiv:2412.05333 [hep-ph]

    Pith/arXiv arXiv 2024

  10. [10]

    A. J. Wildridge, J. P. Rodgers, E. M. Colbert, Y . yao, A. W. Jung, and M. Liu, Bumblebee: Foundation Model for Particle Physics Discovery, in38th conference on Neural Information Processing Systems(2024) arXiv:2412.07867 [hep-ex]

    Pith/arXiv arXiv 2024

  11. [11]

    Bardhan, R

    J. Bardhan, R. Agrawal, A. Tilak, C. Neeraj, and S. Mitra, HEP-JEPA: A foundation model for collider physics using joint embedding predictive architecture, (2025), arXiv:2502.03933 [cs.LG]

    Pith/arXiv arXiv 2025

  12. [12]

    J. Birk, A. Hallin, and G. Kasieczka, OmniJet-α: the first cross- task foundation model for particle physics, Mach. Learn. Sci. Tech.5, 035031 (2024), arXiv:2403.05618 [hep-ph]

    Pith/arXiv arXiv 2024

  13. [13]

    H. Qu, C. Li, and S. Qian, Particle transformer for jet tagging, in Proceedings of the 39th International Conference on Machine Learning, Proceedings of Machine Learning Research, V ol. 162, edited by K. Chaudhuri, S. Jegelka, L. Song, C. Szepes- vari, G. Niu, and S. Sabato (PMLR, 2022) pp. 18281–18292

    2022

  14. [14]

    Mikuni and B

    V . Mikuni and B. Nachman, Solving key challenges in collider physics with foundation models, Phys. Rev. D111, L051504 (2025), arXiv:2404.16091 [hep-ph]

    arXiv 2025

  15. [15]

    Mikuni and B

    V . Mikuni and B. Nachman, Method to simultaneously facil- itate all jet physics tasks, Phys. Rev. D111, 054015 (2025), arXiv:2502.14652 [hep-ph]

    arXiv 2025

  16. [16]

    Bhimji, C

    W. Bhimji, C. Harris, V . Mikuni, and B. Nachman, Om- niLearned: A Foundation Model Framework for All Tasks In- volving Jet Physics, (2025), arXiv:2510.24066 [hep-ph]

    arXiv 2025

  17. [17]

    J. Ho, B. R. Roberts, S. Han, and H. Wang, Pretrained Event Classification Model for High Energy Physics Analysis (2024), arXiv:2412.10665 [hep-ph]

    Pith/arXiv arXiv 2024

  18. [18]

    Harris, J

    P. Harris, J. Krupa, M. Kagan, B. Maier, and N. Wood- ward, Resimulation-based self-supervised learning for pretrain- ing physics foundation models, Phys. Rev. D111, 032010 (2025), arXiv:2403.07066 [hep-ph]

    Pith/arXiv arXiv 2025

  19. [19]

    Z. Hao, R. Kansal, A. Gandrakota, C. Sun, N. Jennifer, J. Duarte, and M. Spiropulu, Rino: Renormalization group in- variance with no labels (2025), arXiv:2509.07486 [hep-ex]

    arXiv 2025

  20. [20]

    Amram, L

    O. Amram, L. Anzalone, J. Birk, D. A. Faroughy, A. Hallin, G. Kasieczka, M. Kr ¨amer, I. Pang, H. Reyes-Gonzalez, and D. Shih, Aspen Open Jets: unlocking LHC data for foundation models in particle physics, Mach. Learn. Sci. Tech.6, 030601 (2025), arXiv:2412.10504 [hep-ph]

    arXiv 2025

  21. [21]

    I. Pang, D. A. Faroughy, D. Shih, R. Das, and G. Kasieczka, SURFing to the Fundamental Limit of Jet Tagging, (2025), arXiv:2511.15779 [hep-ph]

    arXiv 2025

  22. [22]

    A. Lv, K. Zhang, S. Xie, Q. Tu, Y . Chen, J.-R. Wen, and R. Yan, An analysis and mitigation of the reversal curse (2024), arXiv:2311.07468 [cs.CL]

    Pith/arXiv arXiv 2024

  23. [23]

    BehnamGhader, V

    P. BehnamGhader, V . Adlakha, M. Mosbach, D. Bahdanau, N. Chapados, and S. Reddy, LLM2Vec: Large language models are secretly powerful text encoders (2024), arXiv:2404.05961 [cs.CL]

    Pith/arXiv arXiv 2024

  24. [24]

    L. G. G. Charpentier and D. Samuel, GPT or BERT: why not both? (2024), arXiv:2410.24159 [cs.CL]

    Pith/arXiv arXiv 2024

  25. [25]

    X. Yu, B. Guo, S. Luo, J. Wang, T. Ji, and Y . Wu, Antlm: Bridging causal and masked language models (2024), arXiv:2412.03275 [cs.CL]

    Pith/arXiv arXiv 2024

  26. [26]

    H. Liu, X. Geng, L. Lee, I. Mordatch, S. Levine, S. Narang, and P. Abbeel, Towards better few-shot and finetuning per- formance with forgetful causal language models (2023), arXiv:2210.13432 [cs.CL]

    Pith/arXiv arXiv 2023

  27. [27]

    Gisserot-Boukhlef, N

    H. Gisserot-Boukhlef, N. Boizard, M. Faysse, D. M. Alves, E. Malherbe, A. F. T. Martins, C. Hudelot, and P. Colombo, Should we still pretrain encoders with masked language model- ing? (2025), arXiv:2507.00994 [cs.CL]

    Pith/arXiv arXiv 2025

  28. [28]

    Paszke, S

    A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, A. Desmaison, A. Kopf, E. Yang, Z. DeVito, M. Raison, A. Te- jani, S. Chilamkurthy, B. Steiner, L. Fang, J. Bai, and S. Chin- tala, Pytorch: An imperative style, high-performance deep learning library, inAdvances in Neural Information Processing ...

    2019

  29. [29]

    Falcon and The PyTorch Lightning team, PyTorch Lightning (2019)

    W. Falcon and The PyTorch Lightning team, PyTorch Lightning (2019)

    2019

  30. [30]

    J. D. Hunter, Matplotlib: A 2d graphics environment, Comput- ing in Science & Engineering9, 90 (2007)

    2007

  31. [31]

    pandas development team, pandas-dev/pandas: Pandas (2020)

    T. pandas development team, pandas-dev/pandas: Pandas (2020)

    2020

  32. [32]

    Wes McKinney, Data Structures for Statistical Computing in Python, inProceedings of the 9th Python in Science Confer- ence, edited by St´efan van der Walt and Jarrod Millman (2010) pp. 56 – 61

    2010

  33. [33]

    C. R. Harris, K. J. Millman, S. J. van der Walt, R. Gommers, P. Virtanen, D. Cournapeau, E. Wieser, J. Taylor, S. Berg, N. J. Smith, R. Kern, M. Picus, S. Hoyer, M. H. van Kerkwijk, M. Brett, A. Haldane, J. F. del R ´ıo, M. Wiebe, P. Peterson, P. G ´erard-Marchant, K. Sheppard, T. Reddy, W. Weckesser, H. Abbasi, C. Gohlke, and T. E. Oliphant, Array progra...

    2020

  34. [34]

    Pivarski, Jim, Elmer, Peter, and Lange, David, Awkward arrays in python, c++, and numba, EPJ Web Conf.245, 05023 (2020)

    2020

  35. [35]

    Pivarski, I

    J. Pivarski, I. Osborne, I. Ifrim, H. Schreiner, A. Hollands, A. Biswas, P. Das, S. Roy Choudhury, N. Smith, M. Goyal, P. Fackeldey, and I. Krommydas, Awkward array (2025)

    2025

  36. [36]

    A. Roy, J. Pivarski, L. Gray, C. Papageorgakis, M. Feickert, J. Duarte, H. Schreiner, R. Kansal, cmoore24 24, 5ean github, K. W. Ho, K. Lieret, P. Fackeldey, and S. Rothman, scikit- hep/fastjet (2025)

    2025

  37. [37]

    Cacciari, G

    M. Cacciari, G. P. Salam, and G. Soyez, FastJet User Manual, Eur. Phys. J. C72, 1896 (2012), arXiv:1111.6097 [hep-ph]

    Pith/arXiv arXiv 2012

  38. [38]

    Huh, vqtorch: PyTorch package for vector quantization, https://github.com/minyoungg/vqtorch(2022)

    M. Huh, vqtorch: PyTorch package for vector quantization, https://github.com/minyoungg/vqtorch(2022)

    2022

  39. [39]

    Pedregosa, G

    F. Pedregosa, G. Varoquaux, A. Gramfort, V . Michel, B. Thirion, O. Grisel, M. Blondel, P. Prettenhofer, R. Weiss, V . Dubourg, J. Vanderplas, A. Passos, D. Cournapeau, M. Brucher, M. Perrot, and E. Duchesnay, Scikit-learn: Ma- chine learning in Python, Journal of Machine Learning Re- search12, 2825 (2011)

    2011

  40. [40]

    Buss, FLC-QU-hep/ranger-lite: RangerLite 1.0.0 (2025)

    T. Buss, FLC-QU-hep/ranger-lite: RangerLite 1.0.0 (2025)

    2025

  41. [41]

    Chopra, H

    S. Chopra, H. Schreiner, E. Rodrigues, J. Eschle, and J. Pi- varski, Vector: Jit-compilable mathematical manipulations of ragged lorentz vectors, Journal of Open Source Software10, 7791 (2025)

    2025

  42. [42]

    H. Qu, C. Li, and S. Qian, Jetclass: A large-scale dataset for deep learning in jet physics, 10.5281/zenodo.6619768 (2022)

    doi:10.5281/zenodo.6619768 2022

  43. [43]

    Kasieczka, T

    G. Kasieczka, T. Plehn, J. Thompson, and M. Russel, Top quark tagging reference dataset, 10.5281/zenodo.2603256 (2019)

    doi:10.5281/zenodo.2603256 2019

  44. [44]

    Alwall, R

    J. Alwall, R. Frederix, S. Frixione, V . Hirschi, F. Maltoni, O. Mattelaer, H. S. Shao, T. Stelzer, P. Torrielli, and M. Zaro, The automated computation of tree-level and next-to-leading order differential cross sections, and their matching to parton shower simulations, JHEP07, 079, arXiv:1405.0301 [hep-ph]

    Pith/arXiv arXiv

  45. [45]

    Sjostrand, S

    T. Sjostrand, S. Mrenna, and P. Z. Skands, A Brief Introduc- tion to PYTHIA 8.1, Comput. Phys. Commun.178, 852 (2008), arXiv:0710.3820 [hep-ph]

    Pith/arXiv arXiv 2008

  46. [46]

    Sj ¨ostrand, S

    T. Sj ¨ostrand, S. Ask, J. R. Christiansen, R. Corke, N. Desai, P. Ilten, S. Mrenna, S. Prestel, C. O. Rasmussen, and P. Z. Skands, An introduction to PYTHIA 8.2, Comput. Phys. Com- mun.191, 159 (2015), arXiv:1410.3012 [hep-ph]

    Pith/arXiv arXiv 2015

  47. [47]

    de Favereau, C

    J. de Favereau, C. Delaere, P. Demin, A. Giammanco, V . Lemaˆıtre, A. Mertens, and M. Selvaggi (DELPHES 3), DELPHES 3, A modular framework for fast simulation of a generic collider experiment, JHEP02, 057, arXiv:1307.6346 [hep-ex]

    Pith/arXiv arXiv

  48. [48]

    Chatrchyanet al.(CMS), The CMS Experiment at the CERN LHC, JINST3, S08004

    S. Chatrchyanet al.(CMS), The CMS Experiment at the CERN LHC, JINST3, S08004

  49. [49]

    Cacciari, G

    M. Cacciari, G. P. Salam, and G. Soyez, The anti-k t jet cluster- ing algorithm, JHEP04, 063, arXiv:0802.1189 [hep-ph]

    Pith/arXiv arXiv

  50. [50]

    Aadet al.(ATLAS), The ATLAS Experiment at the CERN Large Hadron Collider, JINST3, S08003

    G. Aadet al.(ATLAS), The ATLAS Experiment at the CERN Large Hadron Collider, JINST3, S08003

  51. [51]

    van den Oord, O

    A. van den Oord, O. Vinyals, and k. kavukcuoglu, Neural dis- crete representation learning, inAdvances in Neural Informa- tion Processing Systems, V ol. 30, edited by I. Guyon, U. V . Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (Curran Associates, Inc., 2017)

    2017

  52. [52]

    M. Huh, B. Cheung, P. Agrawal, and P. Isola, Straightening out the straight-through estimator: Overcoming optimization chal- lenges in vector quantized networks, inProceedings of the 40th International Conference on Machine Learning, Proceedings of Machine Learning Research, V ol. 202, edited by A. Krause, E. Brunskill, K. Cho, B. Engelhardt, S. Sabato, a...

    2023

  53. [53]

    Buhmann, C

    E. Buhmann, C. Ewen, G. Kasieczka, V . Mikuni, B. Nach- man, and D. Shih, Full phase space resonant anomaly detection, Phys. Rev. D109, 055015 (2024), arXiv:2310.06897 [hep-ph]

    Pith/arXiv arXiv 2024

  54. [54]

    Devlin, M.-W

    J. Devlin, M.-W. Chang, K. Lee, and K. Toutanova, BERT: Pre- training of Deep Bidirectional Transformers for Language Un- derstanding (2019), arXiv:1810.04805 [cs.CL]

    Pith/arXiv arXiv 2019

  55. [55]

    K. He, X. Chen, S. Xie, Y . Li, P. Doll ´ar, and R. Girshick, Masked autoencoders are scalable vision learners, inProceed- ings of the IEEE/CVF Conference on Computer Vision and Pat- tern Recognition (CVPR)(2022) pp. 16000–16009

    2022

  56. [56]

    H. Bao, L. Dong, S. Piao, and F. Wei, BEit: BERT pre-training of image transformers, inInternational Conference on Learning Representations(2022)

    2022

  57. [57]

    R. Wang, D. Chen, Z. Wu, Y . Chen, X. Dai, M. Liu, Y .-G. Jiang, L. Zhou, and L. Yuan, Bevt: Bert pretraining of video transformers, inProceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)(2022) pp. 14733–14743

    2022

  58. [58]

    Baevski, W.-N

    A. Baevski, W.-N. Hsu, Q. Xu, A. Babu, J. Gu, and M. Auli, data2vec: A general framework for self-supervised learning in speech, vision and language, inProceedings of the 39th In- ternational Conference on Machine Learning, Proceedings of Machine Learning Research, V ol. 162, edited by K. Chaud- huri, S. Jegelka, L. Song, C. Szepesvari, G. Niu, and S. Sab...

    2022

  59. [59]

    C. Wei, H. Fan, S. Xie, C.-Y . Wu, A. Yuille, and C. Feicht- enhofer, Masked feature prediction for self-supervised visual pre-training, inProceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)(2022) pp. 14668–14678

    2022

  60. [60]

    Spinner, V

    J. Spinner, V . Bres ´o, P. de Haan, T. Plehn, J. Thaler, and J. Brehmer, Lorentz-Equivariant Geometric Algebra Trans- formers for High-Energy Physics, in38th conference on Neu- ral Information Processing Systems(2024) arXiv:2405.14806 [physics.data-an]

    Pith/arXiv arXiv 2024

  61. [61]

    Brehmer, V

    J. Brehmer, V . Bres´o, P. de Haan, T. Plehn, H. Qu, J. Spinner, and J. Thaler, A Lorentz-equivariant transformer for all of the LHC, SciPost Phys.19, 108 (2025), arXiv:2411.00446 [hep- ph]

    arXiv 2025

  62. [62]

    Tschannen, C

    M. Tschannen, C. Eastwood, and F. Mentzer, Givt: Genera- tive infinite-vocabulary transformers (2024), arXiv:2312.02116 [cs.CV]

    Pith/arXiv arXiv 2024

  63. [63]

    Radford, K

    A. Radford, K. Narasimhan, T. Salimans, and I. Sutskever, Improving language understanding by generative pre-training (2018)

    2018

  64. [64]

    Xiong, Y

    R. Xiong, Y . Yang, D. He, K. Zheng, S. Zheng, C. Xing, H. Zhang, Y . Lan, L. Wang, and T.-Y . Liu, On layer normal- ization in the transformer architecture, inProceedings of the 10 37th International Conference on Machine Learning, ICML’20 (JMLR.org, 2020)

    2020

  65. [65]

    Touvron, M

    H. Touvron, M. Cord, A. Sablayrolles, G. Synnaeve, and H. J ´egou, Going deeper with image transformers, inProceed- ings of the IEEE/CVF International Conference on Computer Vision (ICCV)(2021) pp. 32–42

    2021

  66. [66]

    Darcet, M

    T. Darcet, M. Oquab, J. Mairal, and P. Bojanowski, Vision transformers need registers, inThe Twelfth International Con- ference on Learning Representations(2024)

    2024

  67. [67]

    Shleifer, J

    S. Shleifer, J. Weston, and M. Ott, Normformer: Im- proved transformer pretraining with extra normalization (2021), arXiv:2110.09456 [cs.CL]

    Pith/arXiv arXiv 2021

  68. [68]

    Wright, Ranger - a synergistic opti- mizer.,https://github.com/lessw2020/ Ranger-Deep-Learning-Optimizer(2019)

    L. Wright, Ranger - a synergistic opti- mizer.,https://github.com/lessw2020/ Ranger-Deep-Learning-Optimizer(2019)

    2019

  69. [69]

    Zhang, J

    M. Zhang, J. Lucas, J. Ba, and G. E. Hinton, Lookahead opti- mizer: k steps forward, 1 step back, inAdvances in Neural In- formation Processing Systems, V ol. 32, edited by H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alch ´e-Buc, E. Fox, and R. Garnett (Curran Associates, Inc., 2019)

    2019

  70. [70]

    L. Liu, H. Jiang, P. He, W. Chen, X. Liu, J. Gao, and J. Han, On the variance of the adaptive learning rate and beyond, in International Conference on Learning Representations(2020). Appendix A: Model architecture and training details

    2020

  71. [71]

    Backbone architecture changes Compared to OmniJet-α, which uses post-norm trans- former blocks based on [60], we use pre-norm trans- former [61] blocks in the backbone architecture as this choice is more common in modern transformer architectures. Fur- thermore, we add LayerScale[62] to the architecture, but instead of initializing the per-channel weights...

  72. [72]

    Token prediction head As in the original OmniJet-αwork, we use a linear layer to project the backbone output to the logits of the token-IDs when training on the NTP target only. The token prediction head used for MPM, as well as the token prediction head used for the NTP prediction in the joint pre-training consists of two transformer blocks with the same...

  73. [73]

    Furthermore, we do not apply the causal mask in the backbone in classification mode as this is an unnecessary restriction on the model’s capacity when trained for classifi- cation

    Classification head Compared to OmniJet-αwe adapt the classification head to feature two class-attention blocks [10, 62] followed by a lin- ear layer (instead of two linear layers with a summation over all particles between them) to obtain a more capable classifi- cation head. Furthermore, we do not apply the causal mask in the backbone in classification ...

  74. [74]

    Additionally, we use the same transformer block architecture as in the backbone model, in- stead of the NormFormer [64] architecture that was used in OmniJet-α

    VQ-V AE architecture changes As the hybrid setup with continuous input and token-ID prediction requires on-the-fly decoding of token-IDs during the generation loop, the VQ-V AE has to be trained with a causal mask in the decoder. Additionally, we use the same transformer block architecture as in the backbone model, in- stead of the NormFormer [64] archite...

  75. [75]

    A constant learning rate of 1×10−3 and a batch size of 1000 is used for all trainings except for trainings where we use 10k jets or less, for which we use a batch size of

    Training details All models are trained with the Ranger [65] optimizer, which uses the Lookahead [66] optimizer with RAdam [67] as the inner optimizer. A constant learning rate of 1×10−3 and a batch size of 1000 is used for all trainings except for trainings where we use 10k jets or less, for which we use a batch size of

  76. [100]

    Classification trainings are performed for up to 1 M steps as well, but with early stopping based on the validation loss

    Pre-trainings are performed for a total of 1 M steps, with the final model state being used for evaluation. Classification trainings are performed for up to 1 M steps as well, but with early stopping based on the validation loss. Classification runs are repeated with five different random seeds and the average and standard deviation of the five runs is re...

    2000