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

PQuantML es una biblioteca de código abierto que facilita la compresión de modelos neuronales mediante poda y cuantización de punto fijo, optimizada para el despliegue eficiente en entornos con restricciones de latencia como el procesamiento de datos del LHC.

Lo esencial

¡Hola! Imagina que el Gran Colisionador de Hadrones (LHC) es como una fábrica gigante de partículas que produce un torrente de datos tan enorme que es como intentar beber agua de una manguera de incendios con una pajita. Si intentaran guardar todo lo que pasa, se ahogarían en información.

Para solucionar esto, los científicos usan "guardianes" (llamados disparadores o triggers) que deciden en microsegundos qué datos son interesantes y cuáles son basura. El problema es que estos guardianes viven en chips especiales (FPGAs) que son muy rápidos, pero tienen un espacio de memoria y energía muy limitado.

Aquí es donde entra PQuantML, la herramienta que presenta este paper. Vamos a explicarlo con una analogía sencilla:

🎒 El Problema: La Mochila Demasiado Pesada

Imagina que tienes que enviar un mensajero (un modelo de Inteligencia Artificial) a través de una montaña muy estrecha y peligrosa (el chip de hardware).

• El Mensajero: Es un modelo de IA muy inteligente, pero es un gigante con una mochila llena de libros, herramientas y ropa de repuesto (parámetros y datos).
• El Camino: Es estrecho (latencia baja) y el mensajero no puede cargar mucho peso (recursos limitados del chip).
• El Riesgo: Si el mensajero es demasiado pesado, se queda atascado en la puerta y nunca llega a tiempo. Si es demasiado ligero, no sabe qué hacer y comete errores.

✂️ La Solución: PQuantML (El "Sastre y Empaquetador" Inteligente)

Antes, tenías que contratar a dos personas diferentes: una para cortar la ropa (podar el modelo) y otra para empaquetar las cosas en cajas pequeñas (cuantizar). A veces, estas dos personas no se entendían y el resultado era un desastre.

PQuantML es como un sastre y empaquetador todo en uno que trabaja contigo desde el principio. No solo corta la ropa, sino que te enseña a caminar con menos peso sin que te caigas.

1. La Poda (Pruning): Cortar lo que no sirve

Imagina que tienes un árbol con miles de ramas. Algunas ramas dan frutos deliciosos, pero otras están secas o son solo hojas muertas que no hacen nada.

• PQuantML tiene tijeras mágicas que pueden cortar:
  • Ramas enteras: (Poda estructurada) Quita todo un grupo de hojas de golpe. Es como quitar una rama completa del árbol. Es fácil para el mensajero caminar, pero a veces cortas un poco de fruta buena.
  • Hojas individuales: (Poda no estructurada) Corta solo las hojas secas específicas. Es más preciso, pero deja el árbol con agujeros raros que son difíciles de navegar.
  • Patrones especiales: (Poda N:M) Corta, por ejemplo, 2 hojas de cada grupo de 4. Es un punto medio perfecto: el árbol sigue teniendo estructura, pero pesa menos.

2. La Cuantización (Quantization): Cambiar el idioma

Imagina que el mensajero habla un idioma muy complejo y preciso (números con decimales infinitos, como 3.14159265...). Esto ocupa mucho espacio en su mochila.

• PQuantML le enseña a hablar un idioma más simple (números enteros o con pocos decimales, como 3.14).
• La magia: En lugar de hacerlo al final (cuando el mensajero ya está cansado), PQuantML le enseña este "idioma simple" mientras aprende. Así, el mensajero se acostumbra a pensar de forma sencilla desde el primer día y no pierde su inteligencia.

🚀 ¿Qué lograron con esto?

Los científicos probaron PQuantML en una tarea real: identificar qué tipo de partícula creó una "lluvia" de partículas (llamada jet tagging).

