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

Das Open-Source-Tool PQuantML ermöglicht hardwarebewusste, End-to-End-Neuronale-Netzwerk-Kompression durch kombiniertes Pruning und Fixed-Point-Quantisierung, wobei es in Tests zur Jet-Tagging-Aufgabe signifikante Reduktionen bei gleichzeitiger Beibehaltung der Genauigkeit im Vergleich zu bestehenden Lösungen wie QKeras und HGQ erzielt.

Das Wesentliche

Das Problem: Der überfüllte Bahnhof

Stell dir vor, der Large Hadron Collider (LHC) ist ein riesiger, extrem belebter Bahnhof. Jeden Tag kommen Millionen von Zügen (Teilchenkollisionen) an. Das Problem: Der Bahnhof ist so voll, dass man unmöglich jeden einzelnen Zug, jeden Koffer und jedes Ticket aufheben und in ein riesiges Archiv (die Festplatte) legen kann. Es würde einfach platzen.

Deshalb brauchen die Wissenschaftler einen super-schnellen Schaffner (den "Trigger"), der in Bruchteilen einer Sekunde entscheidet:

1. Ist dieser Zug wichtig? (Behalten!)
2. Ist er nur Langeweile? (Weg damit!)

Früher waren diese Schaffner sehr schnell, aber auch sehr dumm. Sie schauten nur auf einfache Dinge wie "Ist der Koffer schwer?". Heute wollen wir aber kluge KI-Modelle (Neuronale Netze) als Schaffner einsetzen, die Muster erkennen können. Aber KI-Modelle sind normalerweise wie dicke, schwere Riesen, die viel Zeit brauchen, um zu denken. Auf dem Bahnhof (der Hardware/FPGA) gibt es aber nur wenig Platz und wenig Zeit. Wenn der Schaffner zu lange nachdenkt, verpasst er den nächsten Zug.

Die Lösung: PQuantML – Der große Umzug

Hier kommt PQuantML ins Spiel. Es ist wie ein professionelles Umzugsunternehmen für KI-Modelle, das speziell dafür trainiert wurde, diese riesigen Riesen so klein und leicht zu machen, dass sie in eine winzige Schachtel (den Chip) passen, ohne ihre Intelligenz zu verlieren.

PQuantML macht das mit zwei Haupt-Tricks:

1. Der "Ausmisten"-Trick (Pruning)

Stell dir vor, dein KI-Modell ist ein riesiges Bürogebäude mit tausenden Räumen. Die meisten Räume sind aber leer oder werden nie benutzt.

• Was PQuantML tut: Es geht durch das Gebäude und schließt alle leeren Räume. Es entfernt ganze Stockwerke oder sogar ganze Flügel, die niemand braucht.
• Der Clou: Früher hat man das gemacht, nachdem das Gebäude fertig war (was oft dazu führte, dass das Gebäude instabil wurde). PQuantML baut das Gebäude aber während des Trainings schon so, dass es von Anfang an nur die notwendigen Räume hat. Es ist wie ein Architekt, der weiß: "Wir brauchen kein 10. Stockwerk, das spart uns viel Beton."

2. Der "Verdichtungs"-Trick (Quantization)

Stell dir vor, die Mitarbeiter in deinem Büro schreiben ihre Notizen. Normalerweise schreiben sie mit Tinte in feinsten Details (hohe Genauigkeit, aber viel Platz nötig).

• Was PQuantML tut: Es sagt den Mitarbeitern: "Schreibt nur noch mit Bleistift und in groben Strichen."
• Der Clou: Die Notizen sind jetzt viel kleiner und leichter zu transportieren, aber man kann sie immer noch lesen und verstehen. PQuantML sorgt dafür, dass das Modell lernt, mit diesen "grob gestrichelten" Notizen (weniger Rechenleistung) trotzdem genauso klug zu sein wie mit der feinen Tinte.

Warum ist das besonders? (Die "All-in-One"-Box)

Bisher mussten Wissenschaftler zwei verschiedene Werkzeuge benutzen:

1. Ein Werkzeug, um Räume zu entfernen (Pruning).
2. Ein anderes Werkzeug, um die Notizen zu verdichten (Quantization).

Das war wie wenn man erst den Möbelwagen packen und dann erst die Koffer zusammenfalten müsste – sehr umständlich und fehleranfällig.

PQuantML ist wie ein Schweizer Taschenmesser: Es macht beides gleichzeitig und perfekt abgestimmt. Es ist wie ein Dirigent, der die Musik (das Training) so leitet, dass das Orchester (das Modell) lernt, mit weniger Instrumenten (weniger Hardware-Ressourcen) trotzdem einen perfekten Klang (hohe Genauigkeit) zu erzeugen.

