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

PQuantML is an open-source, hardware-aware library that streamlines end-to-end neural network compression through joint or individual pruning and fixed-point quantization, demonstrating significant parameter and bit-width reductions with maintained accuracy on real-time LHC jet tagging tasks compared to existing tools like QKeras and HGQ.

The Gist

Imagine you are a chef trying to cook a massive, gourmet feast (a complex Artificial Intelligence model) for a very specific guest: a high-speed race car driver (the Large Hadron Collider's trigger system).

The problem? The driver has to make a decision in the blink of an eye—microseconds, not seconds. If the chef takes too long to chop vegetables or plate the dish, the race is over, and the data is lost.

PQuantML is the new, magical kitchen toolkit designed to help chefs shrink their giant, slow recipes into tiny, lightning-fast meals without losing the flavor (accuracy).

Here is how it works, broken down into simple concepts:

1. The Problem: The "Too Big to Fit" Dilemma

In the world of particle physics, machines like the Large Hadron Collider (LHC) smash particles together 40 million times a second. This creates a tsunami of data. You can't save it all, so the machine has to decide instantly: "Keep this interesting crash" or "Ignore that boring one."

To make these decisions, they use special computer chips called FPGAs. Think of these chips as tiny, custom-built kitchens with very limited counter space and a very strict time limit.

• Old AI models are like a 50-course tasting menu. They are delicious (accurate), but they take too long to cook and require too much counter space. They don't fit in the tiny kitchen.
• The Goal: We need to turn that 50-course menu into a single, perfect bite-sized appetizer that tastes just as good but cooks in a split second.

2. The Solution: PQuantML (The "Shrink-Ray" Kit)

PQuantML is a software tool that acts like a master chef's shrink-ray. It doesn't just throw away ingredients; it intelligently reorganizes the recipe. It does two main things:

A. Pruning (The "Edit" Button)

Imagine your recipe has 1,000 steps, but 400 of them are just "stir the pot" or "check the salt." They don't actually change the taste.

• Pruning is like a smart editor that reads the recipe and says, "We don't need these 400 steps. Let's cut them out."
• PQuantML can cut out individual ingredients (unstructured pruning) or entire sections of the recipe (structured pruning). This makes the recipe shorter and faster to follow.

B. Quantization (The "Rough Draft" Translator)

Usually, AI models speak in "high-definition" math (using complex decimal numbers like 3.14159265...). This is precise but takes up a lot of space and time to calculate.

• Quantization is like translating that high-definition math into "low-resolution" math (like rounding 3.14159 to just 3.1).
• It's like switching from a 4K movie to a 480p video. The picture is slightly less sharp, but it loads instantly and takes up way less memory. PQuantML teaches the model to be comfortable with this "rougher" math while it is learning, so it doesn't lose its taste.

3. The Secret Sauce: "Training While Compressing"

In the past, people would train a giant model first, and then try to shrink it. This is like baking a massive cake and then trying to slice off pieces to make it fit in a tiny box. The cake often falls apart, and the taste suffers.

PQuantML does something smarter: It shrinks the model while it is learning.

• Imagine a student learning to play the piano. Instead of learning a full symphony and then trying to play it on a toy piano, PQuantML makes them practice on the toy piano while they are learning the notes.
• By the time they are done, they are a master of the toy piano. The model learns to be accurate even with fewer ingredients and simpler math.

4. The Result: A Perfect Fit for the Race Car

The paper tested PQuantML on a task called "Jet Tagging." This is like trying to identify what kind of particle exploded in a split second.

• Before: The models were too big and slow for the FPGA chips.
• After: PQuantML shrunk the models by cutting out unnecessary steps and simplifying the math.
• The Outcome: The models became tiny (using less than half the computer memory) and super fast (running in nanoseconds), but they were still 99% as accurate as the giant versions.

Why This Matters

Think of PQuantML as the bridge between "Smart AI" and "Real-World Speed."

• For Scientists: It means they can use powerful AI to catch rare particles in real-time, rather than having to use slow, simple rules.
• For Everyone Else: It shows how we can make AI run on small, battery-powered devices (like your phone or a smartwatch) without needing a massive server farm.

In short, PQuantML is the tool that teaches big, slow AI models to run fast, light, and efficient, making them ready for the high-speed world of particle physics and beyond.

Technical Summary

1. Problem Statement

High-energy physics (HEP) experiments at the Large Hadron Collider (LHC), such as ATLAS and CMS, face extreme data challenges. Proton-proton collisions occur at 40 MHz, generating hundreds of terabytes per second. To manage this, the Level-1 Trigger (L1T) system must make real-time decisions within microseconds using Field-Programmable Gate Arrays (FPGAs).

• The Challenge: Deploying complex Machine Learning (ML) models on FPGAs for real-time inference is difficult due to strict constraints on latency (microseconds), throughput, and hardware resources (DSPs, LUTs, BRAM).
• The Gap: Existing tools often handle quantization or pruning separately.
  • Libraries like QKeras focus on quantization but lack advanced pruning capabilities.
  • Libraries like HGQ (High-Granularity Quantization) offer fine-grained quantization but require users to implement pruning strategies separately.
  • Post-training compression often leads to significant accuracy degradation.
• The Need: An end-to-end, unified framework that integrates pruning and quantization-aware training (QAT) specifically designed for hardware-aware optimization in HEP real-time systems.

2. Methodology: PQuantML Architecture

PQuantML is an open-source library designed to simplify the training of compressed models for FPGA deployment. It bridges the gap between high-level model definition and low-level hardware synthesis (via hls4ml).

Core Components

• Unified Interface: Supports both PyTorch and Keras/TensorFlow backends.
• Configuration-Driven: Uses YAML configuration files (validated via Pydantic) to define model architectures, pruning strategies, quantization settings, and hyperparameters.
• Two Workflow Modes:
  1. Direct Layer Definition: Users explicitly define layers using PQuantML's custom modules (e.g., PQConv2d, PQDense).
  2. Layer Replacement: Automatically injects compression operators into existing standard Keras/PyTorch models.

Key Techniques

• Quantization:
  • Supports Fixed-Point Quantization (parameterized by integer and fractional bits).
  • Implements High-Granularity Quantization (HGQ), which learns individual bit-widths for weights/activations via gradient-based optimization.
  • Uses Quantization-Aware Training (QAT) to minimize accuracy loss during the training process.
• Pruning:
  • Implements multiple pruning methods with varying granularities:
    • Unstructured: AutoSparse, Continuous Sparsification (CS), DST, PDP, Wanda.
    • Structured: Activation Pruning (AP), MDMM (Modified Differential Method of Multipliers).
    • Semi-Structured (N:M): Supported for specific methods.
  • FITCompress: An algorithm that determines per-layer bit-widths and global sparsity to meet a target compression ratio.
• Training Pipeline:
  • Supports multi-stage training: Pre-training (optional), Training (mask learning), and Fine-tuning (hardening soft masks).
  • Integrates Optuna for automated hyperparameter optimization (Bayesian search) and MLflow for experiment tracking.

3. Key Contributions

1. Unified Framework: PQuantML is the first library to seamlessly integrate structured/unstructured pruning with both fixed-point and high-granularity quantization in a single, hardware-aware workflow.
2. Hardware Awareness: It is designed specifically for FPGA constraints, supporting the generation of bit-accurate models compatible with hls4ml for firmware synthesis.
3. Flexibility & Usability: Provides a declarative configuration system that allows physicists to experiment with complex compression strategies without rewriting training loops.
4. Benchmarking: Demonstrates that combining pruning and quantization yields better trade-offs than using either technique alone or using existing tools like QKeras.

4. Results and Benchmarks

The authors evaluated PQuantML on the Jet Substructure Classification (JSC) task (identifying the particle initiating a jet), a critical real-time LHC application. Two datasets were used: High-Level Features (HLF) and Particle-Level Features (PLF).

Key Findings:

• Accuracy vs. Resources: PQuantML achieved comparable accuracy to existing tools while significantly reducing resource usage.
  • Latency: Reduced inference latency from ~105 ns (QKeras baseline) to ~36–47 ns.
  • Resources:
    • LUTs: Reduced from 60,089 (unpruned) to as low as 3,895 (DST pruning).
    • DSPs: Reduced from 1,175 to as low as 45 (FITCompress).
• Comparison with HGQ: Models trained with PQuantML using HGQ produced results nearly identical to models trained with the native HGQ library, validating PQuantML's implementation.
• Comparison with QKeras: PQuantML outperformed QKeras in latency and resource efficiency when combining pruning with fixed-point quantization, while maintaining similar accuracy (~76%).
• Granularity Impact:
  • Per-tensor quantization offered a good balance.
  • Per-weight quantization significantly reduced LUT/FF usage but slightly increased latency.
  • Structured pruning (DST) proved highly effective, matching QKeras accuracy while drastically cutting resources.

5. Significance and Future Outlook

Significance:
PQuantML lowers the barrier for HEP physicists to adopt advanced model compression techniques. By automating the complex interplay between pruning, quantization, and hardware constraints, it enables the deployment of sophisticated ML models in the L1T systems of future LHC upgrades, where microsecond latency is non-negotiable.

Limitations & Future Work:

• Current Limitations: The TensorFlow backend has limited support for complex architectures (e.g., branching, residual connections). FITCompress is currently PyTorch-only.
• Future Directions:
  • Full integration of FITCompress into TensorFlow.
  • Automation of hyperparameter search space generation.
  • Integration of Knowledge Distillation and Low-Rank Adaptation to further improve the accuracy-efficiency trade-off.
  • Expansion to support other hardware-specific compilation toolchains beyond hls4ml.

In conclusion, PQuantML represents a significant step forward in making ML-driven real-time triggers feasible for next-generation particle physics experiments by providing a robust, end-to-end tool for hardware-aware model compression.