• Resultado: Crearon mensajeros que eran mucho más ligeros (ocupaban menos espacio en el chip) y más rápidos (llegaban antes), pero que seguían siendo tan inteligentes como los gigantes originales.
• Comparación: Lo probaron contra otras herramientas (como QKeras y HGQ). PQuantML logró resultados similares o mejores, pero con la ventaja de que puedes usarlo con un solo botón de configuración, sin tener que ser un experto en programación de chips.

🌟 En resumen

PQuantML es una caja de herramientas mágica que permite tomar modelos de Inteligencia Artificial gigantes y pesados, y convertirlos en mensajeros ágiles y rápidos capaces de correr por los pasillos estrechos de los chips de los aceleradores de partículas, sin perder su capacidad de tomar decisiones inteligentes.

Es como si pudieras convertir un camión de mudanzas en una bicicleta de carreras, pero manteniendo la misma capacidad de carga y velocidad, todo gracias a un diseño inteligente que se adapta a las reglas del camino.

Resumen técnico

1. El Problema: Despliegue de ML en Tiempo Real en HEP

Los experimentos de física de altas energías (HEP) en el Gran Colisionador de Hadrones (LHC), como ATLAS y CMS, enfrentan un desafío de datos extremo. Las colisiones protón-protón ocurren a 40 MHz, generando cientos de terabytes por segundo. Para gestionar esto, se utilizan sistemas de disparo (trigger) multinivel:

• Nivel 1 (L1T): Basado en hardware (FPGAs), toma decisiones en microsegundos.
• Nivel Alto (HLT): Basado en software, refina la selección en ~1 segundo.

El problema central es que los algoritmos avanzados de aprendizaje automático (ML), aunque precisos, suelen ser demasiado lentos y consumir demasiados recursos para implementarse directamente en los FPGAs del L1T, que tienen restricciones estrictas de latencia y presupuesto de recursos (DSP, BRAM, LUT). Las técnicas tradicionales de compresión de modelos a menudo se aplican como pasos posteriores al entrenamiento, lo que puede degradar el rendimiento o no optimizar adecuadamente para el hardware específico. Se necesita una metodología de compresión consciente del hardware (hardware-aware) y end-to-end que integre la poda (pruning) y la cuantización durante el entrenamiento.

2. Metodología: PQuantML

PQuantML es una biblioteca de código abierto diseñada para unificar y simplificar el flujo de trabajo de compresión de modelos neuronales, integrando poda y cuantización en un solo marco de trabajo.

Arquitectura y Diseño

• Interfaz Unificada: Ofrece una API coherente para PyTorch y Keras (TensorFlow).
• Configuración Declarativa: Las estrategias de compresión se definen mediante archivos YAML validados con esquemas Pydantic, lo que facilita la reproducibilidad.
• Capas Conscientes de Compresión: Proporciona implementaciones de capas estándar (convoluciones, densas, activaciones) que soportan poda y cuantización nativamente.
• Dos Flujos de Trabajo:
  1. Definición Directa: Construir el modelo usando las capas específicas de PQuantML.
  2. Sustitución de Capas: Reemplazar automáticamente las capas de un modelo estándar por versiones comprimidas basadas en una configuración.

Técnicas de Compresión Implementadas

• Cuantización Consciente del Entrenamiento (QAT):
  • Soporta cuantización de punto fijo (parámetros $k, i, f$ para signo, enteros y fraccionarios).
  • Integra Cuantización de Alta Granularidad (HGQ): Aprende bit-widths individuales para pesos y activaciones mediante optimización basada en gradientes, minimizando las operaciones efectivas de bits (EBOPs) y regularizando para evitar que los bits crezcan excesivamente.
• Métodos de Poda (Pruning):
  • No Estructurada: Elimina pesos individuales (ej. AutoSparse, CS, DST, PDP, Wanda).
  • Estructurada: Elimina canales, filtros o neuronas completas (ej. Activación Pruning, MDMM).
  • Semi-estructurada (N:M): Patrones regulares que equilibran expresividad y eficiencia de hardware.
