JetFormer: A Scalable and Efficient Transformer for Jet Tagging from Offline Analysis to FPGA Triggers

JetFormer es una arquitectura Transformer escalable y eficiente diseñada para la clasificación de jets en el LHC, capaz de ofrecer un alto rendimiento tanto en análisis offline como en sistemas de disparo (triggers) basados en FPGA mediante optimizaciones de hardware.

Lo esencial

El Gran Filtro de Partículas: La historia de "JetFormer"

Imagina que estás en un concierto de rock masivo en un estadio gigante. Hay miles de personas gritando, instrumentos sonando y gente moviéndose al mismo tiempo. De repente, en medio de todo ese ruido, alguien lanza un pequeño cohete de bengalas que brilla de una forma muy especial. Tu misión es identificar ese cohete de inmediato, entre el caos, para poder grabarlo y estudiarlo. Si esperas a que termine el concierto para revisar las grabaciones, el momento ya habrá pasado.

Eso es exactamente lo que pasa en el Gran Colisionador de Hadrones (LHC).

1. El Problema: Un caos de datos

En el LHC, se chocan partículas a velocidades increíbles. Estos choques crean "jets" (chorros de partículas). Algunos de estos chorros son "tesoros" (como el Bosón de Higgs) y otros son solo "ruido" (partículas comunes).

El problema es que el colisionador produce miles de millones de choques por segundo. No tenemos capacidad para guardar todo. Necesitamos un "portero" o un "filtro" ultra rápido que decida en microsegundos: "Este choque es interesante, guárdalo" o "Este es ruido, deséchalo".

2. La Solución: JetFormer (El "Súper Detective")

Hasta ahora, teníamos dos tipos de detectives:

• Los Detectives de Oficina: Son muy inteligentes y analizan cada detalle con lupa, pero son lentos. Se usan para estudiar los datos después de que el experimento termina (análisis offline).
• Los Guardias de Seguridad: Son muy rápidos y están en la puerta, pero no son muy detallistas. A veces dejan pasar cosas importantes o se equivocan (implementaciones en hardware online).

JetFormer es como un detective que puede transformarse. Es una arquitectura de Inteligencia Artificial (un "Transformer") diseñada para ser lo suficientemente inteligente para el trabajo de oficina, pero lo suficientemente ágil para trabajar en la puerta del estadio.

3. ¿Cómo lo logra? (Las herramientas de magia)

Para que este detective sea tan versátil, los científicos usaron tres trucos de "magia" matemática:

• El Entrenamiento Inteligente (Optimización): En lugar de crear un detective gigante y pesado, usaron un proceso llamado Optuna para encontrar el tamaño perfecto. Es como diseñar un coche que sea tan potente como un Ferrari pero tan ligero como una bicicleta.
• La Poda (Pruning): Imagina que un detective tiene un manual de 1,000 páginas, pero descubre que la mitad de las páginas son instrucciones repetitivas. La "poda" consiste en arrancar esas páginas innecesarias para que el detective pueda leer y decidir mucho más rápido, sin perder su inteligencia.
• La Simplificación Extrema (Cuantización de 1 bit): Normalmente, los cálculos matemáticos son como usar una regla con milímetros y micras (muy precisos pero lentos). La cuantización de 1 bit es como decir: "No me des decimales, solo dime si el valor es positivo o negativo". Es mucho más rápido de procesar para los chips de hardware (FPGAs), permitiendo que el cerebro del detector trabaje a la velocidad de la luz.

4. ¿Por qué es importante?

Gracias a JetFormer, los científicos del CERN podrán capturar señales mucho más raras y valiosas de la naturaleza. Es como haber pasado de tener un colador de cocina para buscar diamantes, a tener un escáner láser de alta velocidad que detecta cada brillo en medio de una tormenta de arena.

En resumen: JetFormer es un cerebro artificial diseñado para ser "el filtro perfecto": lo suficientemente listo para no perderse nada importante y lo suficientemente rápido para no perder ni un solo segundo.

Resumen técnico

Resumen Técnico: JetFormer

1. El Problema (Contexto y Desafío)

