Economical Jet Taggers -- Equivariant, Slim, and Quantized

Este artículo presenta una versión del l-GATr jet tagger esbelta, cuantizada y con reducción de parámetros que logra una reducción de un orden de magnitud en el costo energético con solo una disminución moderada en el rendimiento, allanando el camino para un etiquetado de jets eficiente a nivel de trigger en el LHC.

Lo esencial

Imagina el Gran Colisionador de Hadrones (LHC) como una enorme y rápida fábrica de partículas. Cada segundo, estrella protones entre sí, creando una caótica lluvia de escombros. Los físicos necesitan clasificar estos escombros para encontrar partículas específicas y raras (como el "quark cima") ocultas entre miles de millones de otras ordinarias. Este proceso de clasificación se llama etiquetado de jets (jet tagging).

Durante años, los científicos han utilizado complejos programas informáticos (Aprendizaje Automático) para realizar esta clasificación. Los actuales campeones son los "Transformers": modelos de IA potentes que son increíblemente precisos pero también enormes, lentos y hambrientos de energía. Son como una flota de camiones masivos y voraces en combustible que intentan entregar una sola carta; hacen el trabajo, pero son demasiado grandes y costosos para usarse en el preciso momento en que se están recolectando los datos (el nivel del "disparador" o trigger).

Este artículo plantea una pregunta sencilla: ¿Podemos encoger estos gigantescos camiones hasta convertirlos en diminutos y eficientes motocarros sin perder la capacidad de entregar la carta?

Así es como los autores lo hicieron, utilizando tres estrategias principales:

1. La versión "Delgada" (L-GATr-slim)

El modelo "L-GATr" original es como una navaja suiza que transporta todas las herramientas posibles: escalares, vectores, tensores y más. Sin embargo, los autores se dieron cuenta de que, para la mayoría de los trabajos de física de partículas, solo necesitas dos herramientas: escalares (números) y vectores (flechas con dirección).

• La analogía: Imagina a un chef que insiste en usar una cocina industrial completa con hornos, licuadoras y batidoras solo para hacer un sándwich sencillo. Los autores dijeron: "Vamos a usar solo un cuchillo y una tabla de cortar".
• El resultado: Construyeron una versión "Slim" (delgada) de la IA que elimina las herramientas innecesarias. Funciona tan bien como la versión gigante, pero es mucho más rápida de entrenar y utiliza menos memoria. Es como cambiar un camión de carga pesada por un coche deportivo ágil que realiza el mismo trabajo.

2. La versión "Diminuta" (Etiquetadores Ultra-mini)

Los autores luego se preguntaron: "¿Qué tan pequeño podemos llegar a ser?". Intentaron encoger estos modelos de IA al tamaño de un pequeño coche de juguete (alrededor de 1.000 parámetros, comparado con los millones del original).

• La analogía: Piensa en intentar meter toda la información de una biblioteca en una sola tarjeta postal. Normalmente, pierdes la historia. Pero los autores descubrieron que si organizas la información correctamente (usando reglas específicas "Lorentz-equivariantes" que respetan las leyes de la física), puedes ajustar el conocimiento esencial en un espacio diminuto.
• El resultado: Descubrieron que para modelos muy pequeños, la arquitectura "LLoCa" funciona mejor si se reduce el número de capas, mientras que el "L-GATr-slim" funciona mejor si se reduce el ancho de las capas. Incluso a este tamaño microscópico, superaron a modelos de IA más antiguos que no son conscientes de la física.

3. La versión "Cuantizada" (Matemáticas de baja precisión)

Este es el ahorro de energía más drástico. La IA estándar utiliza matemáticas muy precisas (como medir una distancia hasta la milmillonésima de milímetro). Los autores se dieron cuenta de que, para el etiquetado de jets, no se necesita tanta precisión. Se puede prescindir de ella redondeando los números significativamente.

• La analogía: Imagina que estás contando manzanas en un almacén.
  • IA Estándar: Pesas cada una de las manzanas hasta el microgramo. (Preciso, pero toma mucho tiempo y consume mucha energía de la báscula).
  • IA Cuantizada: Simplemente las cuentas en números enteros. (Rápido, consume casi nada de energía y, para el propósito de saber "cuántas manzanas hay", es perfectamente adecuado).
