Economical Jet Taggers -- Equivariant, Slim, and Quantized

Dit artikel presenteert een slanke, gekwantiseerde en met parameters gereduceerde versie van de L-GATr jet tagger die een reductie in energiekosten van een orde van grootte bereikt met slechts een matige prestatievermindering, wat de weg vrijmaakt voor efficiënte jet tagging op trigger-niveau bij de LHC.

De kern

Stel je de Large Hadron Collider (LHC) voor als een enorme, hogesnelheidsdeeltjesfabriek. Elke seconde laat hij protonen op elkaar botsen, wat een chaotische spray van puin creëert. Natuurkundigen moeten dit puin sorteren om specifieke, zeldzame deeltjes (zoals het "topquark") te vinden die verborgen liggen tussen miljarden gewone deeltjes. Dit sorteerproces wordt jet tagging genoemd.

Jarenlang hebben wetenschappers complexe computerprogramma's (Machine Learning) gebruikt om dit sorteren te doen. De huidige kampioenen zijn "Transformers"—krachtige AI-modellen die ongelooflijk nauwkeurig zijn, maar ook enorm, traag en hongerig naar energie. Ze zijn als een vloot enorme, brandstofverslindende vrachtwagens die proberen een enkele brief te bezorgen; ze krijgen de klus geklaard, maar ze zijn te groot en te duur om op het moment dat de data wordt verzameld (het "trigger"-niveau) te gebruiken.

Dit artikel stelt een simpele vraag: Kunnen we deze gigantische vrachtwagens verkleinen tot kleine, brandstofefficiënte scooters zonder het vermogen te verliezen om de brief te bezorgen?

Zo hebben de auteurs dit aangepakt, met behulp van drie hoofdstrategieën:

1. De "Slanke" versie (L-GATr-slim)

Het originele "L-GATr"-model is als een Zwitsers zakmes dat elk mogelijk gereedschap bij zich draagt: scalairen, vectoren, tensoren en meer. De auteurs realiseerden zich echter dat je voor de meeste deeltjesfysica-taken eigenlijk maar twee gereedschappen nodig hebt: scalairen (getallen) en vectoren (pijlen met een richting).

• De Analogie: Stel je een chef voor die erop staat een volledige industriële keuken met ovens, blenders en mengers te gebruiken om slechts een simpel broodje te maken. De auteurs zeiden: "Laten we gewoon een mes en een snijplank gebruiken."
• Het Resultaat: Ze bouwden een "Slim" versie van de AI die de overbodige gereedschappen verwijdert. Het presteert net zo goed als de gigantische versie, maar is veel sneller te trainen en verbruikt minder geheugen. Het is alsof je wisselt van een zware vrachtwagen naar een wendbare sportauto die dezelfde klus klaart.

2. De "Miniatuur" versie (Ultra-mini Taggers)

De auteurs vroegen zich vervolgens af: "Hoe klein kunnen we gaan?" Ze probeerden deze AI-modellen te verkleinen tot de grootte van een piepkleine speelgoedauto (ongeveer 1.000 parameters, vergeleken met de miljoenen in het origineel).

• De Analogie: Denk aan het proberen om de kennis van een hele bibliotheek in een enkele ansichtkaart te passen. Normaal gesproken verlies je dan het verhaal. Maar de auteurs ontdekten dat als je de informatie correct organiseert (met behulp van specifieke "Lorentz-equivariante" regels die de natuurwetten respecteren), je de essentiële kennis in een minuscule ruimte kunt passen.
• Het Resultaat: Ze ontdekten dat voor zeer kleine modellen de "LLoCa"-architectuur het beste werkt als je het aantal lagen verkleint, terwijl de "L-GATr-slim" het beste werkt als je de breedte van de lagen verkleint. Zelfs op deze microscopische grootte presteren ze nog steeds beter dan oudere, niet-natuurkundig-bewuste AI-modellen.

3. De "Gekwantiseerde" versie (Wiskunde met lage precisie)

Dit is de meest dramatische energiebesparing. Standaard AI gebruikt zeer nauwkeurige wiskunde (zoals het meten van een afstand tot de miljardste millimeter nauwkeurig). De auteurs realiseerden zich dat je voor jet tagging niet zoveel precisie nodig hebt. Je kunt het je veroorloven om getallen aanzienlijk af te ronden.

• De Analogie: Stel je voor dat je appels telt in een magazijn.
  • Standaard AI: Je weegt elke appel tot op de microgram nauwkeurig. (Nauwkeurig, maar het duurt eeuwen en verbruikt veel energie van de weegschaal).
  • Gekwantiseerde AI: Je telt ze gewoon in hele getallen. (Snel, verbruikt bijna geen energie, en voor het doel om te weten "hoeveel appels er zijn", is het volkomen voldoende).