En la física de altas energías, específicamente en el Gran Colisionador de Hadrones (LHC), la identificación de "jets" (chorros de partículas) es crucial para distinguir señales de partículas pesadas (como el bosón de Higgs) del ruido de fondo. Actualmente, existen dos regímenes de trabajo con necesidades opuestas:

• Análisis Offline: Requiere la máxima precisión posible, utilizando modelos de aprendizaje profundo complejos (como Transformers) que son computacionalmente costosos.
• Disparadores Online (Triggers L1): Requieren una latencia ultra baja (sub-microsegundo) para decidir en tiempo real qué datos guardar, operando en hardware limitado como FPGAs.

El desafío reside en que los modelos Transformer actuales son demasiado grandes y complejos para ser implementados en FPGAs, y los modelos ligeros actuales (como MLPs o GNNs simples) carecen de la precisión necesaria.

2. Metodología

Los autores proponen JetFormer, una arquitectura de Transformer basada únicamente en un codificador (encoder-only) diseñada para ser versátil y escalable.

Innovaciones Arquitectónicas:

• Invariancia de Permutación: A diferencia de los Transformers estándar, no utiliza codificación posicional para tratar las partículas como un conjunto desordenado.
• Token de Clase [CLS]: Utiliza un token aprendible para resumir la información global del jet para la clasificación.
• Optimización para Hardware: Sustituyen la Layer Normalization por Batch Normalization (cuyos parámetros son fijos durante la inferencia) y cambian la función de activación SiLU por ReLU para reducir la complejidad computacional en el hardware.

Pipeline de Optimización y Compresión:
Para lograr la viabilidad en FPGAs, implementan un flujo de trabajo de tres etapas:

1. Optimización de Hiperparámetros (HPO): Utilizan el framework Optuna con una estrategia de optimización multiobjetivo para maximizar la precisión y minimizar los FLOPs simultáneamente.
2. Poda Estructurada (Pruning): Eliminan canales y filtros completos utilizando la importancia de Taylor para reducir el costo computacional sin perder precisión.
3. Cuantización de 1 bit: Implementan una técnica de entrenamiento consciente de la cuantización (QAT) para convertir los pesos a valores de $\pm 1$, reduciendo drásticamente el tamaño del modelo.

3. Contribuciones Clave

• Arquitectura Unificada: Un único marco que escala desde variantes diminutas para triggers hasta modelos de gran escala para análisis offline.
• Extensión de Allo: Ampliaron el framework de compilación Allo (basado en MLIR) para soportar operaciones esenciales de Transformers (como log_softmax y operaciones de tokens de clase) en síntesis de hardware (HLS).
• Pipeline de Compresión de Extremo a Extremo: Un método sistemático que integra búsqueda de hiperparámetros, poda y cuantización agresiva.

4. Resultados Principales

• Rendimiento en Offline (Dataset JETCLASS): El modelo de gran escala alcanza una precisión de 0.829, compitiendo directamente con el estado del arte (ParT) con solo un 37.4% menos de FLOPs.
• Eficiencia en Benchmarks (HLS4ML 150P): JetFormer supera a modelos como MLP, Deep Sets e Interaction Networks por un margen de 3-4% en precisión.
• Compresión Extrema:
  • La poda redujo los FLOPs en un 50% con una pérdida de precisión mínima (<0.5%).
  • La cuantización de 1 bit logró reducir el tamaño del modelo entre un 82% y 92%, con una caída de precisión de apenas 1.5-3.5%.
• Implementación en Hardware: Se validó la síntesis de la variante JetFormer-tiny en FPGAs, demostrando que es posible ejecutar inferencia de Transformers en entornos de latencia extremadamente baja.

5. Significado e Impacto

El trabajo de JetFormer es significativo porque cierra la brecha entre la investigación de vanguardia en Deep Learning y la implementación práctica en la física experimental. Al demostrar que un Transformer puede ser lo suficientemente eficiente para vivir en un chip FPGA sin sacrificar la capacidad de descubrimiento científico, sientan las bases para la próxima generación de sistemas de adquisición de datos en colisionadores de partículas, permitiendo una selección de eventos mucho más inteligente y precisa en tiempo real.