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

PQuantML est une nouvelle bibliothèque open source de compression de modèles neuronaux adaptée au matériel, qui simplifie le déploiement de modèles performants sous contraintes de latence en offrant une interface unifiée pour l'élagage et la quantification, comme le démontre son efficacité sur des tâches de classification de jets au LHC.

L'essentiel

🚀 PQuantML : Le "Tuning" pour les Cerveaux d'Ordinateur

Imaginez que vous êtes un ingénieur dans une usine géante (le Grand Collisionneur de Hadrons ou LHC au CERN). Cette usine produit des millions de collisions de particules chaque seconde. C'est comme si vous aviez un torrent d'eau qui dévale une montagne : il y a trop d'eau pour tout stocker dans des réservoirs !

Pour gérer ce déluge, vous avez besoin de gardiens très rapides (les déclencheurs ou triggers) qui doivent décider en une fraction de seconde (des microsecondes) quelles gouttes d'eau sont intéressantes et lesquelles sont juste de l'eau sale à jeter.

Le problème ? Les "cerveaux" (les modèles d'intelligence artificielle) que nous voulons utiliser pour faire ce tri sont trop gros, trop lourds et trop lents pour tenir dans les petits ordinateurs spéciaux (les FPGA) installés sur place.

C'est là qu'intervient PQuantML.

🛠️ Qu'est-ce que PQuantML ?

Pensez à PQuantML comme à un atelier de rénovation ultra-puissant pour les modèles d'intelligence artificielle. Son but est de transformer un "gros éléphant" (un modèle complexe et lent) en un "furet agile" (un modèle rapide, léger et précis) sans perdre sa capacité à chasser (sa précision).

Il le fait grâce à deux techniques principales, que l'on peut imaginer ainsi :

1. La Taille (Le Pruning) 🌳

Imaginez un arbre très touffu. Beaucoup de branches sont mortes ou ne portent pas de fruits.

• Ce que fait PQuantML : Il coupe les branches inutiles.
• La différence : Les anciens outils coupaient au hasard (comme un tondeuse qui fait des trous bizarres dans la pelouse). PQuantML est plus intelligent : il peut couper des branches entières d'un coup (ce qui est plus facile à gérer pour les machines) ou même des motifs précis (comme couper 2 branches sur 4). Cela rend l'arbre plus petit et plus rapide à traverser.

2. La Précision (La Quantification) 📏

Imaginez que vous devez mesurer la hauteur de chaque personne dans une foule.

• Avant : Vous utilisez un mètre ruban ultra-précis capable de mesurer au millimètre près (des nombres à virgule complexes). C'est précis, mais ça prend du temps et de la place pour écrire les chiffres.
• Avec PQuantML : Il vous dit : "Pour notre besoin, on n'a pas besoin du millimètre. On peut arrondir au centimètre, voire au demi-mètre."
• Le résultat : Les chiffres deviennent plus petits (des nombres entiers simples). Les machines FPGA adorent ça : elles calculent beaucoup plus vite avec des chiffres simples. PQuantML apprend au modèle à faire ce "rondissage" pendant qu'il apprend, pour qu'il ne se trompe pas trop.

🤝 La Grande Innovation : Tout en Un

Avant PQuantML, c'était comme si vous deviez embaucher deux artisans séparés :

1. Un pour tailler l'arbre (le Pruning).
2. Un pour changer les règles de mesure (la Quantification).
   Ils ne parlaient pas le même langage, et le résultat était souvent désordonné.

PQuantML est le chef d'orchestre qui fait tout en même temps. Il dit à l'IA : "Coupe cette branche, et en même temps, arrondis ce chiffre". Il apprend au modèle à s'adapter à ces contraintes pendant qu'il étudie, pas après. C'est comme apprendre à conduire une voiture de course sur un circuit étroit dès le premier jour, plutôt que d'essayer de la réduire après l'avoir conduite sur autoroute.

🏁 Le Résultat : Pourquoi c'est génial ?

Les auteurs ont testé leur outil sur un vrai problème du CERN : identifier le type de particule qui a créé un "jet" de débris (un peu comme identifier le type de voiture qui a causé un accident en regardant les débris).

Les résultats montrent que PQuantML :

• Réduit la taille du modèle de façon spectaculaire (moins de mémoire nécessaire).
• Accélère le calcul (moins de temps de réponse, crucial pour ne pas rater les événements rares).
• Garde la précision : Le modèle reste aussi bon que l'original, même s'il est beaucoup plus petit.
• Est compatible : Il se connecte directement aux outils qui transforment le code en circuits électroniques réels.

🎯 En résumé

PQuantML, c'est l'outil qui permet de transformer des modèles d'intelligence artificielle lourds et lents en des versions "compactes" et ultra-rapides, prêtes à être installées dans les petits ordinateurs des accélérateurs de particules.

C'est comme passer d'un camion de déménagement à une moto de course : même destination, même capacité à transporter l'essentiel, mais beaucoup plus rapide et agile pour naviguer dans les embouteillages de données du futur ! 🏎️💨

Résumé technique

1. Problématique

Les expériences de physique des hautes énergies (PHE) au Grand collisionneur de hadrons (LHC), telles qu'ATLAS et CMS, font face à un défi de données extrême. Les collisions proton-proton se produisent à une fréquence de 40 MHz, générant des centaines de téraoctets de données par seconde. Il est impossible de stocker tous les événements en mode hors ligne.

Pour y remédier, le système de déclenchement (trigger) opère en deux niveaux :

• Niveau 1 (L1T) : Un système matériel basé sur des FPGA (Field-Programmable Gate Arrays) qui prend des décisions en quelques microsecondes.
• Niveau haut (HLT) : Un système logiciel plus lent.

Le L1T repose traditionnellement sur des algorithmes simples et rapides. Cependant, l'intégration de modèles d'apprentissage automatique (ML) complexes pour améliorer la précision de sélection est entravée par des contraintes strictes de latence et de ressources matérielles (DSP, LUT, BRAM, FF). Les modèles ML standards sont trop lourds et trop lents pour être déployés directement sur ces FPGA.

Le défi principal réside dans la nécessité de réduire drastiquement la complexité computationnelle et la précision numérique des modèles (par le pruning et la quantification) tout en préservant les performances physiques, le tout sans sacrifier la latence déterministe requise par le matériel.

2. Méthodologie

Pour répondre à ce besoin, les auteurs présentent PQuantML, une bibliothèque open-source et unifiée conçue pour le compression de modèles de bout en bout, spécifiquement adaptée aux contraintes matérielles des FPGA.

L'approche repose sur plusieurs piliers techniques :

• Intégration Unifiée du Pruning et de la Quantification : Contrairement aux outils existants qui traitent souvent la quantification et le pruning séparément, PQuantML les intègre dans un seul cadre. Elle permet d'appliquer ces techniques soit conjointement, soit individuellement, directement durant l'entraînement (Quantization-Aware Training - QAT).
• Pruning (Élagage) Multi-Granularité : La bibliothèque implémente plusieurs méthodes d'élagage avec différentes granularités :
  • Non structuré : Suppression de poids individuels (ex: AutoSparse, DST, Wanda).
  • Structuré : Suppression de canaux entiers ou de filtres (ex: Activation Pruning, MDMM).
  • Semi-structuré (N:M) : Motifs réguliers (ex: 2:4) compatibles avec les accélérateurs matériels.
  • Des méthodes avancées comme MDMM (Modified Differential Method of Multipliers) formulent l'élagage comme un problème d'optimisation sous contraintes.
• Quantification à Haute Granularité (HGQ) : PQuantML étend les fonctionnalités de HGQ (High-Granularity Quantization). Elle permet une quantification en point fixe avec un contrôle fin de la précision (bits entiers et fractionnaires) par tensor, par canal ou par poids. Elle utilise une fonction de perte basée sur les EBOPs (Effective Bit Operations) pour optimiser le compromis entre précision et coût matériel.
• Flux de Travail Automatisé et Optimisation des Hyperparamètres :
  • La configuration se fait via des fichiers YAML validés par Pydantic.
  • L'outil intègre Optuna pour l'optimisation bayésienne des hyperparamètres (recherche de la meilleure configuration de compression).
  • Il utilise MLflow pour le suivi des expériences et la reproductibilité.
• Compatibilité Matérielle (hls4ml) : PQuantML est conçue pour interagir nativement avec hls4ml, l'outil standard de synthèse de haut niveau pour FPGA. Une fois le modèle entraîné et compressé, il est converti en code HLS (High-Level Synthesis) pour le déploiement sur FPGA.

3. Contributions Clés

• Première bibliothèque unifiée : PQuantML est la première bibliothèque open-source offrant une interface unifiée pour combiner le pruning (structuré et non structuré) et la quantification (point fixe et HGQ) dans un flux de travail d'entraînement adapté au matériel.
• Support Multi-Frontend : Elle fonctionne avec les backends PyTorch et Keras/TensorFlow, offrant une interface utilisateur cohérente.
• Méthodes Avancées d'Élagage : Introduction et implémentation de méthodes sophistiquées comme MDMM et FITCompress (détermination automatique de la largeur de bits et du niveau d'élagage par couche).
• Réduction de la Barrière d'Entrée : En automatisant les flux de travail complexes de compression et d'optimisation, PQuantML rend accessible aux physiciens l'adoption de techniques de compression avancées pour les applications temps réel.
• Validation sur des Tâches Réelles : Le framework a été testé sur des tâches critiques de physique, notamment la classification de la sous-structure des jets (jet tagging).

