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

PQuantML è una nuova libreria open-source per la compressione hardware-consapevole di modelli neurali che semplifica i flussi di lavoro end-to-end, integrando potatura e quantizzazione a punto fisso per ottenere riduzioni significative di parametri e larghezza di bit mantenendo l'accuratezza, come dimostrato nel contesto della classificazione della struttura dei getti per l'elaborazione dei dati del LHC.

L'essenziale

Immagina di dover costruire un motore da Formula 1 che deve essere così piccolo e leggero da poter stare nel bagagliaio di una Fiat 500, ma che deve comunque correre alla stessa velocità del motore originale.

Questa è esattamente la sfida che affrontano gli scienziati al CERN (il laboratorio dove si studia la materia, come l'origine dell'universo). Lì, le particelle si scontrano a velocità incredibili, creando una quantità di dati così enorme che sarebbe come cercare di bere un fiume con un cucchiaino: non riescono a salvare tutto. Devono decidere in microsecondi (milionesimi di secondo) quali dati sono interessanti e quali possono essere scartati.

Per farlo, usano dei "filtri" intelligenti basati sull'intelligenza artificiale (Machine Learning). Ma c'è un problema: questi filtri intelligenti sono solitamente molto pesanti e lenti, come un camion che non può passare su un ponte stretto.

Cos'è PQuantML?

PQuantML è il nuovo "kit di strumenti" che gli scienziati hanno creato per trasformare quei "camioni" pesanti in "Fiat 500" veloci e leggere, senza perdere la loro capacità di correre (ovvero, senza perdere la precisione nel riconoscere le particelle).

Ecco come funziona, usando delle metafore semplici:

1. Il Taglio dei Panni (Pruning)

