TrackCore-F: Deploying Transformer-Based Subatomic Particle Tracking on FPGAs

Dit paper beschrijft de ontwikkeling van methoden en tools voor het implementeren van Transformer-gebaseerde subatomaire deeltjessporen op FPGAs, met als doel de beperkingen van bestaande hulpmiddelen en hardwarebronnen te overwinnen door middel van monolithische of gepartitioneerde synthese.

De kern

🚀 De Jacht op de Spookdeeltjes: Hoe we AI op een "Schakelbord" zetten

Stel je voor dat je in een gigantische, donkere danszaal staat (deeltjesversneller LHC). Miljoenen mensen rennen rond, botsen en vliegen in alle richtingen. Je taak is om te kijken wie met wie heeft gedanst, zelfs als ze elkaar maar een fractie van een seconde hebben geraakt. Dit noemen we deeltjes-tracking.

Vroeger deden computers dit met een heel trage, slimme "post-mortem" analyse (na de feiten). Maar nu willen we dit live doen, terwijl de botsingen plaatsvinden. Dat is als proberen de danspartners te herkennen terwijl de muziek nog speelt en de dansvloer trilt.

Hier komt dit paper om de hoek kijken. Het vertelt ons hoe we een heel slimme AI (een Transformer, de technologie achter moderne chatbots) op een heel specifiek stukje hardware zetten: een FPGA.

🧩 Wat is een FPGA? (Het Legobord dat zichzelf herschrijft)

Normaal gesproken draait zware AI op enorme, energievretende videokaarten (GPUs), alsof je een Formule 1-auto gebruikt om boodschappen te doen. Dat is te duur en te groot voor een deeltjesdetector.

Een FPGA is anders. Stel je voor dat het een legendarisch Legobord is.

• Je kunt er een auto van bouwen.
• Je kunt er een vliegtuig van bouwen.
• Je kunt er zelfs een robot van bouwen.
• En het allerbelangrijkste: je kunt het terwijl het draait ombouwen naar precies wat je nodig hebt.

Het paper wil die "robot" bouwen die live meekijkt naar de deeltjesbotsingen. Het is energiezuinig, snel en past in de machine.

🏗️ Het Probleem: De AI is te groot voor het Legobord

De "slimme" AI-modellen die ze gebruiken (genaamd EncCla en EncReg) zijn als een gigantische bibliotheek. Ze bevatten miljoenen regels code (parameters).

• Het Legobord (de FPGA) heeft echter maar een beperkte hoeveelheid Legostenen (geheugen en rekenkracht).
• Je kunt niet de hele bibliotheek in één keer op het bord zetten; het past er niet.

De oplossing van de auteurs:
In plaats van de hele bibliotheek te verplaatsen, nemen ze slechts één hoofdstuk (één laag van de AI) en bouwen ze dat als een speciaal blokje op het Legobord.

• De rest van de AI draait nog op de normale computer.
• Het "hoofdstuk" op het Legobord doet het zware werk.
• Ze sluiten de stukken aan alsof je een stukje van een trein koppelt aan de rest.

Dit noemen ze partitioneren: het AI-model in stukjes hakken en de moeilijkste stukken op het Legobord zetten.

⚡ De Uitdaging: De "Kwaliteit" van de AI

Om dit Legobord niet te overladen, willen ze de "Legostenen" kleiner maken. In de AI-wereld noemen ze dit kwantisatie.

• Normaal: De AI gebruikt heel precieze getallen (zoals 3,14159265...). Dit is als het meten met een microscoop.
• Gekwantiseerd: Ze ronden af naar hele getallen (zoals 3 of 4). Dit is als meten met een liniaal.

Het resultaat (de verrassing):
De auteurs hebben gekeken wat er gebeurt als ze de getallen afkappen.

• Als ze alleen de "gewichten" (de kennis van de AI) afkappen, blijft het nog redelijk goed.
• Maar als ze ook de "activaties" (de actieve gedachten van de AI) afkappen, valt de prestatie in elkaar. De AI wordt plotseling heel dom en maakt veel fouten.
• Vergelijking: Het is alsof je een chef-kok dwingt om met een lepel te koken in plaats van een mes. Hij kan nog wel koken, maar de gerechten worden minder lekker.

