JetFormer: A Scalable and Efficient Transformer for Jet Tagging from Offline Analysis to FPGA Triggers

JetFormer is een schaalbare en efficiënte Transformer-architectuur voor jet-tagging bij de LHC, die zowel hoge nauwkeurigheid biedt voor offline analyses als een lage latentie voor implementatie op FPGA-triggers.

De kern

Stel je voor dat je een superdrukke verkeersregelaar bent bij een gigantisch kruispunt waar miljoenen auto's per seconde voorbijrazen. Je moet in een fractie van een seconde beslissen: is dit een gewone gezinsauto, of is het een gevaarlijke raceauto die we direct moeten volgen? Als je te lang nadenkt, raakt de hele stad in de knoop.

Dit is precies het probleem waar wetenschappers bij de deeltjesversneller (de LHC) mee kampen. Ze botsen deeltjes tegen elkaar aan en krijgen een enorme explosie van "jets" (bundels deeltjes). Ze willen weten wat voor soort deeltje er net is ontstaan, maar de data komt zo snel binnen dat een normale computer het niet bijhoudt.

Dit paper introduceert JetFormer. Hier is de uitleg in begrijpelijke taal:

1. Wat is JetFormer? (De Slimme Verkeersregelaar)

In plaats van een trage, zware computer die elk detail van elke auto bestudeert, is JetFormer een soort "slimme scanner". Het is een AI-model (een Transformer) dat heel snel naar de verzameling deeltjes kijkt en zegt: "Dit is een Higgs-boson!" of "Dit is gewoon ruis."

Het bijzondere is dat JetFormer een "alleskunner" is. Het kan worden ingezet als een super-expert die heel diep nadenkt voor wetenschappelijk onderzoek (offline), maar het kan ook worden omgebouwd tot een razendsnelle, compacte versie voor de directe actie (online).

2. De Magie van de "Compacte Versie" (De LEGO-methode)