Immagina che la tua rete neurale (il cervello dell'IA) sia un abito da sposa ricchissimo di pizzi, perline e strascichi. È bellissimo, ma troppo pesante per correre.

• Cosa fa PQuantML: È come un sarto esperto che ti dice: "Ehi, guarda qui! Questi 30 centimetri di pizzo non servono a nessuno, tagliali via!". Oppure: "Queste perline sono tutte uguali, ne teniamo solo una ogni quattro".
• Il risultato: L'abito è molto più leggero e veloce da indossare, ma sembra ancora lo stesso abito da sposa. In termini tecnici, questo si chiama pruning (potatura): si eliminano i "pesi" inutili del modello.

2. La Riduzione delle Parole (Quantization)

Ora immagina che l'abito sia fatto di parole scritte su fogli di carta. Normalmente, ogni parola è scritta con un inchiostro costosissimo e dettagliatissimo (come i numeri a virgola mobile che usano i computer normali).

• Cosa fa PQuantML: Decide di riscrivere tutto usando un inchiostro più semplice e parole più corte. Invece di dire "Il valore esatto è 3,14159265...", dice semplicemente "È circa 3,14".
• Il risultato: I fogli diventano molto più piccoli e leggeri. Il computer li legge molto più velocemente perché deve fare meno calcoli complessi. Questo si chiama quantizzazione.

3. L'Allenamento Intelligente

Il trucco di PQuantML non è solo tagliare e semplificare dopo aver costruito il modello. È come se il sarto e lo scrittore lavorassero mentre l'abito viene cucito.

• Invece di costruire un abito enorme e poi tagliarlo (che spesso rovinerebbe il vestito), PQuantML insegna al modello a diventare leggero durante l'allenamento. È come se un atleta si allenasse già con zavorre leggere, così quando le toglie, è pronto a correre subito senza perdere forma.

Perché è importante?

Prima di PQuantML, se volevi rendere un modello veloce per i chip dei computer (chiamati FPGA, che sono come piccoli cervelli elettronici), dovevi usare diversi strumenti separati: uno per tagliare, uno per semplificare i numeri, e dovevi fare tutto a mano, come se dovessi costruire un mobile con istruzioni in lingue diverse.

PQuantML è come un fai-da-te tutto-in-uno:

• Ti dà un'unica scatola con tutti gli attrezzi.
• Ti permette di dire: "Voglio tagliare il 50% dei pizzi e usare parole corte" e lui lo fa automaticamente.
• Si assicura che il risultato finale sia perfetto per essere installato sui chip veloci del CERN.

Il Risultato

Grazie a questo strumento, gli scienziati hanno potuto prendere modelli di intelligenza artificiale complessi e renderli:

• Più piccoli: Occupano meno spazio sui chip.
• Più veloci: Rispondono in tempi brevissimi (fondamentali per fermare le particelle che volano via).
• Precisi: Non hanno perso la capacità di riconoscere le particelle importanti.

In sintesi, PQuantML è il "maghetto" che prende un'intelligenza artificiale goffa e pesante e la trasforma in un'atleta agile e veloce, pronta a correre nella gara più veloce dell'universo: quella dei dati del CERN.

Sommario tecnico

1. Il Problema: Sfide nell'Implementazione ML in Fisica delle Alte Energie

Gli esperimenti di fisica delle alte energie (HEP) al Large Hadron Collider (LHC), come ATLAS e CMS, affrontano una sfida dati estrema: le collisioni protone-protone avvengono a 40 MHz, generando centinaia di terabyte al secondo. È impossibile archiviare tutti gli eventi offline.
Per mitigare ciò, si utilizzano sistemi di "trigger" multi-livello:

• L1T (Livello 1): Basato su hardware (FPGA), deve prendere decisioni in microsecondi.
• HLT (High-Level Trigger): Basato su software, opera entro circa 1 secondo.

Attualmente, l'L1T utilizza algoritmi semplici e veloci, mentre gli algoritmi complessi basati sul Machine Learning (ML) sono troppo lenti per essere implementati direttamente nell'hardware di trigger. Tuttavia, i modelli ML offrono approssimazioni rapide di algoritmi costosi. La sfida principale è deployare modelli ML su FPGA rispettando vincoli stringenti di:

• Latenza: Deve essere deterministica e nell'ordine dei microsecondi.
• Risorse Hardware: Limiti su DSP (Digital Signal Processors), BRAM, LUT e Flip-Flops.
• Efficienza: I modelli devono essere ottimizzati per ridurre l'uso delle risorse senza degradare le prestazioni fisiche (accuratezza).

Le tecniche tradizionali di compressione (potatura e quantizzazione) applicate post-training spesso causano una significativa perdita di accuratezza. È necessario un approccio che integri la compressione direttamente durante l'addestramento (compression-aware training).

2. Metodologia: PQuantML

PQuantML è una nuova libreria open-source che fornisce un framework unificato per la compressione end-to-end di modelli neurali, consapevole dell'hardware. Si basa e estende le funzionalità di librerie esistenti come HGQ (High-Granularity Quantization) e QKeras, integrando sistematicamente la potatura (pruning) con la quantizzazione.

Architettura e Design

• Interfaccia Unificata: Supporta sia PyTorch che Keras/TensorFlow. Offre un'interfaccia coerente indipendentemente dal backend.
• Configurazione Dichiarativa: Le strategie di compressione (potatura, quantizzazione, iperparametri) sono definite tramite file YAML validati con schemi Pydantic, garantendo flessibilità e sicurezza dei tipi.
• Pipeline di Addestramento: Gestisce automaticamente flussi di lavoro multi-stadio (pre-addestramento, addestramento, fine-tuning) necessari per diversi metodi di compressione.
• Ottimizzazione degli Iperparametri: Integrazione nativa con Optuna (per la ricerca bayesiana) e MLflow (per il tracciamento degli esperimenti), permettendo l'ottimizzazione automatica di parametri come tassi di potatura, bit-width e schedulazione.

Tecniche di Compressione Implementate

1. Quantizzazione (QAT - Quantization-Aware Training):

   • Utilizza rappresentazioni in punto fisso (fixed-point).
   • Supporta la High-Granularity Quantization (HGQ): apprende bit-width individuali per ogni peso e attivazione tramite ottimizzazione basata sul gradiente.
   • Minimizza un loss composto da accuratezza e costo hardware stimato (EBOPs - Effective Bit Operations).
   • Supporta quantizzazione per tensore, per canale o per peso.

2. Potatura (Pruning):

   • Implementa metodi strutturati e non strutturati con diverse granularità.
   • Metodi Supportati:
     • Strutturati: Activation Pruning (AP), MDMM (basato su moltiplicatori differenziali modificati).
     • Non Strutturati: AutoSparse, Continuous Sparsification (CS), DST, PDP (basato sulla distribuzione), Wanda.
     • Supporto per pattern N:M (semi-strutturati).
   • La potatura viene appresa durante l'addestramento tramite maschere (soft o hard) e schedulazione.

3. Integrazione con HLS:

   • Una volta addestrato, il modello compresso può essere convertito direttamente in codice HLS (High-Level Synthesis) per FPGA utilizzando il pacchetto hls4ml, garantendo che le previsioni siano bit-accurate rispetto al modello originale.

3. Contributi Chiave

• Framework Unificato: PQuantML colma il divario tra librerie focalizzate solo sulla quantizzazione (come QKeras) e quelle che richiedono implementazioni manuali complesse per la potatura.
• Hardware-Aware: Progettato specificamente per i vincoli delle FPGA nei sistemi di trigger LHC, ottimizzando l'uso di DSP e LUT.
• Automazione: Riduce la barriera all'ingresso per i fisici, automatizzando flussi di lavoro complessi di compressione, ottimizzazione degli iperparametri e tracciamento degli esperimenti.
• Flessibilità: Permette di combinare diverse strategie (es. potatura strutturata + quantizzazione HGQ) e supporta sia l'approccio "layer-by-layer" manuale che la sostituzione automatica dei layer in architetture esistenti.

4. Risultati Sperimentali

Il framework è stato valutato sul compito di classificazione della sottostruttura dei jet (Jet Substructure Classification - JSC), un problema critico per l'LHC. Sono stati utilizzati due dataset:

1. HLF (High-Level Features): 16 feature, rete fully-connected.
2. PLF (Particle-Level Features): Rappresentazione 2D delle particelle.

Prestazioni e Confronti:

• Riduzione delle Risorse: PQuantML ha ottenuto riduzioni sostanziali nell'uso delle risorse FPGA rispetto ai modelli non compressi e rispetto a QKeras.
  • Esempio (Dataset HLF OpenML): L'uso di DST (potatura) con quantizzazione per-tensore ha ridotto l'uso di LUT da 60.089 a 3.895 e la latenza da 105 ns a ~47 ns, mantenendo un'accuratezza del 76.3% (simile al modello non compresso).
  • FITCompress ha ridotto ulteriormente l'uso di DSP (fino a 45) con una latenza di 36.2 ns.
• Confronto con HGQ Nativo: I modelli addestrati con PQuantML utilizzando HGQ hanno mostrato prestazioni (accuratezza e uso risorse) quasi identiche a quelle ottenute addestrando direttamente con la libreria HGQ nativa, validando PQuantML come alternativa affidabile.
• Confronto con QKeras: PQuantML supera QKeras quando combina potatura e quantizzazione a bit-width fissi, offrendo un miglior compromesso tra accuratezza e consumo di risorse.

5. Significato e Impatto

PQuantML rappresenta un passo fondamentale verso l'adozione del Machine Learning nei sistemi di trigger in tempo reale dell'LHC.

• Efficienza Operativa: Dimostra che è possibile raggiungere livelli di compressione elevati (riduzione di bit-width e parametri) mantenendo l'accuratezza fisica necessaria per la selezione degli eventi.
• Scalabilità: La capacità di gestire flussi di lavoro end-to-end e l'ottimizzazione automatica degli iperparametri rende la tecnologia accessibile a un pubblico più ampio di ricercatori, non solo esperti di hardware.
• Futuro: Sebbene esistano limitazioni attuali (es. supporto parziale per architetture complesse in TensorFlow e metriche di compressione da affinare), PQuantML stabilisce una base solida per l'integrazione di tecniche avanzate come la distillazione della conoscenza e la fattorizzazione a basso rango nei flussi di lavoro HEP.

In sintesi, PQuantML è uno strumento ingegneristico pratico e un prototipo di ricerca che abilita il deployment di modelli neurali complessi su hardware con vincoli severi, facilitando l'evoluzione dei sistemi di trigger del CERN.