📉 De Resultaten: Hoeveel past er?

Ze hebben het getest op een bordje genaamd ZCU102.

• Ze konden één laag van de AI perfect laten draaien.
• Ze hebben gekeken hoeveel ruimte dit inneemt. Het grootste probleem was het geheugen (BRAM). Het bordje was al 38% vol met alleen maar één laag.
• Conclusie: Ze kunnen maximaal 4 lagen van deze AI op dit specifieke bordje kwijt. Meer dan dat, en het bordje springt uit elkaar (of wordt te traag).

💡 De Grote Les

Het paper zegt eigenlijk: "We kunnen deze super-snelle AI op de deeltjesversneller zetten, maar we moeten slim zijn."

1. Splitsen is slim: Je hoeft niet alles op het bord te zetten. Deel het op.
2. Kwaliteit vs. Snelheid: Als je de AI te veel "verkleint" (kwantisatie) om ruimte te besparen, wordt hij te dom om nuttig te zijn.
3. De toekomst: Dit is de eerste stap. Ze hebben bewezen dat het werkt, en nu moeten ze de gereedschappen verbeteren zodat dit proces in de toekomst automatisch gaat, in plaats van handmatig.

Kortom: Ze zijn bezig met het bouwen van een "live-deeltjesdetector" die slim genoeg is om de chaos van het universum te ordenen, maar klein en zuinig genoeg om in de machine te passen. Een echte wetenschappelijke puzzel!

Technische samenvatting

Probleemstelling

In de Hoge-Energie Fysica (HEP), en specifiek bij de Large Hadron Collider (LHC) en de toekomstige High-Luminosity LHC (HL-LHC), is de reconstructie van deeltjesbanen (tracking) een kritische taak. Traditioneel is dit een post-mortem proces vanwege de hoge rekenvereisten. Hoewel Machine Learning (ML) modellen, zoals Transformatoren, de nauwkeurigheid en snelheid van deze reconstructie hebben verbeterd, vormt de implementatie van deze modellen in real-time (online of pseudo-online) een grote uitdaging.