• Optimización de Hiperparámetros: Integra Optuna para la búsqueda bayesiana de hiperparámetros y MLflow para el seguimiento de experimentos, permitiendo optimizar simultáneamente la precisión y el uso de recursos.

Flujo de Trabajo

El proceso abarca desde la definición del modelo, pasando por un entrenamiento multi-etapa (pre-entrenamiento, entrenamiento con máscara suave, ajuste fino con máscara dura), hasta la conversión final a código HLS (High-Level Synthesis) compatible con hls4ml para su implementación en FPGAs.

3. Contribuciones Clave

1. Integración Unificada: PQuantML es la primera herramienta que combina sistemáticamente poda (en múltiples granularidades) y cuantización (incluyendo HGQ) en un solo framework accesible para físicos.
2. Facilidad de Uso: Abstrae la complejidad de los flujos de trabajo de compresión, permitiendo a los usuarios definir estrategias mediante configuración en lugar de código complejo.
3. Compatibilidad con Hardware: Diseñado específicamente para cumplir con las restricciones de latencia y recursos de los FPGAs del LHC, integrándose nativamente con el ecosistema hls4ml.
4. Validación de HGQ: Demuestra que PQuantML puede replicar y extender las capacidades de la biblioteca HGQ, ofreciendo una alternativa robusta y unificada.

4. Resultados y Evaluación

Los autores evaluaron PQuantML en la tarea de clasificación de subestructura de chorros (jet tagging), un problema crítico para el procesamiento de datos en tiempo real del LHC. Se utilizaron dos conjuntos de datos: características de alto nivel (HLF) y características a nivel de partícula (PLF).

• Rendimiento vs. Recursos:

  • PQuantML logró reducciones significativas en el uso de recursos de FPGA (LUT, DSP, Flip-Flops) y latencia en comparación con modelos sin comprimir.
  • Comparación con QKeras: En el conjunto de datos OpenML, PQuantML (usando poda DST) mantuvo una precisión similar (~76.3%) a QKeras pero redujo el uso de LUT de 5,504 a 3,895 y la latencia de 105 ns a ~47 ns.
  • Comparación con HGQ Nativo: Los modelos entrenados con PQuantML usando HGQ mostraron resultados comparables (precisión y uso de recursos) a los entrenados directamente con la biblioteca HGQ, validando la integración.
  • FITCompress: El algoritmo FITCompress (integrado en PQuantML) logró reducir el uso de DSP a niveles muy bajos (ej. 45 DSP) con una ligera pérdida de precisión, demostrando flexibilidad en la compensación entre recursos y rendimiento.

• Eficiencia de Hardware: Los modelos comprimidos lograron frecuencias máximas (Fmax) superiores y latencias más bajas (ej. 36-47 ns) en FPGAs Xilinx Virtex UltraScale+, cumpliendo con los requisitos de microsegundos del L1T.

5. Significado e Impacto

PQuantML representa un avance significativo para la implementación de ML en entornos de física de altas energías y más allá:

• Barrera de Entrada Reducida: Permite a los físicos adoptar técnicas avanzadas de compresión sin necesidad de ser expertos en ingeniería de hardware o en la implementación manual de algoritmos de poda/cuantización.
• Optimización End-to-End: Cierra la brecha entre el entrenamiento de modelos y su despliegue en hardware, asegurando que las métricas de compresión se traduzcan directamente en beneficios de hardware (menor latencia, menor consumo).
• Escalabilidad: Al soportar múltiples backends (PyTorch, TensorFlow) y técnicas de compresión, se posiciona como una herramienta fundamental para las futuras actualizaciones del LHC (High-Luminosity LHC), donde los requisitos de procesamiento en tiempo real serán aún más estrictos.

En resumen, PQuantML no es solo una biblioteca de compresión, sino un puente esencial que permite trasladar modelos de IA complejos a sistemas de disparo de hardware en tiempo real, manteniendo la precisión física necesaria mientras se maximiza la eficiencia de los recursos.