Das Ergebnis im echten Leben

Die Autoren haben PQuantML an einem echten Test geprüft: Sie mussten Teilchenstrahlen (Jets) in der Physik identifizieren.

• Das Ergebnis: Die KI wurde durch PQuantML so stark verkleinert, dass sie viel weniger Platz auf dem Chip brauchte (bis zu 90% weniger!) und viel schneller reagierte.
• Der Vergleich: Andere bekannte Werkzeuge (wie QKeras oder HGQ) waren entweder nicht so gut beim "Ausmisten" oder nicht so gut beim "Verdichten". PQuantML hat beides kombiniert und war damit der Gewinner: Schnell, klein und trotzdem sehr schlau.

Fazit

PQuantML ist das Werkzeug, das es Wissenschaftlern erlaubt, ihre "dummen, schweren" KI-Modelle in "schlank, schnelle und kluge" Schaffner zu verwandeln, die auf den extremen Hochgeschwindigkeits-Bahnhöfen der Teilchenphysik arbeiten können. Ohne dieses Werkzeug wären viele dieser fortschrittlichen KI-Modelle einfach zu schwer für die Hardware, die wir in Echtzeit brauchen.

Es ist im Grunde die Brücke zwischen der komplexen Welt der KI-Forschung und der harten Realität der Elektronik, die in Millisekunden entscheiden muss.

Technische Zusammenfassung

1. Problemstellung und Motivation

Hochenergiephysik-Experimente am Large Hadron Collider (LHC), wie ATLAS und CMS, stehen vor der Herausforderung extrem hoher Datenraten (ca. 40 MHz Kollisionsrate). Da nicht alle Ereignisse offline gespeichert werden können, sind mehrstufige Trigger-Systeme notwendig. Das Level-1-Trigger-System (L1T) arbeitet auf Hardware-Basis (FPGAs) und muss Entscheidungen innerhalb von Mikrosekunden treffen.

• Herausforderung: Traditionelle Algorithmen sind oft zu einfach, während fortschrittliche Machine-Learning (ML)-Modelle für die Echtzeit-Datenverarbeitung zu rechenintensiv und speicherhungrig für FPGAs sind.
• Ziel: ML-Modelle müssen so komprimiert werden, dass sie strenge Latenzanforderungen (Determinismus) und Hardware-Ressourcenbeschränkungen (DSP, LUT, BRAM) erfüllen, ohne die physikalische Genauigkeit (z. B. bei der Jet-Tagging-Klassifizierung) signifikant zu beeinträchtigen.
• Lücke in bestehenden Tools: Bestehende Bibliotheken wie QKeras oder HGQ konzentrieren sich primär auf Quantisierung. Das gleichzeitige Anwenden von fortschrittlichen Pruning-Methoden (Strukturiertes/Unstrukturiertes Pruning) erfordert oft separate Implementierungen und manuelle Workflows, was den Einsatz in komplexen End-to-End-Pipelines erschwert.

2. Methodik: PQuantML

PQuantML ist eine Open-Source-Bibliothek, die Pruning und Quantisierung in einem einheitlichen, hardwarebewussten Framework integriert. Sie ermöglicht das Training komprimierter Modelle direkt für die FPGA-Bereitstellung.

Kernarchitektur und Features:

• Einheitliche Schnittstelle: Bietet eine konsistente API für PyTorch und Keras (TensorFlow).
• Konfiguration: Alle Strategien (Pruning, Quantisierung, Hyperparameter) werden über YAML-Konfigurationsdateien gesteuert, validiert durch Pydantic-Schemas.
• Hardware-Awareness: Integriert direkt mit hls4ml (High-Level Synthesis), um trainierte Modelle in FPGA-Firmware zu übersetzen.
• Automatisierung: Enthält eine Pipeline für Hyperparameter-Optimierung (via Optuna) und Experiment-Tracking (via MLflow).

Technische Komponenten:

1. Quantisierung (Quantization-Aware Training - QAT):
   • Unterstützt Fixed-Point-Quantisierung (Parameter: Vorzeichen, Ganzzahl- und Bruchteile-Bits).
   • Implementiert High-Granularity Quantization (HGQ): Lernt bitweise Bit-Breiten für Gewichte und Aktivierungen mittels Gradientenoptimierung. Dies minimiert den „Effective Bit Operations" (EBOPs)-Wert, ein Metrik für den Hardware-Kosten-Nutzen.
   • Unterstützt Tensor-, Kanal- und Gewichts-spezifische Quantisierung.