De huidige ML-deployments zijn gedomineerd door Graphics Processing Units (GPUs), maar deze zijn niet altijd geschikt voor on-site implementatie in deeltjesdetectoren vanwege beperkingen in energie-efficiëntie en latentie. Field-Programmable Gate Arrays (FPGA's) bieden een aantrekkelijk alternatief vanwege hun lage latentie en energie-efficiëntie. Echter, het deployen van complexe Transformer-architecturen op FPGA's is nog in een vroeg stadium. Bestaande tools bieden beperkte of geen ondersteuning voor het automatisch partitioneren van grote modellen, en FPGA-resources (zoals geheugen en logische blokken) zijn een significant knelpunt.

Methodologie

Het paper beschrijft een gestructureerde ontwikkelworkflows voor het deployen van vooraf getrainde Transformer-modellen op FPGA's, specifiek gericht op de tracking van subatomaire deeltjes.

1. Doelmodellen: De auteurs focussen op twee Transformer-architecturen uit het TrackFormers-project:
   • EncCla (Encoder-Classifier): Een encoder-only model dat hit-coördinaten mapt naar vooraf gedefinieerde klasse-labels (binning van track-parameters).
   • EncReg (Encoder-Regressor): Een encoder-only model dat direct de parameters van tracks regresseert, gevolgd door clustering (HDBSCAN) om tracks te associëren.
2. Hardware Platform: De implementatie is getest op een ARM Zynq UltraScale+ MPSoC ZCU102 ontwikkelbord, met een ZU9EG-chip (Quad-Core ARM Cortex-A53 + Programmable Logic).
3. Toolchain en Workflow:
   • PyTorch & ONNX: Modellen worden geconverteerd van PyTorch naar het Open Neural Network Exchange (ONNX) formaat om laag-niveau operaties en data-types bloot te leggen.
   • Handmatig Partitioneren: Omdat volledige modellen vaak te groot zijn, wordt het model handmatig opgesplitst in drie delen: het segment voor de encoder, de encoder-laag zelf (die op de FPGA wordt geïmplementeerd), en het resterende model.
   • Vitis HLS: De geselecteerde encoder-laag wordt herschreven in C/C++ en geoptimaliseerd met HLS-pragmas om een hardware IP-kern (Intellectual Property) te genereren.
   • Vivado: De gegenereerde IP wordt geïntegreerd met het MPSoC-systeem, verbonden via AXI4-bussen, en de uiteindelijke bitstream wordt gegenereerd.
4. Quantisatie: Er is onderzocht of het gebruik van kwantisatie (INT8 voor gewichten en activeringen) de prestaties kan verbeteren ten koste van de nauwkeurigheid.

Belangrijkste Bijdragen

• Ontwikkelworkflows voor Transformer-deploy: Het paper presenteert een concrete methodologie voor het deployen van Transformer-modellen op FPGA's, inclusief de noodzaak van handmatig partitioneren en het fusioneren van dataflow tussen software- en hardware-gedeeltes.
• Monolithische Kernel Ontwerp: In plaats van individuele ONNX-operaties te implementeren, tonen de auteurs aan dat het ontwikkelen van een monolithische kern voor een encoder-segment de resource-efficiëntie (CLB en geheugen) verbetert.
• Analyse van Quantisatie-impact: Een gedetailleerde analyse van hoe quantisatie (INT8 vs. INT16) de voorspellingsnauwkeurigheid beïnvloedt, wat essentieel is voor het afwegen van prestaties versus nauwkeurigheid.
• Resource-gebruik Beoordeling: Een realistische inschatting van de hardware-beperkingen op een ZCU102 FPGA voor Transformer-architecturen.

Resultaten

• Nauwkeurigheid en Quantisatie:
  • Het basismodel heeft een voorspellingsnauwkeurigheid van 0,97.
  • Quantisatie heeft een significant negatief effect op de nauwkeurigheid.
  • Het kwantiseren van alleen de gewichten (INT8) met behoud van INT16-activeringen resulteert in een nauwkeurigheid van 0,90.
  • Het kwantiseren van zowel gewichten als activeringen naar INT8 daalt de nauwkeurigheid drastisch naar 0,70.
  • Conclusie: Het kwantiseren van activeringen heeft een grotere negatieve impact dan het kwantiseren van gewichten.
• Resource-gebruik (ZCU102):
  • Voor een enkele encoder-laag is de Block RAM (BRAM) de beperkende factor met een gebruik van 38,04% (347 van 912 BRAMs).
  • De gebruikte LUTs bedragen 24,01% en Flip-Flops 13,53%.
  • Op basis van de BRAM-beperkingen kan er maximaal 4 encoder-lagen op deze specifieke FPGA worden geïmplementeerd voor dit model.
  • Het gebruik van extern DDR-geheugen kan de BRAM-belasting verminderen, maar introduceert extra latentie en data-overdracht.

Betekenis en Conclusie

Het paper demonstreert dat het deployen van Transformer-based tracking op FPGA's haalbaar is, zelfs als volledige deploy niet mogelijk is door resource-beperkingen. Een gedeeltelijke deploy (waarbij slechts een deel van het model op de FPGA draait) wordt beschouwd als waardevol, omdat het toelaat om inferentie uit te voeren op meer toegankelijke hardware met betere latentie dan CPU/GPU-oplossingen.

De studie benadrukt dat optimalisaties (zoals het ontwerp van de kern en het partitioneren) vaak belangrijker zijn dan de HLS-implementatie zelf. Daarnaast waarschuwt het voor de gevaren van agressieve quantisatie, die de nauwkeurigheid van deeltjesreconstructie ernstig kan aantasten. De ontwikkelde methodologie vormt een basis voor toekomstige online tracking-systemen in de Hoge-Energie Fysica, met name voor de HL-LHC, waarbij energie-efficiëntie en lage latentie cruciaal zijn.