Om JetFormer geschikt te maken voor de razendsnelle hardware (FPGA's) die in de versneller zit, hebben de onderzoekers drie slimme trucjes gebruikt:

• Pruning (Snoeien): Denk aan een struik die te groot is. Door de onnodige takjes weg te knippen, blijft de vorm van de struik hetzelfde, maar is hij veel lichter en makkelijker te verplaatsen. De onderzoekers hebben de "onnodige hersencellen" van de AI weggehaald zonder dat hij zijn intelligentie verloor.
• Quantization (De Versimpeling): In plaats van te rekenen met extreem precieze getallen (zoals 3,14159265...), zegt de AI nu gewoon: "Het is ongeveer 3." Dit is alsof je een ingewikkelde wiskundesom oplost door alleen met hele getallen te werken. Het gaat veel sneller en kost veel minder energie, terwijl de uitkomst bijna hetzelfde blijft.
• Hyperparameter Optimization (De Perfecte Receptuur): Ze hebben een computerprogramma gebruikt om de perfecte balans te zoeken tussen "slim zijn" en "snel zijn". Het is alsof je een recept voor een taart blijft aanpassen tot hij precies lekker is, maar ook snel genoeg gebakken kan worden in een kleine oven.

3. Waarom is dit belangrijk?

Vroeger moest je kiezen: of je hebt een model dat heel slim is maar traag (zoals een professor die een boek leest), of een model dat snel is maar dom (zoals een reflex).

JetFormer doorbreekt dit. Het is een architectuur die kan schalen. Je kunt hem "groot" maken voor de wetenschap, of "tiny" maken voor de razendsnelle hardware in de machine. Het is de eerste keer dat een dergelijk geavanceerd AI-systeem zo soepel kan schakelen tussen die twee werelden.

Samenvatting in één zin:

JetFormer is een slimme, aanpasbare AI die de enorme chaos van deeltjesbotsingen kan begrijpen, zowel met de precisie van een microscoop als met de snelheid van een flits.

Technische samenvatting

Technische Samenvatting: JetFormer

1. Het Probleem (De Uitdaging)

In de deeltjesfysica bij de Large Hadron Collider (LHC) worden bij botsingen "jets" geproduceerd: bundels deeltjes die wijzen op de aanwezigheid van zware deeltjes (zoals Higgs-bosonen of top-quarks). Jet tagging is het proces waarbij deze jets worden geclassificeerd.

Er bestaat momenteel een kloof tussen twee toepassingsgebieden:

• Offline analyse: Vereist maximale nauwkeurigheid en maakt gebruik van complexe Transformer-modellen (zoals ParT), maar deze zijn te groot en traag voor real-time gebruik.
• Online triggering (L1T): Moet beslissingen nemen in sub-microseconde tijd op FPGA-hardware (Field Programmable Gate Arrays). Bestaande modellen (zoals MLP's of GNN's) zijn vaak minder nauwkeurig of niet flexibel genoeg om de kracht van Transformers te benutten binnen de strikte hardwarebeperkingen.

2. Methodologie

De auteurs introduceren JetFormer, een encoder-only Transformer-architectuur die specifiek is ontworpen om schaalbaar te zijn van offline naar online omgevingen.

Architecturale aanpassingen voor hardware:

• Encoder-only structuur: Gebaseerd op het BERT-principe, waarbij een speciale [CLS] token wordt gebruikt om een globale representatie van de jet te extraheren.
• Permutatie-invariantie: Er wordt geen positional encoding gebruikt, omdat de volgorde van deeltjes in een jet niet relevant is.
• Hardware-vriendelijke operaties:
  • Layer Normalization is vervangen door Batch Normalization (omdat parameters vooraf berekend kunnen worden).
  • SiLU-activatiefuncties zijn vervangen door ReLU (om complexe exponentiële berekeningen te vermijden).
• Geen expliciete interactietermen: Het model verwerkt deeltjeskenmerken direct zonder de rekenintensieve berekening van alle paren, wat efficiëntie oplevert.

Optimalisatie- en compressie-pipeline:

1. Multi-objective Hyperparameter Optimization (HPO): Gebruik van Optuna om een balans te vinden tussen maximale nauwkeurigheid en minimale FLOPs (Floating Point Operations). Dit resulteerde in de compacte JetFormer-tiny.
2. Structured Pruning: Het systematisch verwijderen van volledige kanalen/filters om de rekenlast met 50% te verminderen met minimaal nauwkeurigheidsverlies.
3. 1-bit Quantization (BitNet): Het reduceren van gewichten naar slechts +1 of -1, waardoor vermenigvuldigingen worden vervangen door eenvoudige optellingen.
4. Allo Framework: Een extensie van het Allo (MLIR-gebaseerde) framework om Transformer-operaties te vertalen naar HLS (High-Level Synthesis) voor FPGA-implementatie.

3. Belangrijkste Bijdragen

• Een verenigd framework: Het overbrugt de kloof tussen high-performance deep learning en real-time hardware-implementatie.
• Efficiënte architectuur: Een model dat de nauwkeurigheid van state-of-the-art modellen benadert, maar aanzienlijk minder rekenkracht vereist.
• Hardware-aware pipeline: Een complete workflow van modelontwerp tot aan de synthese van gecomprimeerde modellen op FPGA's.

4. Resultaten

• Nauwkeurigheid en Efficiëntie (JETCLASS dataset): De volledige JetFormer bereikt een nauwkeurigheid die bijna gelijk is aan het ParT-model (verschil van slechts 0,7%), maar gebruikt 37,4% minder FLOPs.
• Benchmark prestaties (HLS4ML 150P): JetFormer presteert 3–4% beter dan bestaande modellen zoals MLP, Deep Sets en Interaction Networks.
• Compressie-effectiviteit:
  • Pruning: Vermindert de rekenlast met 50% met <0,5% verlies in nauwkeurigheid.
  • 1-bit Quantization: Verkleint de modelgrootte met 82–92%, met een relatief klein verlies in nauwkeurigheid (1,5–3,5%).
• FPGA-implementatie: De geprunede JetFormer-tiny variant is succesvol gesynthetiseerd op FPGA-hardware, wat de haalbaarheid van real-time Transformer-inferentie in de LHC-trigger bewijst.

5. Betekenis (Significance)

JetFormer biedt een praktisch pad voor de volgende generatie deeltjesdetectoren. Door Transformers schaalbaar te maken, kunnen experimenten bij de LHC (zoals de CMS Phase-2 upgrade) geavanceerdere algoritmen direct in de hardware (Level-1 Trigger) implementeren. Dit stelt wetenschappers in staat om zeldzame fysieke fenomenen efficiënter te identificeren in de enorme stroom aan data die tijdens botsingen vrijkomt.