2. Pruning-Methoden:
   • Bietet eine Vielzahl von Algorithmen mit unterschiedlichen Granularitäten:
     • Strukturiert: Entfernt ganze Kanäle oder Filter (z. B. Activation Pruning, MDMM).
     • Unstrukturiert: Entfernt einzelne Gewichte (z. B. AutoSparse, Continuous Sparsification, DST, PDP, Wanda).
     • Semi-strukturiert: Unterstützt N:M-Muster (z. B. 2:4), die für Hardware-Optimierung geeignet sind.
   • FITCompress: Ein Algorithmus, der per-Layer-Bit-Breiten und globale Sparsity-Level bestimmt, um eine Ziel-Kompressionsrate zu erreichen.
3. Workflow:
   • Modelldefinition: Entweder direkte Definition von PQuantML-Layern oder automatischer Ersatz von Layern in bestehenden Modellen.
   • Training: Mehrstufiger Prozess (Pre-Training, Training, Fine-Tuning), der je nach Pruning-Methode variiert (z. B. Soft-Mask zu Hard-Mask Konvertierung).
   • Deployment: Nach dem Training werden Gewichte und Masken angewendet und das Modell via hls4ml in HLS-Code für FPGAs übersetzt.

3. Wichtige Beiträge

• Integration von Pruning und Quantisierung: PQuantML ist das erste Tool, das fortschrittliche Pruning-Strategien nahtlos mit HGQ und QAT in einem einzigen Framework vereint.
• End-to-End-Workflow: Von der Modelldefinition über das Training bis zur FPGA-Synthese (hls4ml) ohne manuelle Bruchstellen.
• Flexibilität: Unterstützung von PyTorch und TensorFlow sowie verschiedene Granularitäten für Pruning (unstrukturiert, strukturiert, N:M).
• Reproduzierbarkeit: Durch Integration von MLflow und Optuna werden Experimente, Hyperparameter und Metriken systematisch protokolliert und optimiert.

4. Ergebnisse und Benchmarks

Die Evaluation erfolgte auf der Aufgabe der Jet-Substruktur-Klassifizierung (Jet Tagging) mit zwei Datensätzen (HLF und PLF) auf einem Xilinx Virtex UltraScale+ FPGA.

• Vergleich mit QKeras:
  • PQuantML (mit DST-Pruning und Fixed-Point-Quantisierung) erreichte eine Genauigkeit von 76,3 % (ähnlich wie QKeras), reduzierte aber die LUT-Nutzung um ca. 30 % (von 5.504 auf 3.895) und die Latenz von 105 ns auf ca. 47 ns.
  • Die DSP-Nutzung wurde ebenfalls signifikant gesenkt.
• Vergleich mit HGQ:
  • Modelle, die mit PQuantML unter Verwendung von HGQ trainiert wurden, erzielten fast identische Genauigkeit und Ressourcennutzung wie Modelle, die direkt mit der nativen HGQ-Bibliothek trainiert wurden. Dies validiert PQuantML als zuverlässige Alternative.
• FITCompress-Ergebnisse:
  • Der FITCompress-Algorithmus ermöglichte eine weitere Reduktion der DSP-Nutzung (bis auf 45 DSPs) bei nur minimalen Genauigkeitsverlusten (76,0 %), was zeigt, dass dynamische Bit-Breiten-Anpassung effektiv ist.
• Zusammenfassung: PQuantML erreicht hohe Kompressionsraten bei gleichzeitiger Beibehaltung der physikalischen Genauigkeit und übertrifft bestehende Tools in der Flexibilität und Effizienz der Ressourcennutzung.

5. Bedeutung und Ausblick

PQuantML senkt die Einstiegshürde für Physiker, um fortschrittliche Komprimierungstechniken in Echtzeit-ML-Workflows für den LHC einzuführen. Es adressiert die kritische Lücke zwischen theoretischer ML-Modellentwicklung und den strengen Hardware-Anforderungen von FPGA-basierten Triggersystemen.

Zukünftige Arbeiten:

• Erweiterung der Unterstützung für TensorFlow (insbesondere für verzweigte Architekturen wie ResNet oder Attention-Blöcke).
• Integration weiterer Komprimierungstechniken wie Knowledge Distillation und Low-Rank-Faktorisierung.
• Verbesserung der Metriken für die Hardware-Kosten (z. B. direkte Zielvorgabe von EBOPs statt nur Bit-Operationen).
• Vollständige Integration von FITCompress in beide Backend-Frameworks.

Fazit: PQuantML stellt einen wichtigen Schritt dar, um ML-Modelle für die nächste Generation von Echtzeit-Trigger-Systemen am LHC effizient, hardwarebewusst und reproduzierbar zu trainieren und bereitzustellen.