• El método: Utilizaron una técnica llamada PARQ (Cuantización Regularizada de Afín por Tramos). Piensa en esto como una regla de redondeo inteligente que empuja suavemente los números para que sean simples (como 0, 1 o -1) durante el proceso de entrenamiento, en lugar de forzarlos abruptamente.
• El resultado: Al cambiar a estos números más "toscos", redujeron el costo energético de ejecutar la IA en 10 veces (un orden de magnitud). La IA se volvió increíblemente rápida y eficiente energéticamente, con solo una mínima caída en la precisión.

El panorama general

Los autores combinaron estas tres estrategias —adelgazar la arquitectura, miniaturizar el tamaño y cuantizar las matemáticas— para crear "Etiquetadores de Jets Económicos".

• ¿Por qué es esto importante? Actualmente, estos modelos de IA potentes son demasiado grandes para ejecutarse en el hardware que decide en tiempo real qué colisiones conservar y cuáles descartar (el "disparador").
• El objetivo: Al hacer que estos modelos sean pequeños, rápidos y eficientes energéticamente, los autores esperan poder ejecutarlos directamente en el hardware del disparador. Esto permitiría que el LHC utilice la IA para tomar decisiones en fracciones de segundo sobre qué colisiones de partículas guardar, permitiendo potencialmente el descubrimiento de nueva física que antes se perdía porque los datos se descartaban demasiado rápido.

En resumen: Tomaron una IA gigante y hambrienta de energía, le pusieron a dieta, la encogieron y le enseñaron a hacer matemáticas con menos decimales, resultando en un motor diminuto y súper eficiente que aún puede reconocer las partículas más importantes del universo.

Resumen técnico

Resumen Técnico: Marcadores de Jets Económicos – Equivariantes, Esbeltos y Cuantizados

Planteamiento del Problema
El aprendizaje automático (ML) moderno ha transformado la identificación de jets (jet tagging) en el Gran Colisionador de Hadrones (LHC), emergiendo los transformadores con equivariancia de Lorentz como arquitecturas de vanguardia. Sin embargo, los modelos líderes como L-GATr son computacionalamente costosos, requiriendo una memoria y un tiempo de entrenamiento significativos. Mientras que las tendencias de la industria favorecen el escalado de redes y conjuntos de datos, la física del LHC enfrenta restricciones específicas, particularmente respecto a los requisitos de memoria y latencia del hardware de disparo (trigger). La clasificación de jets actual aún no desempeña un papel en el disparo, pero los autores argumentan que debería hacerlo. El desafío central abordado es cómo reducir el tamaño y el costo computacional de los modernos marcadores de jets equivariantes minimizando la degradación del rendimiento, permitiendo potencialmente su despliegue en el nivel de disparo.

Metodología
El artículo propone una estrategia de dos vertientes para optimizar la eficiencia de los recursos: el adelgazamiento arquitectónico y la cuantización numérica.

1. Arquitectura L-GATr-slim:
   Los autores introducen una versión simplificada del transformador con equivariancia de Lorentz (L-GATr). El L-GATr estándar utiliza una representación de álgebra geométrica que involucra escalares, pseudoescalares, vectores, axiales-vectores y tensores antisimétricos de rango dos. Los autores observan que los pseudoescalares, los axiales-vectores y los tensores son innecesarios para la mayoría de las aplicaciones del LHC. En consecuencia, L-GATr-slim restringe la representación latente únicamente a escalares y vectores.

   • Capas Lineales: Se extienden para operar sobre representaciones acopladas de escalar y vector, asegurando que los componentes vectoriales compartan un único coeficiente escalar aprendible para mantener la equivariancia de Lorentz.
   • No linealidad: Adapta la Unidad Lineal Puerta (GLU) aplicando la no linealidad al producto interno de dos vectores, multiplicado por el vector de salida.
   • Normalización: Modifica RMSNorm para utilizar el valor absoluto del producto interno de Minkowski para los canales vectoriales.
   • Atención: Construye matrices de atención escalar utilizando una formulación específica que evita el costoso producto exterior utilizado en el L-GATr completo.
   • Implementación: La arquitectura está diseñada para ser compilada con torch.compile para mayor eficiencia.