4. Résultats

Les auteurs ont évalué PQuantML sur la tâche de classification de la sous-structure des jets (Jet Substructure Classification - JSC) en utilisant deux jeux de données (HLF et PLF) et en comparant les performances avec des outils existants comme QKeras et HGQ.

• Réduction des Ressources :
  • Sur le jeu de données HLF, PQuantML (avec élagage DST) a réduit l'utilisation des LUT de 5 504 à 3 895 et la latence de 105 ns à environ 47 ns, tout en maintenant une précision d'environ 76,3 %.
  • L'utilisation de FITCompress a permis de réduire l'utilisation des DSP à 45 (contre 1175 pour le modèle non compressé) avec une latence de 36,2 ns, au prix d'une légère baisse de précision (76,0 %).
• Comparaison avec l'État de l'Art :
  • PQuantML surpasse QKeras en termes de réduction des ressources et de latence pour des niveaux de précision similaires.
  • Les modèles entraînés avec PQuantML utilisant la quantification HGQ obtiennent des résultats comparables (en termes de précision et d'utilisation des ressources) à ceux entraînés directement avec la bibliothèque HGQ native, validant ainsi l'intégration.
• Flexibilité : Les résultats montrent que PQuantML permet d'explorer efficacement le front de Pareto entre la précision du modèle, la consommation de ressources (DSP, LUT) et la latence, offrant des compromis adaptés aux contraintes spécifiques du LHC.

5. Signification

PQuantML représente une avancée significative pour le déploiement de l'intelligence artificielle dans les systèmes de déclenchement du LHC et au-delà.

• Impact Physique : En permettant le déploiement de modèles ML complexes sur le L1T (niveau matériel), PQuantML ouvre la voie à une sélection d'événements plus intelligente et plus précise, potentiellement capable de capturer des phénomènes physiques rares qui échapperaient aux algorithmes traditionnels.
• Efficacité Ingénieriale : La bibliothèque résout le problème de la fragmentation des outils de compression. Elle fournit un pipeline complet, de la définition du modèle à la génération du firmware FPGA, réduisant le temps de développement et les erreurs d'intégration.
• Avenir du ML Temps Réel : En démontrant que l'entraînement conscient de la compression (QAT) couplé à l'élagage structuré et à la quantification fine est viable sur FPGA, PQuantML établit un nouveau standard pour les applications ML à faible latence et à ressources contraintes dans la science des données et l'informatique embarquée.

En résumé, PQuantML est un outil essentiel pour transformer des architectures neuronales génériques en designs matériels optimisés, répondant aux exigences extrêmes de la physique des hautes énergies moderne.