• De Methode: Ze gebruikten een techniek genaamd PARQ (Piecewise-Affine Regularized Quantization). Zie dit als een slimme afrondingsregel die de getallen tijdens het trainingsproces voorzichtig naar eenvoudige waarden stuurt (zoals 0, 1 of -1), in plaats van ze abrupt te dwingen.
• Het Resultaat: Door over te stappen op deze "ruwere" getallen, verminderden ze de energiekosten van het draaien van de AI met 10 keer (een orde van grootte). De AI werd ongelooflijk snel en energiezuinig, met slechts een minimale daling in nauwkeurigheid.

Het Grotere Plaatje

De auteurs combineerden deze drie strategieën — het slanker maken van de architectuur, het verkleinen van de omvang en het kwantiseren van de wiskunde — om "Economische Jet Taggers" te creëren.

• Waarom is dit belangrijk? Momenteel zijn deze krachtige AI-modellen te groot om te draaien op de hardware die in realtime beslist welke botsingen bewaard worden en welke worden weggegooid (de "trigger").
• Het Doel: Door deze modellen klein, snel en energiezuinig te maken, hopen de auteurs ze uiteindelijk direct op de trigger-hardware te kunnen draaien. Dit zou de LHC in staat stellen om AI te gebruiken voor beslissingen in een fractie van een seconde over welke deeltjesbotsingen bewaard moeten worden, wat potentieel nieuwe fysica kan ontdekken die eerder werd gemist omdat de data te snel werd weggegooid.

Kortom: Ze hebben een gigantische, energieverslindende AI een dieet gegeven, hem kleiner gemaakt en hem geleerd om wiskunde te doen met minder decimalen, wat resulteerde in een kleine, super-efficiënte motor die nog steeds de belangrijkste deeltjes in het universum kan herkennen.

Technische samenvatting

Technische Samenvatting: Economische Jet Taggers – Equivariant, Slim en Gekwantiseerd

Probleemstelling
Moderne machine learning (ML) heeft de jet-tagging bij de Large Hadron Collider (LHC) getransformeerd, waarbij Lorentz-equivariante transformers de standaard zijn geworden. Echter, toonaangevende modellen zoals L-GATr zijn computationeel duur en vereisen aanzienlijk geheugen en trainingstijd. Terwijl de industriële trend de nadruk legt op het opschalen van netwerken en datasets, wordt de LHC-fysica geconfronteerd met specifieke beperkingen, met name met betrekking tot de geheugen- en latentievereisten van hardware voor event-triggering. Huidige jet-classificatie speelt nog geen rol in de trigger, maar de auteurs stellen dat dit wel zou moeten. De centrale uitdaging die wordt aangepakt, is hoe de omvang en de computationele kosten van moderne equivariante jet taggers kunnen worden verminderd, terwijl de degradatie van de prestaties tot een minimum wordt beperkt, met het oog op mogelijke implementatie op het trigger-niveau.

Methodologie
Het artikel stelt een tweeledige strategie voor om de efficiëntie van de middelen te optimaliseren: architecturale inkrimping (slimming) en numerieke kwantisering.

1. L-GATr-slim Architectuur:
   De auteurs introduceren een gestroomlijnde versie van de Lorentz-equivariante transformer (L-GATr). Standaard L-GATr maakt gebruik van een geometrische algebra-representatie bestaande uit scalairen, pseudo-scalaren, vectoren, axiale-vectoren en antisymmetrische rang-twee tensoren. De auteurs observeren dat pseudo-scalaren, axiale-vectoren en tensoren voor de meeste LHC-toepassingen overbodig zijn. Daarom beperkt L-GATr-slim de latente representatie tot enkel scalaren en vectoren.

   • Lineaire Lagen: Worden uitgebreid om te opereren op gekoppelde scalaire en vector-representaties, waarbij de vectorcomponenten een enkele leerbare scalaire coëfficiënt delen om Lorentz-equivariantie te behouden.
   • Niet-lineariteit: Past de Gated Linear Unit (GLU) aan door niet-lineariteit toe te passen op het inproduct van twee vectoren, vermenigvuldigd met de vector-output.
   • Normalisatie: Wijzigt RMSNorm om het absolute getal van het Minkowski-inproduct voor vectorkanalen te gebruiken.
   • Attention: Construeert scalaire attention-matrices met een specifieke formulering die het computationeel dure buitenproduct vermijdt dat in de volledige L-GATr wordt gebruikt.
   • Implementatie: De architectuur is ontworpen om te worden gecompileerd met torch.compile voor efficiëntie.