2. Estrategias de Cuantización:
   Los autores aplican tipos de datos de baja precisión y cuantización de pesos para reducir aún más los costos.

   • Cuantización de Tipos de Datos: Las entradas a las capas lineales se cuantizan a int8 (usando cuantización de punto cero) manteniendo bfloat16 para operaciones sensibles a la precisión y para el paso hacia atrás (backward pass). Esto se aplica a las capas ocultas de Transformer, ParT, L-GATr-slim y LLoCa-Transformer.
   • Cuantización de Pesos: Los pesos lineales se cuantizan a valores binarios o ternarios utilizando Cuantización de Gradiente Proximal (PARQ). Este método trata la cuantización como una restricción de regularización, utilizando un operador proximal para actualizar los pesos. Los autores comparan PARQ contra la Estimación de Paso Directo (STE), encontrando que PARQ ofrece mejor estabilidad y rendimiento.
   • Preservación de la Equivariancia: Se toma especial cuidado para asegurar que la cuantización no viole la equivariancia de Lorentz. Para LLoCa, la ortonormalización y las proyecciones de marco permanecen en precisión completa (float32), limitando las operaciones de baja precisión a los invariantes de Lorentz. Para L-GATr-slim, los vectores completos se multiplican por pesos cuantizados, lo que no introduce violaciones adicionales de simetría.

3. Escalado Ultra-Mini:
   Los autores investigan el rendimiento de estas arquitecturas hasta los 1,000 parámetros mediante la reducción del número de bloques o el ancho (canales) de la red.

Resultados Clave
El estudio evalúa las metodologías propuestas en tres tareas: top tagging, regresión de amplitud y generación de eventos.

• Rendimiento vs. Eficiencia (L-GATr-slim):

  • En el conjunto de datos JetClass (clasificación de jets multiclase), L-GATr-slim iguala el rendimiento del L-GATr completo y del LLoCa-Transformer (AUC ~0.9885), pero reduce el tiempo de entrenamiento por un factor de seis (de 166h a 27h en una GPU H100) y el consumo de memoria por un factor de dos.
  • En regresión de amplitud ($Z + 4g$), L-GATr-slim logra el mismo Error Cuadrático Medio (MSE) que el L-GATr completo pero requiere 20 veces menos operaciones de entrenamiento y la mitad del tiempo de entrenamiento.
  • En generación de eventos ($t\bar{t} + nj$), la arquitectura esbelta iguala el rendimiento de log-verosimilitud negativa de los modelos completos.

• Marcadores Ultra-Mini:

  • Al reducir el número de bloques (profundidad), el LLoCa-Transformer supera al L-GATr-slim en tamaños muy pequeños (por ejemplo, 1,000 parámetros).
  • Al mantener fijo el número de bloques (10) y reducir los canales (ancho), L-GATr-slim mantiene una tasa de rechazo de fondo por encima de 1,000 con solo 2 canales vectoriales y 4 escalares, superando a otras arquitecturas de 1,000 parámetros.

• Ganancias de Cuantización:

  • Cuantizar las entradas a int8 y los pesos a valores ternarios reduce el consumo de energía aproximadamente en un orden de magnitud (factor de 10) con solo una pérdida marginal de rendimiento.
  • El LLoCa-Transformer y el L-GATr-slim son robustos a la cuantización, manteniendo un alto rendimiento donde los transformadores estándar podrían degradarse más significativamente.
  • Para el escenario con mayores restricciones de recursos (1 bloque, espacio latente de 16 dimensiones, int8), el LLoCa-Transformer cuantizado (canonicidad global) retiene un rendimiento superior al de los marcadores de grafos previos, a pesar de una reducción de factor de dos en el rechazo de fondo respecto a su contraparte de tamaño completo.

Significancia y Reivindicaciones
El artículo afirma que estas versiones "económicas" de los transformadores equivariantes representan un camino viable hacia la identificación de jets a nivel de disparo (trigger-level) en la Alta Luminosidad del LHC (HL-LHC). Al combinar el adelgazamiento arquitectónico (eliminación de componentes innecesarios del álgebra geométrica) y la cuantización agresiva (PARQ e int8), los autores demuestran que es posible crear marcadores con ~1,000 parámetros que conserven los beneficios motivados por la física de la equivariancia de Lorentz.

Los autores enfatizan que, mientras que el escalado es el estándar de la industria, la física del LHC requiere un enfoque de "reducción de escala consciente de la física" (physics-aware downscaling). Los resultados sugieren que las redes pequeñas, cuantizadas y equivariantes pueden desplegarse en hardware con recursos limitados (como FPGAs) sin sacrificar las simetrías fundamentales que hacen efectivas a estas modelos, abriendo potencialmente nuevas vías para el análisis en tiempo real de la subestructura de jets.