PQuantML: A Tool for End-to-End Hardware-aware Model Compression

PQuantML is een nieuwe open-source bibliotheek voor hardware-bewuste compressie van neurale netwerken die pruning en kwantisatie integreert om modellen met hoge prestaties en lage latentie te trainen, zoals geïllustreerd door toepassing op jet-tagging voor LHC-data.

De kern

🚀 De Grote Uitdaging: De "Super-Snelle" Deurwachter

Stel je voor dat het Large Hadron Collider (LHC) een gigantisch feest is waar elke seconde 40 miljoen deeltjesbotsingen plaatsvinden. Dat is alsof er elke seconde een hele stad vol mensen door een deur moet rennen.

Je kunt niet alles opslaan (je hebt geen ruimte in je kelder voor zoveel foto's). Dus hebben ze een deurwachter (de "trigger"). Deze deurwachter moet in een fractie van een seconde beslissen: "Blijf je hier? Of ga je weg?"

• Het probleem: De deurwachter moet extreem snel zijn (microseconden) en werkt op speciale, kleine computers (FPGA's) die niet veel ruimte of energie hebben.
• De oplossing: Mensen proberen slimme AI (kunstmatige intelligentie) te gebruiken om deze beslissingen te nemen. Maar normale AI-modellen zijn vaak te groot, te traag en te "dik" voor deze kleine deurwachter. Ze passen niet in de koffer en zijn te traag om op tijd te beslissen.

✂️ Enter PQuantML: De Slimme Verkleiner

PQuantML is een nieuw, gratis gereedschap (een bibliotheek) dat wetenschappers helpt om die zware AI-modellen te "slanken" zonder dat ze hun intelligentie verliezen. Het doet dit op twee manieren, die we kunnen vergelijken met het inpakken van een koffer voor een lange reis:

1. Snoeien (Pruning) = De "Onnodige Kleding" wegdoen

Stel je voor dat je een zware winterjas hebt vol met zakken, maar je gaat naar een tropisch eiland. Je draagt die jas niet, en de zakken zijn nutteloos.

• Wat PQuantML doet: Het kijkt naar het AI-model en zegt: "Deze specifieke knoppen, deze zenuwen en deze geheugenplekken worden nooit gebruikt. Laten we ze eruit knippen."
• Het resultaat: Het model wordt lichter en sneller, maar kan nog steeds dezelfde beslissingen nemen. PQuantML kan dit op verschillende manieren doen:
  • Ongeordend snoeien: Willekeurig kleine stukjes wegdoen (zoals een kale plek in je haar).
  • Geordend snoeien: Hele blokken of lagen wegdoen (zoals het verwijderen van een hele mouw). Dit werkt beter op de speciale hardware.

2. Quantisatie = De "Grote Maat" verkleinen

Stel je voor dat je een foto maakt. Normaal gesproken slaat je computer elke pixel op met 32 cijfers (heel precies, maar veel ruimte nodig).

• Wat PQuantML doet: Het zegt: "We hoeven niet tot op de tiende duizendste nauwkeurig te zijn. Laten we de foto opslaan met maar 8 of zelfs 4 cijfers."
• Het resultaat: De foto (het model) wordt veel kleiner en sneller te verwerken, maar je ziet nog steeds dat het een kat is en geen hond. PQuantML zorgt ervoor dat het model tijdens het leren al gewend raakt aan deze "ruwe" cijfers, zodat het niet vergeten wordt hoe het moet werken.

🛠️ Waarom is dit zo speciaal?

Vroeger moesten wetenschappers deze twee stappen (snoeien en verkleinen) handmatig doen, alsof ze een auto zelf moesten demonteren en opnieuw moesten bouwen. Het was lastig en foutgevoelig.

PQuantML is als een automatische carwash en tune-up service in één:

1. Alles-in-één: Je geeft het een configuratiebestand (een recept), en het doet het snoeien én het verkleinen tegelijkertijd.
2. Hardware-bewust: Het weet precies wat de "deurwachter" (de FPGA) aankan. Het zorgt ervoor dat het eindresultaat past in de ruimte en binnen de tijdslimiet valt.
3. Automatisch optimaliseren: Het probeert automatisch duizenden combinaties om de perfecte balans te vinden tussen "hoe klein kan het zijn?" en "hoe goed werkt het nog?".

🏆 De Test: Jet Tagging

De auteurs hebben PQuantML getest op een echte taak: Jet Tagging.

• Wat is dat? Het herkennen van welk type deeltje een "jet" (een straal deeltjes) heeft veroorzaakt. Dit is cruciaal voor de veiligheid en wetenschap op het LHC.
• Het resultaat: PQuantML slaagde erin om de modellen enorm te verkleinen (minder ruimte, minder energie) terwijl ze net zo goed bleven presteren als de zware, onbewerkte modellen.
• Vergelijking: Het deed het beter dan of net zo goed als bestaande tools (zoals QKeras en HGQ), maar dan met meer flexibiliteit en minder gedoe voor de gebruiker.

💡 Conclusie in één zin

PQuantML is de slimme "verkleiner" die ervoor zorgt dat zware, trage AI-modellen net zo snel en efficiënt worden als een racefiets, zodat ze perfect kunnen werken op de kleine, snelle computers die nodig zijn om deeltjesversnellers veilig en slim te laten draaien.

Het maakt het voor fysici makkelijker om de allerbeste AI te gebruiken, zonder zich zorgen te hoeven maken over de technische beperkingen van de hardware.

Technische samenvatting

Probleemstelling

Experimenten in de Hoge-Energie Fysica (HEP), zoals die aan de Large Hadron Collider (LHC) van CERN, staan voor een extreme data-uitdaging. Proton-proton botsingen vinden plaats met een frequentie van 40 MHz, wat resulteert in datastromen van honderden terabytes per seconde. Het is onmogelijk om alle geregistreerde gebeurtenissen offline op te slaan. Om dit op te lossen, gebruiken experimenten zoals ATLAS en CMS multi-level trigger-systemen. Het hardware-gebaseerde Level-1 trigger-systeem (L1T) moet beslissingen nemen binnen microseconden.

Traditioneel gebruikt het L1T snelle, maar simpele algoritmen. Moderne machine learning (ML) modellen kunnen complexe patronen herkennen en zijn nauwkeuriger, maar ze zijn vaak te zwaar voor directe implementatie in FPGA's (Field-Programmable Gate Arrays) vanwege strenge beperkingen op:

• Latentie: Operaties moeten binnen een vast aantal klokcycli voltooid zijn.
• Hardware-begroting: Beperkt gebruik van DSP's (Digital Signal Processors), BRAM (Block RAM), LUT's (Look-Up Tables) en flip-flops.
• Determinisme: De uitvoeringstijd mag niet variëren op basis van de inputcomplexiteit.

Bestaande tools voor modelcompressie (zoals QKeras en HGQ) focussen vaak primair op kwantisatie, terwijl geavanceerde pruning-strategieën vaak handmatig moeten worden geïmplementeerd. Er is een behoefte aan een geïntegreerde, end-to-end oplossing die zowel pruning als kwantisatie combineert en specifiek is ontworpen voor hardware-bewuste training.

Methodologie: PQuantML

PQuantML is een open-source bibliotheek die een unificatie biedt van pruning en kwantisatie binnen een hardware-bewust trainingsframework. De kern van de methologie omvat:

1. Unified Framework: PQuantML integreert pruning en kwantisatie-aware training (QAT) in één workflow. Het ondersteunt zowel PyTorch als TensorFlow/Keras backends.
2. Hardware-Aware Training: In plaats van compressie als een post-processing stap toe te passen, wordt de compressie direct tijdens het trainingsproces geïntegreerd. Dit stelt het netwerk in staat om robuust te worden tegen de verlies van precisie en parameters.
3. Kwantisatie (Quantization):
   • Ondersteunt Fixed-Point Quantization (laag-precisie vaste komma-getallen) in plaats van floating-point.
   • Implementeert High-Granularity Quantization (HGQ), waarbij bit-breedtes per gewicht of activatie kunnen worden geleerd via gradiënt-gebaseerde optimalisatie.
   • Gebruikt een loss-functie die de "Effective Bit Operations" (EBOPs) minimaliseert, wat een schatting is van de hardware-kost, gecombineerd met een L1-regularisatie om bit-breedtes klein te houden.
4. Pruning:
   • Ondersteunt diverse methoden met verschillende granulariteiten:
     • Unstructured: Verwijdert individuele gewichten (bijv. AutoSparse, Continuous Sparsification).
     • Structured: Verwijdert hele kanalen of filters.
     • Semi-structured (N:M): Verwijdert N gewichten uit elke groep van M, wat een balans biedt tussen expressiviteit en hardware-efficiëntie.
   • Specifieke methoden zoals MDMM (Modified Differential Method of Multipliers) gebruiken constraint-optimatie om sparsiteit te forceren zonder heuristieken.
5. Automatisering en Configuratie:
   • Configuratie gebeurt via YAML-bestanden (gevalideerd met Pydantic), waardoor gebruikers strategieën kunnen specificeren zonder code te herschrijven.
   • Integratie met Optuna voor hyperparameter-optimalisatie (Bayesian search) om de beste trade-offs tussen nauwkeurigheid en hardware-gebruik te vinden.
   • Integratie met MLflow voor experiment-tracking en reproduceerbaarheid.
6. Deployment: Geoptimaliseerde modellen kunnen direct worden geconverteerd naar HLS-code (High-Level Synthesis) via het hls4ml framework voor implementatie op FPGA's.

Belangrijkste Bijdragen

• Eerste stabiele release van PQuantML: Een bibliotheek die pruning en kwantisatie naadloos combineert in één interface, specifiek gericht op real-time HEP-toepassingen.
• End-to-End Workflow: Van modeldefinitie tot hardware-deployment, inclusief automatische hyperparameter-tuning en experiment-tracking.
• Flexibiliteit: Ondersteuning voor zowel directe laag-definitie als automatische vervanging van lagen in bestaande modellen.
• Hardware-bewuste optimalisatie: De mogelijkheid om direct te trainen voor specifieke hardware-beperkingen (DSP, LUT, FF) en latentie-eisen.
• Validatie van HGQ-integratie: Bewijs dat PQuantML modellen kan trainen met HGQ-technieken die vergelijkbaar zijn met de native HGQ-bibliotheek, maar dan binnen een breder compressie-ecosysteem.

Resultaten

De auteurs hebben PQuantML geëvalueerd op de taak van jet substructuur classificatie (jet tagging), een kritieke real-time taak voor LHC-data.

• Datasets: Twee datasets werden gebruikt: High-Level Features (HLF) en Particle-Level Features (PLF).
• Vergelijking: PQuantML werd vergeleken met bestaande tools zoals QKeras en HGQ.
• Prestaties:
  • Nauwkeurigheid: PQuantML behaalde vergelijkbare of betere nauwkeurigheid dan QKeras bij hoge compressie. Bijvoorbeeld, op de HLF-dataset behaalde PQuantML (met DST-pruning) 76,3% nauwkeurigheid, vergelijkbaar met QKeras (76,3%), maar met aanzienlijk minder resources.
  • Resource-reductie:
    • LUT-gebruik: Reductie van 5.504 (QKeras) naar 3.895 (PQuantML DSTw), een daling van ~30%.
    • DSP-gebruik: FITCompress in PQuantML reduceerde DSP-gebruik tot 45 (tegenover 1175 voor een niet-gepruned model).
    • Latentie: Significant verbeterd. Bijvoorbeeld, latentie daalde van 105 ns (QKeras) naar ongeveer 47 ns (PQuantML).
  • HGQ-vergelijking: Modellen getraind met PQuantML + HGQ lieten zeer vergelijkbare resultaten zien (nauwkeurigheid en resource-gebruik) als modellen getraind met de native HGQ-bibliotheek, wat de betrouwbaarheid van de PQuantML-implementatie bevestigt.

Betekenis en Toekomstperspectief

PQuantML vult een cruciale kloof in de toolkit voor fysici en ingenieurs die ML-modellen willen deployen op FPGA's voor real-time triggersystemen. Het stelt onderzoekers in staat om complexe compressietechnieken (pruning + kwantisatie) te gebruiken zonder handmatige ingrepen, wat de drempel verlaagt voor het adopteren van geavanceerde ML in hardware-beperkte omgevingen.

Beperkingen en Toekomstig Werk:

• Huidige beperkingen in de TensorFlow-backend voor complexe architecturen (zoals residual connections).
• De FITCompress-methode is momenteel beperkt tot PyTorch.
• Toekomstige plannen omvatten het uitbreiden van ondersteuning voor TensorFlow, het automatiseren van de zoekruimte voor hyperparameters, en het integreren van aanvullende compressietechnieken zoals kennisdistillatie (knowledge distillation) en low-rank factorisatie.

Concluderend biedt PQuantML een gestroomlijnde, flexibele en reproduceerbare aanpak voor het trainen en deployen van geoptimaliseerde neurale netwerken in de strikte omgeving van de LHC-triggersystemen.