2. Kwantiseringsstrategieën:
   De auteurs passen low-precision datatypen en gewichtskwantisering toe om verdere kosten te verlagen.

   • Datatype Kwantisering: Inputs naar lineaire lagen worden gekwantiseerd naar int8 (met gebruik van zero-point kwantisering), terwijl bfloat16 behouden blijft voor precisiegevoelige operaties en de backward pass. Dit wordt toegepast op de verborgen lagen van Transformer, ParT, L-GATr-slim en LLoCa-Transformer.
   • Gewichtskwantisering: Lineaire gewichten worden gekwantiseerd naar binaire of ternaire waarden met behulp van Proximal Gradient Quantization (PARQ). Deze methode behandelt kwantisering als een regularisatie-constraint, waarbij een proximale operator wordt gebruikt om gewichten bij te werken. De auteurs vergelijken PARQ met Straight-Through Estimation (STE) en stellen vast dat PARQ een betere stabiliteit en prestaties biedt.
   • Behoud van Equivariantie: Er wordt speciale zorg besteed aan het waarborgen dat kwantisering de Lorentz-equivariantie niet schendt. Voor LLoCa blijven orthonormalisatie en frame-projecties in volledige precisie (float32), waardoor low-precision operaties beperkt blijven tot Lorentz-invarianten. Voor L-GATr-slim worden volledige vectoren vermenigvuldigd met gekwantiseerde gewichten, wat geen extra schendingen van de symmetrie introduceert.

3. Ultra-Mini Schaling:
   De auteurs onderzoeken de prestaties van deze architecturen tot beneden de 1.000 parameters door het aantal blokken of de breedte (kanalen) van het netwerk te verminderen.

Belangrijkste Resultaten
De studie benchmarkt de voorgestelde methoden op drie taken: top tagging, amplitude regressie en event generatie.

• Prestaties vs. Efficiëntie (L-GATr-slim):

  • Op de JetClass dataset (multi-class jet tagging) evenaart L-GATr-slim de prestaties van de volledige L-GATr en LLoCa-Transformer (AUC ~0.9885), maar vermindert de trainingstijd met een factor zes (van 166u naar 27u op een H100 GPU) en het geheugengebruik met een factor twee.
  • In amplitude regressie ($Z + 4g$) bereikt L-GATr-slim dezelfde Mean Squared Error (MSE) als de volledige L-GATr, maar vereist 20 keer minder trainingsoperaties en de helft van de trainingstijd.
  • In event generatie ($t\bar{t} + nj$) evenaart de slimme architectuur de negative log-likelihood prestaties van de volledige modellen.

• Ultra-Mini Taggers:

  • Wanneer het aantal blokken (diepte) wordt verminderd, presteert de LLoCa-Transformer beter dan L-GATr-slim bij zeer kleine formaten (bijv. 1.000 parameters).
  • Wanneer het aantal kanalen (breedte) gelijk wordt gehouden bij een vast aantal blokken (10), behoudt L-GATr-slim een achtergrond-rejection rate boven de 1.000 met slechts 2 vector- en 4 scalaire kanalen, waarmee het andere 1.000-parameter architecturen overtreft.

• Kwantiseringswinsten:

  • Het kwantiseren van inputs naar int8 en gewichten naar ternaire waarden vermindert het energieverbruik met ongeveer een orde van grootte (factor 10) met slechts marginaal prestatieverlies.
  • De LLoCa-Transformer en L-GATr-slim zijn robuust tegen kwantisering; ze behouden een hoge prestatie waar standaard transformers meer zouden degraderen.
  • Voor het meest hulpbron-beperkte scenario (1 blok, 16-dimensionale latente ruimte, int8), behoudt de gekwantiseerde LLoCa-Transformer (globale canonicalisatie) een prestatie die superieur is aan pre-graph taggers, ondanks een factor-twee reductie in achtergrond-rejection vergeleken met zijn volledige tegenhanger.

Betekenis en Claims
Het artikel beweert dat deze "economische" versies van equivariante transformers een levensvatbaar pad vormen naar trigger-level jet tagging bij de High-Luminosity LHC (HL-LHC). Door het combineren van architecturale inkrimping (het verwijderen van onnodige geometrische algebra componenten) en agressieve kwantisering (PARQ en int8), demonstreren de auteurs dat het mogelijk is om taggers met ~1.000 parameters te creëren die de natuurkunde-gemotiveerde voordelen van Lorentz-equivariantie behouden.

De auteurs benadrukken dat hoewel opschalen de industriestandaard is, de LHC-fysica een "physics-aware downscaling" benadering vereist. De resultaten suggereren dat kleine, gekwantiseerde en equivariante netwerken op hulpbron-beperkte hardware (zoals FPGA's) kunnen worden ingezet zonder de fundamentele symmetrieën op te offeren die hun effectiviteit bepalen, wat potentieel nieuwe wegen opent voor real-time analyse van jet-substructuur.