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

arXiv: https://arxiv.org/abs/2512.04149 · HTML: https://mlsanity.com/paper/2512.04149/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 37 of its 74 references that are in the map.
- First posted on arXiv 3 Dec 2025.

## Builds on (most cited first, 30 of 37)

- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Masked Autoencoders Are Scalable Vision Learners](https://mlsanity.com/paper/2111.06377/) (2021)
- [Neural Discrete Representation Learning](https://mlsanity.com/paper/1711.00937/) (2017)
- [On the Opportunities and Risks of Foundation Models](https://mlsanity.com/paper/2108.07258/) (2021)
- [BEiT: BERT Pre-Training of Image Transformers](https://mlsanity.com/paper/2106.08254/) (2021)
- [On Layer Normalization in the Transformer Architecture](https://mlsanity.com/paper/2002.04745/) (2020)
- [Vision Transformers Need Registers](https://mlsanity.com/paper/2309.16588/) (2023)
- [Going deeper with Image Transformers](https://mlsanity.com/paper/2103.17239/) (2021)
- [data2vec: A General Framework for Self-supervised Learning in Speech, Vision and Language](https://mlsanity.com/paper/2202.03555/) (2022)
- [Masked Feature Prediction for Self-Supervised Visual Pre-Training](https://mlsanity.com/paper/2112.09133/) (2021)
- [LLM2Vec: Large Language Models Are Secretly Powerful Text Encoders](https://mlsanity.com/paper/2404.05961/) (2024)
- [Lookahead Optimizer: k steps forward, 1 step back](https://mlsanity.com/paper/1907.08610/) (2019)
- [BEVT: BERT Pretraining of Video Transformers](https://mlsanity.com/paper/2112.01529/) (2021)
- [Particle Transformer for Jet Tagging](https://mlsanity.com/paper/2202.03772/) (2022)
- [NormFormer: Improved Transformer Pretraining with Extra Normalization](https://mlsanity.com/paper/2110.09456/) (2021)
- [Straightening Out the Straight-Through Estimator: Overcoming Optimization Challenges in Vector Quantized Networks](https://mlsanity.com/paper/2305.08842/) (2023)
- [GIVT: Generative Infinite-Vocabulary Transformers](https://mlsanity.com/paper/2312.02116/) (2023)
- [Masked Particle Modeling on Sets: Towards Self-Supervised High Energy Physics Foundation Models](https://mlsanity.com/paper/2401.13537/) (2024)
- [OmniJet-α: The first cross-task foundation model for particle physics](https://mlsanity.com/paper/2403.05618/) (2024)
- [Lorentz-Equivariant Geometric Algebra Transformers for High-Energy Physics](https://mlsanity.com/paper/2405.14806/) (2024)
- [A Lorentz-Equivariant Transformer for All of the LHC](https://mlsanity.com/paper/2411.00446/) (2024)
- [GPT or BERT: why not both?](https://mlsanity.com/paper/2410.24159/) (2024)
- [An Analysis and Mitigation of the Reversal Curse](https://mlsanity.com/paper/2311.07468/) (2023)
- [Re-Simulation-based Self-Supervised Learning for Pre-Training Foundation Models](https://mlsanity.com/paper/2403.07066/) (2024)
- [Is Tokenization Needed for Masked Particle Modelling?](https://mlsanity.com/paper/2409.12589/) (2024)
- [Should We Still Pretrain Encoders with Masked Language Modeling?](https://mlsanity.com/paper/2507.00994/) (2025)
- [Aspen Open Jets: Unlocking LHC Data for Foundation Models in Particle Physics](https://mlsanity.com/paper/2412.10504/) (2024)
- [HEP-JEPA: A foundation model for collider physics using joint embedding predictive architecture](https://mlsanity.com/paper/2502.03933/) (2025)
- [Learning Symmetry-Independent Jet Representations via Jet-Based Joint Embedding Predictive Architecture](https://mlsanity.com/paper/2412.05333/) (2024)

## Built on by (most cited first, 3 of 3)

- [Neural Scaling Laws for Boosted Jet Tagging](https://mlsanity.com/paper/2602.15781/) (2026)
- [Neural Scaling Laws for Jet Generation](https://mlsanity.com/paper/2605.28940/) (2026)
- [Prompting Particle Physics: Tokenized Multi-modal Foundation Models for Combinatorially Many Tasks](https://mlsanity.com/paper/2609.31862/) (2026)

## Abstract (arXiv)

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-α. 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.